Skip to content

Add support for query source in COPY .. TO statement #858

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 1, 2023
Merged

Add support for query source in COPY .. TO statement #858

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c6a656d
Add support for query source in COPY .. TO statement
aprimadi Apr 27, 2023
c926a68
Fix compile error
aprimadi Apr 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 30 additions & 11 deletions src/ast/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1171,14 +1171,11 @@ pub enum Statement {
source: Box<Query>,
},
Copy {
/// TABLE
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
table_name: ObjectName,
/// COLUMNS
columns: Vec<Ident>,
/// The source of 'COPY TO', or the target of 'COPY FROM'
source: CopySource,
/// If true, is a 'COPY TO' statement. If false is a 'COPY FROM'
to: bool,
/// The source of 'COPY FROM', or the target of 'COPY TO'
/// The target of 'COPY TO', or the source of 'COPY FROM'
target: CopyTarget,
/// WITH options (from PostgreSQL version 9.0)
options: Vec<CopyOption>,
Expand Down Expand Up @@ -1902,17 +1899,25 @@ impl fmt::Display for Statement {
}

Statement::Copy {
table_name,
columns,
source,
to,
target,
options,
legacy_options,
values,
} => {
write!(f, "COPY {table_name}")?;
if !columns.is_empty() {
write!(f, " ({})", display_comma_separated(columns))?;
write!(f, "COPY")?;
match source {
CopySource::Query(query) => write!(f, " ({query})")?,
CopySource::Table {
table_name,
columns,
} => {
write!(f, " {table_name}")?;
if !columns.is_empty() {
write!(f, " ({})", display_comma_separated(columns))?;
}
}
}
write!(f, " {} {}", if *to { "TO" } else { "FROM" }, target)?;
if !options.is_empty() {
Expand Down Expand Up @@ -3663,6 +3668,20 @@ impl fmt::Display for SqliteOnConflict {
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum CopySource {
Table {
/// The name of the table to copy from.
table_name: ObjectName,
/// A list of column names to copy. Empty list means that all columns
/// are copied.
columns: Vec<Ident>,
},
Query(Box<Query>),
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
Expand Down
26 changes: 22 additions & 4 deletions src/parser.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4000,13 +4000,32 @@ impl<'a> Parser<'a> {

/// Parse a copy statement
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
let table_name = self.parse_object_name()?;
let columns = self.parse_parenthesized_column_list(Optional, false)?;
let source;
if self.consume_token(&Token::LParen) {
source = CopySource::Query(Box::new(self.parse_query()?));
self.expect_token(&Token::RParen)?;
} else {
let table_name = self.parse_object_name()?;
let columns = self.parse_parenthesized_column_list(Optional, false)?;
source = CopySource::Table {
table_name,
columns,
};
}
let to = match self.parse_one_of_keywords(&[Keyword::FROM, Keyword::TO]) {
Some(Keyword::FROM) => false,
Some(Keyword::TO) => true,
_ => self.expected("FROM or TO", self.peek_token())?,
};
if !to {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should remove this error in the parser per https://github.com/sqlparser-rs/sqlparser-rs#syntax-vs-semantics (if downstream crates want to error that is fine).

I don't feel super strongly on this point

what do you think @ankrgyl ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I argue that this is syntactic as the syntax is described here and is different for COPY .. FROM vs COPY .. TO. https://www.postgresql.org/docs/current/sql-copy.html.

Looking at DuckDB syntax graph also there seems to be difference syntactically between COPY .. FROM vs COPY .. TO. https://duckdb.org/docs/sql/statements/copy

But yeah let's wait on what @ankrgyl think about it. I'll go with whatever the consensus decide since I'm not an expert in various SQL dialects.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As I said above, I think this way is fine (and kudos to you for covering it with tests). I agree let's see if @ankrgyl has any thoughts

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the really fast review @alamb -- must be tough reviewing so many repos ❤️

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤯 most days

// Use a separate if statement to prevent Rust compiler from complaining about
// "if statement in this position is unstable: https://github.com/rust-lang/rust/issues/53667"
if let CopySource::Query(_) = source {
return Err(ParserError::ParserError(
"COPY ... FROM does not support query as a source".to_string(),
));
}
}
let target = if self.parse_keyword(Keyword::STDIN) {
CopyTarget::Stdin
} else if self.parse_keyword(Keyword::STDOUT) {
Expand Down Expand Up @@ -4037,8 +4056,7 @@ impl<'a> Parser<'a> {
vec![]
};
Ok(Statement::Copy {
table_name,
columns,
source,
to,
target,
options,
Expand Down
139 changes: 113 additions & 26 deletions tests/sqlparser_postgres.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -691,8 +691,10 @@ fn test_copy_from() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: false,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -707,8 +709,10 @@ fn test_copy_from() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: false,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -723,8 +727,10 @@ fn test_copy_from() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: false,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -745,8 +751,10 @@ fn test_copy_to() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -761,8 +769,10 @@ fn test_copy_to() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -777,8 +787,10 @@ fn test_copy_to() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand Down Expand Up @@ -816,8 +828,10 @@ fn parse_copy_from() {
assert_eq!(
pg_and_generic().one_statement_parses_to(sql, ""),
Statement::Copy {
table_name: ObjectName(vec!["table".into()]),
columns: vec!["a".into(), "b".into()],
source: CopySource::Table {
table_name: ObjectName(vec!["table".into()]),
columns: vec!["a".into(), "b".into()],
},
to: false,
target: CopyTarget::File {
filename: "file.csv".into()
Expand Down Expand Up @@ -845,14 +859,25 @@ fn parse_copy_from() {
);
}

#[test]
fn parse_copy_from_error() {
let res = pg().parse_sql_statements("COPY (SELECT 42 AS a, 'hello' AS b) FROM 'query.csv'");
assert_eq!(
ParserError::ParserError("COPY ... FROM does not support query as a source".to_string()),
res.unwrap_err()
);
}

#[test]
fn parse_copy_to() {
let stmt = pg().verified_stmt("COPY users TO 'data.csv'");
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -867,8 +892,10 @@ fn parse_copy_to() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["country".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["country".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::Stdout,
options: vec![CopyOption::Delimiter('|')],
Expand All @@ -882,8 +909,10 @@ fn parse_copy_to() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["country".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["country".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::Program {
command: "gzip > /usr1/proj/bray/sql/country_data.gz".into(),
Expand All @@ -893,6 +922,58 @@ fn parse_copy_to() {
values: vec![],
}
);

let stmt = pg().verified_stmt("COPY (SELECT 42 AS a, 'hello' AS b) TO 'query.csv'");
assert_eq!(
stmt,
Statement::Copy {
source: CopySource::Query(Box::new(Query {
with: None,
body: Box::new(SetExpr::Select(Box::new(Select {
distinct: false,
top: None,
projection: vec![
SelectItem::ExprWithAlias {
expr: Expr::Value(number("42")),
alias: Ident {
value: "a".into(),
quote_style: None,
},
},
SelectItem::ExprWithAlias {
expr: Expr::Value(Value::SingleQuotedString("hello".into())),
alias: Ident {
value: "b".into(),
quote_style: None,
},
}
],
into: None,
from: vec![],
lateral_views: vec![],
selection: None,
group_by: vec![],
having: None,
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
qualify: None,
}))),
order_by: vec![],
limit: None,
offset: None,
fetch: None,
locks: vec![],
})),
to: true,
target: CopyTarget::File {
filename: "query.csv".into(),
},
options: vec![],
legacy_options: vec![],
values: vec![],
}
)
}

#[test]
Expand All @@ -901,8 +982,10 @@ fn parse_copy_from_before_v9_0() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: false,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -928,8 +1011,10 @@ fn parse_copy_from_before_v9_0() {
assert_eq!(
pg_and_generic().one_statement_parses_to(sql, ""),
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: false,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand All @@ -954,8 +1039,10 @@ fn parse_copy_to_before_v9_0() {
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
source: CopySource::Table {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
},
to: true,
target: CopyTarget::File {
filename: "data.csv".to_string(),
Expand Down