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

View file

@ -2233,7 +2233,12 @@ impl<'a> Parser<'a> {
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> { pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
let table_name = self.parse_object_name()?; let table_name = self.parse_object_name()?;
let columns = self.parse_parenthesized_column_list(Optional)?; 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; let mut filename = None;
// check whether data has to be copied form table or std in. // check whether data has to be copied form table or std in.
if !self.parse_keyword(Keyword::STDIN) { if !self.parse_keyword(Keyword::STDIN) {
@ -2271,6 +2276,7 @@ impl<'a> Parser<'a> {
filename, filename,
delimiter, delimiter,
csv_header, csv_header,
to,
}) })
} }

View file

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