1use super::{CopyDirection, CopyEndpoint, CopyFormat, CopyStatement};
9use crate::{assignment::columns::ColumnCatalogError, ast::ColumnDef, SQLError};
10use std::collections::BTreeSet;
11pub 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 ©.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}