diff --git a/src/datanode/src/sql/create.rs b/src/datanode/src/sql/create.rs index bf011a94f9..2f2d3bd367 100644 --- a/src/datanode/src/sql/create.rs +++ b/src/datanode/src/sql/create.rs @@ -22,7 +22,7 @@ use common_telemetry::tracing::info; use common_telemetry::tracing::log::error; use datatypes::schema::SchemaBuilder; use snafu::{ensure, OptionExt, ResultExt}; -use sql::ast::TableConstraint; +use sql::ast::{ColumnOption, TableConstraint}; use sql::statements::column_def_to_schema; use sql::statements::create::CreateTable; use store_api::storage::consts::TIME_INDEX_NAME; @@ -135,6 +135,30 @@ impl SqlHandler { .map(|(k, v)| (v, k)) .collect::>(); + let pk_map = stmt + .columns + .iter() + .filter(|col| { + col.options.iter().any(|options| match options.option { + ColumnOption::Unique { is_primary } => is_primary, + _ => false, + }) + }) + .map(|col| col.name.value.clone()) + .collect::>(); + + ensure!( + pk_map.len() < 2, + InvalidPrimaryKeySnafu { + msg: "Multiple definitions of primary key found" + } + ); + + if let Some(pk) = pk_map.first() { + // # Safety: Both pk_map and col_map are collected from stmt.columns + primary_keys.push(*col_map.get(pk).unwrap()); + } + for c in stmt.constraints { match c { TableConstraint::Unique { @@ -156,6 +180,12 @@ impl SqlHandler { .fail(); } } else if is_primary { + if !primary_keys.is_empty() { + return InvalidPrimaryKeySnafu { + msg: "Multiple definitions of primary key found", + } + .fail(); + } for col in columns { primary_keys.push(*col_map.get(&col.value).context( KeyColumnNotFoundSnafu { @@ -246,6 +276,21 @@ mod tests { } } + #[tokio::test] + pub async fn test_create_with_inline_primary_key() { + let handler = create_mock_sql_handler().await; + let parsed_stmt = sql_to_statement( + r#"CREATE TABLE demo_table (timestamp BIGINT TIME INDEX, value DOUBLE, host STRING PRIMARY KEY) engine=mito with(regions=1);"#, + ); + let c = handler + .create_to_request(42, parsed_stmt, &TableReference::bare("demo_table")) + .unwrap(); + assert_eq!("demo_table", c.table_name); + assert_eq!(42, c.id); + assert!(!c.create_if_not_exists); + assert_eq!(vec![2], c.primary_key_indices); + } + #[tokio::test] pub async fn test_create_to_request() { let handler = create_mock_sql_handler().await; @@ -269,6 +314,37 @@ mod tests { assert_eq!(4, c.schema.column_schemas().len()); } + #[tokio::test] + pub async fn test_multiple_primary_key_definitions() { + let handler = create_mock_sql_handler().await; + let parsed_stmt = sql_to_statement( + r#"create table demo_table ( + timestamp BIGINT TIME INDEX, + value DOUBLE, + host STRING PRIMARY KEY, + PRIMARY KEY(host)) engine=mito with(regions=1);"#, + ); + let error = handler + .create_to_request(42, parsed_stmt, &TableReference::bare("demo_table")) + .unwrap_err(); + assert_matches!(error, Error::InvalidPrimaryKey { .. }); + } + + #[tokio::test] + pub async fn test_multiple_inline_primary_key_definitions() { + let handler = create_mock_sql_handler().await; + let parsed_stmt = sql_to_statement( + r#"create table demo_table ( + timestamp BIGINT TIME INDEX, + value DOUBLE PRIMARY KEY, + host STRING PRIMARY KEY) engine=mito with(regions=1);"#, + ); + let error = handler + .create_to_request(42, parsed_stmt, &TableReference::bare("demo_table")) + .unwrap_err(); + assert_matches!(error, Error::InvalidPrimaryKey { .. }); + } + #[tokio::test] pub async fn test_primary_key_not_specified() { let handler = create_mock_sql_handler().await; diff --git a/tests/cases/standalone/create/create.result b/tests/cases/standalone/create/create.result index e3975593be..e398d95754 100644 --- a/tests/cases/standalone/create/create.result +++ b/tests/cases/standalone/create/create.result @@ -88,3 +88,33 @@ DROP TABLE test2; Affected Rows: 1 +CREATE TABLE test_pk (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE); + +Affected Rows: 0 + +DESC TABLE test_pk; + ++-----------+---------+------+---------+---------------+ +| Field | Type | Null | Default | Semantic Type | ++-----------+---------+------+---------+---------------+ +| timestamp | Int64 | NO | | TIME INDEX | +| host | String | YES | | PRIMARY KEY | +| value | Float64 | YES | | VALUE | ++-----------+---------+------+---------+---------------+ + +DROP TABLE test_pk; + +Affected Rows: 1 + +CREATE TABLE test_multiple_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE, PRIMARY KEY(host)); + +Error: 1004(InvalidArguments), Invalid primary key: Multiple definitions of primary key found + +CREATE TABLE test_multiple_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE, PRIMARY KEY(host), PRIMARY KEY(host)); + +Error: 1004(InvalidArguments), Invalid primary key: Multiple definitions of primary key found + +CREATE TABLE test_multiple_inline_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE PRIMARY KEY); + +Error: 1004(InvalidArguments), Invalid primary key: Multiple definitions of primary key found + diff --git a/tests/cases/standalone/create/create.sql b/tests/cases/standalone/create/create.sql index 8f5501f92a..f9bcbfea11 100644 --- a/tests/cases/standalone/create/create.sql +++ b/tests/cases/standalone/create/create.sql @@ -35,3 +35,16 @@ DROP TABLE times; DROP TABLE test1; DROP TABLE test2; + +CREATE TABLE test_pk (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE); + +DESC TABLE test_pk; + +DROP TABLE test_pk; + +CREATE TABLE test_multiple_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE, PRIMARY KEY(host)); + +CREATE TABLE test_multiple_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE, PRIMARY KEY(host), PRIMARY KEY(host)); + +CREATE TABLE test_multiple_inline_pk_definitions (timestamp BIGINT TIME INDEX, host STRING PRIMARY KEY, value DOUBLE PRIMARY KEY); +