uqa_sql/
copy.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `PostgreSQL` `COPY` statement envelopes and stream codecs.
8//!
9//! `COPY FROM STDIN` and `COPY TO STDOUT` carry their byte stream outside the
10//! SQL statement. The regular [`crate::compile`] entry point therefore cannot
11//! execute one by itself; callers use [`compile_copy`] to validate the SQL
12//! envelope and pair it with the stream through the engine COPY API.
13
14use std::collections::BTreeSet;
15
16use pg_query::NodeEnum;
17
18use crate::SQLError;
19
20mod codec;
21
22pub use codec::{decode_copy_input, encode_copy_result, encode_copy_result_with_engine};
23
24/// COPY stream direction.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum CopyDirection {
27    From,
28    To,
29}
30
31/// COPY data encoding implemented by the embedded stream API.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum CopyFormat {
34    Text,
35    Csv,
36    Binary,
37}
38
39/// `PostgreSQL`'s three `HEADER` states.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum CopyHeader {
42    False,
43    True,
44    Match,
45}
46
47/// Parsed COPY options after format-dependent defaults have been applied.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct CopyOptions {
50    pub format: CopyFormat,
51    pub delimiter: u8,
52    pub null: String,
53    pub header: CopyHeader,
54    pub quote: u8,
55    pub escape: u8,
56}
57
58/// The relation or query supplying COPY rows.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum CopyTarget {
61    Relation {
62        name: String,
63        qualifier: String,
64        columns: Vec<String>,
65    },
66    Query(String),
67}
68
69/// Where `PostgreSQL` expects the COPY stream to be connected.
70#[derive(Debug, Clone, PartialEq, Eq)]
71pub enum CopyEndpoint {
72    Stdio,
73    File(String),
74    Program(String),
75}
76
77/// Fully validated COPY statement envelope.
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct CopyStatement {
80    pub direction: CopyDirection,
81    pub target: CopyTarget,
82    pub endpoint: CopyEndpoint,
83    pub options: CopyOptions,
84}
85
86/// One decoded COPY field. `None` is SQL NULL; strings remain in `PostgreSQL`'s
87/// textual input representation so the target column's normal input coercion
88/// remains authoritative.
89pub type CopyInputField = Option<String>;
90
91/// Parse exactly one `PostgreSQL` COPY statement.
92pub fn compile_copy(sql: &str) -> Result<CopyStatement, SQLError> {
93    let parsed = pg_query::parse(sql)?;
94    if parsed.protobuf.stmts.len() != 1 {
95        return Err(SQLError::Routine {
96            sqlstate: "42601".into(),
97            message: "COPY stream API accepts exactly one COPY statement".into(),
98        });
99    }
100    let node = parsed.protobuf.stmts[0]
101        .stmt
102        .as_deref()
103        .and_then(|node| node.node.as_ref())
104        .ok_or_else(|| SQLError::Internal("parser returned an empty COPY statement".into()))?;
105    let NodeEnum::CopyStmt(statement) = node else {
106        return Err(SQLError::Routine {
107            sqlstate: "42601".into(),
108            message: "COPY stream API requires a COPY statement".into(),
109        });
110    };
111    let direction = if statement.is_from {
112        CopyDirection::From
113    } else {
114        CopyDirection::To
115    };
116    if statement.where_clause.is_some() {
117        return Err(SQLError::Unsupported(
118            "COPY FROM WHERE is not implemented".into(),
119        ));
120    }
121    let target = match (&statement.relation, &statement.query) {
122        (Some(relation), None) => {
123            if !relation.catalogname.is_empty() {
124                return Err(SQLError::Unsupported(
125                    "COPY across databases is not supported".into(),
126                ));
127            }
128            let name = if relation.schemaname.is_empty() {
129                quote_relation_component(&relation.relname)
130            } else {
131                format!(
132                    "{}.{}",
133                    quote_relation_component(&relation.schemaname),
134                    quote_relation_component(&relation.relname)
135                )
136            };
137            let columns = statement
138                .attlist
139                .iter()
140                .map(copy_column_name)
141                .collect::<Result<Vec<_>, _>>()?;
142            CopyTarget::Relation {
143                name,
144                qualifier: relation.relname.clone(),
145                columns,
146            }
147        }
148        (None, Some(query)) if direction == CopyDirection::To => CopyTarget::Query(
149            query
150                .deparse()
151                .map_err(|error| SQLError::Parse(error.to_string()))?,
152        ),
153        (None, Some(_)) => {
154            return Err(SQLError::Routine {
155                sqlstate: "42601".into(),
156                message: "COPY FROM does not accept a query".into(),
157            });
158        }
159        _ => {
160            return Err(SQLError::Internal(
161                "COPY statement has neither one relation nor one query".into(),
162            ));
163        }
164    };
165    let endpoint = if statement.is_program {
166        CopyEndpoint::Program(statement.filename.clone())
167    } else if statement.filename.is_empty() {
168        CopyEndpoint::Stdio
169    } else {
170        CopyEndpoint::File(statement.filename.clone())
171    };
172    let options = compile_copy_options(&statement.options)?;
173    if direction == CopyDirection::To && options.header == CopyHeader::Match {
174        return Err(copy_option_error(
175            "COPY HEADER MATCH is only valid for COPY FROM",
176        ));
177    }
178    Ok(CopyStatement {
179        direction,
180        target,
181        endpoint,
182        options,
183    })
184}
185
186fn quote_relation_component(component: &str) -> String {
187    crate::expr::quote_ident(component)
188}
189
190fn copy_column_name(node: &pg_query::Node) -> Result<String, SQLError> {
191    match node.node.as_ref() {
192        Some(NodeEnum::String(value)) => Ok(value.sval.clone()),
193        other => Err(SQLError::Internal(format!(
194            "COPY column list contains malformed node {other:?}"
195        ))),
196    }
197}
198
199fn compile_copy_options(nodes: &[pg_query::Node]) -> Result<CopyOptions, SQLError> {
200    let raw = nodes
201        .iter()
202        .map(|node| {
203            let Some(NodeEnum::DefElem(option)) = node.node.as_ref() else {
204                return Err(SQLError::Internal(
205                    "COPY option list contains a malformed node".into(),
206                ));
207            };
208            Ok((
209                option.defname.to_ascii_lowercase(),
210                copy_option_value(option)?,
211            ))
212        })
213        .collect::<Result<Vec<_>, SQLError>>()?;
214    let mut names = BTreeSet::new();
215    for (name, _) in &raw {
216        if !names.insert(name.clone()) {
217            return Err(SQLError::Routine {
218                sqlstate: "42601".into(),
219                message: format!("COPY option \"{name}\" specified more than once"),
220            });
221        }
222    }
223    let format = raw.iter().find(|(name, _)| name == "format").map_or(
224        Ok(CopyFormat::Text),
225        |(_, value)| match value.to_ascii_lowercase().as_str() {
226            "text" => Ok(CopyFormat::Text),
227            "csv" => Ok(CopyFormat::Csv),
228            "binary" => Ok(CopyFormat::Binary),
229            other => Err(copy_option_error(format!(
230                "COPY format \"{other}\" not recognized"
231            ))),
232        },
233    )?;
234    let mut options = CopyOptions {
235        format,
236        delimiter: if format == CopyFormat::Csv {
237            b','
238        } else {
239            b'\t'
240        },
241        null: if format == CopyFormat::Csv {
242            String::new()
243        } else {
244            "\\N".into()
245        },
246        header: CopyHeader::False,
247        quote: b'"',
248        escape: b'"',
249    };
250    let quote_explicit = names.contains("quote");
251    let escape_explicit = names.contains("escape");
252    for (name, value) in raw {
253        match name.as_str() {
254            "format" => {}
255            "delimiter" => options.delimiter = copy_single_byte_option("DELIMITER", &value)?,
256            "null" => options.null = value,
257            "header" => {
258                options.header = if value.eq_ignore_ascii_case("match") {
259                    CopyHeader::Match
260                } else if copy_boolean_option("HEADER", &value)? {
261                    CopyHeader::True
262                } else {
263                    CopyHeader::False
264                };
265            }
266            "quote" => options.quote = copy_single_byte_option("QUOTE", &value)?,
267            "escape" => options.escape = copy_single_byte_option("ESCAPE", &value)?,
268            "encoding" if matches!(value.to_ascii_lowercase().as_str(), "utf8" | "utf-8") => {}
269            "encoding" => {
270                return Err(SQLError::Unsupported(format!(
271                    "COPY ENCODING {value} is not implemented"
272                )));
273            }
274            other => {
275                return Err(SQLError::Unsupported(format!(
276                    "COPY option `{other}` is not implemented"
277                )));
278            }
279        }
280    }
281    validate_copy_options(&options, quote_explicit, escape_explicit)?;
282    Ok(options)
283}
284
285fn copy_option_value(option: &pg_query::protobuf::DefElem) -> Result<String, SQLError> {
286    match option
287        .arg
288        .as_deref()
289        .and_then(|argument| argument.node.as_ref())
290    {
291        None => Ok("true".into()),
292        Some(NodeEnum::String(value)) => Ok(value.sval.clone()),
293        Some(NodeEnum::Integer(value)) => Ok(value.ival.to_string()),
294        Some(NodeEnum::Float(value)) => Ok(value.fval.clone()),
295        Some(NodeEnum::Boolean(value)) => Ok(value.boolval.to_string()),
296        other => Err(copy_option_error(format!(
297            "COPY option \"{}\" requires a scalar value, got {other:?}",
298            option.defname
299        ))),
300    }
301}
302
303fn copy_boolean_option(name: &str, value: &str) -> Result<bool, SQLError> {
304    match value.to_ascii_lowercase().as_str() {
305        "true" | "on" | "yes" | "1" => Ok(true),
306        "false" | "off" | "no" | "0" => Ok(false),
307        _ => Err(copy_option_error(format!(
308            "COPY {name} requires a Boolean value"
309        ))),
310    }
311}
312
313fn copy_single_byte_option(name: &str, value: &str) -> Result<u8, SQLError> {
314    if value.len() != 1 {
315        return Err(copy_option_error(format!(
316            "COPY {name} must be a single one-byte character"
317        )));
318    }
319    let byte = value.as_bytes()[0];
320    if matches!(byte, 0 | b'\n' | b'\r') {
321        return Err(copy_option_error(format!(
322            "COPY {name} cannot be newline or carriage return"
323        )));
324    }
325    Ok(byte)
326}
327
328fn validate_copy_options(
329    options: &CopyOptions,
330    quote_explicit: bool,
331    escape_explicit: bool,
332) -> Result<(), SQLError> {
333    if options.null.contains(['\n', '\r']) {
334        return Err(copy_option_error(
335            "COPY null representation cannot use newline or carriage return",
336        ));
337    }
338    if options.format == CopyFormat::Text && options.delimiter == b'\\' {
339        return Err(copy_option_error("COPY delimiter cannot be backslash"));
340    }
341    if options.format != CopyFormat::Csv && (quote_explicit || escape_explicit) {
342        return Err(copy_option_error(
343            "COPY quote or escape available only in CSV mode",
344        ));
345    }
346    if options.format == CopyFormat::Binary && options.header != CopyHeader::False {
347        return Err(copy_option_error(
348            "COPY HEADER is not available in binary mode",
349        ));
350    }
351    if options.format == CopyFormat::Csv && options.delimiter == options.quote {
352        return Err(copy_option_error(
353            "COPY delimiter and quote must be different",
354        ));
355    }
356    if options.null.as_bytes().contains(&options.delimiter) {
357        return Err(copy_option_error(
358            "COPY delimiter character must not appear in the NULL specification",
359        ));
360    }
361    if options.format == CopyFormat::Csv && options.null.as_bytes().contains(&options.quote) {
362        return Err(copy_option_error(
363            "CSV quote character must not appear in the NULL specification",
364        ));
365    }
366    Ok(())
367}
368
369fn copy_option_error(message: impl Into<String>) -> SQLError {
370    SQLError::Routine {
371        sqlstate: "22023".into(),
372        message: message.into(),
373    }
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    #[test]
381    fn copy_parser_preserves_relation_query_and_format_contracts() {
382        let from = compile_copy(
383            "COPY app.items (bucket, payload) FROM STDIN WITH (FORMAT csv, HEADER MATCH)",
384        )
385        .unwrap();
386        assert_eq!(from.direction, CopyDirection::From);
387        assert_eq!(from.endpoint, CopyEndpoint::Stdio);
388        assert_eq!(from.options.format, CopyFormat::Csv);
389        assert_eq!(from.options.header, CopyHeader::Match);
390        assert_eq!(
391            from.target,
392            CopyTarget::Relation {
393                name: "app.items".into(),
394                qualifier: "items".into(),
395                columns: vec!["bucket".into(), "payload".into()]
396            }
397        );
398
399        let to = compile_copy("COPY (SELECT * FROM ONLY items ORDER BY id) TO STDOUT").unwrap();
400        assert_eq!(to.direction, CopyDirection::To);
401        let CopyTarget::Query(query) = to.target else {
402            panic!("expected COPY query target");
403        };
404        assert!(query.contains("ONLY items"));
405
406        let quoted =
407            compile_copy("COPY \"copy.schema\".\"Odd.Table\" (\"Odd Column\") FROM STDIN").unwrap();
408        assert_eq!(
409            quoted.target,
410            CopyTarget::Relation {
411                name: "\"copy.schema\".\"Odd.Table\"".into(),
412                qualifier: "Odd.Table".into(),
413                columns: vec!["Odd Column".into()],
414            }
415        );
416
417        for invalid in [
418            r"COPY items FROM STDIN WITH (FORMAT text, DELIMITER E'\\')",
419            "COPY items FROM STDIN WITH (FORMAT text, QUOTE '\"')",
420            "COPY items FROM STDIN WITH (FORMAT binary, ESCAPE '\"')",
421            "COPY items FROM STDIN WITH (FORMAT binary, HEADER)",
422            "COPY items FROM STDIN WITH (DELIMITER '|', NULL 'a|b')",
423            "COPY items FROM STDIN WITH (FORMAT csv, NULL 'a\"b')",
424        ] {
425            assert_eq!(compile_copy(invalid).unwrap_err().sqlstate(), Some("22023"));
426        }
427    }
428
429    #[test]
430    fn text_codec_distinguishes_null_escapes_and_end_marker() {
431        let options = compile_copy("COPY t FROM STDIN").unwrap().options;
432        let columns = vec!["a".into(), "b".into()];
433        let rows = decode_copy_input(
434            b"\\N\tempty\\tvalue\n\\\\N\tline\\nvalue\n\\.\nignored\tx\n",
435            &options,
436            &columns,
437        )
438        .unwrap();
439        assert_eq!(rows.len(), 2);
440        assert_eq!(rows[0][0], None);
441        assert_eq!(rows[0][1].as_deref(), Some("empty\tvalue"));
442        assert_eq!(rows[1][0].as_deref(), Some("\\N"));
443        assert_eq!(rows[1][1].as_deref(), Some("line\nvalue"));
444
445        let custom = compile_copy("COPY t FROM STDIN WITH (DELIMITER '|')")
446            .unwrap()
447            .options;
448        let escaped_delimiter =
449            decode_copy_input(b"left\\|right|tail\n", &custom, &columns).unwrap();
450        assert_eq!(
451            escaped_delimiter[0],
452            vec![Some("left|right".into()), Some("tail".into())]
453        );
454        assert_eq!(
455            decode_copy_input(b"\\000\tvalue\n", &options, &columns)
456                .unwrap_err()
457                .sqlstate(),
458            Some("22021")
459        );
460
461        let with_header = compile_copy("COPY t FROM STDIN WITH (HEADER MATCH)")
462            .unwrap()
463            .options;
464        assert_eq!(
465            decode_copy_input(b"a\tb\n1\t2\n", &with_header, &columns).unwrap(),
466            vec![vec![Some("1".into()), Some("2".into())]]
467        );
468    }
469
470    #[test]
471    fn csv_codec_preserves_quoted_empty_and_embedded_newline() {
472        let options = compile_copy("COPY t FROM STDIN WITH (FORMAT csv)")
473            .unwrap()
474            .options;
475        let columns = vec!["a".into(), "b".into()];
476        let rows =
477            decode_copy_input(b",\"\"\n\"a,b\",\"line\nvalue\"\n", &options, &columns).unwrap();
478        assert_eq!(rows[0], vec![None, Some(String::new())]);
479        assert_eq!(rows[1][0].as_deref(), Some("a,b"));
480        assert_eq!(rows[1][1].as_deref(), Some("line\nvalue"));
481
482        let permissive = decode_copy_input(b"\"a\"b,c\n\"a\" ,b\n", &options, &columns).unwrap();
483        assert_eq!(permissive[0], vec![Some("ab".into()), Some("c".into())]);
484        assert_eq!(permissive[1], vec![Some("a ".into()), Some("b".into())]);
485    }
486}
uqa_sql/copy.rs

uqa_sql/
copy.rs
