Add support for COPY TO (apache#441)

matthewmturner · web-flow · commit b68e9a380191 · 2022-03-22T18:21:44.000-04:00
* Start adding COPY TO

* Fix statement and add tests

* Merge copy statements

* Remove extra line

* Clippy

* Cleanup
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -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,14 +1145,21 @@ impl fmt::Display for Statement {
                 delimiter,
                 filename,
                 csv_header,
+                to,
             } => {
                 write!(f, "COPY {}", table_name)?;
                 if !columns.is_empty() {
                     write!(f, " ({})", display_comma_separated(columns))?;
                 }
 
                 if let Some(name) = filename {
-                    write!(f, " FROM {}", name)?;
+                    if *to {
+                        write!(f, " TO {}", name)?
+                    } else {
+                        write!(f, " FROM {}", name)?;
+                    }
+                } else if *to {
+                    write!(f, " TO stdin ")?
                 } else {
                     write!(f, " FROM stdin ")?;
                 }
diff --git a/src/parser.rs b/src/parser.rs
@@ -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,
         })
     }
 
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -402,7 +402,7 @@ PHP	&#8369; 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
         }
     )
 }
