Add support for COPY TO (#441)

* Start adding COPY TO

* Fix statement and add tests

* Merge copy statements

* Remove extra line

* Clippy

* Cleanup
This commit is contained in:
Matthew Turner 2022-03-22 18:21:44 -04:00 committed by GitHub
parent 12a3e97ef3
commit b68e9a3801
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 82 additions and 4 deletions

View file

@ -749,6 +749,8 @@ pub enum Statement {
delimiter: Option<Ident>,
/// CSV HEADER
csv_header: bool,
/// If true, is a 'COPY TO' statement. If false is a 'COPY FROM'
to: bool,
},
/// UPDATE
Update {
@ -1143,6 +1145,7 @@ impl fmt::Display for Statement {
delimiter,
filename,
csv_header,
to,
} => {
write!(f, "COPY {}", table_name)?;
if !columns.is_empty() {
@ -1150,7 +1153,13 @@ impl fmt::Display for Statement {
}
if let Some(name) = filename {
if *to {
write!(f, " TO {}", name)?
} else {
write!(f, " FROM {}", name)?;
}
} else if *to {
write!(f, " TO stdin ")?
} else {
write!(f, " FROM stdin ")?;
}

View file

@ -2233,7 +2233,12 @@ impl<'a> Parser<'a> {
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
let table_name = self.parse_object_name()?;
let columns = self.parse_parenthesized_column_list(Optional)?;
self.expect_keywords(&[Keyword::FROM])?;
let to_or_from = self.expect_one_of_keywords(&[Keyword::FROM, Keyword::TO])?;
let to: bool = match to_or_from {
Keyword::TO => true,
Keyword::FROM => false,
_ => unreachable!("something wrong while parsing copy statment :("),
};
let mut filename = None;
// check whether data has to be copied form table or std in.
if !self.parse_keyword(Keyword::STDIN) {
@ -2271,6 +2276,7 @@ impl<'a> Parser<'a> {
filename,
delimiter,
csv_header,
to,
})
}

View file

@ -402,7 +402,7 @@ PHP ₱ USD $
}
#[test]
fn test_copy() {
fn test_copy_from() {
let stmt = pg().verified_stmt("COPY users FROM 'data.csv'");
assert_eq!(
stmt,
@ -415,7 +415,8 @@ fn test_copy() {
}),
values: vec![],
delimiter: None,
csv_header: false
csv_header: false,
to: false
}
);
@ -435,6 +436,7 @@ fn test_copy() {
quote_style: Some('\'')
}),
csv_header: false,
to: false
}
);
@ -454,6 +456,67 @@ fn test_copy() {
quote_style: Some('\'')
}),
csv_header: true,
to: false
}
)
}
#[test]
fn test_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![],
filename: Some(Ident {
value: "data.csv".to_string(),
quote_style: Some('\'')
}),
values: vec![],
delimiter: None,
csv_header: false,
to: true
}
);
let stmt = pg().verified_stmt("COPY users TO 'data.csv' DELIMITER ','");
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
filename: Some(Ident {
value: "data.csv".to_string(),
quote_style: Some('\'')
}),
values: vec![],
delimiter: Some(Ident {
value: ",".to_string(),
quote_style: Some('\'')
}),
csv_header: false,
to: true
}
);
let stmt = pg().verified_stmt("COPY users TO 'data.csv' DELIMITER ',' CSV HEADER");
assert_eq!(
stmt,
Statement::Copy {
table_name: ObjectName(vec!["users".into()]),
columns: vec![],
filename: Some(Ident {
value: "data.csv".to_string(),
quote_style: Some('\'')
}),
values: vec![],
delimiter: Some(Ident {
value: ",".to_string(),
quote_style: Some('\'')
}),
csv_header: true,
to: true
}
)
}