feat: support Create Table ... Like (#3372)

* feat: support `Create Table ... Like` * fix: `check_permission` for `Create Table ... Like` * style: renaming `name` -> `table_name` & `target` -> `source_name` and make `Create Table ... Like` testcase more complicated * rebase * avoid _ fn Signed-off-by: tison <[email protected]> --------- Signed-off-by: tison <[email protected]> Co-authored-by: tison <[email protected]>
GreptimeTeam · Mar 2, 2024 · 00cbbc9 · 00cbbc9
1 parent 7d30c24
commit 00cbbc9
Show file tree

Hide file tree

Showing 13 changed files with 179 additions and 44 deletions.
diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
@@ -471,6 +471,10 @@ pub fn check_permission(
         Statement::CreateTable(stmt) => {
             validate_param(&stmt.name, query_ctx)?;
         }
+        Statement::CreateTableLike(stmt) => {
+            validate_param(&stmt.table_name, query_ctx)?;
+            validate_param(&stmt.source_name, query_ctx)?;
+        }
         Statement::DropTable(drop_stmt) => {
             validate_param(drop_stmt.table_name(), query_ctx)?;
         }

diff --git a/src/operator/src/expr_factory.rs b/src/operator/src/expr_factory.rs
@@ -181,7 +181,7 @@ pub(crate) async fn create_external_expr(
     Ok(expr)
 }
 
-/// Convert `CreateTable` statement to `CreateExpr` gRPC request.
+/// Convert `CreateTable` statement to [`CreateTableExpr`] gRPC request.
 pub fn create_to_expr(create: &CreateTable, query_ctx: QueryContextRef) -> Result<CreateTableExpr> {
     let (catalog_name, schema_name, table_name) =
         table_idents_to_full_name(&create.name, &query_ctx)

diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs
@@ -153,6 +153,10 @@ impl StatementExecutor {
                 let _ = self.create_table(stmt, query_ctx).await?;
                 Ok(Output::AffectedRows(0))
             }
+            Statement::CreateTableLike(stmt) => {
+                let _ = self.create_table_like(stmt, query_ctx).await?;
+                Ok(Output::AffectedRows(0))
+            }
             Statement::CreateExternalTable(stmt) => {
                 let _ = self.create_external_table(stmt, query_ctx).await?;
                 Ok(Output::AffectedRows(0))
@@ -174,7 +178,6 @@ impl StatementExecutor {
                 let table_name = TableName::new(catalog, schema, table);
                 self.truncate_table(table_name).await
             }
-
             Statement::CreateDatabase(stmt) => {
                 self.create_database(
                     query_ctx.current_catalog(),
@@ -183,7 +186,6 @@ impl StatementExecutor {
                 )
                 .await
             }
-
             Statement::ShowCreateTable(show) => {
                 let (catalog, schema, table) =
                     table_idents_to_full_name(&show.table_name, &query_ctx)

diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs
@@ -21,6 +21,7 @@ use catalog::CatalogManagerRef;
 use chrono::Utc;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_catalog::format_full_table_name;
+use common_error::ext::BoxedError;
 use common_meta::cache_invalidator::Context;
 use common_meta::ddl::ExecutorContext;
 use common_meta::key::schema_name::{SchemaNameKey, SchemaNameValue};
@@ -34,11 +35,13 @@ use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::RawSchema;
 use lazy_static::lazy_static;
 use partition::partition::{PartitionBound, PartitionDef};
+use query::sql::create_table_stmt;
 use regex::Regex;
 use session::context::QueryContextRef;
+use session::table_name::table_idents_to_full_name;
 use snafu::{ensure, IntoError, OptionExt, ResultExt};
 use sql::statements::alter::AlterTable;
-use sql::statements::create::{CreateExternalTable, CreateTable, Partitions};
+use sql::statements::create::{CreateExternalTable, CreateTable, CreateTableLike, Partitions};
 use table::dist_table::DistTable;
 use table::metadata::{self, RawTableInfo, RawTableMeta, TableId, TableInfo, TableType};
 use table::requests::{AlterKind, AlterTableRequest, TableOptions};
@@ -54,6 +57,7 @@ use crate::error::{
     UnrecognizedTableOptionSnafu,
 };
 use crate::expr_factory;
+use crate::statement::show::create_partitions_stmt;
 
 lazy_static! {
     static ref NAME_PATTERN_REG: Regex = Regex::new(&format!("^{NAME_PATTERN}$")).unwrap();
@@ -71,6 +75,46 @@ impl StatementExecutor {
             .await
     }
 
+    #[tracing::instrument(skip_all)]
+    pub async fn create_table_like(
+        &self,
+        stmt: CreateTableLike,
+        ctx: QueryContextRef,
+    ) -> Result<TableRef> {
+        let (catalog, schema, table) = table_idents_to_full_name(&stmt.source_name, &ctx)
+            .map_err(BoxedError::new)
+            .context(error::ExternalSnafu)?;
+        let table_ref = self
+            .catalog_manager
+            .table(&catalog, &schema, &table)
+            .await
+            .context(CatalogSnafu)?
+            .context(TableNotFoundSnafu { table_name: &table })?;
+        let partitions = self
+            .partition_manager
+            .find_table_partitions(table_ref.table_info().table_id())
+            .await
+            .context(error::FindTablePartitionRuleSnafu { table_name: table })?;
+
+        let quote_style = ctx.quote_style();
+        let mut create_stmt = create_table_stmt(&table_ref.table_info(), quote_style)
+            .context(error::ParseQuerySnafu)?;
+        create_stmt.name = stmt.table_name;
+        create_stmt.if_not_exists = false;
+
+        let partitions = create_partitions_stmt(partitions)?.and_then(|mut partitions| {
+            if !partitions.column_list.is_empty() {
+                partitions.set_quote(quote_style);
+                Some(partitions)
+            } else {
+                None
+            }
+        });
+
+        let create_expr = &mut expr_factory::create_to_expr(&create_stmt, ctx.clone())?;
+        self.create_table_inner(create_expr, partitions, &ctx).await
+    }
+
     #[tracing::instrument(skip_all)]
     pub async fn create_external_table(
         &self,
@@ -160,7 +204,7 @@ impl StatementExecutor {
 
         table_info.ident.table_id = table_id;
 
-        let table_info = Arc::new(table_info.try_into().context(error::CreateTableInfoSnafu)?);
+        let table_info = Arc::new(table_info.try_into().context(CreateTableInfoSnafu)?);
         create_table.table_id = Some(api::v1::TableId { id: table_id });
 
         let table = DistTable::table(table_info);

diff --git a/src/operator/src/statement/show.rs b/src/operator/src/statement/show.rs
@@ -76,7 +76,7 @@ impl StatementExecutor {
     }
 }
 
-fn create_partitions_stmt(partitions: Vec<PartitionInfo>) -> Result<Option<Partitions>> {
+pub(crate) fn create_partitions_stmt(partitions: Vec<PartitionInfo>) -> Result<Option<Partitions>> {
     if partitions.is_empty() {
         return Ok(None);
     }

diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs
@@ -37,6 +37,7 @@ use object_store::ObjectStore;
 use once_cell::sync::Lazy;
 use regex::Regex;
 use session::context::QueryContextRef;
+pub use show_create_table::create_table_stmt;
 use snafu::{ensure, OptionExt, ResultExt};
 use sql::statements::create::Partitions;
 use sql::statements::show::{ShowDatabases, ShowKind, ShowTables, ShowVariables};
@@ -258,16 +259,9 @@ pub fn show_create_table(
     let table_info = table.table_info();
     let table_name = &table_info.name;
 
-    // Default to double quote and fallback to back quote
-    let quote_style = if query_ctx.sql_dialect().is_delimited_identifier_start('"') {
-        '"'
-    } else if query_ctx.sql_dialect().is_delimited_identifier_start('\'') {
-        '\''
-    } else {
-        '`'
-    };
+    let quote_style = query_ctx.quote_style();
 
-    let mut stmt = show_create_table::create_table_stmt(&table_info, quote_style)?;
+    let mut stmt = create_table_stmt(&table_info, quote_style)?;
     stmt.partitions = partitions.map(|mut p| {
         p.set_quote(quote_style);
         p

diff --git a/src/session/src/context.rs b/src/session/src/context.rs
@@ -172,6 +172,17 @@ impl QueryContext {
             session.set_timezone(tz.as_ref().clone())
         }
     }
+
+    /// Default to double quote and fallback to back quote
+    pub fn quote_style(&self) -> char {
+        if self.sql_dialect().is_delimited_identifier_start('"') {
+            '"'
+        } else if self.sql_dialect().is_delimited_identifier_start('\'') {
+            '\''
+        } else {
+            '`'
+        }
+    }
 }
 
 impl QueryContextBuilder {

diff --git a/src/sql/src/parser.rs b/src/sql/src/parser.rs
@@ -75,19 +75,23 @@ impl<'a> ParserContext<'a> {
     }
 
     pub fn parse_table_name(sql: &'a str, dialect: &dyn Dialect) -> Result<ObjectName> {
-        let mut parser = Parser::new(dialect)
+        let parser = Parser::new(dialect)
             .with_options(ParserOptions::new().with_trailing_commas(true))
             .try_with_sql(sql)
             .context(SyntaxSnafu)?;
+        ParserContext { parser, sql }.intern_parse_table_name()
+    }
 
-        let raw_table_name = parser.parse_object_name().context(error::UnexpectedSnafu {
-            sql,
-            expected: "a table name",
-            actual: parser.peek_token().to_string(),
-        })?;
-        let table_name = Self::canonicalize_object_name(raw_table_name);
-
-        Ok(table_name)
+    pub(crate) fn intern_parse_table_name(&mut self) -> Result<ObjectName> {
+        let raw_table_name = self
+            .parser
+            .parse_object_name()
+            .context(error::UnexpectedSnafu {
+                sql: self.sql,
+                expected: "a table name",
+                actual: self.parser.peek_token().to_string(),
+            })?;
+        Ok(Self::canonicalize_object_name(raw_table_name))
     }
 
     pub fn parse_function(sql: &'a str, dialect: &dyn Dialect) -> Result<Expr> {

diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs
@@ -32,7 +32,7 @@ use crate::error::{
 };
 use crate::parser::ParserContext;
 use crate::statements::create::{
-    CreateDatabase, CreateExternalTable, CreateTable, Partitions, TIME_INDEX,
+    CreateDatabase, CreateExternalTable, CreateTable, CreateTableLike, Partitions, TIME_INDEX,
 };
 use crate::statements::get_data_type_by_alias_name;
 use crate::statements::statement::Statement;
@@ -66,15 +66,7 @@ impl<'a> ParserContext<'a> {
         let if_not_exists =
             self.parser
                 .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
-        let raw_table_name = self
-            .parser
-            .parse_object_name()
-            .context(error::UnexpectedSnafu {
-                sql: self.sql,
-                expected: "a table name",
-                actual: self.peek_token_as_string(),
-            })?;
-        let table_name = Self::canonicalize_object_name(raw_table_name);
+        let table_name = self.intern_parse_table_name()?;
         let (columns, constraints) = self.parse_columns()?;
         let engine = self.parse_table_engine(common_catalog::consts::FILE_ENGINE)?;
         let options = self
@@ -136,15 +128,16 @@ impl<'a> ParserContext<'a> {
             self.parser
                 .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
 
-        let raw_table_name = self
-            .parser
-            .parse_object_name()
-            .context(error::UnexpectedSnafu {
-                sql: self.sql,
-                expected: "a table name",
-                actual: self.peek_token_as_string(),
-            })?;
-        let table_name = Self::canonicalize_object_name(raw_table_name);
+        let table_name = self.intern_parse_table_name()?;
+
+        if self.parser.parse_keyword(Keyword::LIKE) {
+            let source_name = self.intern_parse_table_name()?;
+
+            return Ok(Statement::CreateTableLike(CreateTableLike {
+                table_name,
+                source_name,
+            }));
+        }
 
         let (columns, constraints) = self.parse_columns()?;
 
@@ -739,6 +732,23 @@ mod tests {
     use crate::dialect::GreptimeDbDialect;
     use crate::parser::ParseOptions;
 
+    #[test]
+    fn test_parse_create_table_like() {
+        let sql = "CREATE TABLE t1 LIKE t2";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(1, stmts.len());
+        match &stmts[0] {
+            Statement::CreateTableLike(c) => {
+                assert_eq!(c.table_name.to_string(), "t1");
+                assert_eq!(c.source_name.to_string(), "t2");
+            }
+            _ => unreachable!(),
+        }
+    }
+
     #[test]
     fn test_validate_external_table_options() {
         let sql = "CREATE EXTERNAL TABLE city (

diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs
@@ -219,6 +219,14 @@ pub struct CreateExternalTable {
     pub engine: String,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut)]
+pub struct CreateTableLike {
+    /// Table name
+    pub table_name: ObjectName,
+    /// The table that is designated to be imitated by `Like`
+    pub source_name: ObjectName,
+}
+
 #[cfg(test)]
 mod tests {
     use crate::dialect::GreptimeDbDialect;

diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs
@@ -19,7 +19,9 @@ use sqlparser_derive::{Visit, VisitMut};
 use super::show::ShowVariables;
 use crate::error::{ConvertToDfStatementSnafu, Error};
 use crate::statements::alter::AlterTable;
-use crate::statements::create::{CreateDatabase, CreateExternalTable, CreateTable};
+use crate::statements::create::{
+    CreateDatabase, CreateExternalTable, CreateTable, CreateTableLike,
+};
 use crate::statements::delete::Delete;
 use crate::statements::describe::DescribeTable;
 use crate::statements::drop::DropTable;
@@ -45,6 +47,8 @@ pub enum Statement {
     CreateTable(CreateTable),
     // CREATE EXTERNAL TABLE
     CreateExternalTable(CreateExternalTable),
+    // CREATE TABLE ... LIKE
+    CreateTableLike(CreateTableLike),
     // DROP TABLE
     DropTable(DropTable),
     // CREATE DATABASE

diff --git a/tests/cases/standalone/common/create/create.result b/tests/cases/standalone/common/create/create.result
@@ -143,3 +143,43 @@ DROP TABLE neg_default_value;
 
 Affected Rows: 0
 
+CREATE TABLE test_like_1 (PK STRING PRIMARY KEY, i INTEGER DEFAULT 7, j TIMESTAMP TIME INDEX);
+
+Affected Rows: 0
+
+CREATE TABLE test_like_2 LIKE test_like_1;
+
+Affected Rows: 0
+
+CREATE TABLE test_like_2 LIKE test_like_1;
+
+Error: 4000(TableAlreadyExists), Table already exists: `greptime.public.test_like_2`
+
+DESC TABLE test_like_1;
+
++--------+----------------------+-----+------+---------+---------------+
+| Column | Type                 | Key | Null | Default | Semantic Type |
++--------+----------------------+-----+------+---------+---------------+
+| pk     | String               | PRI | YES  |         | TAG           |
+| i      | Int32                |     | YES  | 7       | FIELD         |
+| j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
++--------+----------------------+-----+------+---------+---------------+
+
+DESC TABLE test_like_2;
+
++--------+----------------------+-----+------+---------+---------------+
+| Column | Type                 | Key | Null | Default | Semantic Type |
++--------+----------------------+-----+------+---------+---------------+
+| pk     | String               | PRI | YES  |         | TAG           |
+| i      | Int32                |     | YES  | 7       | FIELD         |
+| j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
++--------+----------------------+-----+------+---------+---------------+
+
+DROP TABLE test_like_1;
+
+Affected Rows: 0
+
+DROP TABLE test_like_2;
+
+Affected Rows: 0
+
diff --git a/tests/cases/standalone/common/create/create.sql b/tests/cases/standalone/common/create/create.sql
@@ -57,3 +57,17 @@ CREATE TABLE neg_default_value(i INT DEFAULT -1024, ts TIMESTAMP TIME INDEX);
 desc TABLE neg_default_value;
 
 DROP TABLE neg_default_value;
+
+CREATE TABLE test_like_1 (PK STRING PRIMARY KEY, i INTEGER DEFAULT 7, j TIMESTAMP TIME INDEX);
+
+CREATE TABLE test_like_2 LIKE test_like_1;
+
+CREATE TABLE test_like_2 LIKE test_like_1;
+
+DESC TABLE test_like_1;
+
+DESC TABLE test_like_2;
+
+DROP TABLE test_like_1;
+
+DROP TABLE test_like_2;