feat: drop column for alter table (#562)

* feat: drop column for alter table

* refactor: rename RemoveColumns to DropColumns

* test: alter table

* chore: error msg

Co-authored-by: Ruihang Xia <waynestxia@gmail.com>

* fix: test_parse_alter_drop_column

Co-authored-by: Ruihang Xia <waynestxia@gmail.com>

src/api/greptime/v1/admin.proto

@@ -51,6 +51,7 @@ message AlterExpr {
   string table_name = 3;
   oneof kind {
     AddColumns add_columns = 4;
+    DropColumns drop_columns = 5;
   }
 }
 
@@ -58,11 +59,19 @@ message AddColumns {
   repeated AddColumn add_columns = 1;
 }
 
+message DropColumns {
+  repeated DropColumn drop_columns = 1;
+}
+
 message AddColumn {
   ColumnDef column_def = 1;
   bool is_key = 2;
 }
 
+message DropColumn {
+  string name = 1;
+}
+
 message CreateDatabaseExpr {
   //TODO(hl): maybe rename to schema_name?
   string database_name = 1;

src/datanode/src/server/grpc/ddl.rs

@@ -17,7 +17,7 @@ use std::sync::Arc;
 use api::helper::ColumnDataTypeWrapper;
 use api::result::AdminResultBuilder;
 use api::v1::alter_expr::Kind;
-use api::v1::{AdminResult, AlterExpr, ColumnDef, CreateExpr};
+use api::v1::{AdminResult, AlterExpr, ColumnDef, CreateExpr, DropColumns};
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_error::prelude::{ErrorExt, StatusCode};
 use common_query::Output;
@@ -191,6 +191,19 @@ fn alter_expr_to_request(expr: AlterExpr) -> Result<Option<AlterTableRequest>> {
             };
             Ok(Some(request))
         }
+        Some(Kind::DropColumns(DropColumns { drop_columns })) => {
+            let alter_kind = AlterKind::DropColumns {
+                names: drop_columns.into_iter().map(|c| c.name).collect(),
+            };
+
+            let request = AlterTableRequest {
+                catalog_name: expr.catalog_name,
+                schema_name: expr.schema_name,
+                table_name: expr.table_name,
+                alter_kind,
+            };
+            Ok(Some(request))
+        }
         None => Ok(None),
     }
 }

src/datanode/src/sql/alter.rs

@@ -72,6 +72,9 @@ impl SqlHandler {
                     is_key: false,
                 }],
             },
