Parse COPY FROM STDIN payload as rows

Author: LucaCappelletti94

src/ast/mod.rs

@@ -3650,8 +3650,9 @@ pub enum Statement {
         options: Vec<CopyOption>,
         /// WITH options (before PostgreSQL version 9.0)
         legacy_options: Vec<CopyLegacyOption>,
-        /// VALUES a vector of values to be copied
-        values: Vec<Option<String>>,
+        /// The inline payload of `COPY ... FROM STDIN`, one inner vector per
+        /// row, where `None` is the `\N` null field.
+        values: Vec<Vec<Option<String>>>,
     },
     /// ```sql
     /// COPY INTO <table> | <location>
@@ -5349,17 +5350,19 @@ impl fmt::Display for Statement {
                 }
                 if !values.is_empty() {
                     writeln!(f, ";")?;
-                    let mut delim = "";
-                    for v in values {
-                        write!(f, "{delim}")?;
-                        delim = "\t";
-                        if let Some(v) = v {
-                            write!(f, "{v}")?;
-                        } else {
-                            write!(f, "\\N")?;
+                    for row in values {
+                        let mut delim = "";
+                        for field in row {
+                            write!(f, "{delim}")?;
+                            delim = "\t";
+                            match field {
+                                Some(field) => write!(f, "{field}")?,
+                                None => write!(f, "\\N")?,
+                            }
                         }
+                        writeln!(f)?;
                     }
-                    write!(f, "\n\\.")?;
+                    write!(f, "\\.")?;
                 }
                 Ok(())
             }

src/parser/mod.rs

@@ -12353,29 +12353,50 @@ impl<'a> Parser<'a> {
 
     /// Parse a tab separated values in
     /// COPY payload
-    pub fn parse_tsv(&mut self) -> Vec<Option<String>> {
+    pub fn parse_tsv(&mut self) -> Vec<Vec<Option<String>>> {
         self.parse_tab_value()
     }
 
-    /// Parse a single tab-separated value row used by `COPY` payload parsing.
-    pub fn parse_tab_value(&mut self) -> Vec<Option<String>> {
-        let mut values = vec![];
+    /// Parse the tab separated payload rows of a `COPY ... FROM STDIN`.
+    ///
+    /// A tab ends a field, a newline ends a row, and `\N` is a null field.
+    pub fn parse_tab_value(&mut self) -> Vec<Vec<Option<String>>> {
+        /// Take the field accumulated so far, `None` when it was `\N`.
+        fn take_field(content: &mut String, is_null: &mut bool) -> Option<String> {
+            if core::mem::take(is_null) {
+                content.clear();
+                None
+            } else {
+                Some(core::mem::take(content))
+            }
+        }
+
+        // The newline closing the `COPY ... ;` header opens the payload, so it
+        // is not a row terminator. Every later newline is.
+        if self.peek_token_no_skip().token == Token::Whitespace(Whitespace::Newline) {
+            self.next_token_no_skip();
+        }
+
+        let mut rows = vec![];
+        let mut row: Vec<Option<String>> = vec![];
         let mut content = String::new();
+        let mut is_null = false;
         while let Some(t) = self.next_token_no_skip().map(|t| &t.token) {
             match t {
                 Token::Whitespace(Whitespace::Tab) => {
-                    values.push(Some(core::mem::take(&mut content)));
+                    row.push(take_field(&mut content, &mut is_null));
                 }
                 Token::Whitespace(Whitespace::Newline) => {
-                    values.push(Some(core::mem::take(&mut content)));
+                    row.push(take_field(&mut content, &mut is_null));
+                    rows.push(core::mem::take(&mut row));
                 }
                 Token::Backslash => {
                     if self.consume_token(&Token::Period) {
-                        return values;
+                        break;
                     }
                     if let Token::Word(w) = self.next_token().token {
                         if w.value == "N" {
-                            values.push(None);
+                            is_null = true;
                         }
                     }
                 }
@@ -12384,7 +12405,11 @@ impl<'a> Parser<'a> {
                 }
             }
         }
-        values
+        if !row.is_empty() || !content.is_empty() || is_null {
+            row.push(take_field(&mut content, &mut is_null));
+            rows.push(row);
+        }
+        rows
     }
 
     /// Parse a literal value (numbers, strings, date/time, booleans)
@@ -21337,7 +21362,6 @@ mod tests {
                     unit: None
                 }))
             );
-
             test_parse_data_type!(
                 dialect,
                 "CHAR(20 CHARACTERS)",

tests/sqlparser_postgres.rs

@@ -9663,3 +9663,32 @@ fn parse_right_deep_join_chain() {
     // NATURAL JOIN followed by a constrained join must stay left-associative.
     pg().verified_stmt("SELECT * FROM t0 NATURAL JOIN t1 INNER JOIN t2 ON true");
 }
+
+#[test]
+fn parse_copy_from_stdin_payload_rows() {
+    fn payload(sql: &str) -> Vec<Vec<Option<String>>> {
+        let Statement::Copy { values, .. } = pg().verified_stmt(sql) else {
+            panic!("expected a COPY statement");
+        };
+        values
+    }
+
+    // Rows keep their boundaries, and `\N` is a null field rather than an extra one.
+    assert_eq!(
+        payload("COPY t (a, b) FROM STDIN;\n1\t\\N\n2\ty\n\\."),
+        vec![
+            vec![Some("1".to_string()), None],
+            vec![Some("2".to_string()), Some("y".to_string())],
+        ]
+    );
+    // An empty field is not a null one.
+    assert_eq!(
+        payload("COPY t (a, b) FROM STDIN;\n\t\\N\n\\."),
+        vec![vec![Some(String::new()), None]]
+    );
+    // A row of empty fields is still a row.
+    assert_eq!(
+        payload("COPY t (a) FROM STDIN;\n\n\\."),
+        vec![vec![Some(String::new())]]
+    );
+}