Capture the COPY ... FROM STDIN payload verbatim

Author: LucaCappelletti94

src/ast/mod.rs

@@ -3687,8 +3687,11 @@ 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`, as written, without the
+        /// `\.` terminator. Interpreting it (field separators, escapes and the
+        /// `\N` null marker) is left to the caller. `None` when the statement
+        /// carries no inline payload.
+        payload: Option<String>,
     },
     /// ```sql
     /// COPY INTO <table> | <location>
@@ -5367,7 +5370,7 @@ impl fmt::Display for Statement {
                 target,
                 options,
                 legacy_options,
-                values,
+                payload,
             } => {
                 write!(f, "COPY")?;
                 match source {
@@ -5389,19 +5392,10 @@ impl fmt::Display for Statement {
                 if !legacy_options.is_empty() {
                     write!(f, " {}", display_separated(legacy_options, " "))?;
                 }
-                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")?;
-                        }
-                    }
-                    write!(f, "\n\\.")?;
+                if let Some(payload) = payload {
+                    // The payload starts on the line after the command and ends
+                    // at the `\.` terminator.
+                    write!(f, ";\n{payload}\\.")?;
                 }
                 Ok(())
             }

src/ast/spans.rs

@@ -347,7 +347,7 @@ impl Spanned for Statement {
                 target: _,
                 options: _,
                 legacy_options: _,
-                values: _,
+                payload: _,
             } => source.span(),
             Statement::CopyIntoSnowflake {
                 into: _,

src/parser/mod.rs

@@ -12134,20 +12134,20 @@ impl<'a> Parser<'a> {
         while let Some(opt) = self.maybe_parse(|parser| parser.parse_copy_legacy_option())? {
             legacy_options.push(opt);
         }
-        let values =
+        let payload =
             if matches!(target, CopyTarget::Stdin) && self.peek_token_ref().token != Token::EOF {
                 self.expect_token(&Token::SemiColon)?;
-                self.parse_tsv()
+                self.parse_copy_payload()
             } else {
-                vec![]
+                None
             };
         Ok(Statement::Copy {
             source,
             to,
             target,
             options,
             legacy_options,
-            values,
+            payload,
         })
     }
 
@@ -12488,40 +12488,50 @@ impl<'a> Parser<'a> {
         Ok(s.chars().next().unwrap())
     }
 
-    /// Parse a tab separated values in
-    /// COPY payload
-    pub fn parse_tsv(&mut self) -> 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![];
-        let mut content = String::new();
-        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)));
-                }
-                Token::Whitespace(Whitespace::Newline) => {
-                    values.push(Some(core::mem::take(&mut content)));
-                }
+    /// Parse the inline payload of `COPY ... FROM STDIN` as written.
+    ///
+    /// The payload is data rather than SQL, so it is captured uninterpreted.
+    /// It starts on the line following the command and runs up to the `\.`
+    /// terminator, which is consumed but not returned. Field
+    /// separators, escapes and the `\N` null marker are left to the caller.
+    ///
+    /// Returns `None` when nothing but the statement delimiter follows, so a
+    /// bare `COPY t FROM STDIN;` is not given a payload it never carried.
+    fn parse_copy_payload(&mut self) -> Option<String> {
+        use core::fmt::Write;
+
+        let mut payload = String::new();
+        // The newline ending the command line is a delimiter, not data.
+        if self.peek_nth_token_no_skip_ref(0).token == Token::Whitespace(Whitespace::Newline) {
+            let _ = self.next_token_no_skip();
+        }
+        while let Some(token) = self.next_token_no_skip().map(|t| &t.token) {
+            match token {
                 Token::Backslash => {
-                    if self.consume_token(&Token::Period) {
-                        return values;
-                    }
-                    if let Token::Word(w) = self.next_token().token {
-                        if w.value == "N" {
-                            values.push(None);
-                        }
+                    // `\.` terminates the payload, every other backslash is data.
+                    if self.peek_nth_token_no_skip_ref(0).token == Token::Period {
+                        let _ = self.next_token_no_skip();
+                        return Some(payload);
                     }
+                    payload.push('\\');
                 }
-                _ => {
-                    content.push_str(&t.to_string());
+                // `Display` for a single line comment appends a newline that the
+                // tokenizer left in the stream as a separate token, so render it
+                // here instead to avoid duplicating it.
+                Token::Whitespace(Whitespace::SingleLineComment { prefix, comment }) => {
+                    payload.push_str(prefix);
+                    payload.push_str(comment);
+                }
+                // Rendering into the buffer avoids a temporary `String` per
+                // token, and a payload can hold an entire table.
+                token => {
+                    let _ = write!(payload, "{token}");
                 }
             }
         }
-        values
+        // An unterminated payload keeps whatever data it holds, but an empty one
+        // means the statement had no inline data at all.
+        (!payload.is_empty()).then_some(payload)
     }
 
     /// Parse a literal value (numbers, strings, date/time, booleans)

tests/sqlparser_postgres.rs

@@ -1443,7 +1443,10 @@ Kwara & Kogi
 PHP	₱ USD $
 \N  Some other value
 \\."#;
-    pg_and_generic().one_statement_parses_to(sql, "");
+    // The payload survives printing unchanged, so the only difference from the
+    // input is the normalized `STDIN` keyword.
+    let canonical = sql.replacen("FROM stdin;", "FROM STDIN;", 1);
+    pg_and_generic().one_statement_parses_to(sql, &canonical);
 }
 
 #[test]
@@ -1460,7 +1463,7 @@ fn parse_copy_from_stdin_without_semicolon() {
             target: CopyTarget::Stdin,
             options: vec![],
             legacy_options: vec![CopyLegacyOption::Null("null".into())],
-            values: vec![],
+            payload: None,
         }
     );
 }