+            AlterTableOperation::DropColumn { name } => AlterKind::DropColumns {
+                names: vec![name.value.clone()],
+            },
         };
         Ok(AlterTableRequest {
             catalog_name: Some(catalog_name),

src/datanode/src/tests/instance_test.rs

@@ -179,8 +179,7 @@ async fn assert_query_result(instance: &Instance, sql: &str, ts: i64, host: &str
     }
 }
 
-#[tokio::test(flavor = "multi_thread")]
-async fn test_execute_insert() {
+async fn setup_test_instance() -> Instance {
     common_telemetry::init_default_ut_logging();
 
     let (opts, _guard) = test_util::create_tmp_dir_and_datanode_opts("execute_insert");
@@ -195,6 +194,12 @@ async fn test_execute_insert() {
     .await
     .unwrap();
 
+    instance
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_execute_insert() {
+    let instance = setup_test_instance().await;
     let output = instance
         .execute_sql(
             r#"insert into demo(host, cpu, memory, ts) values
@@ -479,6 +484,7 @@ async fn test_alter_table() {
         .await
         .unwrap();
 
+    // Add column
     let output = instance
         .execute_sql("alter table demo add my_tag string null")
         .await
@@ -498,7 +504,10 @@ async fn test_alter_table() {
         .unwrap();
     assert!(matches!(output, Output::AffectedRows(1)));
 
-    let output = instance.execute_sql("select * from demo").await.unwrap();
+    let output = instance
+        .execute_sql("select * from demo order by ts")
+        .await
+        .unwrap();
     let expected = vec![
         "+-------+-----+--------+---------------------+--------+",
         "| host  | cpu | memory | ts                  | my_tag |",
@@ -509,6 +518,51 @@ async fn test_alter_table() {
         "+-------+-----+--------+---------------------+--------+",
     ];
     check_output_stream(output, expected).await;
+
+    // Drop a column
+    let output = instance
+        .execute_sql("alter table demo drop column memory")
+        .await
+        .unwrap();
+    assert!(matches!(output, Output::AffectedRows(0)));
+
+    let output = instance
+        .execute_sql("select * from demo order by ts")
+        .await
+        .unwrap();
+    let expected = vec![
+        "+-------+-----+---------------------+--------+",
+        "| host  | cpu | ts                  | my_tag |",
+        "+-------+-----+---------------------+--------+",
+        "| host1 | 1.1 | 1970-01-01 00:00:01 |        |",
+        "| host2 | 2.2 | 1970-01-01 00:00:02 | hello  |",
+        "| host3 | 3.3 | 1970-01-01 00:00:03 |        |",
+        "+-------+-----+---------------------+--------+",
+    ];
+    check_output_stream(output, expected).await;
+
+    // insert a new row
+    let output = instance
+        .execute_sql("insert into demo(host, cpu, ts, my_tag) values ('host4', 400, 4000, 'world')")
+        .await
+        .unwrap();
+    assert!(matches!(output, Output::AffectedRows(1)));
+
+    let output = instance
+        .execute_sql("select * from demo order by ts")
+        .await
+        .unwrap();
+    let expected = vec![
+        "+-------+-----+---------------------+--------+",
+        "| host  | cpu | ts                  | my_tag |",
+        "+-------+-----+---------------------+--------+",
+        "| host1 | 1.1 | 1970-01-01 00:00:01 |        |",
+        "| host2 | 2.2 | 1970-01-01 00:00:02 | hello  |",
+        "| host3 | 3.3 | 1970-01-01 00:00:03 |        |",
+        "| host4 | 400 | 1970-01-01 00:00:04 | world  |",
+        "+-------+-----+---------------------+--------+",
+    ];
+    check_output_stream(output, expected).await;
 }
 
 async fn test_insert_with_default_value_for_type(type_name: &str) {

src/mito/src/engine.rs

@@ -952,7 +952,7 @@ mod tests {
             catalog_name: None,
             schema_name: None,
             table_name: TABLE_NAME.to_string(),
-            alter_kind: AlterKind::RemoveColumns {
+            alter_kind: AlterKind::DropColumns {
                 names: vec![String::from("memory"), String::from("my_field")],
             },
         };

src/mito/src/table.rs

@@ -481,7 +481,7 @@ fn create_alter_operation(
         AlterKind::AddColumns { columns } => {
             create_add_columns_operation(table_name, columns, table_meta)
         }
-        AlterKind::RemoveColumns { names } => Ok(AlterOperation::DropColumns {
+        AlterKind::DropColumns { names } => Ok(AlterOperation::DropColumns {
             names: names.to_vec(),
         }),
     }

src/sql/src/parsers/alter_parser.rs

@@ -41,9 +41,19 @@ impl<'a> ParserContext<'a> {
                 let column_def = parser.parse_column_def()?;
                 AlterTableOperation::AddColumn { column_def }
             }
+        } else if parser.parse_keyword(Keyword::DROP) {
+            if parser.parse_keyword(Keyword::COLUMN) {
+                let name = self.parser.parse_identifier()?;
+                AlterTableOperation::DropColumn { name }
+            } else {
+                return Err(ParserError::ParserError(format!(
+                    "expect keyword COLUMN after ALTER TABLE DROP, found {}",
+                    parser.peek_token()
+                )));
+            }
         } else {
             return Err(ParserError::ParserError(format!(
-                "expect ADD or DROP after ALTER TABLE, found {}",
+                "expect keyword ADD or DROP after ALTER TABLE, found {}",
                 parser.peek_token()
             )));
         };
@@ -89,4 +99,35 @@ mod tests {
             _ => unreachable!(),
         }
     }
+
+    #[test]
+    fn test_parse_alter_drop_column() {
+        let sql = "ALTER TABLE my_metric_1 DROP a";
+        let result = ParserContext::create_with_dialect(sql, &GenericDialect {}).unwrap_err();
+        assert!(result
+            .to_string()
+            .contains("expect keyword COLUMN after ALTER TABLE DROP"));
+
+        let sql = "ALTER TABLE my_metric_1 DROP COLUMN a";
+        let mut result = ParserContext::create_with_dialect(sql, &GenericDialect {}).unwrap();
+        assert_eq!(1, result.len());
+
+        let statement = result.remove(0);
+        assert_matches!(statement, Statement::Alter { .. });
+        match statement {
+            Statement::Alter(alter_table) => {
+                assert_eq!("my_metric_1", alter_table.table_name().0[0].value);
+
+                let alter_operation = alter_table.alter_operation();
+                assert_matches!(alter_operation, AlterTableOperation::DropColumn { .. });
+                match alter_operation {
+                    AlterTableOperation::DropColumn { name } => {
+                        assert_eq!("a", name.value);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
 }

src/sql/src/statements/alter.rs

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use api::v1::{alter_expr, AddColumn, AlterExpr};
-use sqlparser::ast::{ColumnDef, ObjectName, TableConstraint};
+use api::v1::{alter_expr, AddColumn, AlterExpr, DropColumn};
+use sqlparser::ast::{ColumnDef, Ident, ObjectName, TableConstraint};
 
 use crate::error::UnsupportedAlterTableStatementSnafu;
 use crate::statements::{sql_column_def_to_grpc_column_def, table_idents_to_full_name};
@@ -47,7 +47,8 @@ pub enum AlterTableOperation {
     AddConstraint(TableConstraint),
     /// `ADD [ COLUMN ] <column_def>`
     AddColumn { column_def: ColumnDef },
-    // TODO(hl): support remove column
+    /// `DROP COLUMN <name>`
+    DropColumn { name: Ident },
 }
 
 /// Convert `AlterTable` statement to `AlterExpr` for gRPC
@@ -72,6 +73,11 @@ impl TryFrom<AlterTable> for AlterExpr {
                     }],
                 })
             }
+            AlterTableOperation::DropColumn { name } => {
+                alter_expr::Kind::DropColumns(api::v1::DropColumns {
+                    drop_columns: vec![DropColumn { name: name.value }],
+                })
+            }
         };
         let expr = AlterExpr {
             catalog_name: Some(catalog),

src/table/src/metadata.rs

@@ -130,7 +130,7 @@ impl TableMeta {
     ) -> Result<TableMetaBuilder> {
         match alter_kind {
             AlterKind::AddColumns { columns } => self.add_columns(table_name, columns),
-            AlterKind::RemoveColumns { names } => self.remove_columns(table_name, names),
+            AlterKind::DropColumns { names } => self.remove_columns(table_name, names),
         }
     }
 
@@ -576,7 +576,7 @@ mod tests {
         // Add more columns so we have enough candidate columns to remove.
         let meta = add_columns_to_meta(&meta);
 
-        let alter_kind = AlterKind::RemoveColumns {
+        let alter_kind = AlterKind::DropColumns {
             names: vec![String::from("col2"), String::from("my_field")],
         };
         let new_meta = meta
@@ -625,7 +625,7 @@ mod tests {
             .unwrap();
 
         // Remove columns in reverse order to test whether timestamp index is valid.
-        let alter_kind = AlterKind::RemoveColumns {
+        let alter_kind = AlterKind::DropColumns {
             names: vec![String::from("col3"), String::from("col1")],
         };
         let new_meta = meta
@@ -685,7 +685,7 @@ mod tests {
             .build()
             .unwrap();
 
-        let alter_kind = AlterKind::RemoveColumns {
+        let alter_kind = AlterKind::DropColumns {
             names: vec![String::from("unknown")],
         };
 
@@ -708,7 +708,7 @@ mod tests {
             .unwrap();
 
         // Remove column in primary key.
-        let alter_kind = AlterKind::RemoveColumns {
+        let alter_kind = AlterKind::DropColumns {
             names: vec![String::from("col1")],
         };
 
@@ -719,7 +719,7 @@ mod tests {
         assert_eq!(StatusCode::InvalidArguments, err.status_code());
 
         // Remove timestamp column.
-        let alter_kind = AlterKind::RemoveColumns {
+        let alter_kind = AlterKind::DropColumns {
             names: vec![String::from("ts")],
         };
 

src/table/src/requests.rs

@@ -77,7 +77,7 @@ pub struct AddColumnRequest {
 #[derive(Debug)]
 pub enum AlterKind {
     AddColumns { columns: Vec<AddColumnRequest> },
-    RemoveColumns { names: Vec<String> },
+    DropColumns { names: Vec<String> },
 }
 
 /// Drop table request