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 <wenymedia@gmail.com>

* fix: pass clippy check

---------

Co-authored-by: Weny Xu <wenymedia@gmail.com>
This commit is contained in:
sarailQAQ
2024-06-04 21:11:41 +08:00
committed by GitHub
parent c0aed1d267
commit 98c19ed0fa
7 changed files with 207 additions and 54 deletions

View File

@@ -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);

View File

@@ -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) => {

View File

@@ -643,18 +643,44 @@ impl StatementExecutor {
drop_if_exists: bool,
query_context: QueryContextRef,
) -> Result<Output> {
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<Output> {
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)]

View File

@@ -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<Statement> {
@@ -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
))
)

View File

@@ -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<ObjectName>,
/// 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<ObjectName>, 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())

View File

@@ -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

View File

@@ -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;