@@ -1489,11 +1492,11 @@ fn parse_copy_from_stdin_without_semicolon_variants() {
             Statement::Copy {
                 to: false,
                 target: CopyTarget::Stdin,
-                values,
+                payload,
                 ..
             } => {
                 assert!(
-                    values.is_empty(),
+                    payload.is_none(),
                     "expected no inline COPY payload for `{sql}`"
                 );
             }
@@ -1518,7 +1521,7 @@ fn test_copy_from() {
             },
             options: vec![],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1536,7 +1539,7 @@ fn test_copy_from() {
             },
             options: vec![],
             legacy_options: vec![CopyLegacyOption::Delimiter(',')],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1557,7 +1560,7 @@ fn test_copy_from() {
                 CopyLegacyOption::Delimiter(','),
                 CopyLegacyOption::Csv(vec![CopyLegacyCsvOption::Header,])
             ],
-            values: vec![],
+            payload: None,
         }
     );
 }
@@ -1578,7 +1581,7 @@ fn test_copy_to() {
             },
             options: vec![],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1596,7 +1599,7 @@ fn test_copy_to() {
             },
             options: vec![],
             legacy_options: vec![CopyLegacyOption::Delimiter(',')],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1617,7 +1620,7 @@ fn test_copy_to() {
                 CopyLegacyOption::Delimiter(','),
                 CopyLegacyOption::Csv(vec![CopyLegacyCsvOption::Header,])
             ],
-            values: vec![],
+            payload: None,
         }
     )
 }
@@ -1671,7 +1674,7 @@ fn parse_copy_from() {
                 CopyOption::Encoding("utf8".into()),
             ],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 }
@@ -1701,7 +1704,7 @@ fn parse_copy_to() {
             },
             options: vec![],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1717,7 +1720,7 @@ fn parse_copy_to() {
             target: CopyTarget::Stdout,
             options: vec![CopyOption::Delimiter('|')],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1736,7 +1739,7 @@ fn parse_copy_to() {
             },
             options: vec![],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1806,7 +1809,7 @@ fn parse_copy_to() {
             },
             options: vec![],
             legacy_options: vec![],
-            values: vec![],
+            payload: None,
         }
     )
 }
