fix: add new column as primary key can't work (#2876)

GreptimeTeam · Dec 5, 2023 · f9e7762 · f9e7762
1 parent 0b421b5
commit f9e7762
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs
@@ -383,13 +383,23 @@ pub fn sql_column_def_to_grpc_column_def(col: &ColumnDef) -> Result<api::v1::Col
         .context(ConvertToGrpcDataTypeSnafu)?
         .to_parts();
 
+    let is_primary_key = col
+        .options
+        .iter()
+        .any(|o| matches!(o.option, ColumnOption::Unique { is_primary: true }));
+
+    let semantic_type = if is_primary_key {
+        SemanticType::Tag
+    } else {
+        SemanticType::Field
+    };
+
     Ok(api::v1::ColumnDef {
         name,
         data_type: datatype as i32,
         is_nullable,
         default_constraint: default_constraint.unwrap_or_default(),
-        // TODO(#1308): support adding new primary key columns
-        semantic_type: SemanticType::Field as _,
+        semantic_type: semantic_type as _,
         comment: String::new(),
         datatype_extension: datatype_ext,
     })
@@ -826,6 +836,7 @@ mod tests {
         assert!(grpc_column_def.is_nullable); // nullable when options are empty
         assert_eq!(ColumnDataType::Float64 as i32, grpc_column_def.data_type);
         assert!(grpc_column_def.default_constraint.is_empty());
+        assert_eq!(grpc_column_def.semantic_type, SemanticType::Field as i32);
 
         // test not null
         let column_def = ColumnDef {
@@ -840,6 +851,20 @@ mod tests {
 
         let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def).unwrap();
         assert!(!grpc_column_def.is_nullable);
+
+        // test primary key
+        let column_def = ColumnDef {
+            name: "col".into(),
+            data_type: SqlDataType::Double,
+            collation: None,
+            options: vec![ColumnOptionDef {
+                name: None,
+                option: ColumnOption::Unique { is_primary: true },
+            }],
+        };
+
+        let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def).unwrap();
+        assert_eq!(grpc_column_def.semantic_type, SemanticType::Tag as i32);
     }
 
     #[test]

diff --git a/tests/cases/standalone/common/alter/add_col.result b/tests/cases/standalone/common/alter/add_col.result
@@ -59,7 +59,7 @@ DESC TABLE test;
 | i      | Int32                |     | YES  |         | FIELD         |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
 | k      | Int32                |     | YES  |         | FIELD         |
-| host   | String               |     | YES  |         | FIELD         |
+| host   | String               | PRI | YES  |         | TAG           |
 +--------+----------------------+-----+------+---------+---------------+
 
 ALTER TABLE test ADD COLUMN idc STRING default 'idc' PRIMARY KEY;
@@ -83,8 +83,8 @@ DESC TABLE test;
 | i      | Int32                |     | YES  |         | FIELD         |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
 | k      | Int32                |     | YES  |         | FIELD         |
-| host   | String               |     | YES  |         | FIELD         |
-| idc    | String               |     | YES  | idc     | FIELD         |
+| host   | String               | PRI | YES  |         | TAG           |
+| idc    | String               | PRI | YES  | idc     | TAG           |
 +--------+----------------------+-----+------+---------+---------------+
 
 ALTER TABLE test ADD COLUMN "IdC" STRING default 'idc' PRIMARY KEY;
@@ -99,9 +99,9 @@ DESC TABLE test;
 | i      | Int32                |     | YES  |         | FIELD         |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
 | k      | Int32                |     | YES  |         | FIELD         |
-| host   | String               |     | YES  |         | FIELD         |
-| idc    | String               |     | YES  | idc     | FIELD         |
-| IdC    | String               |     | YES  | idc     | FIELD         |
+| host   | String               | PRI | YES  |         | TAG           |
+| idc    | String               | PRI | YES  | idc     | TAG           |
+| IdC    | String               | PRI | YES  | idc     | TAG           |
 +--------+----------------------+-----+------+---------+---------------+
 
 DROP TABLE test;