fix: can't adding new columns as primary key (#2310)

src/datanode/src/sql/alter.rs

@@ -17,7 +17,7 @@ use common_query::Output;
 use common_telemetry::logging::info;
 use snafu::prelude::*;
 use sql::statements::alter::{AlterTable, AlterTableOperation};
-use sql::statements::column_def_to_schema;
+use sql::statements::{column_def_to_schema, has_primary_key_option};
 use table::engine::TableReference;
 use table::metadata::TableId;
 use table::requests::{AddColumnRequest, AlterKind, AlterTableRequest};
@@ -106,8 +106,7 @@ impl SqlHandler {
                 columns: vec![AddColumnRequest {
                     column_schema: column_def_to_schema(column_def, false)
                         .context(error::ParseSqlSnafu)?,
-                    // FIXME(dennis): supports adding key column
-                    is_key: false,
+                    is_key: has_primary_key_option(column_def),
                     location: location.clone(),
                 }],
             },
@@ -172,11 +171,42 @@ mod tests {
         assert_matches!(alter_kind, AlterKind::AddColumns { .. });
         match alter_kind {
             AlterKind::AddColumns { columns } => {
-                let new_column = &columns[0].column_schema;
+                let request = &columns[0];
+                let new_column = &request.column_schema;
 
                 assert_eq!(new_column.name, "tagk_i");
                 assert!(new_column.is_nullable());
                 assert_eq!(new_column.data_type, ConcreteDataType::string_datatype());
+                assert!(!request.is_key);
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_alter_to_request_with_adding_key_column() {
+        let alter_table = parse_sql("ALTER TABLE my_metric_1 ADD tagk_i STRING PRIMARY KEY;");
+        let req = SqlHandler::alter_to_request(
+            alter_table,
+            TableReference::full("greptime", "public", "my_metric_1"),
+            1,
+        )
+        .unwrap();
+        assert_eq!(req.catalog_name, "greptime");
+        assert_eq!(req.schema_name, "public");
+        assert_eq!(req.table_name, "my_metric_1");
+
+        let alter_kind = req.alter_kind;
+        assert_matches!(alter_kind, AlterKind::AddColumns { .. });
+        match alter_kind {
+            AlterKind::AddColumns { columns } => {
+                let request = &columns[0];
+                let new_column = &request.column_schema;
+
+                assert_eq!(new_column.name, "tagk_i");
+                assert!(new_column.is_nullable());
+                assert_eq!(new_column.data_type, ConcreteDataType::string_datatype());
+                assert!(request.is_key);
             }
             _ => unreachable!(),
         }

src/datanode/src/sql/create.rs

@@ -21,9 +21,9 @@ use common_telemetry::tracing::info;
 use datatypes::schema::RawSchema;
 use session::context::QueryContextRef;
 use snafu::{ensure, OptionExt, ResultExt};
-use sql::ast::{ColumnOption, TableConstraint};
-use sql::statements::column_def_to_schema;
+use sql::ast::TableConstraint;
 use sql::statements::create::{CreateTable, TIME_INDEX};
+use sql::statements::{column_def_to_schema, has_primary_key_option};
 use sql::util::to_lowercase_options_map;
 use table::engine::TableReference;
 use table::metadata::TableId;
@@ -131,12 +131,7 @@ impl SqlHandler {
         let pk_map = stmt
             .columns
             .iter()
-            .filter(|col| {
-                col.options.iter().any(|options| match options.option {
-                    ColumnOption::Unique { is_primary } => is_primary,
-                    _ => false,
-                })
-            })
+            .filter(|c| has_primary_key_option(c))
             .map(|col| col.name.value.clone())
             .collect::<Vec<_>>();
 

src/sql/src/statements.rs

@@ -270,6 +270,17 @@ fn parse_column_default_constraint(
     }
 }
 
+/// Return true when the `ColumnDef` options contain primary key
+pub fn has_primary_key_option(column_def: &ColumnDef) -> bool {
+    column_def
+        .options
+        .iter()
+        .any(|options| match options.option {
+            ColumnOption::Unique { is_primary } => is_primary,
+            _ => false,
+        })
+}
+
 // TODO(yingwen): Make column nullable by default, and checks invalid case like
 // a column is not nullable but has a default value null.
 /// Create a `ColumnSchema` from `ColumnDef`.
@@ -748,6 +759,28 @@ mod tests {
         assert!(!grpc_column_def.is_nullable);
     }
 
+    #[test]
+    pub fn test_has_primary_key_option() {
+        let column_def = ColumnDef {
+            name: "col".into(),
+            data_type: SqlDataType::Double,
+            collation: None,
+            options: vec![],
+        };
+        assert!(!has_primary_key_option(&column_def));
+
+        let column_def = ColumnDef {
+            name: "col".into(),
+            data_type: SqlDataType::Double,
+            collation: None,
+            options: vec![ColumnOptionDef {
+                name: None,
+                option: ColumnOption::Unique { is_primary: true },
+            }],
+        };
+        assert!(has_primary_key_option(&column_def));
+    }
+
     #[test]
     pub fn test_column_def_to_schema() {
         let column_def = ColumnDef {

tests/cases/standalone/common/alter/add_col.sql

@@ -1,9 +1,25 @@
 CREATE TABLE test(i INTEGER, j TIMESTAMP TIME INDEX);
 
+DESC TABLE test;
+
 INSERT INTO test VALUES (1, 1), (2, 2);
 
 ALTER TABLE test ADD COLUMN k INTEGER;
 
 SELECT * FROM test;
 
+DESC TABLE test;
+
+ALTER TABLE test ADD COLUMN host STRING PRIMARY KEY;
+
+SELECT * FROM test;
+
+DESC TABLE test;
+
+ALTER TABLE test ADD COLUMN idc STRING default "idc" PRIMARY KEY;
+
+SELECT * FROM test;
+
+DESC TABLE test;
+
 DROP TABLE test;