From 98c19ed0fa0d3111bb13093296fffe2ca3c0255e Mon Sep 17 00:00:00 2001 From: sarailQAQ <57243380+sarailQAQ@users.noreply.github.com> Date: Tue, 4 Jun 2024 21:11:41 +0800 Subject: [PATCH] feat: implement drop multiple tables (#4085) * feat: implement drop multiple tables * fix: pass fmt and clippy checks * add: drop multiple sqlness test * update: accept review suggestions * update: accept reviem suggestion Co-authored-by: Weny Xu * fix: pass clippy check --------- Co-authored-by: Weny Xu --- src/frontend/src/instance.rs | 4 +- src/operator/src/statement.rs | 15 +++-- src/operator/src/statement/ddl.rs | 59 ++++++++++++------- src/sql/src/parsers/drop_parser.rs | 51 ++++++++++------ src/sql/src/statements/drop.rs | 42 ++++++++++--- .../standalone/common/drop/drop_table.result | 55 +++++++++++++++++ .../standalone/common/drop/drop_table.sql | 35 +++++++++++ 7 files changed, 207 insertions(+), 54 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index c04770313a..4a1dcfbf29 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -543,7 +543,9 @@ pub fn check_permission( validate_param(&stmt.source_name, query_ctx)?; } Statement::DropTable(drop_stmt) => { - validate_param(drop_stmt.table_name(), query_ctx)?; + for table_name in drop_stmt.table_names() { + validate_param(table_name, query_ctx)?; + } } Statement::ShowTables(stmt) => { validate_db_permission!(stmt, query_ctx); diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 649af286a4..52c1d3b7b7 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -185,12 +185,15 @@ impl StatementExecutor { } Statement::Alter(alter_table) => self.alter_table(alter_table, query_ctx).await, Statement::DropTable(stmt) => { - let (catalog, schema, table) = - table_idents_to_full_name(stmt.table_name(), &query_ctx) - .map_err(BoxedError::new) - .context(error::ExternalSnafu)?; - let table_name = TableName::new(catalog, schema, table); - self.drop_table(table_name, stmt.drop_if_exists(), query_ctx) + let mut table_names = Vec::with_capacity(stmt.table_names().len()); + for table_name_stmt in stmt.table_names() { + let (catalog, schema, table) = + table_idents_to_full_name(table_name_stmt, &query_ctx) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + table_names.push(TableName::new(catalog, schema, table)); + } + self.drop_tables(&table_names[..], stmt.drop_if_exists(), query_ctx.clone()) .await } Statement::DropDatabase(stmt) => { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 2cce8f38d7..251c5529bf 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -643,18 +643,44 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - if let Some(table) = self - .catalog_manager - .table( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - ) + // Reserved for grpc call + self.drop_tables(&[table_name], drop_if_exists, query_context) .await - .context(CatalogSnafu)? - { - let table_id = table.table_info().table_id(); - self.drop_table_procedure(&table_name, table_id, drop_if_exists, query_context) + } + + #[tracing::instrument(skip_all)] + pub async fn drop_tables( + &self, + table_names: &[TableName], + drop_if_exists: bool, + query_context: QueryContextRef, + ) -> Result { + let mut tables = Vec::with_capacity(table_names.len()); + for table_name in table_names { + if let Some(table) = self + .catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .await + .context(CatalogSnafu)? + { + tables.push(table.table_info().table_id()); + } else if drop_if_exists { + // DROP TABLE IF EXISTS meets table not found - ignored + continue; + } else { + return TableNotFoundSnafu { + table_name: table_name.to_string(), + } + .fail(); + } + } + + for (table_name, table_id) in table_names.iter().zip(tables.into_iter()) { + self.drop_table_procedure(table_name, table_id, drop_if_exists, query_context.clone()) .await?; // Invalidates local cache ASAP. @@ -668,17 +694,8 @@ impl StatementExecutor { ) .await .context(error::InvalidateTableCacheSnafu)?; - - 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)) - } else { - TableNotFoundSnafu { - table_name: table_name.to_string(), - } - .fail() } + Ok(Output::new_with_affected_rows(0)) } #[tracing::instrument(skip_all)] diff --git a/src/sql/src/parsers/drop_parser.rs b/src/sql/src/parsers/drop_parser.rs index 6730e3887e..a59f0d9d71 100644 --- a/src/sql/src/parsers/drop_parser.rs +++ b/src/sql/src/parsers/drop_parser.rs @@ -68,22 +68,29 @@ impl<'a> ParserContext<'a> { let _ = self.parser.next_token(); let if_exists = self.parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]); - let raw_table_ident = - self.parse_object_name() - .with_context(|_| error::UnexpectedSnafu { - sql: self.sql, - expected: "a table name", - actual: self.peek_token_as_string(), - })?; - let table_ident = Self::canonicalize_object_name(raw_table_ident); - ensure!( - !table_ident.0.is_empty(), - InvalidTableNameSnafu { - name: table_ident.to_string() + let mut table_names = Vec::with_capacity(1); + loop { + let raw_table_ident = + self.parse_object_name() + .with_context(|_| error::UnexpectedSnafu { + sql: self.sql, + expected: "a table name", + actual: self.peek_token_as_string(), + })?; + let table_ident = Self::canonicalize_object_name(raw_table_ident); + ensure!( + !table_ident.0.is_empty(), + InvalidTableNameSnafu { + name: table_ident.to_string() + } + ); + table_names.push(table_ident); + if !self.parser.consume_token(&Token::Comma) { + break; } - ); + } - Ok(Statement::DropTable(DropTable::new(table_ident, if_exists))) + Ok(Statement::DropTable(DropTable::new(table_names, if_exists))) } fn parse_drop_database(&mut self) -> Result { @@ -122,7 +129,10 @@ mod tests { let mut stmts = result.unwrap(); assert_eq!( stmts.pop().unwrap(), - Statement::DropTable(DropTable::new(ObjectName(vec![Ident::new("foo")]), false)) + Statement::DropTable(DropTable::new( + vec![ObjectName(vec![Ident::new("foo")])], + false + )) ); let sql = "DROP TABLE IF EXISTS foo"; @@ -131,7 +141,10 @@ mod tests { let mut stmts = result.unwrap(); assert_eq!( stmts.pop().unwrap(), - Statement::DropTable(DropTable::new(ObjectName(vec![Ident::new("foo")]), true)) + Statement::DropTable(DropTable::new( + vec![ObjectName(vec![Ident::new("foo")])], + true + )) ); let sql = "DROP TABLE my_schema.foo"; @@ -141,7 +154,7 @@ mod tests { assert_eq!( stmts.pop().unwrap(), Statement::DropTable(DropTable::new( - ObjectName(vec![Ident::new("my_schema"), Ident::new("foo")]), + vec![ObjectName(vec![Ident::new("my_schema"), Ident::new("foo")])], false )) ); @@ -153,11 +166,11 @@ mod tests { assert_eq!( stmts.pop().unwrap(), Statement::DropTable(DropTable::new( - ObjectName(vec![ + vec![ObjectName(vec![ Ident::new("my_catalog"), Ident::new("my_schema"), Ident::new("foo") - ]), + ])], false )) ) diff --git a/src/sql/src/statements/drop.rs b/src/sql/src/statements/drop.rs index 304e52ee52..2048e2c74b 100644 --- a/src/sql/src/statements/drop.rs +++ b/src/sql/src/statements/drop.rs @@ -20,22 +20,23 @@ use sqlparser_derive::{Visit, VisitMut}; /// DROP TABLE statement. #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)] pub struct DropTable { - table_name: ObjectName, + table_names: Vec, + /// drop table if exists drop_if_exists: bool, } impl DropTable { /// Creates a statement for `DROP TABLE` - pub fn new(table_name: ObjectName, if_exists: bool) -> Self { + pub fn new(table_names: Vec, if_exists: bool) -> Self { Self { - table_name, + table_names, drop_if_exists: if_exists, } } - pub fn table_name(&self) -> &ObjectName { - &self.table_name + pub fn table_names(&self) -> &[ObjectName] { + &self.table_names } pub fn drop_if_exists(&self) -> bool { @@ -49,8 +50,14 @@ impl Display for DropTable { if self.drop_if_exists() { f.write_str(" IF EXISTS")?; } - let table_name = self.table_name(); - write!(f, r#" {table_name}"#) + let table_names = self.table_names(); + for (i, table_name) in table_names.iter().enumerate() { + if i > 0 { + f.write_str(",")?; + } + write!(f, " {}", table_name)?; + } + Ok(()) } } @@ -206,6 +213,27 @@ DROP TABLE test"#, } } + let sql = r"drop table test1, test2;"; + let stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, stmts.len()); + assert_matches!(&stmts[0], Statement::DropTable { .. }); + + match &stmts[0] { + Statement::DropTable(set) => { + let new_sql = format!("\n{}", set); + assert_eq!( + r#" +DROP TABLE test1, test2"#, + &new_sql + ); + } + _ => { + unreachable!(); + } + } + let sql = r"drop table if exists test;"; let stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) diff --git a/tests/cases/standalone/common/drop/drop_table.result b/tests/cases/standalone/common/drop/drop_table.result index d554a33a1b..0f1c079938 100644 --- a/tests/cases/standalone/common/drop/drop_table.result +++ b/tests/cases/standalone/common/drop/drop_table.result @@ -20,3 +20,58 @@ DROP TABLE IF EXISTS foo; Affected Rows: 0 +DROP TABLE IF EXISTS foo, bar; + +Affected Rows: 0 + +create table foo ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +Affected Rows: 0 + +DROP TABLE foo, bar; + +Error: 4001(TableNotFound), Table not found: greptime.public.bar + +SHOW TABLES; + ++---------+ +| Tables | ++---------+ +| foo | +| numbers | ++---------+ + +DROP TABLE IF EXISTS foo, bar; + +Affected Rows: 0 + +create table foo ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +Affected Rows: 0 + +create table bar ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +Affected Rows: 0 + +DROP TABLE foo, bar; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/drop/drop_table.sql b/tests/cases/standalone/common/drop/drop_table.sql index 65a092ce40..be671809f9 100644 --- a/tests/cases/standalone/common/drop/drop_table.sql +++ b/tests/cases/standalone/common/drop/drop_table.sql @@ -11,3 +11,38 @@ create table foo ( DROP TABLE IF EXISTS foo; DROP TABLE IF EXISTS foo; + +DROP TABLE IF EXISTS foo, bar; + +create table foo ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +DROP TABLE foo, bar; + +SHOW TABLES; + +DROP TABLE IF EXISTS foo, bar; + +create table foo ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +create table bar ( + host string, + ts timestamp DEFAULT '2024-06-01 00:00:00+00:00', + cpu double default 0, + TIME INDEX (ts), + PRIMARY KEY(host) +) engine=mito; + +DROP TABLE foo, bar; +