From af486ec0d0140675cafe3010ad51395ed511f445 Mon Sep 17 00:00:00 2001
From: Eugene Tolbakov <ev.tolbakov@gmail.com>
Date: Mon, 27 May 2024 04:31:58 +0100
Subject: [PATCH] feat(opertor): check if a database is in use before dropping
 it (#4035)

feat(opertor): check if database is in use before dropping it
---
 src/operator/src/error.rs                        |  8 ++++++++
 src/operator/src/statement/ddl.rs                | 16 ++++++++++------
 .../standalone/common/catalog/schema.result      |  8 ++++++++
 tests/cases/standalone/common/catalog/schema.sql |  4 ++++
 .../common/create/upper_case_table_name.result   |  6 +++++-
 .../common/create/upper_case_table_name.sql      |  2 ++
 6 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs
index 132c3fc36f..6e77b53d3e 100644
--- a/src/operator/src/error.rs
+++ b/src/operator/src/error.rs
@@ -256,6 +256,13 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Schema `{name}` is in use"))]
+    SchemaInUse {
+        name: String,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Table occurs error"))]
     Table {
         #[snafu(implicit)]
@@ -727,6 +734,7 @@ impl ErrorExt for Error {
             | Error::IllegalPrimaryKeysDef { .. }
             | Error::SchemaNotFound { .. }
             | Error::SchemaExists { .. }
+            | Error::SchemaInUse { .. }
             | Error::ColumnNotFound { .. }
             | Error::BuildRegex { .. }
             | Error::InvalidSchema { .. }
diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs
index 28a674fa8f..67c4a4251b 100644
--- a/src/operator/src/statement/ddl.rs
+++ b/src/operator/src/statement/ddl.rs
@@ -68,9 +68,9 @@ use crate::error::{
     CreateLogicalTablesSnafu, CreateTableInfoSnafu, DdlWithMultiCatalogsSnafu,
     DdlWithMultiSchemasSnafu, DeserializePartitionSnafu, EmptyDdlExprSnafu, FlowNotFoundSnafu,
     InvalidPartitionColumnsSnafu, InvalidPartitionRuleSnafu, InvalidTableNameSnafu,
-    InvalidViewNameSnafu, InvalidViewStmtSnafu, ParseSqlValueSnafu, Result, SchemaNotFoundSnafu,
-    SubstraitCodecSnafu, TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
-    UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
+    InvalidViewNameSnafu, InvalidViewStmtSnafu, ParseSqlValueSnafu, Result, SchemaInUseSnafu,
+    SchemaNotFoundSnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu, TableMetadataManagerSnafu,
+    TableNotFoundSnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
 };
 use crate::expr_factory;
 use crate::statement::show::create_partitions_stmt;
@@ -667,10 +667,14 @@ impl StatementExecutor {
             .await
             .context(CatalogSnafu)?
         {
-            self.drop_database_procedure(catalog, schema, drop_if_exists, query_context)
-                .await?;
+            if schema == query_context.current_schema() {
+                SchemaInUseSnafu { name: schema }.fail()
+            } else {
+                self.drop_database_procedure(catalog, schema, drop_if_exists, query_context)
+                    .await?;
 
-            Ok(Output::new_with_affected_rows(0))
+                Ok(Output::new_with_affected_rows(0))
+            }
         } else if drop_if_exists {
             // DROP TABLE IF EXISTS meets table not found - ignored
             Ok(Output::new_with_affected_rows(0))
diff --git a/tests/cases/standalone/common/catalog/schema.result b/tests/cases/standalone/common/catalog/schema.result
index 600415d312..abdef456c5 100644
--- a/tests/cases/standalone/common/catalog/schema.result
+++ b/tests/cases/standalone/common/catalog/schema.result
@@ -120,6 +120,14 @@ SHOW TABLES FROM public WHERE Tables = 'numbers';
 
 DROP SCHEMA test_public_schema;
 
+Error: 1004(InvalidArguments), Schema `test_public_schema` is in use
+
+USE public;
+
+Affected Rows: 0
+
+DROP SCHEMA test_public_schema;
+
 Affected Rows: 0
 
 SELECT * FROM test_public_schema.hello;
diff --git a/tests/cases/standalone/common/catalog/schema.sql b/tests/cases/standalone/common/catalog/schema.sql
index 240c495ea4..0d1d1f0138 100644
--- a/tests/cases/standalone/common/catalog/schema.sql
+++ b/tests/cases/standalone/common/catalog/schema.sql
@@ -40,6 +40,10 @@ SHOW TABLES FROM public WHERE Tables = 'numbers';
 
 DROP SCHEMA test_public_schema;
 
+USE public;
+
+DROP SCHEMA test_public_schema;
+
 SELECT * FROM test_public_schema.hello;
 
 USE public;
diff --git a/tests/cases/standalone/common/create/upper_case_table_name.result b/tests/cases/standalone/common/create/upper_case_table_name.result
index 33e186a63d..6f46741216 100644
--- a/tests/cases/standalone/common/create/upper_case_table_name.result
+++ b/tests/cases/standalone/common/create/upper_case_table_name.result
@@ -100,9 +100,13 @@ Affected Rows: 0
 
 drop database upper_case_table_name;
 
-Affected Rows: 0
+Error: 1004(InvalidArguments), Schema `upper_case_table_name` is in use
 
 use public;
 
 Affected Rows: 0
 
+drop database upper_case_table_name;
+
+Affected Rows: 0
+
diff --git a/tests/cases/standalone/common/create/upper_case_table_name.sql b/tests/cases/standalone/common/create/upper_case_table_name.sql
index 76625e0159..03ee09ac13 100644
--- a/tests/cases/standalone/common/create/upper_case_table_name.sql
+++ b/tests/cases/standalone/common/create/upper_case_table_name.sql
@@ -43,3 +43,5 @@ drop table abcdefge;
 drop database upper_case_table_name;
 
 use public;
+
+drop database upper_case_table_name;