From 9161796dfa756c74df45f801d6b0e2b881326c12 Mon Sep 17 00:00:00 2001 From: Jiachun Feng Date: Mon, 20 Feb 2023 16:43:50 +0800 Subject: [PATCH] feat: export the data from a table to parquet files (#1000) * feat: copy table parser * feat: coopy table * chore: minor fix * chore: give stmt a more clearer name * chore: unified naming * chore: minor change * chore: add a todo * chore: end up with an empty file when occur an empty table * feat: format with copy table * feat: with options * chore: by cr * chore: default 5M rows per segment * Update src/datanode/src/sql/copy_table.rs Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> * Update src/datanode/src/sql/copy_table.rs Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> * Update src/datanode/src/error.rs Co-authored-by: fys <40801205+Fengys123@users.noreply.github.com> --------- Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> Co-authored-by: fys <40801205+Fengys123@users.noreply.github.com> --- src/datanode/src/error.rs | 39 +++++++- src/datanode/src/instance/sql.rs | 62 +++++++----- src/datanode/src/sql.rs | 27 ++--- src/datanode/src/sql/copy_table.rs | 147 ++++++++++++++++++++++++++++ src/frontend/src/instance.rs | 6 +- src/query/src/datafusion/planner.rs | 3 +- src/sql/src/error.rs | 11 ++- src/sql/src/parser.rs | 2 + src/sql/src/parsers.rs | 1 + src/sql/src/parsers/copy_parser.rs | 132 +++++++++++++++++++++++++ src/sql/src/statements.rs | 1 + src/sql/src/statements/copy.rs | 62 ++++++++++++ src/sql/src/statements/statement.rs | 3 + src/table/src/requests.rs | 9 ++ 14 files changed, 463 insertions(+), 42 deletions(-) create mode 100644 src/datanode/src/sql/copy_table.rs create mode 100644 src/sql/src/parsers/copy_parser.rs create mode 100644 src/sql/src/statements/copy.rs diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 052ee7075c7..d02d3200631 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -16,6 +16,7 @@ use std::any::Any; use common_error::prelude::*; use common_recordbatch::error::Error as RecordBatchError; +use datafusion::parquet; use datatypes::prelude::ConcreteDataType; use storage::error::Error as StorageError; use table::error::Error as TableError; @@ -355,6 +356,38 @@ pub enum Error { #[snafu(backtrace)] source: query::error::Error, }, + + #[snafu(display("Failed to copy data from table: {}, source: {}", table_name, source))] + CopyTable { + table_name: String, + #[snafu(backtrace)] + source: TableError, + }, + + #[snafu(display("Failed to execute table scan, source: {}", source))] + TableScanExec { + #[snafu(backtrace)] + source: common_query::error::Error, + }, + + #[snafu(display("Failed to write parquet file, source: {}", source))] + WriteParquet { + source: parquet::errors::ParquetError, + backtrace: Backtrace, + }, + + #[snafu(display("Failed to poll stream, source: {}", source))] + PollStream { + source: datatypes::arrow::error::ArrowError, + backtrace: Backtrace, + }, + + #[snafu(display("Failed to write object into path: {}, source: {}", path, source))] + WriteObject { + path: String, + backtrace: Backtrace, + source: object_store::Error, + }, } pub type Result = std::result::Result; @@ -417,7 +450,9 @@ impl ErrorExt for Error { | MissingRequiredField { .. } | IncorrectInternalState { .. } => StatusCode::Internal, - InitBackend { .. } => StatusCode::StorageUnavailable, + InitBackend { .. } | WriteParquet { .. } | PollStream { .. } | WriteObject { .. } => { + StatusCode::StorageUnavailable + } OpenLogStore { source } => source.status_code(), StartScriptManager { source } => source.status_code(), OpenStorageEngine { source } => source.status_code(), @@ -426,6 +461,8 @@ impl ErrorExt for Error { TableIdProviderNotFound { .. } => StatusCode::Unsupported, BumpTableId { source, .. } => source.status_code(), ColumnDefaultValue { source, .. } => source.status_code(), + CopyTable { source, .. } => source.status_code(), + TableScanExec { source, .. } => source.status_code(), } } diff --git a/src/datanode/src/instance/sql.rs b/src/datanode/src/instance/sql.rs index 28441efa550..79710e9c158 100644 --- a/src/datanode/src/instance/sql.rs +++ b/src/datanode/src/instance/sql.rs @@ -31,7 +31,7 @@ use snafu::prelude::*; use sql::ast::ObjectName; use sql::statements::statement::Statement; use table::engine::TableReference; -use table::requests::{CreateDatabaseRequest, DropTableRequest}; +use table::requests::{CopyTableRequest, CreateDatabaseRequest, DropTableRequest}; use crate::error::{self, BumpTableIdSnafu, ExecuteSqlSnafu, Result, TableIdProviderNotFoundSnafu}; use crate::instance::Instance; @@ -57,10 +57,10 @@ impl Instance { .await .context(ExecuteSqlSnafu) } - QueryStatement::Sql(Statement::Insert(i)) => { + QueryStatement::Sql(Statement::Insert(insert)) => { let requests = self .sql_handler - .insert_to_requests(self.catalog_manager.clone(), *i, query_ctx.clone()) + .insert_to_requests(self.catalog_manager.clone(), *insert, query_ctx.clone()) .await?; match requests { @@ -86,14 +86,14 @@ impl Instance { } } } - QueryStatement::Sql(Statement::Delete(d)) => { - let request = SqlRequest::Delete(*d); + QueryStatement::Sql(Statement::Delete(delete)) => { + let request = SqlRequest::Delete(*delete); self.sql_handler.execute(request, query_ctx).await } - QueryStatement::Sql(Statement::CreateDatabase(c)) => { + QueryStatement::Sql(Statement::CreateDatabase(create_database)) => { let request = CreateDatabaseRequest { - db_name: c.name.to_string(), - create_if_not_exists: c.if_not_exists, + db_name: create_database.name.to_string(), + create_if_not_exists: create_database.if_not_exists, }; info!("Creating a new database: {}", request.db_name); @@ -103,7 +103,7 @@ impl Instance { .await } - QueryStatement::Sql(Statement::CreateTable(c)) => { + QueryStatement::Sql(Statement::CreateTable(create_table)) => { let table_id = self .table_id_provider .as_ref() @@ -111,15 +111,15 @@ impl Instance { .next_table_id() .await .context(BumpTableIdSnafu)?; - let _engine_name = c.engine.clone(); + let _engine_name = create_table.engine.clone(); // TODO(hl): Select table engine by engine_name - let name = c.name.clone(); + let name = create_table.name.clone(); let (catalog, schema, table) = table_idents_to_full_name(&name, query_ctx.clone())?; let table_ref = TableReference::full(&catalog, &schema, &table); - let request = self - .sql_handler - .create_to_request(table_id, c, &table_ref)?; + let request = + self.sql_handler + .create_to_request(table_id, create_table, &table_ref)?; let table_id = request.id; info!("Creating table: {table_ref}, table id = {table_id}",); @@ -148,27 +148,27 @@ impl Instance { .execute(SqlRequest::DropTable(req), query_ctx) .await } - QueryStatement::Sql(Statement::ShowDatabases(stmt)) => { + QueryStatement::Sql(Statement::ShowDatabases(show_databases)) => { self.sql_handler - .execute(SqlRequest::ShowDatabases(stmt), query_ctx) + .execute(SqlRequest::ShowDatabases(show_databases), query_ctx) .await } - QueryStatement::Sql(Statement::ShowTables(stmt)) => { + QueryStatement::Sql(Statement::ShowTables(show_tables)) => { self.sql_handler - .execute(SqlRequest::ShowTables(stmt), query_ctx) + .execute(SqlRequest::ShowTables(show_tables), query_ctx) .await } - QueryStatement::Sql(Statement::Explain(stmt)) => { + QueryStatement::Sql(Statement::Explain(explain)) => { self.sql_handler - .execute(SqlRequest::Explain(Box::new(stmt)), query_ctx) + .execute(SqlRequest::Explain(Box::new(explain)), query_ctx) .await } - QueryStatement::Sql(Statement::DescribeTable(stmt)) => { + QueryStatement::Sql(Statement::DescribeTable(describe_table)) => { self.sql_handler - .execute(SqlRequest::DescribeTable(stmt), query_ctx) + .execute(SqlRequest::DescribeTable(describe_table), query_ctx) .await } - QueryStatement::Sql(Statement::ShowCreateTable(_stmt)) => { + QueryStatement::Sql(Statement::ShowCreateTable(_show_create_table)) => { unimplemented!("SHOW CREATE TABLE is unimplemented yet"); } QueryStatement::Sql(Statement::Use(ref schema)) => { @@ -182,6 +182,22 @@ impl Instance { Ok(Output::RecordBatches(RecordBatches::empty())) } + QueryStatement::Sql(Statement::Copy(copy_table)) => { + let (catalog_name, schema_name, table_name) = + table_idents_to_full_name(copy_table.table_name(), query_ctx.clone())?; + let file_name = copy_table.file_name().to_string(); + + let req = CopyTableRequest { + catalog_name, + schema_name, + table_name, + file_name, + }; + + self.sql_handler + .execute(SqlRequest::CopyTable(req), query_ctx) + .await + } } } diff --git a/src/datanode/src/sql.rs b/src/datanode/src/sql.rs index beea10e72a1..34667260d9a 100644 --- a/src/datanode/src/sql.rs +++ b/src/datanode/src/sql.rs @@ -31,6 +31,7 @@ use crate::error::{self, ExecuteSqlSnafu, GetTableSnafu, Result, TableNotFoundSn use crate::instance::sql::table_idents_to_full_name; mod alter; +mod copy_table; mod create; mod delete; mod drop_table; @@ -48,6 +49,7 @@ pub enum SqlRequest { DescribeTable(DescribeTable), Explain(Box), Delete(Delete), + CopyTable(CopyTableRequest), } // Handler to execute SQL except query @@ -81,31 +83,30 @@ impl SqlHandler { SqlRequest::CreateDatabase(req) => self.create_database(req, query_ctx.clone()).await, SqlRequest::Alter(req) => self.alter(req).await, SqlRequest::DropTable(req) => self.drop_table(req).await, - SqlRequest::Delete(stmt) => self.delete(query_ctx.clone(), stmt).await, - SqlRequest::ShowDatabases(stmt) => { - show_databases(stmt, self.catalog_manager.clone()).context(ExecuteSqlSnafu) + SqlRequest::Delete(req) => self.delete(query_ctx.clone(), req).await, + SqlRequest::CopyTable(req) => self.copy_table(req).await, + SqlRequest::ShowDatabases(req) => { + show_databases(req, self.catalog_manager.clone()).context(ExecuteSqlSnafu) } - SqlRequest::ShowTables(stmt) => { - show_tables(stmt, self.catalog_manager.clone(), query_ctx.clone()) + SqlRequest::ShowTables(req) => { + show_tables(req, self.catalog_manager.clone(), query_ctx.clone()) .context(ExecuteSqlSnafu) } - SqlRequest::DescribeTable(stmt) => { + SqlRequest::DescribeTable(req) => { let (catalog, schema, table) = - table_idents_to_full_name(stmt.name(), query_ctx.clone())?; + table_idents_to_full_name(req.name(), query_ctx.clone())?; let table = self .catalog_manager .table(&catalog, &schema, &table) .context(error::CatalogSnafu)? .with_context(|| TableNotFoundSnafu { - table_name: stmt.name().to_string(), + table_name: req.name().to_string(), })?; describe_table(table).context(ExecuteSqlSnafu) } - SqlRequest::Explain(stmt) => { - explain(stmt, self.query_engine.clone(), query_ctx.clone()) - .await - .context(ExecuteSqlSnafu) - } + SqlRequest::Explain(req) => explain(req, self.query_engine.clone(), query_ctx.clone()) + .await + .context(ExecuteSqlSnafu), }; if let Err(e) = &result { error!(e; "{query_ctx}"); diff --git a/src/datanode/src/sql/copy_table.rs b/src/datanode/src/sql/copy_table.rs new file mode 100644 index 00000000000..5827c41c3f7 --- /dev/null +++ b/src/datanode/src/sql/copy_table.rs @@ -0,0 +1,147 @@ +// Copyright 2023 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::pin::Pin; + +use common_query::physical_plan::SessionContext; +use common_query::Output; +use common_recordbatch::adapter::DfRecordBatchStreamAdapter; +use datafusion::parquet::arrow::ArrowWriter; +use datafusion::parquet::basic::{Compression, Encoding}; +use datafusion::parquet::file::properties::WriterProperties; +use datafusion::physical_plan::RecordBatchStream; +use futures::TryStreamExt; +use object_store::services::fs::Builder; +use object_store::ObjectStore; +use snafu::ResultExt; +use table::engine::TableReference; +use table::requests::CopyTableRequest; + +use crate::error::{self, Result}; +use crate::sql::SqlHandler; + +impl SqlHandler { + pub(crate) async fn copy_table(&self, req: CopyTableRequest) -> Result { + let table_ref = TableReference { + catalog: &req.catalog_name, + schema: &req.schema_name, + table: &req.table_name, + }; + let table = self.get_table(&table_ref)?; + + let stream = table + .scan(None, &[], None) + .await + .with_context(|_| error::CopyTableSnafu { + table_name: table_ref.to_string(), + })?; + + let stream = stream + .execute(0, SessionContext::default().task_ctx()) + .context(error::TableScanExecSnafu)?; + let stream = Box::pin(DfRecordBatchStreamAdapter::new(stream)); + + let accessor = Builder::default().build().unwrap(); + let object_store = ObjectStore::new(accessor); + + let mut parquet_writer = ParquetWriter::new(req.file_name, stream, object_store); + // TODO(jiachun): + // For now, COPY is implemented synchronously. + // When copying large table, it will be blocked for a long time. + // Maybe we should make "copy" runs in background? + // Like PG: https://www.postgresql.org/docs/current/sql-copy.html + let rows = parquet_writer.flush().await?; + + Ok(Output::AffectedRows(rows)) + } +} + +type DfRecordBatchStream = Pin>; + +struct ParquetWriter { + file_name: String, + stream: DfRecordBatchStream, + object_store: ObjectStore, + max_row_group_size: usize, + max_rows_in_segment: usize, +} + +impl ParquetWriter { + pub fn new(file_name: String, stream: DfRecordBatchStream, object_store: ObjectStore) -> Self { + Self { + file_name, + stream, + object_store, + // TODO(jiachun): make these configurable: WITH (max_row_group_size=xxx, max_rows_in_segment=xxx) + max_row_group_size: 4096, + max_rows_in_segment: 5000000, // default 5M rows per segment + } + } + + pub async fn flush(&mut self) -> Result { + let schema = self.stream.as_ref().schema(); + let writer_props = WriterProperties::builder() + .set_compression(Compression::ZSTD) + .set_encoding(Encoding::PLAIN) + .set_max_row_group_size(self.max_row_group_size) + .build(); + let mut total_rows = 0; + loop { + let mut buf = vec![]; + let mut arrow_writer = + ArrowWriter::try_new(&mut buf, schema.clone(), Some(writer_props.clone())) + .context(error::WriteParquetSnafu)?; + + let mut rows = 0; + let mut end_loop = true; + // TODO(hl & jiachun): Since OpenDAL's writer is async and ArrowWriter requires a `std::io::Write`, + // here we use a Vec to buffer all parquet bytes in memory and write to object store + // at a time. Maybe we should find a better way to brige ArrowWriter and OpenDAL's object. + while let Some(batch) = self + .stream + .try_next() + .await + .context(error::PollStreamSnafu)? + { + arrow_writer + .write(&batch) + .context(error::WriteParquetSnafu)?; + rows += batch.num_rows(); + if rows >= self.max_rows_in_segment { + end_loop = false; + break; + } + } + + let start_row_num = total_rows + 1; + total_rows += rows; + arrow_writer.close().context(error::WriteParquetSnafu)?; + + // if rows == 0, we just end up with an empty file. + // + // file_name like: + // "file_name_1_1000000" (row num: 1 ~ 1000000), + // "file_name_1000001_xxx" (row num: 1000001 ~ xxx) + let file_name = format!("{}_{}_{}", self.file_name, start_row_num, total_rows); + let object = self.object_store.object(&file_name); + object.write(buf).await.context(error::WriteObjectSnafu { + path: object.path(), + })?; + + if end_loop { + return Ok(total_rows); + } + } + } +} diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 127597bc7d9..63a27e10885 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -396,7 +396,8 @@ impl Instance { | Statement::Insert(_) | Statement::Delete(_) | Statement::Alter(_) - | Statement::DropTable(_) => self.sql_handler.do_statement_query(stmt, query_ctx).await, + | Statement::DropTable(_) + | Statement::Copy(_) => self.sql_handler.do_statement_query(stmt, query_ctx).await, Statement::Use(db) => self.handle_use(db, query_ctx), Statement::ShowCreateTable(_) => NotSupportedSnafu { feat: format!("{stmt:?}"), @@ -588,6 +589,9 @@ pub fn check_permission( Statement::Delete(delete) => { validate_param(delete.table_name(), query_ctx)?; } + Statement::Copy(stmd) => { + validate_param(stmd.table_name(), query_ctx)?; + } } Ok(()) } diff --git a/src/query/src/datafusion/planner.rs b/src/query/src/datafusion/planner.rs index 15035fbe367..a7493caf9c1 100644 --- a/src/query/src/datafusion/planner.rs +++ b/src/query/src/datafusion/planner.rs @@ -103,7 +103,8 @@ where | Statement::Alter(_) | Statement::Insert(_) | Statement::DropTable(_) - | Statement::Use(_) => unreachable!(), + | Statement::Use(_) + | Statement::Copy(_) => unreachable!(), } } } diff --git a/src/sql/src/error.rs b/src/sql/src/error.rs index dc140d3e526..b3a4328886c 100644 --- a/src/sql/src/error.rs +++ b/src/sql/src/error.rs @@ -135,6 +135,9 @@ pub enum Error { timestamp: Timestamp, target_unit: TimeUnit, }, + + #[snafu(display("Unsupported format option: {}", name))] + UnsupportedCopyFormatOption { name: String }, } impl ErrorExt for Error { @@ -156,12 +159,14 @@ impl ErrorExt for Error { InvalidColumnOption { .. } | InvalidDatabaseName { .. } | ColumnTypeMismatch { .. } - | InvalidTableName { .. } => StatusCode::InvalidArguments, + | InvalidTableName { .. } + | InvalidSqlValue { .. } + | TimestampOverflow { .. } + | UnsupportedCopyFormatOption { .. } => StatusCode::InvalidArguments, + UnsupportedAlterTableStatement { .. } => StatusCode::InvalidSyntax, SerializeColumnDefaultConstraint { source, .. } => source.status_code(), ConvertToGrpcDataType { source, .. } => source.status_code(), - InvalidSqlValue { .. } => StatusCode::InvalidArguments, - TimestampOverflow { .. } => StatusCode::InvalidArguments, } } diff --git a/src/sql/src/parser.rs b/src/sql/src/parser.rs index 47759279ab7..8ad5fcff0df 100644 --- a/src/sql/src/parser.rs +++ b/src/sql/src/parser.rs @@ -117,6 +117,8 @@ impl<'a> ParserContext<'a> { Ok(Statement::Use(database_name.value)) } + Keyword::COPY => self.parse_copy(), + // todo(hl) support more statements. _ => self.unsupported(self.peek_token_as_string()), } diff --git a/src/sql/src/parsers.rs b/src/sql/src/parsers.rs index f356cf70a3a..fb9dfc75de3 100644 --- a/src/sql/src/parsers.rs +++ b/src/sql/src/parsers.rs @@ -13,6 +13,7 @@ // limitations under the License. mod alter_parser; +pub(crate) mod copy_parser; pub(crate) mod create_parser; pub(crate) mod delete_parser; pub(crate) mod insert_parser; diff --git a/src/sql/src/parsers/copy_parser.rs b/src/sql/src/parsers/copy_parser.rs new file mode 100644 index 00000000000..5ee2caca516 --- /dev/null +++ b/src/sql/src/parsers/copy_parser.rs @@ -0,0 +1,132 @@ +// Copyright 2023 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 snafu::ResultExt; +use sqlparser::ast::Value; +use sqlparser::keywords::Keyword; + +use crate::error::{self, Result}; +use crate::parser::ParserContext; +use crate::statements::copy::{CopyTable, Format}; +use crate::statements::statement::Statement; + +// COPY tbl TO 'output.parquet'; +impl<'a> ParserContext<'a> { + pub(crate) fn parse_copy(&mut self) -> Result { + self.parser.next_token(); + let copy_table = self.parse_copy_table()?; + Ok(Statement::Copy(copy_table)) + } + + fn parse_copy_table(&mut self) -> Result { + let table_name = + self.parser + .parse_object_name() + .with_context(|_| error::UnexpectedSnafu { + sql: self.sql, + expected: "a table name", + actual: self.peek_token_as_string(), + })?; + + self.parser + .expect_keyword(Keyword::TO) + .context(error::SyntaxSnafu { sql: self.sql })?; + + let file_name = + self.parser + .parse_literal_string() + .with_context(|_| error::UnexpectedSnafu { + sql: self.sql, + expected: "a file name", + actual: self.peek_token_as_string(), + })?; + + let options = self + .parser + .parse_options(Keyword::WITH) + .context(error::SyntaxSnafu { sql: self.sql })?; + + // default format is parquet + let mut format = Format::Parquet; + for option in options { + if option.name.value.to_uppercase() == "FORMAT" { + if let Value::SingleQuotedString(fmt_str) = option.value { + format = Format::try_from(fmt_str)?; + } + } + } + + Ok(CopyTable::new(table_name, file_name, format)) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches::assert_matches; + + use sqlparser::dialect::GenericDialect; + + use super::*; + + #[test] + fn test_parse_copy_table() { + let sql0 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet'"; + let sql1 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'parquet')"; + let result0 = ParserContext::create_with_dialect(sql0, &GenericDialect {}).unwrap(); + let result1 = ParserContext::create_with_dialect(sql1, &GenericDialect {}).unwrap(); + + for mut result in vec![result0, result1] { + assert_eq!(1, result.len()); + + let statement = result.remove(0); + assert_matches!(statement, Statement::Copy { .. }); + match statement { + Statement::Copy(copy_table) => { + let (catalog, schema, table) = + if let [catalog, schema, table] = ©_table.table_name().0[..] { + ( + catalog.value.clone(), + schema.value.clone(), + table.value.clone(), + ) + } else { + unreachable!() + }; + + assert_eq!("catalog0", catalog); + assert_eq!("schema0", schema); + assert_eq!("tbl", table); + + let file_name = copy_table.file_name(); + assert_eq!("tbl_file.parquet", file_name); + + let format = copy_table.format(); + assert_eq!(Format::Parquet, *format); + } + _ => unreachable!(), + } + } + } + + #[test] + fn test_parse_copy_table_with_unsupopoted_format() { + let sql = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'unknow_format')"; + let result = ParserContext::create_with_dialect(sql, &GenericDialect {}); + assert!(result.is_err()); + assert_matches!( + result.err().unwrap(), + error::Error::UnsupportedCopyFormatOption { .. } + ); + } +} diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 5a87ecb46a6..593b746cb12 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod alter; +pub mod copy; pub mod create; pub mod delete; pub mod describe; diff --git a/src/sql/src/statements/copy.rs b/src/sql/src/statements/copy.rs new file mode 100644 index 00000000000..47c62db2339 --- /dev/null +++ b/src/sql/src/statements/copy.rs @@ -0,0 +1,62 @@ +// Copyright 2023 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 sqlparser::ast::ObjectName; + +use crate::error::{self, Result}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CopyTable { + table_name: ObjectName, + file_name: String, + format: Format, +} + +impl CopyTable { + pub(crate) fn new(table_name: ObjectName, file_name: String, format: Format) -> Self { + Self { + table_name, + file_name, + format, + } + } + + pub fn table_name(&self) -> &ObjectName { + &self.table_name + } + + pub fn file_name(&self) -> &str { + &self.file_name + } + + pub fn format(&self) -> &Format { + &self.format + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Format { + Parquet, +} + +impl TryFrom for Format { + type Error = error::Error; + + fn try_from(name: String) -> Result { + match name.to_uppercase().as_str() { + "PARQUET" => Ok(Format::Parquet), + _ => error::UnsupportedCopyFormatOptionSnafu { name }.fail(), + } + } +} diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs index 558f9a4b835..51e5c1e8bc1 100644 --- a/src/sql/src/statements/statement.rs +++ b/src/sql/src/statements/statement.rs @@ -13,6 +13,7 @@ // limitations under the License. use crate::statements::alter::AlterTable; +use crate::statements::copy::CopyTable; use crate::statements::create::{CreateDatabase, CreateTable}; use crate::statements::delete::Delete; use crate::statements::describe::DescribeTable; @@ -51,6 +52,8 @@ pub enum Statement { // EXPLAIN QUERY Explain(Explain), Use(String), + // COPY + Copy(CopyTable), } /// Comment hints from SQL. diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index f505f65c632..dc4cbc574c7 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -106,3 +106,12 @@ pub struct DeleteRequest { /// The key is the column name, and the value is the column value. pub key_column_values: HashMap, } + +/// Copy table request +#[derive(Debug)] +pub struct CopyTableRequest { + pub catalog_name: String, + pub schema_name: String, + pub table_name: String, + pub file_name: String, +}