Skip to main content

rupg_analyze/ddl/
copy.rs

1//! The relation of `COPY`, as `DoCopy`, `CopyGetAttnums` and `BeginCopyTo` in `copy.c` and `copyto.c` check it: the name, the columns and the kind of the relation. The session reads the rows with a query of the columns, as `DoCopy` does for a table with row security.
2
3use rupg_common::{Error, Result, SqlState};
4use rupg_pgcatalog::builtin;
5use rupg_sql::nodes::{
6    Alias, ColumnRef, CopyStmt, InsertStmt, Node, OverridingKind, ParamRef, RangeSubselect,
7    RangeVar, ResTarget, SelectStmt,
8};
9
10use super::lock::{lookup, relkind};
11use super::relkind_detail;
12use crate::walk::{self, Walker};
13use crate::{Analyzer, Env, Params};
14
15/// The relation of a `COPY ... TO` and the query that reads its rows.
16#[derive(Debug)]
17pub struct CopyRelation {
18    /// The name of the relation, for the errors.
19    pub name: String,
20    /// The names of the user columns of the relation, in the order of the catalog.
21    pub columns: Vec<String>,
22    /// The columns that `COPY` sends, as indexes into `columns`.
23    pub selected: Vec<usize>,
24    /// `SELECT` of the selected columns `FROM ONLY` the relation. For a partitioned table, the query reads the partitions, and `tableoid` comes after the selected columns.
25    pub query: Node,
26    /// The partitions of a partitioned table that have rows, in the order of `find_all_inheritors`. `DoCopyTo` sends the rows of each partition in this order.
27    pub partitions: Option<Vec<u32>>,
28}
29
30/// The table of a `COPY ... FROM`, as `DoCopy` and `BeginCopyFrom` open it.
31#[derive(Debug, Clone)]
32pub struct CopyTarget {
33    /// The OID of the relation.
34    pub relid: u32,
35    /// The name of the relation, for the errors.
36    pub name: String,
37    /// `relkind`. `CopyFrom` checks it after `CopyInResponse`.
38    pub kind: u8,
39    /// The names of the user columns of the relation, in the order of the catalog.
40    pub columns: Vec<String>,
41    /// The columns that the data gives, as indexes into `columns`.
42    pub selected: Vec<usize>,
43}
44
45/// The checks of `table_openrv` and `CopyGetAttnums` for the relation of `COPY`.
46///
47/// # Errors
48///
49/// The errors of `RangeVarGetRelid` for the name, `42809` for an index or a composite type, `42703` and `42701` for the list of columns.
50pub fn copy_target(stmt: &CopyStmt, rv: &RangeVar, env: &dyn Env) -> Result<CopyTarget> {
51    let params = Params::default();
52    let an = Analyzer::new(env, &params);
53    let relid = lookup(&an, rv)?;
54    let name = rv.relname.as_deref().unwrap_or_default().to_string();
55    let kind = relkind(env, relid);
56    if matches!(kind, b'i' | b'I' | b'c') {
57        return Err(Error::new(
58            SqlState::WRONG_OBJECT_TYPE,
59            format!("cannot open relation \"{name}\""),
60        )
61        .with_detail(relkind_detail(kind)));
62    }
63    let columns = relation_columns(env, relid);
64    let names: Vec<String> =
65        stmt.attlist.iter().flatten().filter_map(string).map(str::to_string).collect();
66    let mut selected = attnums(&columns, &names, Some(&name))?;
67    // `CopyGetAttnums`: the default list has no generated column, and the list of the statement cannot name one.
68    let generated: Vec<bool> =
69        env.catalog().and_then(|c| c.relation(relid)).map_or_else(Vec::new, |rel| {
70            rel.columns.iter().filter(|c| !c.dropped).map(|c| c.generated != 0).collect()
71        });
72    let is_generated = |i: usize| generated.get(i).copied().unwrap_or(false);
73    if let Some(condition) = &stmt.whereClause {
74        let mut refs = WhereColumns { relname: &name, columns: &columns, used: Vec::new() };
75        walk::one(&mut refs, Some(condition));
76        let mut used = refs.used;
77        used.sort_unstable();
78        if let Some(&i) = used.iter().find(|&&i| is_generated(i)) {
79            return Err(Error::new(
80                SqlState::INVALID_COLUMN_REFERENCE,
81                "generated columns are not supported in COPY FROM WHERE conditions",
82            )
83            .with_detail(format!("Column \"{}\" is a generated column.", columns[i])));
84        }
85    }
86    if names.is_empty() {
87        selected.retain(|&i| !is_generated(i));
88    } else if let Some(&i) = selected.iter().find(|&&i| is_generated(i)) {
89        return Err(Error::new(
90            SqlState::INVALID_COLUMN_REFERENCE,
91            format!("column \"{}\" is a generated column", columns[i]),
92        )
93        .with_detail("Generated columns cannot be used in COPY."));
94    }
95    Ok(CopyTarget { relid, name, kind, columns, selected })
96}
97
98/// `pull_varattnos` for the `WHERE` of `COPY ... FROM`, on the raw tree: the columns of the relation that the condition names, as indexes into `columns`. A whole-row reference gives all the columns.
99struct WhereColumns<'a> {
100    relname: &'a str,
101    columns: &'a [String],
102    used: Vec<usize>,
103}
104
105impl Walker for WhereColumns<'_> {
106    fn node(&mut self, node: &Node) -> bool {
107        let Node::ColumnRef(cref) = node else { return walk::walk(self, node) };
108        let fields: Vec<Option<&str>> = cref.fields.iter().flatten().map(string).collect();
109        let position = |name: &str| self.columns.iter().position(|c| c == name);
110        match fields.as_slice() {
111            [Some(name)] if position(name).is_some() => self.used.extend(position(name)),
112            [Some(rel)] | [Some(rel), None] if *rel == self.relname => {
113                self.used.extend(0..self.columns.len());
114            }
115            [Some(rel), Some(name)] if *rel == self.relname => self.used.extend(position(name)),
116            _ => {}
117        }
118        false
119    }
120
121    fn select(&mut self, stmt: &SelectStmt) -> bool {
122        walk::select(self, stmt)
123    }
124}
125
126/// The checks of `DoCopy` and `BeginCopyTo` for the relation of `COPY ... TO`: `table_openrv`, `CopyGetAttnums` and the kind of the relation.
127///
128/// # Errors
129///
130/// The errors of `RangeVarGetRelid` for the name, `42809` for an index or a relation that is not a table, `42703` and `42701` for the list of columns.
131pub fn copy_relation(stmt: &CopyStmt, rv: &RangeVar, env: &dyn Env) -> Result<CopyRelation> {
132    let CopyTarget { relid, name, kind, columns, selected } = copy_target(stmt, rv, env)?;
133    let mut partitions = None;
134    match kind {
135        b'r' => {}
136        b'p' => partitions = Some(copy_partitions(env, relid, &name)?),
137        b'm' => {
138            let rel = env.catalog().and_then(|c| c.relation(relid));
139            if rel.and_then(|r| r.view.as_ref()).is_some_and(|v| !v.populated) {
140                return Err(Error::new(
141                    SqlState::FEATURE_NOT_SUPPORTED,
142                    format!("cannot copy from unpopulated materialized view \"{name}\""),
143                )
144                .with_hint("Use the REFRESH MATERIALIZED VIEW command."));
145            }
146        }
147        b'f' => {
148            return Err(Error::new(
149                SqlState::WRONG_OBJECT_TYPE,
150                format!("cannot copy from foreign table \"{name}\""),
151            )
152            .with_hint("Try the COPY (SELECT ...) TO variant."));
153        }
154        b'v' => {
155            return Err(Error::new(
156                SqlState::WRONG_OBJECT_TYPE,
157                format!("cannot copy from view \"{name}\""),
158            )
159            .with_hint("Try the COPY (SELECT ...) TO variant."));
160        }
161        b'S' => {
162            return Err(Error::new(
163                SqlState::WRONG_OBJECT_TYPE,
164                format!("cannot copy from sequence \"{name}\""),
165            ));
166        }
167        _ => {
168            return Err(Error::new(
169                SqlState::WRONG_OBJECT_TYPE,
170                format!("cannot copy from non-table relation \"{name}\""),
171            ));
172        }
173    }
174    let tableoid = partitions.as_ref().map(|_| "tableoid");
175    let targets = selected
176        .iter()
177        .map(|&i| columns[i].as_str())
178        .chain(tableoid)
179        .map(|column| {
180            let field = Some(Node::String(column.into()));
181            let column = ColumnRef { fields: vec![field], location: -1 };
182            let target = ResTarget { val: Some(column.into()), location: -1, ..Default::default() };
183            Some(target.into())
184        })
185        .collect();
186    let from = RangeVar { inh: partitions.is_some(), alias: None, location: -1, ..rv.clone() };
187    let select = SelectStmt {
188        targetList: targets,
189        fromClause: vec![Some(from.into())],
190        ..Default::default()
191    };
192    Ok(CopyRelation { name, columns, selected, query: select.into(), partitions })
193}
194
195/// `find_all_inheritors` of `BeginCopyTo` for the partitioned table `root`: the partitions of each level after the level before it, each level in the order of the OIDs, with no partitioned tables, because they have no rows.
196///
197/// # Errors
198///
199/// `42809` for a foreign table in the partitions.
200fn copy_partitions(env: &dyn Env, root: u32, name: &str) -> Result<Vec<u32>> {
201    let Some(catalog) = env.catalog() else { return Ok(Vec::new()) };
202    let mut all = vec![root];
203    let mut at = 0;
204    while let Some(&oid) = all.get(at) {
205        at += 1;
206        all.extend(catalog.partitions(oid).iter().map(|rel| rel.oid));
207    }
208    let mut out = Vec::new();
209    for oid in all {
210        match relkind(env, oid) {
211            b'p' => {}
212            b'f' => {
213                let child = catalog.relation(oid).map_or("", |rel| rel.name.as_str());
214                return Err(Error::new(
215                    SqlState::WRONG_OBJECT_TYPE,
216                    format!("cannot copy from foreign table \"{child}\""),
217                )
218                .with_detail(format!(
219                    "Partition \"{child}\" is a foreign table in partitioned table \"{name}\"."
220                ))
221                .with_hint("Try the COPY (SELECT ...) TO variant."));
222            }
223            _ => out.push(oid),
224        }
225    }
226    Ok(out)
227}
228
229/// The `INSERT` that writes the rows of a `COPY ... FROM`, as `CopyFrom` fills a row: the parameter `$n` is the value of the selected column `n`, and the other columns get their defaults. `defaults` is true for a selected column whose field is the marker of `DEFAULT`. Such a column also gets its default, and the parameters skip it. `OVERRIDING SYSTEM VALUE` lets the data give a value to an identity column `GENERATED ALWAYS`, as `COPY` does.
230#[must_use]
231pub fn copy_insert(target: &CopyTarget, rv: &RangeVar, defaults: &[bool]) -> Node {
232    let cols: Vec<_> = target
233        .selected
234        .iter()
235        .enumerate()
236        .filter(|&(k, _)| !defaults.get(k).copied().unwrap_or(false))
237        .map(|(_, &i)| {
238            let name = Some(target.columns[i].as_str().into());
239            Some(ResTarget { name, location: -1, ..Default::default() }.into())
240        })
241        .collect();
242    // With no columns, the statement is `INSERT ... DEFAULT VALUES`.
243    let select = (!cols.is_empty()).then(|| values(cols.len()).into());
244    let insert = InsertStmt {
245        relation: Some(Box::new(RangeVar { alias: None, location: -1, ..rv.clone() })),
246        cols,
247        selectStmt: select,
248        r#override: OverridingKind::OVERRIDING_SYSTEM_VALUE,
249        ..Default::default()
250    };
251    insert.into()
252}
253
254/// The query of the condition `WHERE` of a `COPY ... FROM`. The parameter `$n` is the value of the column `n` of the row, and the query gives the value of the condition, as `ExecQual` gives it for the row of `CopyFrom`.
255#[must_use]
256pub fn copy_filter(target: &CopyTarget, rv: &RangeVar, condition: &Node) -> Node {
257    let colnames =
258        target.columns.iter().map(|name| Some(Node::String(name.as_str().into()))).collect();
259    let alias = Alias { aliasname: rv.relname.clone(), colnames };
260    let from = RangeSubselect {
261        lateral: false,
262        subquery: Some(values(target.columns.len()).into()),
263        alias: Some(Box::new(alias)),
264    };
265    let target = ResTarget { val: Some(condition.clone()), location: -1, ..Default::default() };
266    let select = SelectStmt {
267        targetList: vec![Some(target.into())],
268        fromClause: vec![Some(from.into())],
269        ..Default::default()
270    };
271    select.into()
272}
273
274/// The `INSERT` that writes a row of a table again after `ALTER COLUMN ... TYPE`, as `ATRewriteTable` evaluates `tab->newvals`: the parameter `$n` is the old value of the selected column `n`, and `transforms` gives the expression of `USING` of a selected column, or `None` for a column that keeps its old value. The expressions refer to the columns by their names, through `FROM (VALUES ($1, ...)) AS relname(columns)`. The other columns get their defaults.
275#[must_use]
276pub fn rewrite_insert(target: &CopyTarget, rv: &RangeVar, transforms: &[Option<&Node>]) -> Node {
277    let names: Vec<&str> = target.selected.iter().map(|&i| target.columns[i].as_str()).collect();
278    let cols = names
279        .iter()
280        .map(|&name| {
281            Some(ResTarget { name: Some(name.into()), location: -1, ..Default::default() }.into())
282        })
283        .collect();
284    let targets = names
285        .iter()
286        .enumerate()
287        .map(|(k, &name)| {
288            let value = match transforms.get(k).copied().flatten() {
289                Some(expr) => expr.clone(),
290                None => {
291                    let field = Some(Node::String(name.into()));
292                    ColumnRef { fields: vec![field], location: -1 }.into()
293                }
294            };
295            Some(ResTarget { val: Some(value), location: -1, ..Default::default() }.into())
296        })
297        .collect();
298    let colnames = names.iter().map(|&name| Some(Node::String(name.into()))).collect();
299    let alias = Alias { aliasname: rv.relname.clone(), colnames };
300    let from = RangeSubselect {
301        lateral: false,
302        subquery: Some(values(names.len()).into()),
303        alias: Some(Box::new(alias)),
304    };
305    let select = SelectStmt {
306        targetList: targets,
307        fromClause: vec![Some(from.into())],
308        ..Default::default()
309    };
310    let insert = InsertStmt {
311        relation: Some(Box::new(RangeVar { alias: None, location: -1, ..rv.clone() })),
312        cols,
313        selectStmt: Some(select.into()),
314        r#override: OverridingKind::OVERRIDING_SYSTEM_VALUE,
315        ..Default::default()
316    };
317    insert.into()
318}
319
320/// `VALUES ($1, ..., $count)`.
321fn values(count: usize) -> SelectStmt {
322    let params = (1..=count)
323        .map(|n| {
324            let number = i32::try_from(n).unwrap_or(i32::MAX);
325            Some(ParamRef { number, location: -1 }.into())
326        })
327        .collect();
328    SelectStmt { valuesLists: vec![Some(Node::List(params))], ..Default::default() }
329}
330
331/// `CopyGetAttnums`: the indexes of the columns that a list of names gives, or of all the columns when the list is empty. `rel` is the name of the relation for the error of a column that does not exist.
332///
333/// # Errors
334///
335/// `42703` for a name that is not a column and `42701` for a name that the list gives twice.
336pub fn attnums(columns: &[String], names: &[String], rel: Option<&str>) -> Result<Vec<usize>> {
337    if names.is_empty() {
338        return Ok((0..columns.len()).collect());
339    }
340    let mut out = Vec::with_capacity(names.len());
341    for name in names {
342        let Some(at) = columns.iter().position(|c| c == name) else {
343            let message = match rel {
344                Some(rel) => format!("column \"{name}\" of relation \"{rel}\" does not exist"),
345                None => format!("column \"{name}\" does not exist"),
346            };
347            return Err(Error::new(SqlState::UNDEFINED_COLUMN, message));
348        };
349        if out.contains(&at) {
350            return Err(Error::new(
351                SqlState::DUPLICATE_COLUMN,
352                format!("column \"{name}\" specified more than once"),
353            ));
354        }
355        out.push(at);
356    }
357    Ok(out)
358}
359
360/// The names of the user columns of a relation of the catalog or of the built-in catalog.
361pub(super) fn relation_columns(env: &dyn Env, relid: u32) -> Vec<String> {
362    if let Some(rel) = env.catalog().and_then(|c| c.relation(relid)) {
363        return rel.columns.iter().filter(|c| !c.dropped).map(|c| c.name.clone()).collect();
364    }
365    builtin::attributes(relid)
366        .iter()
367        .filter(|a| a.num > 0 && !a.dropped)
368        .map(|a| a.name.to_string())
369        .collect()
370}
371
372/// The text of a `String` node.
373fn string(node: &Node) -> Option<&str> {
374    match node {
375        Node::String(s) => Some(s),
376        _ => None,
377    }
378}