fix: describe distribute table (#988)

* fix: describe distribute table

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 { .. }

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]

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,
 

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)
 }