fix: describe distribute table (GreptimeTeam#988)

MichaelScofield · paomian · commit fa455c430df0 · 2023-10-19T15:50:57.000+08:00
* fix: describe distribute table
diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs
@@ -192,8 +192,8 @@ pub enum Error {
     #[snafu(display("Specified timestamp key or primary key column not found: {}", name))]
     KeyColumnNotFound { name: String, backtrace: Backtrace },
 
-    #[snafu(display("Invalid primary key: {}", msg))]
-    InvalidPrimaryKey { msg: String, backtrace: Backtrace },
+    #[snafu(display("Illegal primary keys definition: {}", msg))]
+    IllegalPrimaryKeysDef { msg: String, backtrace: Backtrace },
 
     #[snafu(display(
         "Constraint in CREATE TABLE statement is not supported yet: {}",
@@ -378,7 +378,7 @@ impl ErrorExt for Error {
             | Error::InvalidSql { .. }
             | Error::NotSupportSql { .. }
             | Error::KeyColumnNotFound { .. }
-            | Error::InvalidPrimaryKey { .. }
+            | Error::IllegalPrimaryKeysDef { .. }
             | Error::MissingTimestampColumn { .. }
             | Error::CatalogNotFound { .. }
             | Error::SchemaNotFound { .. }
diff --git a/src/datanode/src/sql/create.rs b/src/datanode/src/sql/create.rs
@@ -32,7 +32,7 @@ use table::requests::*;
 
 use crate::error::{
     self, CatalogNotFoundSnafu, CatalogSnafu, ConstraintNotSupportedSnafu, CreateSchemaSnafu,
-    CreateTableSnafu, InsertSystemCatalogSnafu, InvalidPrimaryKeySnafu, KeyColumnNotFoundSnafu,
+    CreateTableSnafu, IllegalPrimaryKeysDefSnafu, InsertSystemCatalogSnafu, KeyColumnNotFoundSnafu,
     RegisterSchemaSnafu, Result, SchemaExistsSnafu, SchemaNotFoundSnafu,
 };
 use crate::sql::SqlHandler;
@@ -159,8 +159,8 @@ impl SqlHandler {
 
         ensure!(
             pk_map.len() < 2,
-            InvalidPrimaryKeySnafu {
-                msg: "Multiple definitions of primary key found"
+            IllegalPrimaryKeysDefSnafu {
+                msg: "not allowed to inline multiple primary keys in columns options"
             }
         );
 
@@ -191,8 +191,8 @@ impl SqlHandler {
                         }
                     } else if is_primary {
                         if !primary_keys.is_empty() {
-                            return InvalidPrimaryKeySnafu {
-                                msg: "Multiple definitions of primary key found",
+                            return IllegalPrimaryKeysDefSnafu {
+                                msg: "found definitions of primary keys in multiple places",
                             }
                             .fail();
                         }
@@ -223,7 +223,7 @@ impl SqlHandler {
 
         ensure!(
             !primary_keys.iter().any(|index| *index == ts_index),
-            InvalidPrimaryKeySnafu {
+            IllegalPrimaryKeysDefSnafu {
                 msg: "time index column can't be included in primary key"
             }
         );
@@ -337,7 +337,7 @@ mod tests {
         let error = handler
             .create_to_request(42, parsed_stmt, &TableReference::bare("demo_table"))
             .unwrap_err();
-        assert_matches!(error, Error::InvalidPrimaryKey { .. });
+        assert_matches!(error, Error::IllegalPrimaryKeysDef { .. });
     }
 
     #[tokio::test]
@@ -352,7 +352,7 @@ mod tests {
         let error = handler
             .create_to_request(42, parsed_stmt, &TableReference::bare("demo_table"))
             .unwrap_err();
-        assert_matches!(error, Error::InvalidPrimaryKey { .. });
+        assert_matches!(error, Error::IllegalPrimaryKeysDef { .. });
     }
 
     #[tokio::test]
@@ -410,7 +410,7 @@ mod tests {
         let error = handler
             .create_to_request(42, create_table, &TableReference::full("c", "s", "demo"))
             .unwrap_err();
-        assert_matches!(error, Error::InvalidPrimaryKey { .. });
+        assert_matches!(error, Error::IllegalPrimaryKeysDef { .. });
     }
 
     #[tokio::test]
diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs
@@ -342,6 +342,9 @@ pub enum Error {
         #[snafu(backtrace)]
         source: query::error::Error,
     },
+
+    #[snafu(display("Illegal primary keys definition: {}", msg))]
+    IllegalPrimaryKeysDef { msg: String, backtrace: Backtrace },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -352,7 +355,8 @@ impl ErrorExt for Error {
             Error::ParseAddr { .. }
             | Error::InvalidSql { .. }
             | Error::InvalidInsertRequest { .. }
-            | Error::ColumnValuesNumberMismatch { .. } => StatusCode::InvalidArguments,
+            | Error::ColumnValuesNumberMismatch { .. }
+            | Error::IllegalPrimaryKeysDef { .. } => StatusCode::InvalidArguments,
 
             Error::NotSupported { .. } => StatusCode::Unsupported,
 
diff --git a/src/frontend/src/expr_factory.rs b/src/frontend/src/expr_factory.rs
@@ -22,13 +22,14 @@ use datanode::instance::sql::table_idents_to_full_name;
 use datatypes::schema::ColumnSchema;
 use session::context::QueryContextRef;
 use snafu::{ensure, ResultExt};
-use sql::ast::{ColumnDef, TableConstraint};
+use sql::ast::{ColumnDef, ColumnOption, TableConstraint};
 use sql::statements::column_def_to_schema;
 use sql::statements::create::{CreateTable, TIME_INDEX};
 
 use crate::error::{
     self, BuildCreateExprOnInsertionSnafu, ColumnDataTypeSnafu,
-    ConvertColumnDefaultConstraintSnafu, InvalidSqlSnafu, ParseSqlSnafu, Result,
+    ConvertColumnDefaultConstraintSnafu, IllegalPrimaryKeysDefSnafu, InvalidSqlSnafu,
+    ParseSqlSnafu, Result,
 };
 
 pub type CreateExprFactoryRef = Arc<dyn CreateExprFactory + Send + Sync>;
@@ -88,7 +89,7 @@ pub(crate) fn create_to_expr(
         desc: "".to_string(),
         column_defs: columns_to_expr(&create.columns, &time_index)?,
         time_index,
-        primary_keys: find_primary_keys(&create.constraints)?,
+        primary_keys: find_primary_keys(&create.columns, &create.constraints)?,
         create_if_not_exists: create.if_not_exists,
         // TODO(LFC): Fill in other table options.
         table_options: HashMap::from([("engine".to_string(), create.engine.clone())]),
@@ -98,8 +99,32 @@ pub(crate) fn create_to_expr(
     Ok(expr)
 }
 
-fn find_primary_keys(constraints: &[TableConstraint]) -> Result<Vec<String>> {
-    let primary_keys = constraints
+fn find_primary_keys(
+    columns: &[ColumnDef],
+    constraints: &[TableConstraint],
+) -> Result<Vec<String>> {
+    let columns_pk = columns
+        .iter()
+        .filter_map(|x| {
+            if x.options
+                .iter()
+                .any(|o| matches!(o.option, ColumnOption::Unique { is_primary: true }))
+            {
+                Some(x.name.value.clone())
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<String>>();
+
+    ensure!(
+        columns_pk.len() <= 1,
+        IllegalPrimaryKeysDefSnafu {
+            msg: "not allowed to inline multiple primary keys in columns options"
+        }
+    );
+
+    let constraints_pk = constraints
         .iter()
         .filter_map(|constraint| match constraint {
             TableConstraint::Unique {
@@ -111,6 +136,17 @@ fn find_primary_keys(constraints: &[TableConstraint]) -> Result<Vec<String>> {
         })
         .flatten()
         .collect::<Vec<String>>();
+
+    ensure!(
+        columns_pk.is_empty() || constraints_pk.is_empty(),
+        IllegalPrimaryKeysDefSnafu {
+            msg: "found definitions of primary keys in multiple places"
+        }
+    );
+
+    let mut primary_keys = Vec::with_capacity(columns_pk.len() + constraints_pk.len());
+    primary_keys.extend(columns_pk);
+    primary_keys.extend(constraints_pk);
     Ok(primary_keys)
 }
 
diff --git a/tests/cases/distributed/create/create.result b/tests/cases/distributed/create/create.result
diff --git a/tests/cases/distributed/create/create.sql b/tests/cases/distributed/create/create.sql
diff --git a/tests/cases/standalone/common/create/create.result b/tests/cases/standalone/common/create/create.result
@@ -108,13 +108,13 @@ 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
+Error: 1004(InvalidArguments), Illegal primary keys definition: found definitions of primary keys in multiple places
 
 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
+Error: 1004(InvalidArguments), Illegal primary keys definition: found definitions of primary keys in multiple places
 
 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
+Error: 1004(InvalidArguments), Illegal primary keys definition: not allowed to inline multiple primary keys in columns options
 
diff --git a/tests/cases/standalone/common/create/create.sql b/tests/cases/standalone/common/create/create.sql
diff --git a/tests/cases/standalone/order/order_variable_size_payload.sql b/tests/cases/standalone/order/order_variable_size_payload.sql
@@ -123,3 +123,5 @@ DROP table test7;
 DROP table test8;
 
 DROP TABLE DirectReports;
+
+-- TODO(LFC): Seems creating distributed table has some column schema related issues, look into "order_variable_size_payload" test case.
