feat: Add column supports at first or after the existing columns (#1621)

* feat: Add column supports at first or after the existing columns

* Update src/common/query/Cargo.toml

---------

Co-authored-by: dennis zhuang <killme2008@gmail.com>

src/sql/Cargo.toml

@@ -10,6 +10,7 @@ common-base = { path = "../common/base" }
 common-catalog = { path = "../common/catalog" }
 common-datasource = { path = "../common/datasource" }
 common-error = { path = "../common/error" }
+common-query = { path = "../common/query" }
 common-time = { path = "../common/time" }
 datafusion-sql.workspace = true
 datatypes = { path = "../datatypes" }

src/sql/src/parsers/alter_parser.rs

@@ -12,9 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_query::AddColumnLocation;
 use snafu::ResultExt;
 use sqlparser::keywords::Keyword;
 use sqlparser::parser::ParserError;
+use sqlparser::tokenizer::Token;
 
 use crate::error::{self, Result};
 use crate::parser::ParserContext;
@@ -41,7 +43,25 @@ impl<'a> ParserContext<'a> {
             } else {
                 let _ = parser.parse_keyword(Keyword::COLUMN);
                 let column_def = parser.parse_column_def()?;
-                AlterTableOperation::AddColumn { column_def }
+                let location = if parser.parse_keyword(Keyword::FIRST) {
+                    Some(AddColumnLocation::First)
+                } else if let Token::Word(word) = parser.peek_token().token {
+                    if word.value.to_ascii_uppercase() == "AFTER" {
+                        parser.next_token();
+                        let name = parser.parse_identifier()?;
+                        Some(AddColumnLocation::After {
+                            column_name: name.value,
+                        })
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                };
+                AlterTableOperation::AddColumn {
+                    column_def,
+                    location,
+                }
             }
         } else if parser.parse_keyword(Keyword::DROP) {
             if parser.parse_keyword(Keyword::COLUMN) {
@@ -98,13 +118,90 @@ mod tests {
                 let alter_operation = alter_table.alter_operation();
                 assert_matches!(alter_operation, AlterTableOperation::AddColumn { .. });
                 match alter_operation {
-                    AlterTableOperation::AddColumn { column_def } => {
+                    AlterTableOperation::AddColumn {
+                        column_def,
+                        location,
+                    } => {
                         assert_eq!("tagk_i", column_def.name.value);
                         assert_eq!(DataType::String, column_def.data_type);
                         assert!(column_def
                             .options
                             .iter()
                             .any(|o| matches!(o.option, ColumnOption::Null)));
+                        assert_eq!(&None, location);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[test]
+    fn test_parse_alter_add_column_with_first() {
+        let sql = "ALTER TABLE my_metric_1 ADD tagk_i STRING Null FIRST;";
+        let mut result = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}).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::AddColumn { .. });
+                match alter_operation {
+                    AlterTableOperation::AddColumn {
+                        column_def,
+                        location,
+                    } => {
+                        assert_eq!("tagk_i", column_def.name.value);
+                        assert_eq!(DataType::String, column_def.data_type);
+                        assert!(column_def
+                            .options
+                            .iter()
+                            .any(|o| matches!(o.option, ColumnOption::Null)));
+                        assert_eq!(&Some(AddColumnLocation::First), location);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[test]
+    fn test_parse_alter_add_column_with_after() {
+        let sql = "ALTER TABLE my_metric_1 ADD tagk_i STRING Null AFTER ts;";
+        let mut result = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}).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::AddColumn { .. });
+                match alter_operation {
+                    AlterTableOperation::AddColumn {
+                        column_def,
+                        location,
+                    } => {
+                        assert_eq!("tagk_i", column_def.name.value);
+                        assert_eq!(DataType::String, column_def.data_type);
+                        assert!(column_def
+                            .options
+                            .iter()
+                            .any(|o| matches!(o.option, ColumnOption::Null)));
+                        assert_eq!(
+                            &Some(AddColumnLocation::After {
+                                column_name: "ts".to_string()
+                            }),
+                            location
+                        );
                     }
                     _ => unreachable!(),
                 }

src/sql/src/statements.rs

@@ -28,7 +28,10 @@ pub mod tql;
 use std::str::FromStr;
 
 use api::helper::ColumnDataTypeWrapper;
+use api::v1::add_column::location::LocationType;
+use api::v1::add_column::Location;
 use common_base::bytes::Bytes;
+use common_query::AddColumnLocation;
 use common_time::Timestamp;
 use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, COMMENT_KEY};
@@ -397,6 +400,22 @@ pub fn concrete_data_type_to_sql_data_type(data_type: &ConcreteDataType) -> Resu
     }
 }
 
+pub fn sql_location_to_grpc_add_column_location(
+    location: &Option<AddColumnLocation>,
+) -> Option<api::v1::add_column::Location> {
+    match location {
+        Some(AddColumnLocation::First) => Some(Location {
+            location_type: LocationType::First.into(),
+            after_cloumn_name: "".to_string(),
+        }),
+        Some(AddColumnLocation::After { column_name }) => Some(Location {
+            location_type: LocationType::After.into(),
+            after_cloumn_name: column_name.to_string(),
+        }),
+        None => None,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::assert_matches::assert_matches;

src/sql/src/statements/alter.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_query::AddColumnLocation;
 use sqlparser::ast::{ColumnDef, Ident, ObjectName, TableConstraint};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -41,8 +42,11 @@ impl AlterTable {
 pub enum AlterTableOperation {
     /// `ADD <table_constraint>`
     AddConstraint(TableConstraint),
-    /// `ADD [ COLUMN ] <column_def>`
-    AddColumn { column_def: ColumnDef },
+    /// `ADD [ COLUMN ] <column_def> [location]`
+    AddColumn {
+        column_def: ColumnDef,
+        location: Option<AddColumnLocation>,
+    },
     /// `DROP COLUMN <name>`
     DropColumn { name: Ident },
     /// `RENAME <new_table_name>`