Skip to content

Commit 0ff863b

Browse files
aprimadiaprimadi
authored
Add support for query source in COPY .. TO statement (apache#858)
* Add support for query source in COPY .. TO statement * Fix compile error
1 parent 0113bbd commit 0ff863b

File tree

3 files changed

+165
-41
lines changed

3 files changed

+165
-41
lines changed

‎src/ast/mod.rs

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1173,14 +1173,11 @@ pub enum Statement {
11731173
source: Box<Query>,
11741174
},
11751175
Copy {
1176-
/// TABLE
1177-
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
1178-
table_name: ObjectName,
1179-
/// COLUMNS
1180-
columns: Vec<Ident>,
1176+
/// The source of 'COPY TO', or the target of 'COPY FROM'
1177+
source: CopySource,
11811178
/// If true, is a 'COPY TO' statement. If false is a 'COPY FROM'
11821179
to: bool,
1183-
/// The source of 'COPY FROM', or the target of 'COPY TO'
1180+
/// The target of 'COPY TO', or the source of 'COPY FROM'
11841181
target: CopyTarget,
11851182
/// WITH options (from PostgreSQL version 9.0)
11861183
options: Vec<CopyOption>,
@@ -1926,17 +1923,25 @@ impl fmt::Display for Statement {
19261923
}
19271924

19281925
Statement::Copy {
1929-
table_name,
1930-
columns,
1926+
source,
19311927
to,
19321928
target,
19331929
options,
19341930
legacy_options,
19351931
values,
19361932
} => {
1937-
write!(f, "COPY {table_name}")?;
1938-
if !columns.is_empty() {
1939-
write!(f, " ({})", display_comma_separated(columns))?;
1933+
write!(f, "COPY")?;
1934+
match source {
1935+
CopySource::Query(query) => write!(f, " ({query})")?,
1936+
CopySource::Table {
1937+
table_name,
1938+
columns,
1939+
} => {
1940+
write!(f, " {table_name}")?;
1941+
if !columns.is_empty() {
1942+
write!(f, " ({})", display_comma_separated(columns))?;
1943+
}
1944+
}
19401945
}
19411946
write!(f, " {} {}", if *to { "TO" } else { "FROM" }, target)?;
19421947
if !options.is_empty() {
@@ -3750,6 +3755,20 @@ impl fmt::Display for SqliteOnConflict {
37503755
}
37513756
}
37523757

3758+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
3759+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
3760+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
3761+
pub enum CopySource {
3762+
Table {
3763+
/// The name of the table to copy from.
3764+
table_name: ObjectName,
3765+
/// A list of column names to copy. Empty list means that all columns
3766+
/// are copied.
3767+
columns: Vec<Ident>,
3768+
},
3769+
Query(Box<Query>),
3770+
}
3771+
37533772
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
37543773
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
37553774
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

‎src/parser.rs

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4015,13 +4015,32 @@ impl<'a> Parser<'a> {
40154015

40164016
/// Parse a copy statement
40174017
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
4018-
let table_name = self.parse_object_name()?;
4019-
let columns = self.parse_parenthesized_column_list(Optional, false)?;
4018+
let source;
4019+
if self.consume_token(&Token::LParen) {
4020+
source = CopySource::Query(Box::new(self.parse_query()?));
4021+
self.expect_token(&Token::RParen)?;
4022+
} else {
4023+
let table_name = self.parse_object_name()?;
4024+
let columns = self.parse_parenthesized_column_list(Optional, false)?;
4025+
source = CopySource::Table {
4026+
table_name,
4027+
columns,
4028+
};
4029+
}
40204030
let to = match self.parse_one_of_keywords(&[Keyword::FROM, Keyword::TO]) {
40214031
Some(Keyword::FROM) => false,
40224032
Some(Keyword::TO) => true,
40234033
_ => self.expected("FROM or TO", self.peek_token())?,
40244034
};
4035+
if !to {
4036+
// Use a separate if statement to prevent Rust compiler from complaining about
4037+
// "if statement in this position is unstable: https://github.com/rust-lang/rust/issues/53667"
4038+
if let CopySource::Query(_) = source {
4039+
return Err(ParserError::ParserError(
4040+
"COPY ... FROM does not support query as a source".to_string(),
4041+
));
4042+
}
4043+
}
40254044
let target = if self.parse_keyword(Keyword::STDIN) {
40264045
CopyTarget::Stdin
40274046
} else if self.parse_keyword(Keyword::STDOUT) {
@@ -4052,8 +4071,7 @@ impl<'a> Parser<'a> {
40524071
vec![]
40534072
};
40544073
Ok(Statement::Copy {
4055-
table_name,
4056-
columns,
4074+
source,
40574075
to,
40584076
target,
40594077
options,

‎tests/sqlparser_postgres.rs

Lines changed: 113 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -691,8 +691,10 @@ fn test_copy_from() {
691691
assert_eq!(
692692
stmt,
693693
Statement::Copy {
694-
table_name: ObjectName(vec!["users".into()]),
695-
columns: vec![],
694+
source: CopySource::Table {
695+
table_name: ObjectName(vec!["users".into()]),
696+
columns: vec![],
697+
},
696698
to: false,
697699
target: CopyTarget::File {
698700
filename: "data.csv".to_string(),
@@ -707,8 +709,10 @@ fn test_copy_from() {
707709
assert_eq!(
708710
stmt,
709711
Statement::Copy {
710-
table_name: ObjectName(vec!["users".into()]),
711-
columns: vec![],
712+
source: CopySource::Table {
713+
table_name: ObjectName(vec!["users".into()]),
714+
columns: vec![],
715+
},
712716
to: false,
713717
target: CopyTarget::File {
714718
filename: "data.csv".to_string(),
@@ -723,8 +727,10 @@ fn test_copy_from() {
723727
assert_eq!(
724728
stmt,
725729
Statement::Copy {
726-
table_name: ObjectName(vec!["users".into()]),
727-
columns: vec![],
730+
source: CopySource::Table {
731+
table_name: ObjectName(vec!["users".into()]),
732+
columns: vec![],
733+
},
728734
to: false,
729735
target: CopyTarget::File {
730736
filename: "data.csv".to_string(),
@@ -745,8 +751,10 @@ fn test_copy_to() {
745751
assert_eq!(
746752
stmt,
747753
Statement::Copy {
748-
table_name: ObjectName(vec!["users".into()]),
749-
columns: vec![],
754+
source: CopySource::Table {
755+
table_name: ObjectName(vec!["users".into()]),
756+
columns: vec![],
757+
},
750758
to: true,
751759
target: CopyTarget::File {
752760
filename: "data.csv".to_string(),
@@ -761,8 +769,10 @@ fn test_copy_to() {
761769
assert_eq!(
762770
stmt,
763771
Statement::Copy {
764-
table_name: ObjectName(vec!["users".into()]),
765-
columns: vec![],
772+
source: CopySource::Table {
773+
table_name: ObjectName(vec!["users".into()]),
774+
columns: vec![],
775+
},
766776
to: true,
767777
target: CopyTarget::File {
768778
filename: "data.csv".to_string(),
@@ -777,8 +787,10 @@ fn test_copy_to() {
777787
assert_eq!(
778788
stmt,
779789
Statement::Copy {
780-
table_name: ObjectName(vec!["users".into()]),
781-
columns: vec![],
790+
source: CopySource::Table {
791+
table_name: ObjectName(vec!["users".into()]),
792+
columns: vec![],
793+
},
782794
to: true,
783795
target: CopyTarget::File {
784796
filename: "data.csv".to_string(),
@@ -816,8 +828,10 @@ fn parse_copy_from() {
816828
assert_eq!(
817829
pg_and_generic().one_statement_parses_to(sql, ""),
818830
Statement::Copy {
819-
table_name: ObjectName(vec!["table".into()]),
820-
columns: vec!["a".into(), "b".into()],
831+
source: CopySource::Table {
832+
table_name: ObjectName(vec!["table".into()]),
833+
columns: vec!["a".into(), "b".into()],
834+
},
821835
to: false,
822836
target: CopyTarget::File {
823837
filename: "file.csv".into()
@@ -845,14 +859,25 @@ fn parse_copy_from() {
845859
);
846860
}
847861

862+
#[test]
863+
fn parse_copy_from_error() {
864+
let res = pg().parse_sql_statements("COPY (SELECT 42 AS a, 'hello' AS b) FROM 'query.csv'");
865+
assert_eq!(
866+
ParserError::ParserError("COPY ... FROM does not support query as a source".to_string()),
867+
res.unwrap_err()
868+
);
869+
}
870+
848871
#[test]
849872
fn parse_copy_to() {
850873
let stmt = pg().verified_stmt("COPY users TO 'data.csv'");
851874
assert_eq!(
852875
stmt,
853876
Statement::Copy {
854-
table_name: ObjectName(vec!["users".into()]),
855-
columns: vec![],
877+
source: CopySource::Table {
878+
table_name: ObjectName(vec!["users".into()]),
879+
columns: vec![],
880+
},
856881
to: true,
857882
target: CopyTarget::File {
858883
filename: "data.csv".to_string(),
@@ -867,8 +892,10 @@ fn parse_copy_to() {
867892
assert_eq!(
868893
stmt,
869894
Statement::Copy {
870-
table_name: ObjectName(vec!["country".into()]),
871-
columns: vec![],
895+
source: CopySource::Table {
896+
table_name: ObjectName(vec!["country".into()]),
897+
columns: vec![],
898+
},
872899
to: true,
873900
target: CopyTarget::Stdout,
874901
options: vec![CopyOption::Delimiter('|')],
@@ -882,8 +909,10 @@ fn parse_copy_to() {
882909
assert_eq!(
883910
stmt,
884911
Statement::Copy {
885-
table_name: ObjectName(vec!["country".into()]),
886-
columns: vec![],
912+
source: CopySource::Table {
913+
table_name: ObjectName(vec!["country".into()]),
914+
columns: vec![],
915+
},
887916
to: true,
888917
target: CopyTarget::Program {
889918
command: "gzip > /usr1/proj/bray/sql/country_data.gz".into(),
@@ -893,6 +922,58 @@ fn parse_copy_to() {
893922
values: vec![],
894923
}
895924
);
925+
926+
let stmt = pg().verified_stmt("COPY (SELECT 42 AS a, 'hello' AS b) TO 'query.csv'");
927+
assert_eq!(
928+
stmt,
929+
Statement::Copy {
930+
source: CopySource::Query(Box::new(Query {
931+
with: None,
932+
body: Box::new(SetExpr::Select(Box::new(Select {
933+
distinct: false,
934+
top: None,
935+
projection: vec![
936+
SelectItem::ExprWithAlias {
937+
expr: Expr::Value(number("42")),
938+
alias: Ident {
939+
value: "a".into(),
940+
quote_style: None,
941+
},
942+
},
943+
SelectItem::ExprWithAlias {
944+
expr: Expr::Value(Value::SingleQuotedString("hello".into())),
945+
alias: Ident {
946+
value: "b".into(),
947+
quote_style: None,
948+
},
949+
}
950+
],
951+
into: None,
952+
from: vec![],
953+
lateral_views: vec![],
954+
selection: None,
955+
group_by: vec![],
956+
having: None,
957+
cluster_by: vec![],
958+
distribute_by: vec![],
959+
sort_by: vec![],
960+
qualify: None,
961+
}))),
962+
order_by: vec![],
963+
limit: None,
964+
offset: None,
965+
fetch: None,
966+
locks: vec![],
967+
})),
968+
to: true,
969+
target: CopyTarget::File {
970+
filename: "query.csv".into(),
971+
},
972+
options: vec![],
973+
legacy_options: vec![],
974+
values: vec![],
975+
}
976+
)
896977
}
897978

898979
#[test]
@@ -901,8 +982,10 @@ fn parse_copy_from_before_v9_0() {
901982
assert_eq!(
902983
stmt,
903984
Statement::Copy {
904-
table_name: ObjectName(vec!["users".into()]),
905-
columns: vec![],
985+
source: CopySource::Table {
986+
table_name: ObjectName(vec!["users".into()]),
987+
columns: vec![],
988+
},
906989
to: false,
907990
target: CopyTarget::File {
908991
filename: "data.csv".to_string(),
@@ -928,8 +1011,10 @@ fn parse_copy_from_before_v9_0() {
9281011
assert_eq!(
9291012
pg_and_generic().one_statement_parses_to(sql, ""),
9301013
Statement::Copy {
931-
table_name: ObjectName(vec!["users".into()]),
932-
columns: vec![],
1014+
source: CopySource::Table {
1015+
table_name: ObjectName(vec!["users".into()]),
1016+
columns: vec![],
1017+
},
9331018
to: false,
9341019
target: CopyTarget::File {
9351020
filename: "data.csv".to_string(),
@@ -954,8 +1039,10 @@ fn parse_copy_to_before_v9_0() {
9541039
assert_eq!(
9551040
stmt,
9561041
Statement::Copy {
957-
table_name: ObjectName(vec!["users".into()]),
958-
columns: vec![],
1042+
source: CopySource::Table {
1043+
table_name: ObjectName(vec!["users".into()]),
1044+
columns: vec![],
1045+
},
9591046
to: true,
9601047
target: CopyTarget::File {
9611048
filename: "data.csv".to_string(),

0 commit comments

Comments
 (0)