mirror of
https://github.com/apache/datafusion-sqlparser-rs.git
synced 2025-08-04 06:18:17 +00:00
Add support for query source in COPY .. TO statement (#858)
* Add support for query source in COPY .. TO statement * Fix compile error
This commit is contained in:
Armin Primadi
GitHub
parent
0113bbd924
commit
0ff863b2c7
3 changed files with 165 additions and 41 deletions
|
|
@ -1173,14 +1173,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>,
|
||||
|
|
@ -1926,17 +1923,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() {
|
||||
|
|
@ -3750,6 +3755,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))]
|
||||
|
|
|
|||
|
|
@ -4015,13 +4015,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 {
|
||||
// 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) {
|
||||
|
|
@ -4052,8 +4071,7 @@ impl<'a> Parser<'a> {
|
|||
vec![]
|
||||
};
|
||||
Ok(Statement::Copy {
|
||||
table_name,
|
||||
columns,
|
||||
source,
|
||||
to,
|
||||
target,
|
||||
options,
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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()
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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('|')],
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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]
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
@ -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(),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue