rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-27 14:16:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: impl COPY a query resultset to external file (#5250)

Browse Source 
* feat: impl COPY a query resultset to external file

* chore: add more tests for parse `copy_table_to`

* chore: add more tests for parse `copy_table_to`
This commit is contained in:
Kould
2025-01-06 04:32:51 +00:00
committed by GitHub
parent 513569ed5d
commit 2ad50332cb
9 changed files with 296 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/frontend/src/instance.rs
+2 -1
View File
@@ -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(_)
src/operator/src/statement.rs
+30 -2
View File
@@ -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<CopyQueryToRequest> {
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<CopyTableRequest> {
let direction = match stmt {
CopyTable::To(_) => CopyDirection::Export,
src/operator/src/statement/copy_query_to.rs
+43
View File
@@ -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<usize> {
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
}
}
src/operator/src/statement/copy_table_to.rs
+25 -9
View File
@@ -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<String, String>,
) -> Result<usize> {
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
}
}
src/sql/src/parsers/copy_parser.rs
+115 -2
View File
@@ -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<Statement> {
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<CopyQueryTo> {
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<u64>)> {
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::<HashMap<_, _>>(),
stmt.arg.with.to_str_map()
);
assert_eq!(
[("foo", "Bar"), ("one", "two")]
.into_iter()
.collect::<HashMap<_, _>>(),
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())
}
}
}
src/sql/src/statements/copy.rs
+56
View File
@@ -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<Statement>,
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<u64>,
}
#[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<String> {
@@ -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!();
}
}
}
}
src/table/src/requests.rs
+7
View File
@@ -321,6 +321,13 @@ pub struct CopyDatabaseRequest {
pub time_range: Option<TimestampRange>,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct CopyQueryToRequest {
pub location: String,
pub with: HashMap<String, String>,
pub connection: HashMap<String, String>,
}
#[cfg(test)]
mod tests {
use std::time::Duration;
tests/cases/standalone/common/copy/copy_to_fs.result
+12
View File
@@ -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
tests/cases/standalone/common/copy/copy_to_fs.sql
+6
View File
@@ -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;