From a6887605637a51e35c48aaceb248f383bfde87de Mon Sep 17 00:00:00 2001
From: LFC <990479+MichaelScofield@users.noreply.github.com>
Date: Fri, 15 Sep 2023 18:07:32 +0800
Subject: [PATCH] fix: validate partition columns (#2393)

fix: partition column must belong to primary keys or equals to time index
---
 src/frontend/src/error.rs         | 14 ++++++++-
 src/frontend/src/statement/ddl.rs | 50 +++++++++++++++++++++++++++++--
 tests-integration/src/grpc.rs     |  4 +--
 3 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs
index db2e335602a3..efe79e9f0ca8 100644
--- a/src/frontend/src/error.rs
+++ b/src/frontend/src/error.rs
@@ -689,6 +689,17 @@ pub enum Error {
 
     #[snafu(display("Invalid region request, reason: {}", reason))]
     InvalidRegionRequest { reason: String },
+
+    #[snafu(display(
+        "Invalid partition columns when creating table '{}', reason: {}",
+        table,
+        reason
+    ))]
+    InvalidPartitionColumns {
+        table: String,
+        reason: String,
+        location: Location,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -714,7 +725,8 @@ impl ErrorExt for Error {
             | Error::UnsupportedFormat { .. }
             | Error::EmptyData { .. }
             | Error::ColumnNoneDefaultValue { .. }
-            | Error::IncompleteGrpcRequest { .. } => StatusCode::InvalidArguments,
+            | Error::IncompleteGrpcRequest { .. }
+            | Error::InvalidPartitionColumns { .. } => StatusCode::InvalidArguments,
 
             Error::NotSupported { .. } => StatusCode::Unsupported,
 
diff --git a/src/frontend/src/statement/ddl.rs b/src/frontend/src/statement/ddl.rs
index 5de29d25d786..fa48679f7593 100644
--- a/src/frontend/src/statement/ddl.rs
+++ b/src/frontend/src/statement/ddl.rs
@@ -33,7 +33,7 @@ use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::RawSchema;
 use partition::partition::{PartitionBound, PartitionDef};
 use session::context::QueryContextRef;
-use snafu::{OptionExt, ResultExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use sql::ast::Value as SqlValue;
 use sql::statements::alter::AlterTable;
 use sql::statements::create::{CreateExternalTable, CreateTable, Partitions};
@@ -46,8 +46,9 @@ use table::TableRef;
 use super::StatementExecutor;
 use crate::error::{
     self, AlterExprToRequestSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu,
-    DeserializePartitionSnafu, ParseSqlSnafu, Result, SchemaNotFoundSnafu,
-    TableMetadataManagerSnafu, TableNotFoundSnafu, UnrecognizedTableOptionSnafu,
+    DeserializePartitionSnafu, InvalidPartitionColumnsSnafu, ParseSqlSnafu, Result,
+    SchemaNotFoundSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
+    UnrecognizedTableOptionSnafu,
 };
 use crate::{expr_factory, MAX_VALUE};
 
@@ -101,6 +102,8 @@ impl StatementExecutor {
 
         let (partitions, partition_cols) = parse_partitions(create_table, partitions)?;
 
+        validate_partition_columns(create_table, &partition_cols)?;
+
         let mut table_info = create_table_info(create_table, partition_cols, schema_opts)?;
 
         let resp = self
@@ -346,6 +349,22 @@ impl StatementExecutor {
     }
 }
 
+fn validate_partition_columns(
+    create_table: &CreateTableExpr,
+    partition_cols: &[String],
+) -> Result<()> {
+    ensure!(
+        partition_cols
+            .iter()
+            .all(|col| &create_table.time_index == col || create_table.primary_keys.contains(col)),
+        InvalidPartitionColumnsSnafu {
+            table: &create_table.table_name,
+            reason: "partition column must belongs to primary keys or equals to time index"
+        }
+    );
+    Ok(())
+}
+
 fn parse_partitions(
     create_table: &CreateTableExpr,
     partitions: Option<Partitions>,
@@ -531,6 +550,31 @@ mod test {
     use super::*;
     use crate::expr_factory;
 
+    #[test]
+    fn test_validate_partition_columns() {
+        let create_table = CreateTableExpr {
+            table_name: "my_table".to_string(),
+            time_index: "ts".to_string(),
+            primary_keys: vec!["a".to_string(), "b".to_string()],
+            ..Default::default()
+        };
+
+        assert!(validate_partition_columns(&create_table, &[]).is_ok());
+        assert!(validate_partition_columns(&create_table, &["ts".to_string()]).is_ok());
+        assert!(validate_partition_columns(&create_table, &["a".to_string()]).is_ok());
+        assert!(
+            validate_partition_columns(&create_table, &["b".to_string(), "a".to_string()]).is_ok()
+        );
+
+        assert_eq!(
+            validate_partition_columns(&create_table, &["a".to_string(), "c".to_string()])
+                .unwrap_err()
+                .to_string(),
+            "Invalid partition columns when creating table 'my_table', \
+            reason: partition column must belongs to primary keys or equals to time index",
+        );
+    }
+
     #[tokio::test]
     async fn test_parse_partitions() {
         common_telemetry::init_default_ut_logging();
diff --git a/tests-integration/src/grpc.rs b/tests-integration/src/grpc.rs
index 4bb7b2be8e41..3e01cfb339b0 100644
--- a/tests-integration/src/grpc.rs
+++ b/tests-integration/src/grpc.rs
@@ -347,7 +347,7 @@ CREATE TABLE {table_name} (
                             .collect(),
                         ..Default::default()
                     }),
-                    semantic_type: SemanticType::Tag as i32,
+                    semantic_type: SemanticType::Field as i32,
                     datatype: ColumnDataType::String as i32,
                     ..Default::default()
                 },
@@ -421,7 +421,7 @@ CREATE TABLE {table_name} (
                 },
                 Column {
                     column_name: "b".to_string(),
-                    semantic_type: SemanticType::Tag as i32,
+                    semantic_type: SemanticType::Field as i32,
                     values: Some(Values {
                         string_values: b,
                         ..Default::default()