feat(query): Support show virtual columns

databendlabs · Dec 15, 2023 · 2f7f5ae · 2f7f5ae
1 parent 8f4a558
commit 2f7f5ae
Show file tree

Hide file tree

Showing 26 changed files with 472 additions and 84 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/meta/api/src/schema_api_impl.rs b/src/meta/api/src/schema_api_impl.rs
@@ -1308,15 +1308,19 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
                 get_pb_value(self, &req.name_ident).await?;
 
             if old_virtual_column_opt.is_some() {
-                return Err(KVAppError::AppError(AppError::VirtualColumnAlreadyExists(
-                    VirtualColumnAlreadyExists::new(
-                        req.name_ident.table_id,
-                        format!(
-                            "create virtual column with tenant: {} table_id: {}",
-                            req.name_ident.tenant, req.name_ident.table_id
+                if req.if_not_exists {
+                    return Ok(CreateVirtualColumnReply {});
+                } else {
+                    return Err(KVAppError::AppError(AppError::VirtualColumnAlreadyExists(
+                        VirtualColumnAlreadyExists::new(
+                            req.name_ident.table_id,
+                            format!(
+                                "create virtual column with tenant: {} table_id: {}",
+                                req.name_ident.tenant, req.name_ident.table_id
+                            ),
                         ),
-                    ),
-                )));
+                    )));
+                }
             }
             let virtual_column_meta = VirtualColumnMeta {
                 table_id: req.name_ident.table_id,
@@ -1370,7 +1374,16 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
             trials.next().unwrap()?;
 
             let (seq, old_virtual_column_meta) =
-                get_virtual_column_by_id_or_err(self, &req.name_ident, ctx).await?;
+                match get_virtual_column_by_id_or_err(self, &req.name_ident, ctx).await {
+                    Ok((seq, old_virtual_column_meta)) => (seq, old_virtual_column_meta),
+                    Err(err) => {
+                        if req.if_exists {
+                            return Ok(UpdateVirtualColumnReply {});
+                        } else {
+                            return Err(err);
+                        }
+                    }
+                };
 
             let virtual_column_meta = VirtualColumnMeta {
                 table_id: req.name_ident.table_id,
@@ -1423,7 +1436,13 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
         loop {
             trials.next().unwrap()?;
 
-            let (_, _) = get_virtual_column_by_id_or_err(self, &req.name_ident, ctx).await?;
+            if let Err(err) = get_virtual_column_by_id_or_err(self, &req.name_ident, ctx).await {
+                if req.if_exists {
+                    return Ok(DropVirtualColumnReply {});
+                } else {
+                    return Err(err);
+                }
+            }
 
             // Drop virtual column by deleting this record:
             // (tenant, table_id) -> virtual_column_meta

diff --git a/src/meta/api/src/schema_api_test_suite.rs b/src/meta/api/src/schema_api_test_suite.rs
@@ -5651,6 +5651,7 @@ impl SchemaApiTestSuite {
         {
             info!("--- create virtual column");
             let req = CreateVirtualColumnReq {
+                if_not_exists: false,
                 name_ident: name_ident.clone(),
                 virtual_columns: vec!["variant:k1".to_string(), "variant[1]".to_string()],
             };
@@ -5659,6 +5660,7 @@ impl SchemaApiTestSuite {
 
             info!("--- create virtual column again");
             let req = CreateVirtualColumnReq {
+                if_not_exists: false,
                 name_ident: name_ident.clone(),
                 virtual_columns: vec!["variant:k1".to_string(), "variant[1]".to_string()],
             };
@@ -5693,6 +5695,7 @@ impl SchemaApiTestSuite {
         {
             info!("--- update virtual column");
             let req = UpdateVirtualColumnReq {
+                if_exists: false,
                 name_ident: name_ident.clone(),
                 virtual_columns: vec!["variant:k2".to_string(), "variant[2]".to_string()],
             };
@@ -5718,6 +5721,7 @@ impl SchemaApiTestSuite {
         {
             info!("--- drop virtual column");
             let req = DropVirtualColumnReq {
+                if_exists: false,
                 name_ident: name_ident.clone(),
             };
 
@@ -5738,6 +5742,7 @@ impl SchemaApiTestSuite {
         {
             info!("--- update virtual column after drop");
             let req = UpdateVirtualColumnReq {
+                if_exists: false,
                 name_ident: name_ident.clone(),
                 virtual_columns: vec!["variant:k3".to_string(), "variant[3]".to_string()],
             };

diff --git a/src/meta/app/src/schema/virtual_column.rs b/src/meta/app/src/schema/virtual_column.rs
@@ -56,6 +56,7 @@ pub struct VirtualColumnMeta {
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
 pub struct CreateVirtualColumnReq {
+    pub if_not_exists: bool,
     pub name_ident: VirtualColumnNameIdent,
     pub virtual_columns: Vec<String>,
 }
@@ -75,6 +76,7 @@ pub struct CreateVirtualColumnReply {}
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
 pub struct UpdateVirtualColumnReq {
+    pub if_exists: bool,
     pub name_ident: VirtualColumnNameIdent,
     pub virtual_columns: Vec<String>,
 }
@@ -94,6 +96,7 @@ pub struct UpdateVirtualColumnReply {}
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
 pub struct DropVirtualColumnReq {
+    pub if_exists: bool,
     pub name_ident: VirtualColumnNameIdent,
 }
 

diff --git a/src/query/ast/src/ast/format/ast_format.rs b/src/query/ast/src/ast/format/ast_format.rs
@@ -1839,6 +1839,32 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         self.children.push(node);
     }
 
+    fn visit_show_virtual_columns(&mut self, stmt: &'ast ShowVirtualColumnsStmt) {
+        let mut children = Vec::new();
+        if let Some(database) = &stmt.database {
+            let database_name = format!("Database {}", database);
+            let database_format_ctx = AstFormatContext::new(database_name);
+            let database_node = FormatTreeNode::new(database_format_ctx);
+            children.push(database_node);
+        }
+
+        if let Some(table) = &stmt.database {
+            let table_name = format!("Table {}", table);
+            let table_format_ctx = AstFormatContext::new(table_name);
+            let table_node = FormatTreeNode::new(table_format_ctx);
+            children.push(table_node);
+        }
+
+        if let Some(limit) = &stmt.limit {
+            self.visit_show_limit(limit);
+            children.push(self.children.pop().unwrap());
+        }
+        let name = "ShowVirtualColumns".to_string();
+        let format_ctx = AstFormatContext::with_children(name, children.len());
+        let node = FormatTreeNode::with_children(format_ctx, children);
+        self.children.push(node);
+    }
+
     fn visit_show_users(&mut self) {
         let name = "ShowUsers".to_string();
         let format_ctx = AstFormatContext::new(name);

diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs
@@ -157,6 +157,7 @@ pub enum Statement {
     AlterVirtualColumn(AlterVirtualColumnStmt),
     DropVirtualColumn(DropVirtualColumnStmt),
     RefreshVirtualColumn(RefreshVirtualColumnStmt),
+    ShowVirtualColumns(ShowVirtualColumnsStmt),
 
     // User
     ShowUsers,
@@ -469,6 +470,7 @@ impl Display for Statement {
             Statement::AlterVirtualColumn(stmt) => write!(f, "{stmt}")?,
             Statement::DropVirtualColumn(stmt) => write!(f, "{stmt}")?,
             Statement::RefreshVirtualColumn(stmt) => write!(f, "{stmt}")?,
+            Statement::ShowVirtualColumns(stmt) => write!(f, "{stmt}")?,
             Statement::ShowUsers => write!(f, "SHOW USERS")?,
             Statement::ShowRoles => write!(f, "SHOW ROLES")?,
             Statement::CreateUser(stmt) => write!(f, "{stmt}")?,

diff --git a/src/query/ast/src/ast/statements/virtual_column.rs b/src/query/ast/src/ast/statements/virtual_column.rs
@@ -19,9 +19,11 @@ use crate::ast::write_comma_separated_list;
 use crate::ast::write_dot_separated_list;
 use crate::ast::Expr;
 use crate::ast::Identifier;
+use crate::ast::ShowLimit;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct CreateVirtualColumnStmt {
+    pub if_not_exists: bool,
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
@@ -31,7 +33,11 @@ pub struct CreateVirtualColumnStmt {
 
 impl Display for CreateVirtualColumnStmt {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "CREATE VIRTUAL COLUMN (")?;
+        write!(f, "CREATE VIRTUAL COLUMN ")?;
+        if self.if_not_exists {
+            write!(f, "IF NOT EXISTS ")?;
+        }
+        write!(f, "(")?;
         write_comma_separated_list(f, &self.virtual_columns)?;
         write!(f, ") FOR ")?;
         write_dot_separated_list(
@@ -47,6 +53,7 @@ impl Display for CreateVirtualColumnStmt {
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct AlterVirtualColumnStmt {
+    pub if_exists: bool,
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
@@ -56,7 +63,11 @@ pub struct AlterVirtualColumnStmt {
 
 impl Display for AlterVirtualColumnStmt {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ALTER VIRTUAL COLUMN (")?;
+        write!(f, "ALTER VIRTUAL COLUMN ")?;
+        if self.if_exists {
+            write!(f, "IF EXISTS ")?;
+        }
+        write!(f, "(")?;
         write_comma_separated_list(f, &self.virtual_columns)?;
         write!(f, ") FOR ")?;
         write_dot_separated_list(
@@ -72,14 +83,19 @@ impl Display for AlterVirtualColumnStmt {
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct DropVirtualColumnStmt {
+    pub if_exists: bool,
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
 }
 
 impl Display for DropVirtualColumnStmt {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "DROP VIRTUAL COLUMN FOR ")?;
+        write!(f, "DROP VIRTUAL COLUMN ")?;
+        if self.if_exists {
+            write!(f, "IF EXISTS ")?;
+        }
+        write!(f, "FOR ")?;
         write_dot_separated_list(
             f,
             self.catalog
@@ -112,3 +128,33 @@ impl Display for RefreshVirtualColumnStmt {
         Ok(())
     }
 }
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct ShowVirtualColumnsStmt {
+    pub catalog: Option<Identifier>,
+    pub database: Option<Identifier>,
+    pub table: Option<Identifier>,
+    pub limit: Option<ShowLimit>,
+}
+
+impl Display for ShowVirtualColumnsStmt {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "SHOW VIRTUAL COLUMNS")?;
+        if let Some(table) = &self.table {
+            write!(f, "FROM {}", table)?;
+        }
+        if let Some(database) = &self.database {
+            write!(f, " FROM ")?;
+            if let Some(catalog) = &self.catalog {
+                write!(f, "{catalog}.",)?;
+            }
+            write!(f, "{database}")?;
+        }
+
+        if let Some(limit) = &self.limit {
+            write!(f, " {limit}")?;
+        }
+
+        Ok(())
+    }
+}
diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -958,10 +958,11 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
 
     let create_virtual_column = map(
         rule! {
-            CREATE ~ VIRTUAL ~ COLUMN ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" ~ FOR ~ #dot_separated_idents_1_to_3
+            CREATE ~ VIRTUAL ~ COLUMN ~ ( IF ~ ^NOT ~ ^EXISTS )? ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" ~ FOR ~ #dot_separated_idents_1_to_3
         },
-        |(_, _, _, _, virtual_columns, _, _, (catalog, database, table))| {
+        |(_, _, _, opt_if_not_exists, _, virtual_columns, _, _, (catalog, database, table))| {
             Statement::CreateVirtualColumn(CreateVirtualColumnStmt {
+                if_not_exists: opt_if_not_exists.is_some(),
                 catalog,
                 database,
                 table,
@@ -972,10 +973,11 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
 
     let alter_virtual_column = map(
         rule! {
-            ALTER ~ VIRTUAL ~ COLUMN ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" ~ FOR ~ #dot_separated_idents_1_to_3
+            ALTER ~ VIRTUAL ~ COLUMN ~ ( IF ~ ^EXISTS )? ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" ~ FOR ~ #dot_separated_idents_1_to_3
         },
-        |(_, _, _, _, virtual_columns, _, _, (catalog, database, table))| {
+        |(_, _, _, opt_if_exists, _, virtual_columns, _, _, (catalog, database, table))| {
             Statement::AlterVirtualColumn(AlterVirtualColumnStmt {
+                if_exists: opt_if_exists.is_some(),
                 catalog,
                 database,
                 table,
@@ -986,10 +988,11 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
 
     let drop_virtual_column = map(
         rule! {
-            DROP ~ VIRTUAL ~ COLUMN ~ FOR ~ #dot_separated_idents_1_to_3
+            DROP ~ VIRTUAL ~ COLUMN ~ ( IF ~ ^EXISTS )? ~ FOR ~ #dot_separated_idents_1_to_3
         },
-        |(_, _, _, _, (catalog, database, table))| {
+        |(_, _, _, opt_if_exists, _, (catalog, database, table))| {
             Statement::DropVirtualColumn(DropVirtualColumnStmt {
+                if_exists: opt_if_exists.is_some(),
                 catalog,
                 database,
                 table,
@@ -1010,6 +1013,26 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
         },
     );
 
+    let show_virtual_columns = map(
+        rule! {
+            SHOW ~ VIRTUAL ~ COLUMNS ~ (( FROM | IN ) ~ #ident)? ~ (( FROM | IN ) ~ ^#dot_separated_idents_1_to_2)? ~ #show_limit?
+        },
+        |(_, _, _, opt_table, opt_db, limit)| {
+            let table = opt_table.map(|(_, table)| table);
+            let (catalog, database) = match opt_db {
+                Some((_, (Some(c), d))) => (Some(c), Some(d)),
+                Some((_, (None, d))) => (None, Some(d)),
+                _ => (None, None),
+            };
+            Statement::ShowVirtualColumns(ShowVirtualColumnsStmt {
+                catalog,
+                database,
+                table,
+                limit,
+            })
+        },
+    );
+
     let show_users = value(Statement::ShowUsers, rule! { SHOW ~ USERS });
     let create_user = map(
         rule! {
@@ -1739,6 +1762,7 @@ pub fn statement(i: Input) -> IResult<StatementWithFormat> {
             | #alter_virtual_column: "`ALTER VIRTUAL COLUMN (expr, ...) FOR [<database>.]<table>`"
             | #drop_virtual_column: "`DROP VIRTUAL COLUMN FOR [<database>.]<table>`"
             | #refresh_virtual_column: "`REFRESH VIRTUAL COLUMN FOR [<database>.]<table>`"
+            | #show_virtual_columns : "`SHOW VIRTUAL COLUMNS FROM <table> [FROM|IN <catalog>.<database>] [<show_limit>]`"
         ),
         rule!(
             #show_users : "`SHOW USERS`"

diff --git a/src/query/ast/src/visitors/visitor.rs b/src/query/ast/src/visitors/visitor.rs
@@ -578,6 +578,8 @@ pub trait Visitor<'ast>: Sized {
 
     fn visit_refresh_virtual_column(&mut self, _stmt: &'ast RefreshVirtualColumnStmt) {}
 
+    fn visit_show_virtual_columns(&mut self, _stmt: &'ast ShowVirtualColumnsStmt) {}
+
     fn visit_show_users(&mut self) {}
 
     fn visit_create_user(&mut self, _stmt: &'ast CreateUserStmt) {}

diff --git a/src/query/ast/src/visitors/visitor_mut.rs b/src/query/ast/src/visitors/visitor_mut.rs
@@ -592,6 +592,8 @@ pub trait VisitorMut: Sized {
 
     fn visit_refresh_virtual_column(&mut self, _stmt: &mut RefreshVirtualColumnStmt) {}
 
+    fn visit_show_virtual_columns(&mut self, _stmt: &mut ShowVirtualColumnsStmt) {}
+
     fn visit_show_users(&mut self) {}
 
     fn visit_create_user(&mut self, _stmt: &mut CreateUserStmt) {}

diff --git a/src/query/ast/src/visitors/walk.rs b/src/query/ast/src/visitors/walk.rs
@@ -436,6 +436,7 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
         Statement::AlterVirtualColumn(stmt) => visitor.visit_alter_virtual_column(stmt),
         Statement::DropVirtualColumn(stmt) => visitor.visit_drop_virtual_column(stmt),
         Statement::RefreshVirtualColumn(stmt) => visitor.visit_refresh_virtual_column(stmt),
+        Statement::ShowVirtualColumns(stmt) => visitor.visit_show_virtual_columns(stmt),
         Statement::ShowUsers => visitor.visit_show_users(),
         Statement::ShowRoles => visitor.visit_show_roles(),
         Statement::CreateUser(stmt) => visitor.visit_create_user(stmt),

diff --git a/src/query/ast/src/visitors/walk_mut.rs b/src/query/ast/src/visitors/walk_mut.rs
@@ -441,6 +441,7 @@ pub fn walk_statement_mut<V: VisitorMut>(visitor: &mut V, statement: &mut Statem
         Statement::AlterVirtualColumn(stmt) => visitor.visit_alter_virtual_column(stmt),
         Statement::DropVirtualColumn(stmt) => visitor.visit_drop_virtual_column(stmt),
         Statement::RefreshVirtualColumn(stmt) => visitor.visit_refresh_virtual_column(stmt),
+        Statement::ShowVirtualColumns(stmt) => visitor.visit_show_virtual_columns(stmt),
         Statement::ShowUsers => visitor.visit_show_users(),
         Statement::ShowRoles => visitor.visit_show_roles(),
         Statement::CreateUser(stmt) => visitor.visit_create_user(stmt),