fix: Sql Inline Primary Key definition (#957)

* fix: invalid inline primary key syntax

* fix: format

* fix: clippy fix

* fix: added sqlness tests

* fix: throw exception when multiple inline pk defined

* fix: pr comments

* fix: add ending blank line for create.sql

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::<HashMap<_, _>>();
 
+        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::<Vec<_>>();
+
+        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;

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
+

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);
+