diff --git a/src/ast/mod.rs b/src/ast/mod.rs index eb5bc6cc75..72194057f0 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -3692,8 +3692,10 @@ pub enum Statement { options: Vec, /// WITH options (before PostgreSQL version 9.0) legacy_options: Vec, - /// VALUES a vector of values to be copied - values: Vec>, + /// The inline payload of `COPY ... FROM STDIN`, captured verbatim + /// from the source text. `None` when the statement carries no + /// inline data. + payload: Option, }, /// ```sql /// COPY INTO | @@ -5378,7 +5380,7 @@ impl fmt::Display for Statement { target, options, legacy_options, - values, + payload, } => { write!(f, "COPY")?; match source { @@ -5400,19 +5402,15 @@ impl fmt::Display for Statement { if !legacy_options.is_empty() { write!(f, " {}", display_separated(legacy_options, " "))?; } - if !values.is_empty() { + if let Some(payload) = payload { + // The payload starts on the line after the command and ends + // at the `\.` terminator line, which is only written when + // present in the source. 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, "{}", payload.data)?; + if payload.terminated { + write!(f, "\\.")?; } - write!(f, "\n\\.")?; } Ok(()) } @@ -9555,6 +9553,20 @@ impl fmt::Display for CopyTarget { } } +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +/// The inline payload of a `COPY ... FROM STDIN` statement, captured +/// verbatim from the source text. +pub struct CopyPayload { + /// The payload as written, starting on the line after the command and + /// running up to the `\.` terminator line. Field separators, escapes and + /// the `\N` null marker are left for the caller to interpret. + pub data: String, + /// Whether the `\.` terminator line was present. + pub terminated: bool, +} + #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 8267c71798..ceab88503a 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -347,7 +347,7 @@ impl Spanned for Statement { target: _, options: _, legacy_options: _, - values: _, + payload: _, } => source.span(), Statement::CopyIntoSnowflake { into: _, diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index 9c1dc08b3a..f9d2aae6b8 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -205,6 +205,10 @@ impl Dialect for GenericDialect { true } + fn supports_copy_from_stdin_data(&self) -> bool { + true + } + fn supports_user_host_grantee(&self) -> bool { true } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index a2920e61eb..65b13b0d36 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1306,6 +1306,12 @@ pub trait Dialect: Debug + Any { false } + /// Returns true if the dialect supports inline data for + /// `COPY ... FROM STDIN`, captured verbatim by the tokenizer. + fn supports_copy_from_stdin_data(&self) -> bool { + false + } + /// Returns true if `#` starts a single line comment only when followed by an ASCII /// space or `!`, so `# comment` and `#! comment` are comments but `#comment` is not. fn requires_hash_comment_space(&self) -> bool { diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs index c060f2a8c1..15d8c6733c 100644 --- a/src/dialect/postgresql.rs +++ b/src/dialect/postgresql.rs @@ -458,6 +458,10 @@ impl Dialect for PostgreSqlDialect { true } + fn supports_copy_from_stdin_data(&self) -> bool { + true + } + fn supports_string_escape_constant(&self) -> bool { true } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 276b252ed2..e0d7fe8fb4 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -12229,23 +12229,32 @@ impl<'a> Parser<'a> { while let Some(opt) = self.maybe_parse(|parser| parser.parse_copy_legacy_option())? { legacy_options.push(opt); } - let values = - if matches!(target, CopyTarget::Stdin) && self.peek_token_ref().token != Token::EOF { - self.expect_token(&Token::SemiColon)?; - self.parse_tsv() - } else { - vec![] - }; + let payload = if matches!(target, CopyTarget::Stdin) { + self.take_copy_payload() + } else { + None + }; Ok(Statement::Copy { source, to, target, options, legacy_options, - values, + payload, }) } + /// Take the raw inline payload of a `COPY ... FROM STDIN` statement, or + /// `None` when the statement carried no inline data. + fn take_copy_payload(&mut self) -> Option { + let Token::CopyData { data, terminated } = &self.peek_token_ref().token else { + return None; + }; + let (data, terminated) = (data.clone(), *terminated); + self.next_token(); + Some(CopyPayload { data, terminated }) + } + /// Parse [Statement::Open] fn parse_open(&mut self) -> Result { self.expect_keyword(Keyword::OPEN)?; @@ -12583,42 +12592,6 @@ impl<'a> Parser<'a> { Ok(s.chars().next().unwrap()) } - /// Parse a tab separated values in - /// COPY payload - pub fn parse_tsv(&mut self) -> Vec> { - self.parse_tab_value() - } - - /// Parse a single tab-separated value row used by `COPY` payload parsing. - pub fn parse_tab_value(&mut self) -> Vec> { - 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))); - } - 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); - } - } - } - _ => { - content.push_str(&t.to_string()); - } - } - } - values - } - /// Parse a literal value (numbers, strings, date/time, booleans) pub fn parse_value(&mut self) -> Result { let next_token = self.next_token(); diff --git a/src/tokenizer.rs b/src/tokenizer.rs index b6dafe12e9..a1fa33f161 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -163,6 +163,15 @@ pub enum Token { SemiColon, /// Backslash `\` used in terminating the COPY payload with `\.` Backslash, + /// The verbatim inline payload of a `COPY ... FROM STDIN` statement, + /// captured raw by the tokenizer. `terminated` records whether the `\.` + /// terminator line was present. + CopyData { + /// The payload text, without the terminator line. + data: String, + /// Whether the `\.` terminator line was present. + terminated: bool, + }, /// Left bracket `[` LBracket, /// Right bracket `]` @@ -339,6 +348,7 @@ impl fmt::Display for Token { Token::Assignment => f.write_str(":="), Token::SemiColon => f.write_str(";"), Token::Backslash => f.write_str("\\"), + Token::CopyData { data, .. } => f.write_str(data), Token::LBracket => f.write_str("["), Token::RBracket => f.write_str("]"), Token::Ampersand => f.write_str("&"), @@ -833,6 +843,80 @@ impl State<'_> { } } +/// Tracks whether the statement being tokenized is a `COPY ... FROM STDIN` statement. +struct CopyStatementTracker { + enabled: bool, + first_token_seen: bool, + starts_with_copy: bool, + prev_was_from: bool, + from_stdin: bool, + paren_depth: u32, +} + +impl CopyStatementTracker { + fn new(enabled: bool) -> Self { + Self { + enabled, + first_token_seen: false, + starts_with_copy: false, + prev_was_from: false, + from_stdin: false, + paren_depth: 0, + } + } + + /// Whether the `;` closing the current statement captures the raw payload. + fn captures_payload(&self) -> bool { + self.enabled && self.starts_with_copy && self.from_stdin + } + + /// Record a token, ignoring whitespace, and reset the statement on `;`. + fn note(&mut self, token: &Token) { + if !self.enabled { + return; + } + match token { + Token::SemiColon => { + self.first_token_seen = false; + self.starts_with_copy = false; + self.prev_was_from = false; + self.from_stdin = false; + self.paren_depth = 0; + } + Token::LParen => { + self.paren_depth += 1; + self.first_token_seen = true; + self.prev_was_from = false; + } + Token::RParen => { + self.paren_depth = self.paren_depth.saturating_sub(1); + self.first_token_seen = true; + self.prev_was_from = false; + } + Token::Whitespace(_) => {} + Token::Word(word) => { + if !self.first_token_seen { + self.first_token_seen = true; + self.starts_with_copy = word.keyword == Keyword::COPY; + } + if self.prev_was_from && word.keyword == Keyword::STDIN && self.paren_depth == 0 { + self.from_stdin = true; + } + self.prev_was_from = word.keyword == Keyword::FROM; + } + _ => { + self.first_token_seen = true; + self.prev_was_from = false; + } + } + } +} + +/// The captured `COPY ... FROM STDIN` payload and the rest of the `;` line. +/// The `Option` is the payload token, absent when the data is empty and +/// unterminated, followed by the tokens of the rest of the line. +type CopyCapture = (Option<(String, bool, Span)>, Vec); + /// Represents how many quote characters enclose a string literal. #[derive(Copy, Clone)] enum NumStringQuoteChars { @@ -964,9 +1048,32 @@ impl<'a> Tokenizer<'a> { }; let mut location = state.location(); + let mut copy = CopyStatementTracker::new(self.dialect.supports_copy_from_stdin_data()); while let Some(token) = self.next_token(&mut state, buf.last().map(|t| &t.token))? { let span = location.span_to(state.location()); + // Capture before tokenizing the rest of the line so both keep their source locations. + let mut payload: Option = None; + if let Token::SemiColon = &token { + if copy.captures_payload() { + payload = Some(self.capture_copy_data(&mut state, Some(&token))?); + } + } + copy.note(&token); + + let (copy_data, rest) = match payload { + Some((copy_data, rest)) => (copy_data, Some(rest)), + None => (None, None), + }; + + // The payload token precedes the `;` so the statement loop owns the delimiter. + if let Some((data, terminated, data_span)) = copy_data { + buf.push(mapper(TokenWithSpan { + token: Token::CopyData { data, terminated }, + span: data_span, + })); + } + // Check if this is a multiline comment hint that should be expanded match &token { Token::Whitespace(Whitespace::MultiLineComment(comment)) @@ -981,6 +1088,12 @@ impl<'a> Tokenizer<'a> { } } + if let Some(rest_tokens) = rest { + for rest_token in rest_tokens { + copy.note(&rest_token.token); + buf.push(mapper(rest_token)); + } + } location = state.location(); } Ok(()) @@ -1030,6 +1143,77 @@ impl<'a> Tokenizer<'a> { Ok(()) } + /// Capture the raw payload after the `;` of a `COPY ... FROM STDIN` statement, returning the rest + /// of that line as tokens to emit after the `;`. The payload token is absent when the data is + /// empty and unterminated, e.g. a bare `COPY t FROM STDIN;` at end of input. + fn capture_copy_data( + &self, + state: &mut State, + prev: Option<&Token>, + ) -> Result { + let mut rest = Vec::new(); + let mut location = state.location(); + let mut prev = prev; + while let Some(token) = self.next_token(state, prev)? { + let span = location.span_to(state.location()); + rest.push(TokenWithSpan { token, span }); + location = state.location(); + prev = Some(&rest.last().unwrap().token); + if matches!(prev, Some(&Token::Whitespace(Whitespace::Newline))) { + break; + } + } + let data_start = state.location(); + let (data, terminated) = self.read_copy_data(state); + let payload = if data.is_empty() && !terminated { + None + } else { + Some((data, terminated, data_start.span_to(state.location()))) + }; + Ok((payload, rest)) + } + + /// Read the raw payload from `state` (at the start of a line) up to a + /// line holding only `\.` (the terminator, which is consumed) or end of + /// input. + fn read_copy_data(&self, state: &mut State) -> (String, bool) { + let mut data = String::new(); + loop { + let mut line = String::new(); + while let Some(&ch) = state.peek() { + if ch == '\n' || ch == '\r' { + break; + } + line.push(ch); + state.next(); + } + if line == "\\." { + // A final terminator without a line end still terminates. + if state.peek() == Some(&'\r') { + state.next(); + } + if state.peek() == Some(&'\n') { + state.next(); + } + return (data, true); + } + data.push_str(&line); + if state.peek() == Some(&'\n') { + data.push('\n'); + state.next(); + } else if state.peek() == Some(&'\r') { + data.push('\r'); + state.next(); + if state.peek() == Some(&'\n') { + data.push('\n'); + state.next(); + } + } else { + return (data, false); + } + } + } + // Tokenize the identifier or keywords in `ch` fn tokenize_identifier_or_keyword( &self, @@ -4694,4 +4878,97 @@ mod tests { tokens, ); } + + #[test] + fn tokenize_copy_from_stdin_payload() { + let sql = "COPY t FROM STDIN; SELECT 1;\nrow1\n\\.\nSELECT 2;"; + let tokens = Tokenizer::new(&PostgreSqlDialect {}, sql) + .tokenize() + .unwrap(); + assert_eq!( + tokens, + vec![ + Token::make_word("COPY", None), + Token::Whitespace(Whitespace::Space), + Token::make_word("t", None), + Token::Whitespace(Whitespace::Space), + Token::make_word("FROM", None), + Token::Whitespace(Whitespace::Space), + Token::make_word("STDIN", None), + Token::CopyData { + data: "row1\n".to_string(), + terminated: true, + }, + Token::SemiColon, + Token::Whitespace(Whitespace::Space), + Token::make_word("SELECT", None), + Token::Whitespace(Whitespace::Space), + Token::Number("1".to_string(), false), + Token::SemiColon, + Token::Whitespace(Whitespace::Newline), + Token::make_word("SELECT", None), + Token::Whitespace(Whitespace::Space), + Token::Number("2".to_string(), false), + Token::SemiColon, + ] + ); + } + + #[test] + fn tokenize_copy_from_stdin_payload_span() { + let sql = "COPY t FROM STDIN;\nrow1\n\\."; + let tokens = Tokenizer::new(&PostgreSqlDialect {}, sql) + .tokenize_with_location() + .unwrap(); + let payload = tokens + .iter() + .find(|t| matches!(t.token, Token::CopyData { .. })) + .unwrap(); + assert_eq!( + payload.span, + Span::new(Location::new(2, 1), Location::new(3, 3)) + ); + assert_eq!( + payload.token, + Token::CopyData { + data: "row1\n".to_string(), + terminated: true, + } + ); + } + + #[test] + fn tokenize_copy_from_stdin_payload_malformed_semicolon_line() { + // An unterminated string on the `;` line fails the capture. + let sql = "COPY t FROM STDIN; 'unterminated"; + let err = Tokenizer::new(&PostgreSqlDialect {}, sql) + .tokenize() + .unwrap_err(); + assert_eq!( + err, + TokenizerError { + message: "Unterminated string literal".to_string(), + location: Location::new(1, 20), + } + ); + } + + #[test] + fn tokenize_copy_from_stdin_payload_not_captured() { + // A `FROM STDIN` nested in parens, or a statement not starting with `COPY`, + // keeps the data out of a payload token. + for sql in [ + "COPY (SELECT x FROM stdin) TO STDOUT;\nrow1\n\\.\n", + "SELECT * FROM stdin;\nrow1\n\\.\n", + "COPY t TO STDOUT;\nrow1\n\\.\n", + ] { + let tokens = Tokenizer::new(&PostgreSqlDialect {}, sql) + .tokenize() + .unwrap(); + assert!( + tokens.iter().all(|t| !matches!(t, Token::CopyData { .. })), + "{sql}" + ); + } + } } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 0a3866d29f..f888609e95 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -24,9 +24,11 @@ mod test_utils; use helpers::attached_token::AttachedToken; use sqlparser::ast::*; -use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect}; +use sqlparser::dialect::{ + Dialect, GenericDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect, SnowflakeDialect, +}; use sqlparser::parser::{Parser, ParserError}; -use sqlparser::tokenizer::{Location, Span}; +use sqlparser::tokenizer::{Location, Span, Token, Tokenizer}; use test_utils::*; #[test] @@ -1441,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] @@ -1458,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, } ); } @@ -1487,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}`" ); } @@ -1516,7 +1521,7 @@ fn test_copy_from() { }, options: vec![], legacy_options: vec![], - values: vec![], + payload: None, } ); @@ -1534,7 +1539,7 @@ fn test_copy_from() { }, options: vec![], legacy_options: vec![CopyLegacyOption::Delimiter(',')], - values: vec![], + payload: None, } ); @@ -1555,7 +1560,7 @@ fn test_copy_from() { CopyLegacyOption::Delimiter(','), CopyLegacyOption::Csv(vec![CopyLegacyCsvOption::Header,]) ], - values: vec![], + payload: None, } ); } @@ -1576,7 +1581,7 @@ fn test_copy_to() { }, options: vec![], legacy_options: vec![], - values: vec![], + payload: None, } ); @@ -1594,7 +1599,7 @@ fn test_copy_to() { }, options: vec![], legacy_options: vec![CopyLegacyOption::Delimiter(',')], - values: vec![], + payload: None, } ); @@ -1615,7 +1620,7 @@ fn test_copy_to() { CopyLegacyOption::Delimiter(','), CopyLegacyOption::Csv(vec![CopyLegacyCsvOption::Header,]) ], - values: vec![], + payload: None, } ) } @@ -1669,7 +1674,7 @@ fn parse_copy_from() { CopyOption::Encoding("utf8".into()), ], legacy_options: vec![], - values: vec![], + payload: None, } ); } @@ -1699,7 +1704,7 @@ fn parse_copy_to() { }, options: vec![], legacy_options: vec![], - values: vec![], + payload: None, } ); @@ -1715,7 +1720,7 @@ fn parse_copy_to() { target: CopyTarget::Stdout, options: vec![CopyOption::Delimiter('|')], legacy_options: vec![], - values: vec![], + payload: None, } ); @@ -1734,7 +1739,7 @@ fn parse_copy_to() { }, options: vec![], legacy_options: vec![], - values: vec![], + payload: None, } ); @@ -1804,7 +1809,7 @@ fn parse_copy_to() { }, options: vec![], legacy_options: vec![], - values: vec![], + payload: None, } ) } @@ -1835,7 +1840,7 @@ fn parse_copy_from_before_v9_0() { CopyLegacyCsvOption::ForceNotNull(vec!["column".into()]), ]), ], - values: vec![], + payload: None, } ); @@ -1861,7 +1866,7 @@ fn parse_copy_from_before_v9_0() { CopyLegacyCsvOption::Escape('\\'), ]), ], - values: vec![], + payload: None, } ); } @@ -1892,7 +1897,7 @@ fn parse_copy_to_before_v9_0() { CopyLegacyCsvOption::ForceQuote(vec!["column".into()]), ]), ], - values: vec![], + payload: None, } ) } @@ -10265,3 +10270,204 @@ fn parse_geometric_typed_literal_error_propagates() { ParserError::ParserError("Expected: a value, found: (".to_string()) ); } + +fn assert_copy_stdin_payload(stmt: &Statement, data: &str, terminated: bool, label: &str) { + match stmt { + Statement::Copy { + to: false, + target: CopyTarget::Stdin, + payload: Some(payload), + .. + } => { + assert_eq!(payload.data, data, "{label}"); + assert_eq!(payload.terminated, terminated, "{label}"); + } + other => panic!("expected a COPY t FROM STDIN payload for {label}, got {other:?}"), + } +} + +#[test] +fn parse_copy_from_stdin_payload_verbatim() { + // The payload is data rather than SQL, so quotes, backslashes, tabs and + // control characters stay in it verbatim. + let data_lines = [ + "O'Brien", + "\"x", + "'a''b'", + "0\x1f1", + "E'\\x41'", + "1\tfoo\\.bar", + ]; + for data_line in data_lines { + let sql = format!("COPY t FROM STDIN;\n{data_line}\n\\."); + let stmt = pg_and_generic().verified_stmt(&sql); + assert_copy_stdin_payload(&stmt, &format!("{data_line}\n"), true, data_line); + } +} + +#[test] +fn parse_copy_from_stdin_payload_crlf() { + let sql = "COPY t FROM STDIN;\r\nrow1\r\nrow2\r\n\\.\r\n"; + // The line end following the `;` prints as a single line feed. + let canonical = "COPY t FROM STDIN;\nrow1\r\nrow2\r\n\\."; + let stmt = pg_and_generic().one_statement_parses_to(sql, canonical); + assert_copy_stdin_payload(&stmt, "row1\r\nrow2\r\n", true, "crlf payload"); +} + +#[test] +fn parse_copy_from_stdin_payload_terminator_before_next_statement() { + let statements = pg_and_generic() + .parse_sql_statements("COPY t FROM STDIN;\nrow1\n\\.\nSELECT 1;") + .unwrap(); + assert_eq!(statements.len(), 2); + assert_copy_stdin_payload( + &statements[0], + "row1\n", + true, + "payload before the terminator", + ); + assert_eq!(statements[1].to_string(), "SELECT 1"); +} + +#[test] +fn parse_copy_from_stdin_payload_then_statement_on_same_line() { + // The statement after the `;` on the command line parses after the COPY, + // while the data still belongs to the COPY. + let statements = pg_and_generic() + .parse_sql_statements("COPY t FROM STDIN; SELECT 'after';\nrow1\n\\.\n") + .unwrap(); + assert_eq!(statements.len(), 2); + assert_copy_stdin_payload( + &statements[0], + "row1\n", + true, + "payload with a same-line SELECT", + ); + assert_eq!(statements[1].to_string(), "SELECT 'after'"); +} + +#[test] +fn parse_copy_from_stdin_blank_trailing_line() { + let stmt = pg_and_generic().verified_stmt("COPY t FROM STDIN;\n\n"); + assert_copy_stdin_payload(&stmt, "\n", false, "blank trailing line"); +} + +#[test] +fn parse_copy_from_stdin_unterminated_payload() { + let stmt = pg_and_generic().verified_stmt("COPY t FROM STDIN;\nrow1"); + assert_copy_stdin_payload(&stmt, "row1", false, "unterminated payload"); +} + +#[test] +fn parse_copy_from_stdin_empty_terminated_payload() { + let stmt = pg_and_generic().verified_stmt("COPY t FROM STDIN;\n\\."); + assert_copy_stdin_payload(&stmt, "", true, "empty terminated payload"); +} + +#[test] +fn parse_copy_from_stdin_bare_semicolon() { + let stmt = pg_and_generic().one_statement_parses_to("COPY t FROM STDIN;", "COPY t FROM STDIN"); + match stmt { + Statement::Copy { + to: false, + target: CopyTarget::Stdin, + payload: None, + .. + } => {} + other => panic!("expected a bare COPY t FROM STDIN, got {other:?}"), + } +} + +#[test] +fn parse_copy_from_stdin_payload_with_columns_and_options() { + let sql = "COPY t (a, b) FROM STDIN (FORMAT CSV);\n1\t2\n\\."; + let stmt = pg_and_generic().verified_stmt(sql); + match stmt { + Statement::Copy { + source: + CopySource::Table { + table_name, + columns, + }, + to: false, + target: CopyTarget::Stdin, + options, + legacy_options, + payload: Some(payload), + } => { + assert_eq!(table_name, ObjectName::from(vec!["t".into()])); + assert_eq!(columns, vec!["a".into(), "b".into()]); + assert_eq!(options, vec![CopyOption::Format("CSV".into())]); + assert!(legacy_options.is_empty()); + assert_eq!(payload.data, "1\t2\n"); + assert!(payload.terminated); + } + other => panic!("expected COPY t (a, b) FROM STDIN, got {other:?}"), + } +} + +#[test] +fn parse_copy_from_stdin_payload_dialect_gating() { + // Dialects without `supports_copy_from_stdin_data` never capture the + // payload, so `COPY ... FROM STDIN` parses without inline data. + let dialects = TestedDialects::new(vec![ + Box::new(MySqlDialect {}), + Box::new(SnowflakeDialect {}), + ]); + for dialect in &dialects.dialects { + let tokens = Tokenizer::new(dialect.as_ref(), "COPY t FROM STDIN;\nrow1\n\\.\n") + .tokenize() + .unwrap(); + assert!( + tokens + .iter() + .all(|token| !matches!(token, Token::CopyData { .. })), + "expected no CopyData token for {dialect:?}" + ); + } + dialects.one_statement_parses_to("COPY t FROM STDIN;", "COPY t FROM STDIN"); +} + +#[test] +fn parse_copy_from_stdin_payload_comment_on_semicolon_line() { + // A single-line comment ends the `;` line, so the payload starts on the + // next line instead of being tokenized as SQL. + let stmt = pg().one_statement_parses_to( + "COPY t FROM STDIN; -- note\nO'Brien\n\\.\n", + "COPY t FROM STDIN;\nO'Brien\n\\.", + ); + assert_copy_stdin_payload(&stmt, "O'Brien\n", true, "comment on the semicolon line"); +} + +#[test] +fn parse_copy_from_stdin_payload_cr_only_line_ends() { + // The line end following the `;` prints as a single line feed. + let sql = "COPY t FROM STDIN;\rrow1\r\\."; + let canonical = "COPY t FROM STDIN;\nrow1\r\\."; + let stmt = pg_and_generic().one_statement_parses_to(sql, canonical); + assert_copy_stdin_payload(&stmt, "row1\r", true, "CR-only line ends"); +} + +#[test] +fn display_copy_payload_propagates_formatter_errors() { + use core::fmt::{self, Write}; + + struct RejectingWriter(char); + + impl Write for RejectingWriter { + fn write_str(&mut self, value: &str) -> fmt::Result { + if value.contains(self.0) { + Err(fmt::Error) + } else { + Ok(()) + } + } + } + + // `q` fails the data write and `\` fails the terminator write. + let stmt = pg_and_generic().verified_stmt("COPY t FROM STDIN;\nq1\n\\."); + for rejected in ['q', '\\'] { + let mut writer = RejectingWriter(rejected); + assert_eq!(write!(writer, "{stmt}"), Err(fmt::Error)); + } +}