From 2614576dbf140824cb659bcc7026b21bd1b82f0a Mon Sep 17 00:00:00 2001
From: Poonai <rbalajis25@gmail.com>
Date: Mon, 7 Feb 2022 03:33:28 -0800
Subject: [PATCH] Add support for `FROM <filename>`, `DELIMITER`, and `CSV
 HEADER` options for `COPY` command (#409)

* add additional options for copy command

Signed-off-by: password <rbalajis25@gmail.com>

* cargo fmt

Signed-off-by: password <rbalajis25@gmail.com>
---
 src/ast/mod.rs              | 28 ++++++++++++++++--
 src/keywords.rs             |  1 +
 src/parser.rs               | 37 ++++++++++++++++++++++--
 tests/sqlparser_postgres.rs | 57 +++++++++++++++++++++++++++++++++++++
 4 files changed, 118 insertions(+), 5 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 7511a155..3c769f76 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -670,6 +670,12 @@ pub enum Statement {
         columns: Vec<Ident>,
         /// VALUES a vector of values to be copied
         values: Vec<Option<String>>,
+        /// file name of the data to be copied from
+        filename: Option<Ident>,
+        /// delimiter character
+        delimiter: Option<Ident>,
+        /// CSV HEADER
+        csv_header: bool,
     },
     /// UPDATE
     Update {
@@ -1043,13 +1049,28 @@ impl fmt::Display for Statement {
                 table_name,
                 columns,
                 values,
+                delimiter,
+                filename,
+                csv_header,
             } => {
                 write!(f, "COPY {}", table_name)?;
                 if !columns.is_empty() {
                     write!(f, " ({})", display_comma_separated(columns))?;
                 }
-                write!(f, " FROM stdin; ")?;
+
+                if let Some(name) = filename {
+                    write!(f, " FROM {}", name)?;
+                } else {
+                    write!(f, " FROM stdin ")?;
+                }
+                if let Some(delimiter) = delimiter {
+                    write!(f, " DELIMITER {}", delimiter)?;
+                }
+                if *csv_header {
+                    write!(f, " CSV HEADER")?;
+                }
                 if !values.is_empty() {
+                    write!(f, ";")?;
                     writeln!(f)?;
                     let mut delim = "";
                     for v in values {
@@ -1062,7 +1083,10 @@ impl fmt::Display for Statement {
                         }
                     }
                 }
-                write!(f, "\n\\.")
+                if filename.is_none() {
+                    write!(f, "\n\\.")?;
+                }
+                Ok(())
             }
             Statement::Update {
                 table,
diff --git a/src/keywords.rs b/src/keywords.rs
index 7689fc05..d23c587e 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -173,6 +173,7 @@ define_keywords!(
     DEFAULT,
     DELETE,
     DELIMITED,
+    DELIMITER,
     DENSE_RANK,
     DEREF,
     DESC,
diff --git a/src/parser.rs b/src/parser.rs
index fcadb96b..5faa2ca4 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -2087,13 +2087,44 @@ 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, Keyword::STDIN])?;
-        self.expect_token(&Token::SemiColon)?;
-        let values = self.parse_tsv();
+        self.expect_keywords(&[Keyword::FROM])?;
+        let mut filename = None;
+        // check whether data has to be copied form table or std in.
+        if !self.parse_keyword(Keyword::STDIN) {
+            filename = Some(self.parse_identifier()?)
+        }
+        // parse copy options.
+        let mut delimiter = None;
+        let mut csv_header = false;
+        loop {
+            if let Some(keyword) = self.parse_one_of_keywords(&[Keyword::DELIMITER, Keyword::CSV]) {
+                match keyword {
+                    Keyword::DELIMITER => {
+                        delimiter = Some(self.parse_identifier()?);
+                        continue;
+                    }
+                    Keyword::CSV => {
+                        self.expect_keyword(Keyword::HEADER)?;
+                        csv_header = true
+                    }
+                    _ => unreachable!("something wrong while parsing copy statment :("),
+                }
+            }
+            break;
+        }
+        // copy the values from stdin if there is no file to be copied from.
+        let mut values = vec![];
+        if filename.is_none() {
+            self.expect_token(&Token::SemiColon)?;
+            values = self.parse_tsv();
+        }
         Ok(Statement::Copy {
             table_name,
             columns,
             values,
+            filename,
+            delimiter,
+            csv_header,
         })
     }
 
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 60d9c1cb..2efa2f20 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -404,6 +404,63 @@ PHP	₱ USD $
     //assert_eq!(sql, ast.to_string());
 }
 
+#[test]
+fn test_copy() {
+    let stmt = pg().verified_stmt("COPY users FROM '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
+        }
+    );
+
+    let stmt = pg().verified_stmt("COPY users FROM '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,
+        }
+    );
+
+    let stmt = pg().verified_stmt("COPY users FROM '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,
+        }
+    )
+}
+
 #[test]
 fn parse_set() {
     let stmt = pg_and_generic().verified_stmt("SET a = b");