feat: impl alter table in distributed mode (#572)

Cargo.toml

@@ -15,7 +15,7 @@ members = [
     "src/common/recordbatch",
     "src/common/runtime",
     "src/common/substrait",
-    "src/common/insert",
+    "src/common/grpc-expr",
     "src/common/telemetry",
     "src/common/time",
     "src/datanode",

src/api/src/lib.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod column_def;
 pub mod error;
 pub mod helper;
 pub mod prometheus;

src/api/src/v1.rs

@@ -21,4 +21,5 @@ pub mod codec {
     tonic::include_proto!("greptime.v1.codec");
 }
 
+mod column_def;
 pub mod meta;

src/client/Cargo.toml

@@ -13,7 +13,7 @@ common-error = { path = "../common/error" }
 common-grpc = { path = "../common/grpc" }
 common-query = { path = "../common/query" }
 common-recordbatch = { path = "../common/recordbatch" }
-common-insert = { path = "../common/insert" }
+common-grpc-expr = { path = "../common/grpc-expr" }
 common-time = { path = "../common/time" }
 datafusion = { git = "https://github.com/apache/arrow-datafusion.git", branch = "arrow2", features = [
     "simd",

src/client/src/database.rs

@@ -23,7 +23,7 @@ use api::v1::{
 };
 use common_error::status_code::StatusCode;
 use common_grpc::{AsExecutionPlan, DefaultAsPlanImpl};
-use common_insert::column_to_vector;
+use common_grpc_expr::column_to_vector;
 use common_query::Output;
 use common_recordbatch::{RecordBatch, RecordBatches};
 use datafusion::physical_plan::ExecutionPlan;

src/client/src/error.rs

@@ -103,7 +103,7 @@ pub enum Error {
     #[snafu(display("Failed to convert column to vector, source: {}", source))]
     ColumnToVector {
         #[snafu(backtrace)]
-        source: common_insert::error::Error,
+        source: common_grpc_expr::error::Error,
     },
 }
 

src/cmd/src/error.rs

@@ -97,10 +97,7 @@ mod tests {
     #[test]
     fn test_start_node_error() {
         fn throw_datanode_error() -> StdResult<datanode::error::Error> {
-            datanode::error::MissingFieldSnafu {
-                field: "test_field",
-            }
-            .fail()
+            datanode::error::MissingNodeIdSnafu {}.fail()
         }
 
         let e = throw_datanode_error()

src/common/insert/Cargo.toml → src/common/grpc-expr/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "common-insert"
+name = "common-grpc-expr"
 version = "0.1.0"
 edition = "2021"
 license = "Apache-2.0"
@@ -11,6 +11,7 @@ common-base = { path = "../base" }
 common-error = { path = "../error" }
 common-telemetry = { path = "../telemetry" }
 common-time = { path = "../time" }
+common-catalog = { path = "../catalog" }
 common-query = { path = "../query" }
 datatypes = { path = "../../datatypes" }
 snafu = { version = "0.7", features = ["backtraces"] }

src/common/grpc-expr/src/alter.rs

@@ -0,0 +1,234 @@
+// Copyright 2022 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use api::v1::alter_expr::Kind;
+use api::v1::{AlterExpr, CreateExpr, DropColumns};
+use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef};
+use snafu::{ensure, OptionExt, ResultExt};
+use table::metadata::TableId;
+use table::requests::{AddColumnRequest, AlterKind, AlterTableRequest, CreateTableRequest};
+
+use crate::error::{
+    ColumnNotFoundSnafu, CreateSchemaSnafu, InvalidColumnDefSnafu, MissingFieldSnafu,
+    MissingTimestampColumnSnafu, Result,
+};
+
+/// Convert an [`AlterExpr`] to an optional [`AlterTableRequest`]
+pub fn alter_expr_to_request(expr: AlterExpr) -> Result<Option<AlterTableRequest>> {
+    match expr.kind {
+        Some(Kind::AddColumns(add_columns)) => {
+            let add_column_requests = add_columns
+                .add_columns
+                .into_iter()
+                .map(|ac| {
+                    let column_def = ac.column_def.context(MissingFieldSnafu {
+                        field: "column_def",
+                    })?;
+
+                    let schema =
+                        column_def
+                            .try_as_column_schema()
+                            .context(InvalidColumnDefSnafu {
+                                column: &column_def.name,
+                            })?;
+                    Ok(AddColumnRequest {
+                        column_schema: schema,
+                        is_key: ac.is_key,
+                    })
+                })
+                .collect::<Result<Vec<_>>>()?;
+
+            let alter_kind = AlterKind::AddColumns {
+                columns: add_column_requests,
+            };
+
+            let request = AlterTableRequest {
+                catalog_name: expr.catalog_name,
+                schema_name: expr.schema_name,
+                table_name: expr.table_name,
+                alter_kind,
+            };
+            Ok(Some(request))
+        }
+        Some(Kind::DropColumns(DropColumns { drop_columns })) => {
+            let alter_kind = AlterKind::DropColumns {
+                names: drop_columns.into_iter().map(|c| c.name).collect(),
+            };
+
+            let request = AlterTableRequest {
+                catalog_name: expr.catalog_name,
+                schema_name: expr.schema_name,
+                table_name: expr.table_name,
+                alter_kind,
+            };
+            Ok(Some(request))
+        }
+        None => Ok(None),
+    }
+}
+
+pub fn create_table_schema(expr: &CreateExpr) -> Result<SchemaRef> {
+    let column_schemas = expr
+        .column_defs
+        .iter()
+        .map(|x| {
+            x.try_as_column_schema()
+                .context(InvalidColumnDefSnafu { column: &x.name })
+        })
+        .collect::<Result<Vec<ColumnSchema>>>()?;
+
+    ensure!(
+        column_schemas
+            .iter()
+            .any(|column| column.name == expr.time_index),
+        MissingTimestampColumnSnafu {
+            msg: format!("CreateExpr: {:?}", expr)
+        }
+    );
+
+    let column_schemas = column_schemas
+        .into_iter()
+        .map(|column_schema| {
+            if column_schema.name == expr.time_index {
+                column_schema.with_time_index(true)
+            } else {
+                column_schema
+            }
+        })
+        .collect::<Vec<_>>();
+
+    Ok(Arc::new(
+        SchemaBuilder::try_from(column_schemas)
+            .context(CreateSchemaSnafu)?
+            .build()
+            .context(CreateSchemaSnafu)?,
+    ))
+}
+
+pub fn create_expr_to_request(table_id: TableId, expr: CreateExpr) -> Result<CreateTableRequest> {
+    let schema = create_table_schema(&expr)?;
+    let primary_key_indices = expr
+        .primary_keys
+        .iter()
+        .map(|key| {
+            schema
+                .column_index_by_name(key)
+                .context(ColumnNotFoundSnafu {
+                    column_name: key,
+                    table_name: &expr.table_name,
+                })
+        })
+        .collect::<Result<Vec<usize>>>()?;
+
+    let catalog_name = expr
+        .catalog_name
+        .unwrap_or_else(|| DEFAULT_CATALOG_NAME.to_string());
+    let schema_name = expr
+        .schema_name
+        .unwrap_or_else(|| DEFAULT_SCHEMA_NAME.to_string());
+
+    let region_ids = if expr.region_ids.is_empty() {
+        vec![0]
+    } else {
+        expr.region_ids
+    };
+
+    Ok(CreateTableRequest {
+        id: table_id,
+        catalog_name,
+        schema_name,
+        table_name: expr.table_name,
+        desc: expr.desc,
+        schema,
+        region_numbers: region_ids,
+        primary_key_indices,
+        create_if_not_exists: expr.create_if_not_exists,
+        table_options: expr.table_options,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use api::v1::{AddColumn, AddColumns, ColumnDataType, ColumnDef, DropColumn};
+    use datatypes::prelude::ConcreteDataType;
+
+    use super::*;
+
+    #[test]
+    fn test_alter_expr_to_request() {
+        let expr = AlterExpr {
+            catalog_name: None,
+            schema_name: None,
+            table_name: "monitor".to_string(),
+
+            kind: Some(Kind::AddColumns(AddColumns {
+                add_columns: vec![AddColumn {
+                    column_def: Some(ColumnDef {
+                        name: "mem_usage".to_string(),
+                        datatype: ColumnDataType::Float64 as i32,
+                        is_nullable: false,
+                        default_constraint: None,
+                    }),
+                    is_key: false,
+                }],
+            })),
+        };
+
+        let alter_request = alter_expr_to_request(expr).unwrap().unwrap();
+        assert_eq!(None, alter_request.catalog_name);
+        assert_eq!(None, alter_request.schema_name);
+        assert_eq!("monitor".to_string(), alter_request.table_name);
+        let add_column = match alter_request.alter_kind {
+            AlterKind::AddColumns { mut columns } => columns.pop().unwrap(),
+            _ => unreachable!(),
+        };
+
+        assert!(!add_column.is_key);
+        assert_eq!("mem_usage", add_column.column_schema.name);
+        assert_eq!(
+            ConcreteDataType::float64_datatype(),
+            add_column.column_schema.data_type
+        );
+    }
+
+    #[test]
+    fn test_drop_column_expr() {
+        let expr = AlterExpr {
+            catalog_name: Some("test_catalog".to_string()),
+            schema_name: Some("test_schema".to_string()),
+            table_name: "monitor".to_string(),
+
+            kind: Some(Kind::DropColumns(DropColumns {
+                drop_columns: vec![DropColumn {
+                    name: "mem_usage".to_string(),
+                }],
+            })),
+        };
+
+        let alter_request = alter_expr_to_request(expr).unwrap().unwrap();
+        assert_eq!(Some("test_catalog".to_string()), alter_request.catalog_name);
+        assert_eq!(Some("test_schema".to_string()), alter_request.schema_name);
+        assert_eq!("monitor".to_string(), alter_request.table_name);
+
+        let mut drop_names = match alter_request.alter_kind {
+            AlterKind::DropColumns { names } => names,
+            _ => unreachable!(),
+        };
+        assert_eq!(1, drop_names.len());
+        assert_eq!("mem_usage".to_string(), drop_names.pop().unwrap());
+    }
+}