Skip to content

Commit 88331f0

Browse files
sundy-lisundy-li
authored
feat(query): support with clause in copy into (#15343)
* feat(query): support with clause in copy into * feat(query): support with clause in copy into * feat(query): support with clause in copy into
1 parent f26dba0 commit 88331f0

File tree

6 files changed

+56
-6
lines changed

6 files changed

+56
-6
lines changed

‎src/query/ast/src/ast/statements/copy.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ use crate::ast::Hint;
4040
use crate::ast::Identifier;
4141
use crate::ast::Query;
4242
use crate::ast::TableRef;
43+
use crate::ast::With;
4344

4445
/// CopyIntoTableStmt is the parsed statement of `COPY into <table> from <location>`.
4546
///
@@ -50,6 +51,7 @@ use crate::ast::TableRef;
5051
/// ```
5152
#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
5253
pub struct CopyIntoTableStmt {
54+
pub with: Option<With>,
5355
pub src: CopyIntoTableSource,
5456
pub dst: TableRef,
5557
pub dst_columns: Option<Vec<Identifier>>,
@@ -137,6 +139,9 @@ impl CopyIntoTableStmt {
137139

138140
impl Display for CopyIntoTableStmt {
139141
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142+
if let Some(cte) = &self.with {
143+
write!(f, "WITH {} ", cte)?;
144+
}
140145
write!(f, "COPY")?;
141146
if let Some(hints) = &self.hints {
142147
write!(f, "{} ", hints)?;
@@ -189,6 +194,7 @@ impl Display for CopyIntoTableStmt {
189194
/// CopyIntoLocationStmt is the parsed statement of `COPY into <location> from <table> ...`
190195
#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
191196
pub struct CopyIntoLocationStmt {
197+
pub with: Option<With>,
192198
pub hints: Option<Hint>,
193199
pub src: CopyIntoLocationSource,
194200
pub dst: FileLocation,
@@ -204,6 +210,9 @@ pub struct CopyIntoLocationStmt {
204210

205211
impl Display for CopyIntoLocationStmt {
206212
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
213+
if let Some(cte) = &self.with {
214+
write!(f, "WITH {} ", cte)?;
215+
}
207216
write!(f, "COPY")?;
208217
if let Some(hints) = &self.hints {
209218
write!(f, "{} ", hints)?;

‎src/query/ast/src/parser/copy.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
use nom::branch::alt;
1616
use nom::combinator::map;
1717

18+
use super::query::with;
1819
use crate::ast::CopyIntoLocationOption;
1920
use crate::ast::CopyIntoLocationSource;
2021
use crate::ast::CopyIntoLocationStmt;
@@ -50,14 +51,15 @@ pub fn copy_into_table(i: Input) -> IResult<Statement> {
5051

5152
map(
5253
rule! {
53-
COPY
54+
#with? ~ COPY
5455
~ #hint?
5556
~ INTO ~ #table_ref ~ ( "(" ~ #comma_separated_list1(ident) ~ ")" )?
5657
~ ^FROM ~ ^#copy_into_table_source
5758
~ #copy_into_table_option*
5859
},
59-
|(_copy, opt_hints, _into, dst, dst_columns, _from, src, opts)| {
60+
|(with, _copy, opt_hints, _into, dst, dst_columns, _from, src, opts)| {
6061
let mut copy_stmt = CopyIntoTableStmt {
62+
with,
6163
hints: opt_hints,
6264
src,
6365
dst,
@@ -93,14 +95,15 @@ fn copy_into_location(i: Input) -> IResult<Statement> {
9395

9496
map(
9597
rule! {
96-
COPY
98+
#with? ~ COPY
9799
~ #hint?
98100
~ INTO ~ #file_location
99101
~ ^FROM ~ ^#copy_into_location_source
100102
~ #copy_into_location_option*
101103
},
102-
|(_copy, opt_hints, _into, dst, _from, src, opts)| {
104+
|(with, _copy, opt_hints, _into, dst, _from, src, opts)| {
103105
let mut copy_stmt = CopyIntoLocationStmt {
106+
with,
104107
hints: opt_hints,
105108
src,
106109
dst,

‎src/query/ast/tests/it/testdata/statement.txt

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12185,6 +12185,7 @@ COPY INTO mytable FROM '@~/mybucket/my data.csv' SIZE_LIMIT = 10 PURGE = false F
1218512185
---------- AST ------------
1218612186
CopyIntoTable(
1218712187
CopyIntoTableStmt {
12188+
with: None,
1218812189
src: Location(
1218912190
Stage(
1219012191
"~/mybucket/my data.csv",
@@ -12237,6 +12238,7 @@ COPY INTO mytable FROM '@~/mybucket/data.csv' FILE_FORMAT = (field_delimiter = '
1223712238
---------- AST ------------
1223812239
CopyIntoTable(
1223912240
CopyIntoTableStmt {
12241+
with: None,
1224012242
src: Location(
1224112243
Stage(
1224212244
"~/mybucket/data.csv",
@@ -12303,6 +12305,7 @@ COPY INTO mytable FROM 's3://mybucket/data.csv' FILE_FORMAT = (field_delimiter =
1230312305
---------- AST ------------
1230412306
CopyIntoTable(
1230512307
CopyIntoTableStmt {
12308+
with: None,
1230612309
src: Location(
1230712310
Uri(
1230812311
UriLocation {
@@ -12378,6 +12381,7 @@ COPY INTO mytable FROM 's3://mybucket/data.csv' FILE_FORMAT = (field_delimiter =
1237812381
---------- AST ------------
1237912382
CopyIntoTable(
1238012383
CopyIntoTableStmt {
12384+
with: None,
1238112385
src: Location(
1238212386
Uri(
1238312387
UriLocation {
@@ -12455,6 +12459,7 @@ COPY INTO mytable FROM 's3://mybucket/data.csv' CONNECTION = ( endpoint_url = 'h
1245512459
---------- AST ------------
1245612460
CopyIntoTable(
1245712461
CopyIntoTableStmt {
12462+
with: None,
1245812463
src: Location(
1245912464
Uri(
1246012465
UriLocation {
@@ -12534,6 +12539,7 @@ COPY INTO mytable FROM 's3://mybucket/data.csv' CONNECTION = ( endpoint_url = 'h
1253412539
---------- AST ------------
1253512540
CopyIntoTable(
1253612541
CopyIntoTableStmt {
12542+
with: None,
1253712543
src: Location(
1253812544
Uri(
1253912545
UriLocation {
@@ -12603,6 +12609,7 @@ COPY INTO mytable FROM 'https://127.0.0.1:9900/' PURGE = false FORCE = false DIS
1260312609
---------- AST ------------
1260412610
CopyIntoTable(
1260512611
CopyIntoTableStmt {
12612+
with: None,
1260612613
src: Location(
1260712614
Uri(
1260812615
UriLocation {
@@ -12657,6 +12664,7 @@ COPY INTO mytable FROM 'https://127.0.0.1/' PURGE = false FORCE = false DISABLE_
1265712664
---------- AST ------------
1265812665
CopyIntoTable(
1265912666
CopyIntoTableStmt {
12667+
with: None,
1266012668
src: Location(
1266112669
Uri(
1266212670
UriLocation {
@@ -12719,6 +12727,7 @@ COPY INTO mytable FROM '@my_stage' FILE_FORMAT = (error_on_column_count_mismatch
1271912727
---------- AST ------------
1272012728
CopyIntoTable(
1272112729
CopyIntoTableStmt {
12730+
with: None,
1272212731
src: Location(
1272312732
Stage(
1272412733
"my_stage",
@@ -12786,6 +12795,7 @@ COPY INTO 's3://mybucket/data.csv' FROM mytable FILE_FORMAT = (field_delimiter =
1278612795
---------- AST ------------
1278712796
CopyIntoLocation(
1278812797
CopyIntoLocationStmt {
12798+
with: None,
1278912799
hints: None,
1279012800
src: Table(
1279112801
TableRef {
@@ -12844,6 +12854,7 @@ COPY INTO '@my_stage/my data' FROM mytable SINGLE = false MAX_FILE_SIZE = 0 DETA
1284412854
---------- AST ------------
1284512855
CopyIntoLocation(
1284612856
CopyIntoLocationStmt {
12857+
with: None,
1284712858
hints: None,
1284812859
src: Table(
1284912860
TableRef {
@@ -12886,6 +12897,7 @@ COPY INTO '@my_stage' FROM mytable FILE_FORMAT = (field_delimiter = ',', record_
1288612897
---------- AST ------------
1288712898
CopyIntoLocation(
1288812899
CopyIntoLocationStmt {
12900+
with: None,
1288912901
hints: None,
1289012902
src: Table(
1289112903
TableRef {
@@ -12946,6 +12958,7 @@ COPY INTO mytable FROM 's3://mybucket/data.csv' CONNECTION = ( aws_key_id = 'acc
1294612958
---------- AST ------------
1294712959
CopyIntoTable(
1294812960
CopyIntoTableStmt {
12961+
with: None,
1294912962
src: Location(
1295012963
Uri(
1295112964
UriLocation {
@@ -13023,6 +13036,7 @@ COPY INTO mytable FROM '@external_stage/path/to/file.csv' FILE_FORMAT = (field_d
1302313036
---------- AST ------------
1302413037
CopyIntoTable(
1302513038
CopyIntoTableStmt {
13039+
with: None,
1302613040
src: Location(
1302713041
Stage(
1302813042
"external_stage/path/to/file.csv",
@@ -13088,6 +13102,7 @@ COPY INTO mytable FROM '@external_stage/path/to/dir/' FILE_FORMAT = (field_delim
1308813102
---------- AST ------------
1308913103
CopyIntoTable(
1309013104
CopyIntoTableStmt {
13105+
with: None,
1309113106
src: Location(
1309213107
Stage(
1309313108
"external_stage/path/to/dir/",
@@ -13153,6 +13168,7 @@ COPY INTO mytable FROM '@external_stage/path/to/file.csv' FILE_FORMAT = (field_d
1315313168
---------- AST ------------
1315413169
CopyIntoTable(
1315513170
CopyIntoTableStmt {
13171+
with: None,
1315613172
src: Location(
1315713173
Stage(
1315813174
"external_stage/path/to/file.csv",
@@ -13219,6 +13235,7 @@ COPY INTO mytable FROM 'fs:///path/to/data.csv' FILE_FORMAT = (field_delimiter =
1321913235
---------- AST ------------
1322013236
CopyIntoTable(
1322113237
CopyIntoTableStmt {
13238+
with: None,
1322213239
src: Location(
1322313240
Uri(
1322413241
UriLocation {
@@ -13292,6 +13309,7 @@ COPY INTO books FROM 's3://databend/books.csv' CONNECTION = ( access_key_id = 'R
1329213309
---------- AST ------------
1329313310
CopyIntoTable(
1329413311
CopyIntoTableStmt {
13312+
with: None,
1329513313
src: Location(
1329613314
Uri(
1329713315
UriLocation {
@@ -18953,6 +18971,7 @@ CreatePipe(
1895318971
auto_ingest: true,
1895418972
comments: "This is test pipe 1",
1895518973
copy_stmt: CopyIntoTableStmt {
18974+
with: None,
1895618975
src: Location(
1895718976
Stage(
1895818977
"~/MyStage1",
@@ -19007,6 +19026,7 @@ CreatePipe(
1900719026
auto_ingest: false,
1900819027
comments: "",
1900919028
copy_stmt: CopyIntoTableStmt {
19029+
with: None,
1901019030
src: Location(
1901119031
Stage(
1901219032
"~/mybucket/data.csv",

‎src/query/sql/src/planner/binder/copy_into_location.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ impl<'a> Binder {
5959
}
6060
}
6161
CopyIntoLocationSource::Query(query) => {
62+
if let Some(with) = &stmt.with {
63+
self.add_cte(with, bind_context)?;
64+
}
6265
let select_plan = self
6366
.bind_statement(bind_context, &Statement::Query(query.clone()))
6467
.await?;

‎src/query/sql/src/planner/binder/copy_into_table.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,10 @@ impl<'a> Binder {
9494
.await
9595
}
9696
CopyIntoTableSource::Query(query) => {
97+
if let Some(with) = &stmt.with {
98+
self.add_cte(with, bind_context)?;
99+
}
100+
97101
let mut max_column_position = MaxColumnPosition::new();
98102
query.drive(&mut max_column_position);
99103
self.metadata

‎tests/sqllogictests/suites/base/03_common/03_0028_copy_into_stage.test

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,18 @@ statement ok
4343
CREATE STAGE IF NOT EXISTS hello
4444

4545
statement ok
46-
COPY INTO @hello from (select number from numbers(1)) FILE_FORMAT = (type = parquet)
46+
COPY INTO @hello from (select number from numbers(10)) FILE_FORMAT = (type = parquet)
47+
48+
statement ok
49+
COPY INTO @hello from (select number from numbers(1) where number in (select max(number) from numbers(1000)) ) FILE_FORMAT = (type = parquet)
50+
51+
statement ok
52+
with S as (select number from numbers(1000) where number > 100) COPY INTO @hello from (select number from numbers(1) where number not in (SELECT number FROM S)) FILE_FORMAT = (type = parquet)
53+
54+
query I
55+
select sum($1) from @hello;
56+
----
57+
45
4758

4859
statement ok
4960
CREATE TABLE world(c1 int , c2 int);
@@ -55,7 +66,7 @@ statement ok
5566
DROP STAGE IF EXISTS hello
5667

5768
statement ok
58-
drop table world
69+
drop table world
5970

6071
statement ok
6172
DROP DATABASE db1

0 commit comments

Comments
 (0)