Skip to content

Commit b68e9a3

Browse files
Add support for COPY TO (#441)
* Start adding COPY TO * Fix statement and add tests * Merge copy statements * Remove extra line * Clippy * Cleanup
1 parent 12a3e97 commit b68e9a3

File tree

3 files changed

+82
-4
lines changed

3 files changed

+82
-4
lines changed

src/ast/mod.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -749,6 +749,8 @@ pub enum Statement {
749749
delimiter: Option<Ident>,
750750
/// CSV HEADER
751751
csv_header: bool,
752+
/// If true, is a 'COPY TO' statement. If false is a 'COPY FROM'
753+
to: bool,
752754
},
753755
/// UPDATE
754756
Update {
@@ -1143,14 +1145,21 @@ impl fmt::Display for Statement {
11431145
delimiter,
11441146
filename,
11451147
csv_header,
1148+
to,
11461149
} => {
11471150
write!(f, "COPY {}", table_name)?;
11481151
if !columns.is_empty() {
11491152
write!(f, " ({})", display_comma_separated(columns))?;
11501153
}
11511154

11521155
if let Some(name) = filename {
1153-
write!(f, " FROM {}", name)?;
1156+
if *to {
1157+
write!(f, " TO {}", name)?
1158+
} else {
1159+
write!(f, " FROM {}", name)?;
1160+
}
1161+
} else if *to {
1162+
write!(f, " TO stdin ")?
11541163
} else {
11551164
write!(f, " FROM stdin ")?;
11561165
}

src/parser.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2233,7 +2233,12 @@ impl<'a> Parser<'a> {
22332233
pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
22342234
let table_name = self.parse_object_name()?;
22352235
let columns = self.parse_parenthesized_column_list(Optional)?;
2236-
self.expect_keywords(&[Keyword::FROM])?;
2236+
let to_or_from = self.expect_one_of_keywords(&[Keyword::FROM, Keyword::TO])?;
2237+
let to: bool = match to_or_from {
2238+
Keyword::TO => true,
2239+
Keyword::FROM => false,
2240+
_ => unreachable!("something wrong while parsing copy statment :("),
2241+
};
22372242
let mut filename = None;
22382243
// check whether data has to be copied form table or std in.
22392244
if !self.parse_keyword(Keyword::STDIN) {
@@ -2271,6 +2276,7 @@ impl<'a> Parser<'a> {
22712276
filename,
22722277
delimiter,
22732278
csv_header,
2279+
to,
22742280
})
22752281
}
22762282

tests/sqlparser_postgres.rs

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -402,7 +402,7 @@ PHP ₱ USD $
402402
}
403403

404404
#[test]
405-
fn test_copy() {
405+
fn test_copy_from() {
406406
let stmt = pg().verified_stmt("COPY users FROM 'data.csv'");
407407
assert_eq!(
408408
stmt,
@@ -415,7 +415,8 @@ fn test_copy() {
415415
}),
416416
values: vec![],
417417
delimiter: None,
418-
csv_header: false
418+
csv_header: false,
419+
to: false
419420
}
420421
);
421422

@@ -435,6 +436,7 @@ fn test_copy() {
435436
quote_style: Some('\'')
436437
}),
437438
csv_header: false,
439+
to: false
438440
}
439441
);
440442

@@ -454,6 +456,67 @@ fn test_copy() {
454456
quote_style: Some('\'')
455457
}),
456458
csv_header: true,
459+
to: false
460+
}
461+
)
462+
}
463+
464+
#[test]
465+
fn test_copy_to() {
466+
let stmt = pg().verified_stmt("COPY users TO 'data.csv'");
467+
assert_eq!(
468+
stmt,
469+
Statement::Copy {
470+
table_name: ObjectName(vec!["users".into()]),
471+
columns: vec![],
472+
filename: Some(Ident {
473+
value: "data.csv".to_string(),
474+
quote_style: Some('\'')
475+
}),
476+
values: vec![],
477+
delimiter: None,
478+
csv_header: false,
479+
to: true
480+
}
481+
);
482+
483+
let stmt = pg().verified_stmt("COPY users TO 'data.csv' DELIMITER ','");
484+
assert_eq!(
485+
stmt,
486+
Statement::Copy {
487+
table_name: ObjectName(vec!["users".into()]),
488+
columns: vec![],
489+
filename: Some(Ident {
490+
value: "data.csv".to_string(),
491+
quote_style: Some('\'')
492+
}),
493+
values: vec![],
494+
delimiter: Some(Ident {
495+
value: ",".to_string(),
496+
quote_style: Some('\'')
497+
}),
498+
csv_header: false,
499+
to: true
500+
}
501+
);
502+
503+
let stmt = pg().verified_stmt("COPY users TO 'data.csv' DELIMITER ',' CSV HEADER");
504+
assert_eq!(
505+
stmt,
506+
Statement::Copy {
507+
table_name: ObjectName(vec!["users".into()]),
508+
columns: vec![],
509+
filename: Some(Ident {
510+
value: "data.csv".to_string(),
511+
quote_style: Some('\'')
512+
}),
513+
values: vec![],
514+
delimiter: Some(Ident {
515+
value: ",".to_string(),
516+
quote_style: Some('\'')
517+
}),
518+
csv_header: true,
519+
to: true
457520
}
458521
)
459522
}

0 commit comments

Comments
 (0)