feat: create database if not exists (#1009)

src/datanode/src/error.rs

@@ -222,7 +222,7 @@ pub enum Error {
         source: catalog::error::Error,
     },
 
-    #[snafu(display("Schema already exists, name: {}", name))]
+    #[snafu(display("Schema {} already exists", name))]
     SchemaExists { name: String, backtrace: Backtrace },
 
     #[snafu(display("Failed to convert alter expr to request: {}", source))]

src/datanode/src/instance/grpc.rs

@@ -30,12 +30,16 @@ use crate::error::{self, DecodeLogicalPlanSnafu, ExecuteSqlSnafu, Result};
 use crate::instance::Instance;
 
 impl Instance {
-    pub(crate) async fn handle_create_database(&self, expr: CreateDatabaseExpr) -> Result<Output> {
+    pub(crate) async fn handle_create_database(
+        &self,
+        expr: CreateDatabaseExpr,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
         let req = CreateDatabaseRequest {
             db_name: expr.database_name,
             create_if_not_exists: expr.create_if_not_exists,
         };
-        self.sql_handler().create_database(req).await
+        self.sql_handler.create_database(req, query_ctx).await
     }
 
     pub(crate) async fn execute_logical(&self, plan_bytes: Vec<u8>) -> Result<Output> {
@@ -83,14 +87,14 @@ impl Instance {
         Ok(Output::AffectedRows(affected_rows))
     }
 
-    async fn handle_ddl(&self, request: DdlRequest) -> Result<Output> {
+    async fn handle_ddl(&self, request: DdlRequest, query_ctx: QueryContextRef) -> Result<Output> {
         let expr = request.expr.context(error::MissingRequiredFieldSnafu {
             name: "DdlRequest.expr",
         })?;
         match expr {
             DdlExpr::CreateTable(expr) => self.handle_create(expr).await,
             DdlExpr::Alter(expr) => self.handle_alter(expr).await,
-            DdlExpr::CreateDatabase(expr) => self.handle_create_database(expr).await,
+            DdlExpr::CreateDatabase(expr) => self.handle_create_database(expr, query_ctx).await,
             DdlExpr::DropTable(expr) => self.handle_drop_table(expr).await,
         }
     }
@@ -111,7 +115,7 @@ impl GrpcQueryHandler for Instance {
                     })?;
                 self.handle_query(query, ctx).await
             }
-            GrpcRequest::Ddl(request) => self.handle_ddl(request).await,
+            GrpcRequest::Ddl(request) => self.handle_ddl(request, ctx).await,
         }
     }
 }

src/datanode/src/sql.rs

