feat(ddl): allow drop generated column in table with schema registry (#…

…17689)
risingwavelabs · Jul 23, 2024 · 07bd25b · 07bd25b
1 parent 063db16
commit 07bd25b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 17 deletions.
diff --git a/e2e_test/source_inline/kafka/avro/alter_table.slt b/e2e_test/source_inline/kafka/avro/alter_table.slt
@@ -54,9 +54,9 @@ Caused by these errors (recent errors listed first):
   4: Item not found: Invalid column: bar
 
 
-# TODO: Add support for dropping generated columns for table with schema registry
+# Can't drop non-generated column
 statement error
-ALTER TABLE t DROP COLUMN gen_col;
+ALTER TABLE t DROP COLUMN foo;
 ----
 db error: ERROR: Failed to run the query
 
@@ -65,17 +65,18 @@ Caused by:
 HINT: try `ALTER TABLE .. FORMAT .. ENCODE .. (...)` instead
 
 
-# statement ok
-# ALTER TABLE t DROP COLUMN gen_col;
+# Drop generated column
+statement ok
+ALTER TABLE t DROP COLUMN gen_col;
 
-# # Refresh table schema
-# statement ok
-# ALTER TABLE t REFRESH SCHEMA;
+# Refresh table schema
+statement ok
+ALTER TABLE t REFRESH SCHEMA;
 
-# query ?
-# select * from t
-# ----
-# ABC
+query ?
+select * from t
+----
+ABC
 
 statement ok
 drop table t;
diff --git a/src/frontend/src/handler/alter_table_column.rs b/src/frontend/src/handler/alter_table_column.rs
@@ -131,15 +131,18 @@ pub async fn handle_alter_table_column(
         .clone()
         .map(|source_schema| source_schema.into_v2_with_warning());
 
-    if let Some(source_schema) = &source_schema {
-        if schema_has_schema_registry(source_schema) {
-            return Err(ErrorCode::NotSupported(
+    let fail_if_has_schema_registry = || {
+        if let Some(source_schema) = &source_schema
+            && schema_has_schema_registry(source_schema)
+        {
+            Err(ErrorCode::NotSupported(
                 "alter table with schema registry".to_string(),
                 "try `ALTER TABLE .. FORMAT .. ENCODE .. (...)` instead".to_string(),
-            )
-            .into());
+            ))
+        } else {
+            Ok(())
         }
-    }
+    };
 
     if columns.is_empty() {
         Err(ErrorCode::NotSupported(
@@ -152,6 +155,8 @@ pub async fn handle_alter_table_column(
         AlterTableOperation::AddColumn {
             column_def: new_column,
         } => {
+            fail_if_has_schema_registry()?;
+
             // Duplicated names can actually be checked by `StreamMaterialize`. We do here for
             // better error reporting.
             let new_column_name = new_column.name.real_value();
@@ -189,6 +194,10 @@ pub async fn handle_alter_table_column(
 
             // Check if the column to drop is referenced by any generated columns.
             for column in original_catalog.columns() {
+                if column_name.real_value() == column.name() && !column.is_generated() {
+                    fail_if_has_schema_registry()?;
+                }
+
                 if let Some(expr) = column.generated_expr() {
                     let expr = ExprImpl::from_expr_proto(expr)?;
                     let refs = expr.collect_input_refs(original_catalog.columns().len());