uqa_sql/copy/
codec.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `PostgreSQL` COPY text and CSV stream codecs.
8
9use uqa_core::Value;
10
11use super::{CopyFormat, CopyHeader, CopyInputField, CopyOptions};
12use crate::{ColumnType, SQLError, SQLResult};
13
14/// Decode a complete text or CSV COPY stream into rows.
15pub fn decode_copy_input(
16    bytes: &[u8],
17    options: &CopyOptions,
18    expected_columns: &[String],
19) -> Result<Vec<Vec<CopyInputField>>, SQLError> {
20    if options.format == CopyFormat::Binary {
21        return Err(SQLError::Unsupported(
22            "binary COPY format is not implemented".into(),
23        ));
24    }
25    let input = copy_utf8(bytes)?;
26    let mut rows = match options.format {
27        CopyFormat::Text => decode_text_rows(input, options)?,
28        CopyFormat::Csv => decode_csv_rows(input, options)?,
29        CopyFormat::Binary => unreachable!(),
30    };
31    if options.header != CopyHeader::False {
32        if rows.is_empty() {
33            return Ok(Vec::new());
34        }
35        let header = rows.remove(0);
36        if options.header == CopyHeader::Match {
37            let actual = header
38                .into_iter()
39                .map(|field| field.unwrap_or_default())
40                .collect::<Vec<_>>();
41            if actual != expected_columns {
42                return Err(copy_data_error(
43                    "22P04",
44                    format!(
45                        "wrong number of fields in header line: expected {expected_columns:?}, got {actual:?}"
46                    ),
47                    1,
48                ));
49            }
50        }
51    }
52    let first_data_line = usize::from(options.header != CopyHeader::False) + 1;
53    for (index, row) in rows.iter().enumerate() {
54        if row.len() > expected_columns.len() {
55            return Err(copy_data_error(
56                "22P04",
57                "extra data after last expected column",
58                index + first_data_line,
59            ));
60        }
61        if row.len() < expected_columns.len() {
62            return Err(copy_data_error(
63                "22P04",
64                format!(
65                    "missing data for column \"{}\"",
66                    expected_columns[row.len()]
67                ),
68                index + first_data_line,
69            ));
70        }
71    }
72    Ok(rows)
73}
74
75fn decode_text_rows(
76    input: &str,
77    options: &CopyOptions,
78) -> Result<Vec<Vec<CopyInputField>>, SQLError> {
79    let mut rows = Vec::new();
80    for (line_index, raw) in logical_text_lines(input).into_iter().enumerate() {
81        if raw == "\\." {
82            break;
83        }
84        rows.push(decode_text_row(raw, options, line_index + 1)?);
85    }
86    Ok(rows)
87}
88
89fn decode_text_row(
90    raw: &str,
91    options: &CopyOptions,
92    line: usize,
93) -> Result<Vec<CopyInputField>, SQLError> {
94    let bytes = raw.as_bytes();
95    let mut fields = Vec::new();
96    let mut field_start = 0usize;
97    let mut index = 0usize;
98    while index < bytes.len() {
99        if bytes[index] == options.delimiter {
100            fields.push(decode_text_input_field(
101                &raw[field_start..index],
102                options,
103                line,
104            )?);
105            index += 1;
106            field_start = index;
107            continue;
108        }
109        if bytes[index] != b'\\' {
110            index += 1;
111            continue;
112        }
113        index += 1;
114        if index == bytes.len() {
115            return Err(copy_data_error(
116                "22P04",
117                "unterminated COPY data escape",
118                line,
119            ));
120        }
121        match bytes[index] {
122            b'x' => {
123                index += 1;
124                for _ in 0..2 {
125                    if index >= bytes.len() || hex_value(bytes[index]).is_none() {
126                        break;
127                    }
128                    index += 1;
129                }
130            }
131            b'0'..=b'7' => {
132                index += 1;
133                for _ in 0..2 {
134                    if index >= bytes.len() || !(b'0'..=b'7').contains(&bytes[index]) {
135                        break;
136                    }
137                    index += 1;
138                }
139            }
140            _ => index += 1,
141        }
142    }
143    fields.push(decode_text_input_field(&raw[field_start..], options, line)?);
144    Ok(fields)
145}
146
147fn decode_text_input_field(
148    raw: &str,
149    options: &CopyOptions,
150    line: usize,
151) -> Result<CopyInputField, SQLError> {
152    if raw == options.null {
153        Ok(None)
154    } else {
155        decode_text_field(raw, line).map(Some)
156    }
157}
158
159fn logical_text_lines(input: &str) -> Vec<&str> {
160    if input.is_empty() {
161        return Vec::new();
162    }
163    let mut lines = input.split('\n').collect::<Vec<_>>();
164    if input.ends_with('\n') {
165        lines.pop();
166    }
167    for line in &mut lines {
168        if let Some(without_cr) = line.strip_suffix('\r') {
169            *line = without_cr;
170        }
171    }
172    lines
173}
174
175fn decode_text_field(raw: &str, line: usize) -> Result<String, SQLError> {
176    let bytes = raw.as_bytes();
177    let mut output = Vec::with_capacity(bytes.len());
178    let mut index = 0;
179    while index < bytes.len() {
180        if bytes[index] != b'\\' {
181            output.push(bytes[index]);
182            index += 1;
183            continue;
184        }
185        index += 1;
186        if index == bytes.len() {
187            return Err(copy_data_error(
188                "22P04",
189                "unterminated COPY data escape",
190                line,
191            ));
192        }
193        match bytes[index] {
194            b'b' => output.push(0x08),
195            b'f' => output.push(0x0c),
196            b'n' => output.push(b'\n'),
197            b'r' => output.push(b'\r'),
198            b't' => output.push(b'\t'),
199            b'v' => output.push(0x0b),
200            b'x' => {
201                let mut value = 0u8;
202                let mut digits = 0usize;
203                while digits < 2 && index + 1 < bytes.len() {
204                    let Some(nibble) = hex_value(bytes[index + 1]) else {
205                        break;
206                    };
207                    value = value * 16 + nibble;
208                    index += 1;
209                    digits += 1;
210                }
211                if digits == 0 {
212                    output.push(b'x');
213                } else {
214                    output.push(value);
215                }
216            }
217            digit @ b'0'..=b'7' => {
218                let mut value = digit - b'0';
219                let mut digits = 1usize;
220                while digits < 3 && index + 1 < bytes.len() {
221                    let next = bytes[index + 1];
222                    if !(b'0'..=b'7').contains(&next) {
223                        break;
224                    }
225                    value = value.wrapping_mul(8).wrapping_add(next - b'0');
226                    index += 1;
227                    digits += 1;
228                }
229                output.push(value);
230            }
231            escaped => output.push(escaped),
232        }
233        index += 1;
234    }
235    if output.contains(&0) {
236        return Err(copy_encoding_error(Some(line)));
237    }
238    String::from_utf8(output).map_err(|_| copy_encoding_error(Some(line)))
239}
240
241fn hex_value(byte: u8) -> Option<u8> {
242    match byte {
243        b'0'..=b'9' => Some(byte - b'0'),
244        b'a'..=b'f' => Some(byte - b'a' + 10),
245        b'A'..=b'F' => Some(byte - b'A' + 10),
246        _ => None,
247    }
248}
249
250fn decode_csv_rows(
251    input: &str,
252    options: &CopyOptions,
253) -> Result<Vec<Vec<CopyInputField>>, SQLError> {
254    let bytes = input.as_bytes();
255    let mut rows = Vec::new();
256    let mut row = Vec::new();
257    let mut field = Vec::new();
258    let mut quoted = false;
259    let mut in_quotes = false;
260    let mut after_quote = false;
261    let mut index = 0usize;
262    let mut line = 1usize;
263    while index < bytes.len() {
264        let byte = bytes[index];
265        if in_quotes {
266            if byte == options.escape {
267                if index + 1 < bytes.len()
268                    && matches!(bytes[index + 1], next if next == options.quote || next == options.escape)
269                {
270                    field.push(bytes[index + 1]);
271                    index += 2;
272                    continue;
273                }
274                if options.escape != options.quote {
275                    field.push(byte);
276                    index += 1;
277                    continue;
278                }
279            }
280            if byte == options.quote {
281                in_quotes = false;
282                after_quote = true;
283                index += 1;
284                continue;
285            }
286            if byte == b'\n' {
287                line += 1;
288            }
289            field.push(byte);
290            index += 1;
291            continue;
292        }
293        if after_quote && !matches!(byte, b'\n' | b'\r') && byte != options.delimiter {
294            field.push(byte);
295            index += 1;
296            continue;
297        }
298        if byte == options.delimiter {
299            row.push(csv_field(&field, quoted, options)?);
300            field.clear();
301            quoted = false;
302            after_quote = false;
303            index += 1;
304            continue;
305        }
306        if byte == b'\n' || byte == b'\r' {
307            if byte == b'\r' && index + 1 < bytes.len() && bytes[index + 1] == b'\n' {
308                index += 1;
309            }
310            row.push(csv_field(&field, quoted, options)?);
311            if row.len() == 1 && row[0].as_deref() == Some("\\.") && !quoted {
312                return Ok(rows);
313            }
314            rows.push(std::mem::take(&mut row));
315            field.clear();
316            quoted = false;
317            after_quote = false;
318            index += 1;
319            line += 1;
320            continue;
321        }
322        if byte == options.quote && field.is_empty() && !after_quote {
323            quoted = true;
324            in_quotes = true;
325            index += 1;
326            continue;
327        }
328        field.push(byte);
329        index += 1;
330    }
331    if in_quotes {
332        return Err(copy_data_error(
333            "22P04",
334            "unterminated CSV quoted field",
335            line,
336        ));
337    }
338    if !field.is_empty() || !row.is_empty() || quoted || after_quote {
339        row.push(csv_field(&field, quoted, options)?);
340        if !(row.len() == 1 && row[0].as_deref() == Some("\\.") && !quoted) {
341            rows.push(row);
342        }
343    }
344    Ok(rows)
345}
346
347fn csv_field(
348    field: &[u8],
349    quoted: bool,
350    options: &CopyOptions,
351) -> Result<CopyInputField, SQLError> {
352    let text = copy_utf8(field)?;
353    if !quoted && text == options.null {
354        Ok(None)
355    } else {
356        Ok(Some(text.to_string()))
357    }
358}
359
360fn copy_data_error(sqlstate: &str, message: impl Into<String>, line: usize) -> SQLError {
361    SQLError::Routine {
362        sqlstate: sqlstate.into(),
363        message: format!("{}\nCONTEXT: COPY data, line {line}", message.into()),
364    }
365}
366
367fn copy_utf8(bytes: &[u8]) -> Result<&str, SQLError> {
368    if bytes.contains(&0) {
369        return Err(copy_encoding_error(None));
370    }
371    std::str::from_utf8(bytes).map_err(|_| copy_encoding_error(None))
372}
373
374fn copy_encoding_error(line: Option<usize>) -> SQLError {
375    let context = line.map_or_else(String::new, |line| format!(" at line {line}"));
376    SQLError::Routine {
377        sqlstate: "22021".into(),
378        message: format!("invalid byte sequence for encoding \"UTF8\" in COPY data{context}"),
379    }
380}
381
382/// Encode one SQL result using `PostgreSQL` COPY text or CSV representation.
383pub fn encode_copy_result(result: &SQLResult, options: &CopyOptions) -> Result<Vec<u8>, SQLError> {
384    encode_copy_result_impl(result, options, None)
385}
386
387/// Encode one SQL result while resolving catalog-backed `reg*` text output through the engine hook.
388pub fn encode_copy_result_with_engine(
389    result: &SQLResult,
390    options: &CopyOptions,
391    engine: &dyn crate::expr::EngineHook,
392) -> Result<Vec<u8>, SQLError> {
393    encode_copy_result_impl(result, options, Some(engine))
394}
395
396fn encode_copy_result_impl(
397    result: &SQLResult,
398    options: &CopyOptions,
399    engine: Option<&dyn crate::expr::EngineHook>,
400) -> Result<Vec<u8>, SQLError> {
401    if options.format == CopyFormat::Binary {
402        return Err(SQLError::Unsupported(
403            "binary COPY format is not implemented".into(),
404        ));
405    }
406    let mut output = Vec::new();
407    if options.header != CopyHeader::False {
408        encode_copy_row(
409            result.columns.iter().map(|name| Some(name.as_str())),
410            options,
411            &mut output,
412        );
413    }
414    for row_index in 0..result.rows.len() {
415        let values = result
416            .columns
417            .iter()
418            .enumerate()
419            .map(|(column, _)| {
420                result
421                    .value_at(row_index, column)
422                    .map_or(Ok(None), |value| {
423                        copy_value_text(
424                            value,
425                            result.column_types.get(column).and_then(Option::as_ref),
426                            engine,
427                        )
428                    })
429            })
430            .collect::<Result<Vec<_>, SQLError>>()?;
431        match options.format {
432            CopyFormat::Text => encode_text_value_row(values, options, &mut output),
433            CopyFormat::Csv => encode_csv_value_row(values, options, &mut output),
434            CopyFormat::Binary => unreachable!(),
435        }
436    }
437    Ok(output)
438}
439
440fn copy_value_text(
441    value: &Value,
442    ty: Option<&ColumnType>,
443    engine: Option<&dyn crate::expr::EngineHook>,
444) -> Result<Option<String>, SQLError> {
445    if matches!(value, Value::Null) {
446        return Ok(None);
447    }
448    if let Some(text) = ty
449        .map(|ty| crate::expr::format_regtype_value(value, ty, engine))
450        .transpose()?
451        .flatten()
452    {
453        return Ok(Some(text));
454    }
455    if matches!(ty, Some(ColumnType::Int2Vector | ColumnType::OidVector)) {
456        return Ok(crate::expr::vector_value_to_string(value));
457    }
458    Ok(Some(match value {
459        Value::Bool(true) => "t".into(),
460        Value::Bool(false) => "f".into(),
461        Value::Float(value) if value.is_nan() => "NaN".into(),
462        Value::Float(value) if *value == f64::INFINITY => "Infinity".into(),
463        Value::Float(value) if *value == f64::NEG_INFINITY => "-Infinity".into(),
464        Value::FixedChar(value) => value.clone(),
465        other => crate::expr::value_to_string(other),
466    }))
467}
468
469fn encode_text_value_row(
470    fields: impl IntoIterator<Item = Option<String>>,
471    options: &CopyOptions,
472    output: &mut Vec<u8>,
473) {
474    for (index, field) in fields.into_iter().enumerate() {
475        if index != 0 {
476            output.push(options.delimiter);
477        }
478        match field {
479            None => output.extend_from_slice(options.null.as_bytes()),
480            Some(field) => encode_text_field(&field, options.delimiter, output),
481        }
482    }
483    output.push(b'\n');
484}
485
486fn encode_text_field(field: &str, delimiter: u8, output: &mut Vec<u8>) {
487    for byte in field.bytes() {
488        match byte {
489            0x08 => output.extend_from_slice(b"\\b"),
490            0x0c => output.extend_from_slice(b"\\f"),
491            b'\n' => output.extend_from_slice(b"\\n"),
492            b'\r' => output.extend_from_slice(b"\\r"),
493            b'\t' => output.extend_from_slice(b"\\t"),
494            0x0b => output.extend_from_slice(b"\\v"),
495            b'\\' => output.extend_from_slice(b"\\\\"),
496            escaped if escaped == delimiter => {
497                output.push(b'\\');
498                output.push(escaped);
499            }
500            other => output.push(other),
501        }
502    }
503}
504
505fn encode_csv_value_row(
506    fields: impl IntoIterator<Item = Option<String>>,
507    options: &CopyOptions,
508    output: &mut Vec<u8>,
509) {
510    for (index, field) in fields.into_iter().enumerate() {
511        if index != 0 {
512            output.push(options.delimiter);
513        }
514        match field {
515            None => output.extend_from_slice(options.null.as_bytes()),
516            Some(field) => encode_csv_field(&field, options, output),
517        }
518    }
519    output.push(b'\n');
520}
521
522fn encode_copy_row<'a>(
523    fields: impl IntoIterator<Item = Option<&'a str>>,
524    options: &CopyOptions,
525    output: &mut Vec<u8>,
526) {
527    match options.format {
528        CopyFormat::Text => encode_text_value_row(
529            fields.into_iter().map(|field| field.map(str::to_string)),
530            options,
531            output,
532        ),
533        CopyFormat::Csv => encode_csv_value_row(
534            fields.into_iter().map(|field| field.map(str::to_string)),
535            options,
536            output,
537        ),
538        CopyFormat::Binary => {}
539    }
540}
541
542fn encode_csv_field(field: &str, options: &CopyOptions, output: &mut Vec<u8>) {
543    let needs_quotes = field.is_empty() && options.null.is_empty()
544        || field == options.null
545        || field.bytes().any(|byte| {
546            matches!(byte, b'\n' | b'\r') || byte == options.delimiter || byte == options.quote
547        });
548    if !needs_quotes {
549        output.extend_from_slice(field.as_bytes());
550        return;
551    }
552    output.push(options.quote);
553    for byte in field.bytes() {
554        if byte == options.quote || byte == options.escape {
555            output.push(options.escape);
556        }
557        output.push(byte);
558    }
559    output.push(options.quote);
560}
uqa_sql/copy/codec.rs

uqa_sql/copy/
codec.rs
