Skip to main content

uqa_sql/copy/
stream.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! COPY stream envelope validation and relation-column binding.
8use super::{CopyDirection, CopyEndpoint, CopyFormat, CopyStatement};
9use crate::{assignment::columns::ColumnCatalogError, ast::ColumnDef, SQLError};
10use std::collections::BTreeSet;
11/// Metadata from the relation generation selected before COPY column validation.
12pub trait CopyRelation {
13    fn columns(&self) -> Vec<ColumnDef>;
14    fn is_partitioned(&self) -> bool;
15}
16pub trait CopyCatalog {
17    fn resolve_relation_kind(&self, name: &str)
18        -> Result<Option<(String, &'static str)>, SQLError>;
19    fn table(&self, name: &str) -> Result<Option<Box<dyn CopyRelation + '_>>, ColumnCatalogError>;
20}
21pub fn validate_stream(copy: &CopyStatement, direction: CopyDirection) -> Result<(), SQLError> {
22    ensure_copy_direction(copy, direction)?;
23    ensure_stdio_endpoint(copy)?;
24    ensure_stream_format(copy)
25}
26pub fn relation_columns(
27    catalog: &dyn CopyCatalog,
28    relation: &str,
29    display_name: &str,
30    requested: &[String],
31    reject_partitioned_output: bool,
32) -> Result<(String, Vec<String>), SQLError> {
33    let canonical = match catalog.resolve_relation_kind(relation)? {
34        Some((canonical, "table")) => canonical,
35        Some(_) | None => return Err(SQLError::UnknownTable(relation.to_string())),
36    };
37    let table = catalog
38        .table(&canonical)
39        .map_err(|error| SQLError::Internal(format!("read COPY relation `{canonical}`: {error}")))?
40        .ok_or_else(|| SQLError::UnknownTable(relation.to_string()))?;
41    let definitions = table.columns();
42    let columns = if requested.is_empty() {
43        definitions
44            .into_iter()
45            .filter(|column| column.generated.is_none())
46            .map(|column| column.name)
47            .collect()
48    } else {
49        let mut seen = BTreeSet::new();
50        let mut columns = Vec::with_capacity(requested.len());
51        for requested in requested {
52            if !seen.insert(requested.clone()) {
53                return Err(SQLError::Routine {
54                    sqlstate: "42701".into(),
55                    message: format!("column \"{requested}\" specified more than once"),
56                });
57            }
58            let Some(column) = definitions.iter().find(|column| column.name == *requested) else {
59                return Err(SQLError::Routine {
60                    sqlstate: "42703".into(),
61                    message: format!(
62                        "column \"{requested}\" of relation \"{display_name}\" does not exist"
63                    ),
64                });
65            };
66            if column.generated.is_some() {
67                return Err(SQLError::Routine {
68                        sqlstate: "42P10".into(),
69                        message: format!(
70                            "column \"{requested}\" is a generated column\nDETAIL: Generated columns cannot be used in COPY."
71                        ),
72                    });
73            }
74            columns.push(requested.clone());
75        }
76        columns
77    };
78    if reject_partitioned_output && table.is_partitioned() {
79        return Err(SQLError::Routine {
80                sqlstate: "42809".into(),
81                message: format!(
82                    "cannot copy from partitioned table \"{display_name}\"\nHINT: Try the COPY (SELECT ...) TO variant."
83                ),
84            });
85    }
86    Ok((canonical, columns))
87}
88fn ensure_copy_direction(copy: &CopyStatement, expected: CopyDirection) -> Result<(), SQLError> {
89    if copy.direction == expected {
90        return Ok(());
91    }
92    let expected = match expected {
93        CopyDirection::From => "FROM STDIN",
94        CopyDirection::To => "TO STDOUT",
95    };
96    Err(SQLError::Routine {
97        sqlstate: "42601".into(),
98        message: format!("COPY stream API requires COPY {expected}"),
99    })
100}
101
102fn ensure_stdio_endpoint(copy: &CopyStatement) -> Result<(), SQLError> {
103    match &copy.endpoint {
104        CopyEndpoint::Stdio => Ok(()),
105        CopyEndpoint::File(_) => Err(SQLError::Unsupported(
106            "server-side COPY files are not available through the embedded stream API".into(),
107        )),
108        CopyEndpoint::Program(_) => Err(SQLError::Unsupported(
109            "COPY PROGRAM is not available through the embedded stream API".into(),
110        )),
111    }
112}
113
114fn ensure_stream_format(copy: &CopyStatement) -> Result<(), SQLError> {
115    if copy.options.format == CopyFormat::Binary {
116        Err(SQLError::Unsupported(
117            "binary COPY format is not implemented".into(),
118        ))
119    } else {
120        Ok(())
121    }
122}