@@ -78,7 +78,7 @@ impl SqlHandler {
         let result = match request {
             SqlRequest::Insert(req) => self.insert(req).await,
             SqlRequest::CreateTable(req) => self.create_table(req).await,
-            SqlRequest::CreateDatabase(req) => self.create_database(req).await,
+            SqlRequest::CreateDatabase(req) => self.create_database(req, query_ctx.clone()).await,
             SqlRequest::Alter(req) => self.alter(req).await,
             SqlRequest::DropTable(req) => self.drop_table(req).await,
             SqlRequest::Delete(stmt) => self.delete(query_ctx.clone(), stmt).await,

src/datanode/src/sql/create.rs

@@ -16,11 +16,11 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use catalog::{RegisterSchemaRequest, RegisterTableRequest};
-use common_catalog::consts::DEFAULT_CATALOG_NAME;
 use common_query::Output;
 use common_telemetry::tracing::info;
 use common_telemetry::tracing::log::error;
 use datatypes::schema::SchemaBuilder;
+use session::context::QueryContextRef;
 use snafu::{ensure, OptionExt, ResultExt};
 use sql::ast::{ColumnOption, TableConstraint};
 use sql::statements::column_def_to_schema;
@@ -38,25 +38,35 @@ use crate::error::{
 use crate::sql::SqlHandler;
 
 impl SqlHandler {
-    pub(crate) async fn create_database(&self, req: CreateDatabaseRequest) -> Result<Output> {
+    pub(crate) async fn create_database(
+        &self,
+        req: CreateDatabaseRequest,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        let catalog = query_ctx.current_catalog();
         let schema = req.db_name;
+        if self
+            .catalog_manager
+            .schema(&catalog, &schema)
+            .context(CatalogSnafu)?
+            .is_some()
+        {
+            return if req.create_if_not_exists {
+                Ok(Output::AffectedRows(1))
+            } else {
+                SchemaExistsSnafu { name: schema }.fail()
+            };
+        }
+
         let reg_req = RegisterSchemaRequest {
-            catalog: DEFAULT_CATALOG_NAME.to_string(),
+            catalog,
             schema: schema.clone(),
         };
-        let success = self
-            .catalog_manager
+        self.catalog_manager
             .register_schema(reg_req)
             .await
             .context(RegisterSchemaSnafu)?;
 
-        // FIXME(dennis): looks like register_schema always returns true even
-        // even when the schema already exists.
-        ensure!(
-            success || req.create_if_not_exists,
-            SchemaExistsSnafu { name: schema }
-        );
-
         info!("Successfully created database: {:?}", schema);
         Ok(Output::AffectedRows(1))
     }

src/frontend/src/instance/distributed.rs

@@ -281,7 +281,7 @@ impl DistInstance {
                     database_name: stmt.name.to_string(),
                     create_if_not_exists: stmt.if_not_exists,
                 };
-                Ok(self.handle_create_database(expr).await?)
+                return self.handle_create_database(expr, query_ctx).await;
             }
             Statement::CreateTable(stmt) => {
                 let create_expr = &mut expr_factory::create_to_expr(&stmt, query_ctx)?;
@@ -371,10 +371,30 @@ impl DistInstance {
     }
 
     /// Handles distributed database creation
-    async fn handle_create_database(&self, expr: CreateDatabaseExpr) -> Result<Output> {
+    async fn handle_create_database(
+        &self,
+        expr: CreateDatabaseExpr,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        let catalog = query_ctx.current_catalog();
+        if self
+            .catalog_manager
+            .schema(&catalog, &expr.database_name)
+            .context(CatalogSnafu)?
+            .is_some()
+        {
+            return if expr.create_if_not_exists {
+                Ok(Output::AffectedRows(1))
+            } else {
+                SchemaExistsSnafu {
+                    name: &expr.database_name,
+                }
+                .fail()
+            };
+        }
+
         let key = SchemaKey {
-            // TODO(sunng87): custom catalog
-            catalog_name: DEFAULT_CATALOG_NAME.to_string(),
+            catalog_name: catalog,
             schema_name: expr.database_name,
         };
         let value = SchemaValue {};

src/frontend/src/instance/distributed/grpc.rs

@@ -45,7 +45,7 @@ impl GrpcQueryHandler for DistInstance {
                     err_msg: "Missing 'expr' in DDL request",
                 })?;
                 match expr {
-                    DdlExpr::CreateDatabase(expr) => self.handle_create_database(expr).await,
+                    DdlExpr::CreateDatabase(expr) => self.handle_create_database(expr, ctx).await,
                     DdlExpr::CreateTable(mut expr) => {
                         // TODO(LFC): Support creating distributed table through GRPC interface.
                         // Currently only SQL supports it; how to design the fields in CreateTableExpr?

tests/cases/distributed/catalog/schema.result

@@ -6,6 +6,10 @@ CREATE SCHEMA test_public_schema;
 
 Error: 1004(InvalidArguments), Schema test_public_schema already exists
 
+CREATE SCHEMA IF NOT EXISTS test_public_schema;
+
+Affected Rows: 1
+
 SHOW DATABASES LIKE '%public%';
 
 +--------------------+

tests/cases/distributed/catalog/schema.sql

@@ -2,6 +2,8 @@ CREATE SCHEMA test_public_schema;
 
 CREATE SCHEMA test_public_schema;
 
+CREATE SCHEMA IF NOT EXISTS test_public_schema;
+
 SHOW DATABASES LIKE '%public%';
 
 USE test_public_schema;

tests/cases/standalone/catalog/schema.result

@@ -6,6 +6,10 @@ CREATE SCHEMA test_public_schema;
 
 Error: 1004(InvalidArguments), Schema test_public_schema already exists
 
+CREATE SCHEMA IF NOT EXISTS test_public_schema;
+
+Affected Rows: 1
+
 SHOW DATABASES LIKE '%public%';
 
 +--------------------+

tests/cases/standalone/catalog/schema.sql

@@ -2,6 +2,8 @@ CREATE SCHEMA test_public_schema;
 
 CREATE SCHEMA test_public_schema;
 
+CREATE SCHEMA IF NOT EXISTS test_public_schema;
+
 SHOW DATABASES LIKE '%public%';
 
 USE test_public_schema;