From b1599ad3a56dfc91d6c714c60725087614b54da3 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 3 Sep 2023 22:47:22 -0500 Subject: [PATCH] fix: can't adding new columns as primary key (#2310) --- src/datanode/src/sql/alter.rs | 38 +++++++++++++++++-- src/datanode/src/sql/create.rs | 11 ++---- src/sql/src/statements.rs | 33 ++++++++++++++++ .../cases/standalone/common/alter/add_col.sql | 16 ++++++++ 4 files changed, 86 insertions(+), 12 deletions(-) diff --git a/src/datanode/src/sql/alter.rs b/src/datanode/src/sql/alter.rs index efe1084a8b..934cffe76a 100644 --- a/src/datanode/src/sql/alter.rs +++ b/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!(), } diff --git a/src/datanode/src/sql/create.rs b/src/datanode/src/sql/create.rs index 573ed6bded..8e9e0faf60 100644 --- a/src/datanode/src/sql/create.rs +++ b/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::>(); diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index dba5b3b1fc..ef67c6e81e 100644 --- a/src/sql/src/statements.rs +++ b/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 { diff --git a/tests/cases/standalone/common/alter/add_col.sql b/tests/cases/standalone/common/alter/add_col.sql index 36bf9d68bc..a1c27eac0f 100644 --- a/tests/cases/standalone/common/alter/add_col.sql +++ b/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;