@@ -1837,7 +1840,7 @@ fn parse_copy_from_before_v9_0() {
                     CopyLegacyCsvOption::ForceNotNull(vec!["column".into()]),
                 ]),
             ],
-            values: vec![],
+            payload: None,
         }
     );
 
@@ -1863,7 +1866,7 @@ fn parse_copy_from_before_v9_0() {
                     CopyLegacyCsvOption::Escape('\\'),
                 ]),
             ],
-            values: vec![],
+            payload: None,
         }
     );
 }
@@ -1894,7 +1897,7 @@ fn parse_copy_to_before_v9_0() {
                     CopyLegacyCsvOption::ForceQuote(vec!["column".into()]),
                 ]),
             ],
-            values: vec![],
+            payload: None,
         }
     )
 }
@@ -10126,3 +10129,65 @@ fn parse_bitstring_literal_escaping() {
     pg_and_generic().verified_stmt("SELECT B''''");
     pg_and_generic().verified_stmt("SELECT B'it''s'");
 }
+
+#[test]
+fn parse_copy_from_stdin_payload_is_captured_as_written() {
+    let check = |sql: &str, printed: &str, expected: Option<&str>| match pg_and_generic()
+        .one_statement_parses_to(sql, printed)
+    {
+        Statement::Copy { payload, .. } => {
+            assert_eq!(payload.as_deref(), expected, "for {sql:?}")
+        }
+        other => panic!("expected COPY, got {other:?}"),
+    };
+
+    // Tabs, newlines and the `\N` null marker are all left uninterpreted.
+    let rows = "COPY t (a, b) FROM STDIN;\n1\t\\N\n2\ty\n\\.";
+    check(rows, rows, Some("1\t\\N\n2\ty\n"));
+
+    // Only the keyword casing is normalized.
+    check(
+        "COPY t (a, b) FROM stdin;\n1\t\\N\n2\ty\n\\.",
+        rows,
+        Some("1\t\\N\n2\ty\n"),
+    );
+
+    // An empty payload is still a payload, distinct from a `COPY` that carries
+    // no inline data at all.
+    let empty = "COPY t (a, b) FROM STDIN;\n\\.";
+    check(empty, empty, Some(""));
+
+    // A statement delimiter with nothing behind it is not a payload, so
+    // printing must not invent data and a `\.` terminator.
+    check(
+        "COPY t (a, b) FROM STDIN;",
+        "COPY t (a, b) FROM STDIN",
+        None,
+    );
+    check(
+        "COPY t (a, b) FROM STDIN;\n",
+        "COPY t (a, b) FROM STDIN",
+        None,
+    );
+
+    // A payload cut short keeps the rows it does have, and printing restores
+    // the terminator the input was missing.
+    check(
+        "COPY t (a, b) FROM STDIN;\n1\t2\n",
+        "COPY t (a, b) FROM STDIN;\n1\t2\n\\.",
+        Some("1\t2\n"),
+    );
+
+    // The declared format does not change how the payload is captured.
+    let csv = "COPY t (a, b) FROM STDIN (FORMAT CSV);\na,b\n1,\"x y\"\n\\.";
+    check(csv, csv, Some("a,b\n1,\"x y\"\n"));
+}
+
+#[test]
+fn parse_copy_from_stdin_payload_with_comment_marker() {
+    // A `--` in the payload is data, but the tokenizer still reports it as a
+    // comment, and printing that token appends a newline the tokenizer already
+    // emitted separately.
+    pg_and_generic().verified_stmt("COPY t (a, b) FROM STDIN;\n1\t-- not a comment\n2\ty\n\\.");
+    pg_and_generic().verified_stmt("COPY t (a, b) FROM STDIN;\n1\t/* nor is this */\n\\.");
+}