feat: implement the drop database parser (#3521)

* refactor: refactor drop table parser * feat: implement drop database parser * fix: canonicalize name of create database * test: update sqlness result * Update src/operator/src/statement.rs Co-authored-by: Ruihang Xia <waynestxia@gmail.com> --------- Co-authored-by: Ruihang Xia <waynestxia@gmail.com>
GreptimeTeam · Mar 15, 2024 · a52aede · a52aede
1 parent b6fac61
commit a52aede
Show file tree

Hide file tree

Showing 9 changed files with 136 additions and 9 deletions.
diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
@@ -473,7 +473,8 @@ pub fn check_permission(
         // These are executed by query engine, and will be checked there.
         Statement::Query(_) | Statement::Explain(_) | Statement::Tql(_) | Statement::Delete(_) => {}
         // database ops won't be checked
-        Statement::CreateDatabase(_) | Statement::ShowDatabases(_) => {}
+        Statement::CreateDatabase(_) | Statement::ShowDatabases(_) | Statement::DropDatabase(_) => {
+        }
         // show create table and alter are not supported yet
         Statement::ShowCreateTable(_) | Statement::CreateExternalTable(_) | Statement::Alter(_) => {
         }

diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs
@@ -171,6 +171,13 @@ impl StatementExecutor {
                 let table_name = TableName::new(catalog, schema, table);
                 self.drop_table(table_name, stmt.drop_if_exists()).await
             }
+            Statement::DropDatabase(_stmt) => {
+                // TODO(weny): implement the drop database procedure
+                error::NotSupportedSnafu {
+                    feat: "Drop Database",
+                }
+                .fail()
+            }
             Statement::TruncateTable(stmt) => {
                 let (catalog, schema, table) =
                     table_idents_to_full_name(stmt.table_name(), &query_ctx)

diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs
@@ -119,7 +119,7 @@ impl<'a> ParserContext<'a> {
                 expected: "a database name",
                 actual: self.peek_token_as_string(),
             })?;
-
+        let database_name = Self::canonicalize_object_name(database_name);
         Ok(Statement::CreateDatabase(CreateDatabase {
             name: database_name,
             if_not_exists,
@@ -722,7 +722,7 @@ mod tests {
     use common_catalog::consts::FILE_ENGINE;
     use common_error::ext::ErrorExt;
     use sqlparser::ast::ColumnOption::NotNull;
-    use sqlparser::ast::{BinaryOperator, Value};
+    use sqlparser::ast::{BinaryOperator, ObjectName, Value};
 
     use super::*;
     use crate::dialect::GreptimeDbDialect;
@@ -916,6 +916,18 @@ mod tests {
             }
             _ => unreachable!(),
         }
+
+        let sql = "CREATE DATABASE `fOo`";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        let mut stmts = result.unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::CreateDatabase(CreateDatabase::new(
+                ObjectName(vec![Ident::with_quote('`', "fOo"),]),
+                false
+            ))
+        );
     }
 
     #[test]

diff --git a/src/sql/src/parsers/drop_parser.rs b/src/sql/src/parsers/drop_parser.rs
@@ -13,20 +13,29 @@
 // limitations under the License.
 
 use snafu::{ensure, ResultExt};
-use sqlparser::keywords::Keyword;
+use sqlparser::dialect::keywords::Keyword;
+use sqlparser::tokenizer::Token;
 
 use crate::error::{self, InvalidTableNameSnafu, Result};
 use crate::parser::ParserContext;
-use crate::statements::drop::DropTable;
+use crate::statements::drop::{DropDatabase, DropTable};
 use crate::statements::statement::Statement;
 
 /// DROP statement parser implementation
 impl<'a> ParserContext<'a> {
     pub(crate) fn parse_drop(&mut self) -> Result<Statement> {
         let _ = self.parser.next_token();
-        if !self.matches_keyword(Keyword::TABLE) {
-            return self.unsupported(self.peek_token_as_string());
+        match self.parser.peek_token().token {
+            Token::Word(w) => match w.keyword {
+                Keyword::TABLE => self.parse_drop_table(),
+                Keyword::SCHEMA | Keyword::DATABASE => self.parse_drop_database(),
+                _ => self.unsupported(w.to_string()),
+            },
+            unexpected => self.unsupported(unexpected.to_string()),
         }
+    }
+
+    fn parse_drop_table(&mut self) -> Result<Statement> {
         let _ = self.parser.next_token();
 
         let if_exists = self.parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
@@ -48,6 +57,26 @@ impl<'a> ParserContext<'a> {
 
         Ok(Statement::DropTable(DropTable::new(table_ident, if_exists)))
     }
+
+    fn parse_drop_database(&mut self) -> Result<Statement> {
+        let _ = self.parser.next_token();
+
+        let if_exists = self.parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+        let database_name =
+            self.parser
+                .parse_object_name()
+                .with_context(|_| error::UnexpectedSnafu {
+                    sql: self.sql,
+                    expected: "a database name",
+                    actual: self.peek_token_as_string(),
+                })?;
+        let database_name = Self::canonicalize_object_name(database_name);
+
+        Ok(Statement::DropDatabase(DropDatabase::new(
+            database_name,
+            if_exists,
+        )))
+    }
 }
 
 #[cfg(test)]
@@ -106,4 +135,43 @@ mod tests {
             ))
         )
     }
+
+    #[test]
+    pub fn test_drop_database() {
+        let sql = "DROP DATABASE public";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        let mut stmts = result.unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::DropDatabase(DropDatabase::new(
+                ObjectName(vec![Ident::new("public")]),
+                false
+            ))
+        );
+
+        let sql = "DROP DATABASE IF EXISTS public";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        let mut stmts = result.unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::DropDatabase(DropDatabase::new(
+                ObjectName(vec![Ident::new("public")]),
+                true
+            ))
+        );
+
+        let sql = "DROP DATABASE `fOo`";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        let mut stmts = result.unwrap();
+        assert_eq!(
+            stmts.pop().unwrap(),
+            Statement::DropDatabase(DropDatabase::new(
+                ObjectName(vec![Ident::with_quote('`', "fOo"),]),
+                false
+            ))
+        );
+    }
 }
diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs
@@ -206,6 +206,16 @@ pub struct CreateDatabase {
     pub if_not_exists: bool,
 }
 
+impl CreateDatabase {
+    /// Creates a statement for `CREATE DATABASE`
+    pub fn new(name: ObjectName, if_not_exists: bool) -> Self {
+        Self {
+            name,
+            if_not_exists,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut)]
 pub struct CreateExternalTable {
     /// Table name

diff --git a/src/sql/src/statements/drop.rs b/src/sql/src/statements/drop.rs
@@ -40,3 +40,29 @@ impl DropTable {
         self.drop_if_exists
     }
 }
+
+/// DROP DATABASE statement.
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub struct DropDatabase {
+    name: ObjectName,
+    /// drop table if exists
+    drop_if_exists: bool,
+}
+
+impl DropDatabase {
+    /// Creates a statement for `DROP DATABASE`
+    pub fn new(name: ObjectName, if_exists: bool) -> Self {
+        Self {
+            name,
+            drop_if_exists: if_exists,
+        }
+    }
+
+    pub fn name(&self) -> &ObjectName {
+        &self.name
+    }
+
+    pub fn drop_if_exists(&self) -> bool {
+        self.drop_if_exists
+    }
+}
diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs
@@ -16,6 +16,7 @@ use datafusion_sql::parser::Statement as DfStatement;
 use sqlparser::ast::Statement as SpStatement;
 use sqlparser_derive::{Visit, VisitMut};
 
+use super::drop::DropDatabase;
 use super::show::ShowVariables;
 use crate::error::{ConvertToDfStatementSnafu, Error};
 use crate::statements::alter::AlterTable;
@@ -51,6 +52,8 @@ pub enum Statement {
     CreateTableLike(CreateTableLike),
     // DROP TABLE
     DropTable(DropTable),
+    // DROP DATABASE
+    DropDatabase(DropDatabase),
     // CREATE DATABASE
     CreateDatabase(CreateDatabase),
     /// ALTER TABLE

diff --git a/tests/cases/standalone/common/catalog/schema.result b/tests/cases/standalone/common/catalog/schema.result
@@ -120,7 +120,7 @@ SHOW TABLES FROM public WHERE Tables = 'numbers';
 
 DROP SCHEMA test_public_schema;
 
-Error: 1001(Unsupported), SQL statement is not supported: DROP SCHEMA test_public_schema;, keyword: SCHEMA
+Error: 1001(Unsupported), Not supported: Drop Database
 
 SELECT * FROM test_public_schema.hello;
 

diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result
@@ -447,7 +447,7 @@ Affected Rows: 0
 
 drop schema my_db;
 
-Error: 1001(Unsupported), SQL statement is not supported: drop schema my_db;, keyword: schema
+Error: 1001(Unsupported), Not supported: Drop Database
 
 use information_schema;