diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index be17ee810fc..4f7e78b03d3 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -495,7 +495,8 @@ pub fn check_permission( | Statement::Explain(_) | Statement::Tql(_) | Statement::Delete(_) - | Statement::DeclareCursor(_) => {} + | Statement::DeclareCursor(_) + | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {} // database ops won't be checked Statement::CreateDatabase(_) | Statement::ShowDatabases(_) diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index ad842a40fe2..c72a09d72b4 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -14,6 +14,7 @@ mod admin; mod copy_database; +mod copy_query_to; mod copy_table_from; mod copy_table_to; mod cursor; @@ -57,14 +58,16 @@ use session::context::{Channel, QueryContextRef}; use session::table_name::table_idents_to_full_name; use set::set_query_timeout; use snafu::{ensure, OptionExt, ResultExt}; -use sql::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument}; +use sql::statements::copy::{ + CopyDatabase, CopyDatabaseArgument, CopyQueryToArgument, CopyTable, CopyTableArgument, +}; use sql::statements::set_variables::SetVariables; use sql::statements::show::ShowCreateTableVariant; use sql::statements::statement::Statement; use sql::statements::OptionMap; use sql::util::format_raw_object_name; use sqlparser::ast::ObjectName; -use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest}; +use table::requests::{CopyDatabaseRequest, CopyDirection, CopyQueryToRequest, CopyTableRequest}; use table::table_name::TableName; use table::table_reference::TableReference; use table::TableRef; @@ -165,6 +168,17 @@ impl StatementExecutor { Statement::ShowFlows(stmt) => self.show_flows(stmt, query_ctx).await, + Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(stmt)) => { + let query_output = self + .plan_exec(QueryStatement::Sql(*stmt.query), query_ctx) + .await?; + let req = to_copy_query_request(stmt.arg)?; + + self.copy_query_to(req, query_output) + .await + .map(Output::new_with_affected_rows) + } + Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => { let req = to_copy_table_request(stmt, query_ctx.clone())?; match req.direction { @@ -516,6 +530,20 @@ fn derive_timeout(stmt: &QueryStatement, query_ctx: &QueryContextRef) -> Option< } } +fn to_copy_query_request(stmt: CopyQueryToArgument) -> Result { + let CopyQueryToArgument { + with, + connection, + location, + } = stmt; + + Ok(CopyQueryToRequest { + location, + with: with.into_map(), + connection: connection.into_map(), + }) +} + fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result { let direction = match stmt { CopyTable::To(_) => CopyDirection::Export, diff --git a/src/operator/src/statement/copy_query_to.rs b/src/operator/src/statement/copy_query_to.rs new file mode 100644 index 00000000000..98fc92e23eb --- /dev/null +++ b/src/operator/src/statement/copy_query_to.rs @@ -0,0 +1,43 @@ +// 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 common_datasource::file_format::Format; +use common_query::Output; +use common_telemetry::{debug, tracing}; +use snafu::ResultExt; +use table::requests::CopyQueryToRequest; + +use crate::error; +use crate::error::Result; +use crate::statement::StatementExecutor; + +impl StatementExecutor { + #[tracing::instrument(skip_all)] + pub(crate) async fn copy_query_to( + &self, + req: CopyQueryToRequest, + query_output: Output, + ) -> Result { + let CopyQueryToRequest { + with, + connection, + location, + } = &req; + let format = Format::try_from(with).context(error::ParseFileFormatSnafu)?; + + debug!("Copy query to location: {location}"); + self.copy_to_file(&format, query_output, location, connection) + .await + } +} diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 0e53d79ac60..012bebbf4a7 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use client::OutputData; @@ -22,6 +23,7 @@ use common_datasource::file_format::parquet::stream_to_parquet; use common_datasource::file_format::Format; use common_datasource::object_store::{build_backend, parse_url}; use common_datasource::util::find_dir_and_filename; +use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::SendableRecordBatchStream; use common_telemetry::{debug, tracing}; @@ -135,24 +137,38 @@ impl StatementExecutor { .execute(plan, query_ctx) .await .context(ExecLogicalPlanSnafu)?; + + let CopyTableRequest { + location, + connection, + .. + } = &req; + + debug!("Copy table: {table_id} to location: {location}"); + self.copy_to_file(&format, output, location, connection) + .await + } + + pub(crate) async fn copy_to_file( + &self, + format: &Format, + output: Output, + location: &str, + connection: &HashMap, + ) -> Result { let stream = match output.data { OutputData::Stream(stream) => stream, OutputData::RecordBatches(record_batches) => record_batches.as_stream(), _ => unreachable!(), }; - let (_schema, _host, path) = parse_url(&req.location).context(error::ParseUrlSnafu)?; + let (_schema, _host, path) = parse_url(location).context(error::ParseUrlSnafu)?; let (_, filename) = find_dir_and_filename(&path); let filename = filename.context(error::UnexpectedSnafu { violated: format!("Expected filename, path: {path}"), })?; - let object_store = - build_backend(&req.location, &req.connection).context(error::BuildBackendSnafu)?; - debug!("Copy table: {table_id} to path: {path}"); - let rows_copied = self - .stream_to_file(stream, &format, object_store, &filename) - .await?; - - Ok(rows_copied) + let object_store = build_backend(location, connection).context(error::BuildBackendSnafu)?; + self.stream_to_file(stream, format, object_store, &filename) + .await } } diff --git a/src/sql/src/parsers/copy_parser.rs b/src/sql/src/parsers/copy_parser.rs index 5bf270129f6..ebf8c23ca19 100644 --- a/src/sql/src/parsers/copy_parser.rs +++ b/src/sql/src/parsers/copy_parser.rs @@ -16,11 +16,15 @@ use std::collections::HashMap; use snafu::ResultExt; use sqlparser::keywords::Keyword; +use sqlparser::tokenizer::Token; use sqlparser::tokenizer::Token::Word; use crate::error::{self, Result}; use crate::parser::ParserContext; -use crate::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument}; +use crate::statements::copy::{ + CopyDatabase, CopyDatabaseArgument, CopyQueryTo, CopyQueryToArgument, CopyTable, + CopyTableArgument, +}; use crate::statements::statement::Statement; use crate::util::parse_option_string; @@ -32,7 +36,11 @@ impl ParserContext<'_> { pub(crate) fn parse_copy(&mut self) -> Result { let _ = self.parser.next_token(); let next = self.parser.peek_token(); - let copy = if let Word(word) = next.token + let copy = if next.token == Token::LParen { + let copy_query = self.parse_copy_query_to()?; + + crate::statements::copy::Copy::CopyQueryTo(copy_query) + } else if let Word(word) = next.token && word.keyword == Keyword::DATABASE { let _ = self.parser.next_token(); @@ -126,6 +134,40 @@ impl ParserContext<'_> { } } + fn parse_copy_query_to(&mut self) -> Result { + self.parser + .expect_token(&Token::LParen) + .with_context(|_| error::UnexpectedSnafu { + expected: "'('", + actual: self.peek_token_as_string(), + })?; + let query = self.parse_query()?; + self.parser + .expect_token(&Token::RParen) + .with_context(|_| error::UnexpectedSnafu { + expected: "')'", + actual: self.peek_token_as_string(), + })?; + self.parser + .expect_keyword(Keyword::TO) + .context(error::SyntaxSnafu)?; + let (with, connection, location, limit) = self.parse_copy_parameters()?; + if limit.is_some() { + return error::InvalidSqlSnafu { + msg: "limit is not supported", + } + .fail(); + } + Ok(CopyQueryTo { + query: Box::new(query), + arg: CopyQueryToArgument { + with: with.into(), + connection: connection.into(), + location, + }, + }) + } + fn parse_copy_parameters(&mut self) -> Result<(With, Connection, String, Option)> { let location = self.parser @@ -453,4 +495,75 @@ mod tests { stmt.connection.to_str_map() ); } + + #[test] + fn test_copy_query_to() { + let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + + let Copy(crate::statements::copy::Copy::CopyQueryTo(stmt)) = stmt else { + unreachable!() + }; + + let query = ParserContext::create_with_dialect( + "SELECT * FROM tbl WHERE ts > 10", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0); + + assert_eq!(&query, stmt.query.as_ref()); + assert_eq!( + [("format", "parquet")] + .into_iter() + .collect::>(), + stmt.arg.with.to_str_map() + ); + + assert_eq!( + [("foo", "Bar"), ("one", "two")] + .into_iter() + .collect::>(), + stmt.arg.connection.to_str_map() + ); + } + + #[test] + fn test_invalid_copy_query_to() { + { + let sql = "COPY SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')"; + + assert!(ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err()) + } + { + let sql = "COPY SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')"; + + assert!(ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err()) + } + { + let sql = "COPY (SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')"; + + assert!(ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err()) + } + } } diff --git a/src/sql/src/statements/copy.rs b/src/sql/src/statements/copy.rs index 436d86d3aba..9961162a0ca 100644 --- a/src/sql/src/statements/copy.rs +++ b/src/sql/src/statements/copy.rs @@ -18,12 +18,14 @@ use serde::Serialize; use sqlparser::ast::ObjectName; use sqlparser_derive::{Visit, VisitMut}; +use crate::statements::statement::Statement; use crate::statements::OptionMap; #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub enum Copy { CopyTable(CopyTable), CopyDatabase(CopyDatabase), + CopyQueryTo(CopyQueryTo), } impl Display for Copy { @@ -31,10 +33,32 @@ impl Display for Copy { match self { Copy::CopyTable(s) => s.fmt(f), Copy::CopyDatabase(s) => s.fmt(f), + Copy::CopyQueryTo(s) => s.fmt(f), } } } +#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] +pub struct CopyQueryTo { + pub query: Box, + pub arg: CopyQueryToArgument, +} + +impl Display for CopyQueryTo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "COPY ({}) TO {}", &self.query, &self.arg.location)?; + if !self.arg.with.is_empty() { + let options = self.arg.with.kv_pairs(); + write!(f, " WITH ({})", options.join(", "))?; + } + if !self.arg.connection.is_empty() { + let options = self.arg.connection.kv_pairs(); + write!(f, " CONNECTION ({})", options.join(", "))?; + } + Ok(()) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub enum CopyTable { To(CopyTableArgument), @@ -115,6 +139,13 @@ pub struct CopyTableArgument { pub limit: Option, } +#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] +pub struct CopyQueryToArgument { + pub with: OptionMap, + pub connection: OptionMap, + pub location: String, +} + #[cfg(test)] impl CopyTableArgument { pub fn format(&self) -> Option { @@ -238,4 +269,29 @@ mod tests { } } } + + #[test] + fn test_display_copy_query_to() { + let sql = r"copy (SELECT * FROM tbl WHERE ts > 10) to 's3://my-bucket/data.parquet' + with (format = 'parquet', pattern = '.*parquet.*') + connection(region = 'us-west-2', secret_access_key = '12345678');"; + let stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, stmts.len()); + assert_matches!(&stmts[0], Statement::Copy { .. }); + + match &stmts[0] { + Statement::Copy(copy) => { + let new_sql = format!("{}", copy); + assert_eq!( + r#"COPY (SELECT * FROM tbl WHERE ts > 10) TO s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#, + &new_sql + ); + } + _ => { + unreachable!(); + } + } + } } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index e7de2eb9d6c..05175aa6327 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -321,6 +321,13 @@ pub struct CopyDatabaseRequest { pub time_range: Option, } +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +pub struct CopyQueryToRequest { + pub location: String, + pub with: HashMap, + pub connection: HashMap, +} + #[cfg(test)] mod tests { use std::time::Duration; diff --git a/tests/cases/standalone/common/copy/copy_to_fs.result b/tests/cases/standalone/common/copy/copy_to_fs.result index a3b73ab4adf..13ac66b79e4 100644 --- a/tests/cases/standalone/common/copy/copy_to_fs.result +++ b/tests/cases/standalone/common/copy/copy_to_fs.result @@ -18,6 +18,18 @@ COPY demo TO '${SQLNESS_HOME}/export/demo.json' WITH (format='json', start_time= Affected Rows: 1 +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.parquet'; + +Affected Rows: 1 + +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.csv' WITH (format='csv'); + +Affected Rows: 1 + +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.json' WITH (format='json'); + +Affected Rows: 1 + drop table demo; Affected Rows: 0 diff --git a/tests/cases/standalone/common/copy/copy_to_fs.sql b/tests/cases/standalone/common/copy/copy_to_fs.sql index 82e3ae1fb30..d59ef14dd9c 100644 --- a/tests/cases/standalone/common/copy/copy_to_fs.sql +++ b/tests/cases/standalone/common/copy/copy_to_fs.sql @@ -8,4 +8,10 @@ COPY demo TO '${SQLNESS_HOME}/export/demo.csv' WITH (format='csv', start_time='2 COPY demo TO '${SQLNESS_HOME}/export/demo.json' WITH (format='json', start_time='2022-06-15 07:02:37', end_time='2022-06-15 07:02:38'); +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.parquet'; + +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.csv' WITH (format='csv'); + +COPY (select host, cpu, ts from demo where host = 'host2') TO '${SQLNESS_HOME}/export/demo.json' WITH (format='json'); + drop table demo;