Skip to main content

rupg_analyze/
returning.rs

1//! `transformReturningClause` and `addNSItemForReturning` of `analyze.c`: the list of `RETURNING` of an `INSERT`, an `UPDATE` or a `DELETE`, with the names `old` and `new` of the old and the new row.
2//!
3//! The analysis gives a query that reads one relation: the rows that the statement writes. A row of it has the columns of the table, then the carried values. For an `INSERT` and an `UPDATE` the columns are the new row, and for a `DELETE` they are the old row. A carried value is a value that the query of the statement gives for each row after the columns: a column of a relation of `FROM` or `USING`, or a column of the old row of an `UPDATE`. A column of a row that the statement does not have, as the old row of an `INSERT`, is null.
4
5use rupg_common::{Error, Result, SqlState};
6use rupg_sql::nodes::{Node, ReturningClause, ReturningOptionKind};
7use rupg_types::{Value, oid};
8
9use crate::Analyzer;
10use crate::agg::Kind;
11use crate::coerce::{AtOpt, Context};
12use crate::expr::{Case, CastForm, Expr, ExprKind, Var};
13use crate::from::{Column, FromItem, Relation, TABLE_OID_ATTNUM};
14use crate::insert::WriteKind;
15use crate::merge::carry;
16use crate::select::Query;
17use crate::typename::place;
18
19/// The relation of a column of the old row in [`Shown::targets`], as a `Var` with `varreturningtype` `VAR_RETURNING_OLD`.
20pub const RETURNING_OLD: usize = usize::MAX - 1;
21/// The relation of a column of the new row in [`Shown::targets`], as a `Var` with `varreturningtype` `VAR_RETURNING_NEW`.
22pub const RETURNING_NEW: usize = usize::MAX - 2;
23/// The relation of `merge_action()` in [`Shown::targets`], as a `MergeSupportFunc`.
24pub const MERGE_ACTION: usize = usize::MAX - 3;
25
26/// The list of `RETURNING` as `EXPLAIN` shows it in the `Output` of a `ModifyTable`. A column of `old` or `new` reads [`RETURNING_OLD`] or [`RETURNING_NEW`] with the number of the column in the table, and `merge_action()` reads [`MERGE_ACTION`]. The other columns are as in the targets of the query.
27#[derive(Clone, Debug, PartialEq)]
28pub struct Shown {
29    pub targets: Vec<Expr>,
30    /// `returningOldAlias` and `returningNewAlias`: the names of the old and the new row.
31    pub old: String,
32    pub new: String,
33}
34
35/// The place where a column of `RETURNING` takes its value.
36enum Source {
37    /// The row that the statement writes, which is the relation of the query of `RETURNING`.
38    Row,
39    /// A value that the query of the statement carries for each row: the column of this relation of the statement.
40    Carry(usize),
41    /// A row that the statement does not have.
42    Null,
43    /// The columns of the row of `RETURNING` of a `MERGE` after this number of columns: the old row, the new row or the action.
44    Shift(usize),
45}
46
47impl Analyzer<'_> {
48    /// `transformReturningClause`: the query of `RETURNING` and the values that the query of the statement must carry for it. The first relation of the scope is the target table, and the other relations are the relations of `FROM` or `USING`. The names `old` and `new` stay in the scope only while the analysis reads the list.
49    ///
50    /// # Errors
51    ///
52    /// `42601` for an option that the clause gives twice or for a list with no column, `42712` for a name of an option that the scope has, and the errors of the expressions.
53    pub(crate) fn returning(
54        &mut self,
55        clause: &ReturningClause,
56        kind: WriteKind,
57    ) -> Result<(Query, Vec<Expr>)> {
58        let size = self.scope_size();
59        // `merge_action()` reads the relation of the action, which no name can find.
60        if kind == WriteKind::Merge {
61            let column = Column {
62                name: "merge_action".to_string(),
63                ty: oid::TEXT,
64                typmod: -1,
65                not_null: false,
66            };
67            self.scope.merge_action = Some(self.scope.relations.len());
68            self.scope.relations.push(Relation::new(0, vec![column], None, None, None));
69        }
70        let result = self.returning_list(clause, kind);
71        self.scope.merge_action = None;
72        self.truncate_scope(size);
73        result
74    }
75
76    /// The body of [`Analyzer::returning`], which adds the names `old` and `new` to the scope.
77    fn returning_list(
78        &mut self,
79        clause: &ReturningClause,
80        kind: WriteKind,
81    ) -> Result<(Query, Vec<Expr>)> {
82        let mut names: [Option<(String, Option<usize>)>; 2] = [None, None];
83        for node in &clause.options {
84            let Some(Node::ReturningOption(option)) = node else {
85                return Err(Error::internal("a RETURNING option that is not ReturningOption"));
86            };
87            let (slot, what) = if option.option == ReturningOptionKind::RETURNING_OPTION_OLD {
88                (0, "OLD")
89            } else {
90                (1, "NEW")
91            };
92            let at = place(option.location);
93            if names[slot].is_some() {
94                return Err(Error::new(
95                    SqlState::SYNTAX_ERROR,
96                    format!("{what} cannot be specified multiple times"),
97                )
98                .at_opt(at));
99            }
100            names[slot] = Some((option.value.as_deref().unwrap_or_default().to_string(), at));
101        }
102        // The names `old` and `new` are the default names, unless the scope has a relation with the name.
103        for (slot, name) in names.iter_mut().zip(["old", "new"]) {
104            if slot.is_none() && !self.has_refname(name)? {
105                *slot = Some((name.to_string(), None));
106            }
107        }
108        let mut relations = [None, None];
109        for ((index, slot), old) in relations.iter_mut().zip(&names).zip([true, false]) {
110            let Some((name, at)) = slot else { continue };
111            if self.has_refname(name)? {
112                return Err(Error::new(
113                    SqlState::DUPLICATE_ALIAS,
114                    format!("table name \"{name}\" specified more than once"),
115                )
116                .at_opt(*at));
117            }
118            *index = Some(self.add_returning_item(name, old, kind)?);
119        }
120        let mut targets = self.target_list(&clause.exprs, Kind::Returning)?;
121        if targets.is_empty() {
122            let at = match clause.exprs.first() {
123                Some(Some(Node::ResTarget(rt))) => place(rt.location),
124                _ => None,
125            };
126            return Err(Error::new(
127                SqlState::SYNTAX_ERROR,
128                "RETURNING must have at least one column",
129            )
130            .at_opt(at));
131        }
132        self.mark_origins(&mut targets);
133        self.auto_view_origins(&mut targets);
134        // `resolveTargetListUnknowns`: a column of type `unknown` becomes `text`.
135        for target in &mut targets {
136            if target.expr.ty == oid::UNKNOWN {
137                let expr = std::mem::replace(&mut target.expr, Expr::new(ExprKind::CaseTest, 0));
138                target.expr = self.coerce(expr, oid::TEXT, -1, Context::Implicit, None)?;
139            }
140        }
141        let [old, new] = relations;
142        let mut shown: Vec<Expr> = targets.iter().map(|t| t.expr.clone()).collect();
143        let table = self
144            .scope
145            .relations
146            .first()
147            .ok_or_else(|| Error::internal("RETURNING with no table"))?;
148        let width = table.columns.len();
149        let table_oid = table.oid;
150        let types: Vec<(u32, i32)> = table.columns.iter().map(|c| (c.ty, c.typmod)).collect();
151        let mut relation = Relation::new(table.oid, table.columns.clone(), None, None, None);
152        relation.name.clone_from(&table.name);
153        relation.alias.clone_from(&table.alias);
154        let action = self.scope.merge_action;
155        // The whole old row of a `MERGE` is null for an insert, and the whole new row is null for a delete.
156        let no_row = if kind == WriteKind::Merge {
157            let mut no_row = |name: &str| {
158                let attnum = i16::try_from(3 * width + 1).unwrap_or(i16::MAX);
159                let var = Var { relation: 0, attnum, levels_up: 0 };
160                let action = Expr::new(ExprKind::Var(var), oid::TEXT);
161                let name = Expr::new(ExprKind::Const(Value::Text(name.into())), oid::TEXT);
162                self.make_op(&["="], Some(action), name, None)
163            };
164            Some((no_row("INSERT")?, no_row("DELETE")?))
165        } else {
166            None
167        };
168        // `ReplaceVarsFromTargetList`: the whole row of an automatically updatable view is a row of the columns of the view, which is null when the action has no such row.
169        let view_rows = |expr: &mut Expr, depth: usize| {
170            let ExprKind::Row(fields, CastForm::Implicit) = &expr.kind else { return false };
171            if expr.ty == oid::RECORD {
172                return false;
173            }
174            let first = fields.iter().find_map(|f| {
175                f.find(depth, &mut |e: &Expr, d| match e.kind {
176                    ExprKind::Var(var) if var.levels_up == d => Some(var.relation),
177                    _ => None,
178                })
179            });
180            let Some(relation) = first else { return false };
181            let is_old = Some(relation) == old;
182            if !is_old && Some(relation) != new {
183                return false;
184            }
185            let null = Expr::new(ExprKind::Const(Value::Null), expr.ty);
186            match (&no_row, kind) {
187                (Some((insert, delete)), _) => {
188                    let mut when = if is_old { insert.clone() } else { delete.clone() };
189                    rewrite(&mut when, 0, &mut |e: &mut Expr, _| {
190                        if let ExprKind::Var(v) = &mut e.kind {
191                            v.levels_up = depth;
192                        }
193                    });
194                    let row = std::mem::replace(expr, null.clone());
195                    let case =
196                        Case { arg: None, whens: vec![(when, null)], default: Box::new(row) };
197                    expr.kind = ExprKind::Case(case);
198                }
199                (None, WriteKind::Insert) if is_old => *expr = null,
200                (None, WriteKind::Delete) if !is_old => *expr = null,
201                _ => {}
202            }
203            true
204        };
205        for expr in targets.iter_mut().map(|t| &mut t.expr).chain(&mut shown) {
206            rewrite_until(expr, 0, &mut |e: &mut Expr, depth| view_rows(e, depth));
207        }
208        let source = |relation: usize| {
209            if Some(relation) == old {
210                match kind {
211                    WriteKind::Insert => Source::Null,
212                    WriteKind::Update => Source::Carry(0),
213                    WriteKind::Delete => Source::Row,
214                    WriteKind::Merge => Source::Shift(width),
215                }
216            } else if Some(relation) == new {
217                match kind {
218                    WriteKind::Delete => Source::Null,
219                    WriteKind::Insert | WriteKind::Update => Source::Row,
220                    WriteKind::Merge => Source::Shift(2 * width),
221                }
222            } else if Some(relation) == action {
223                Source::Shift(3 * width)
224            } else if relation == 0 {
225                Source::Row
226            } else {
227                Source::Carry(relation)
228            }
229        };
230        // The carried values of a `MERGE` come after the row, the old row, the new row and the action. A row of a partitioned table has the OID of its partition after the columns, or after the action for a `MERGE`, and a row of a table with children of `INHERITS` has the OID of its table after the columns.
231        let partitioned = self.is_partitioned(table_oid) || self.inherits_target(kind);
232        let base = match kind {
233            WriteKind::Merge => 3 * width + 1 + usize::from(partitioned),
234            _ if partitioned => width + 1,
235            _ => width,
236        };
237        let partition_at = if kind == WriteKind::Merge { 3 * width + 2 } else { width + 1 };
238        let partition_attnum = i16::try_from(partition_at).unwrap_or(i16::MAX);
239        let mut carried: Vec<Expr> = Vec::new();
240        let mut change = |expr: &mut Expr, depth: usize| {
241            let ExprKind::Var(var) = expr.kind else { return };
242            if var.levels_up != depth {
243                return;
244            }
245            match source(var.relation) {
246                Source::Row if partitioned && var.attnum == TABLE_OID_ATTNUM => {
247                    expr.kind = ExprKind::Var(Var {
248                        relation: 0,
249                        attnum: partition_attnum,
250                        levels_up: depth,
251                    });
252                }
253                Source::Row => expr.kind = ExprKind::Var(Var { relation: 0, ..var }),
254                Source::Null => expr.kind = ExprKind::Const(Value::Null),
255                Source::Shift(by) if var.attnum == 0 => {
256                    // The planner of PostgreSQL reads the whole row as a `RowExpr` of the columns. A dropped column has no field.
257                    let fields = types.iter().enumerate().filter(|(_, (ty, _))| *ty != 0);
258                    let fields = fields
259                        .map(|(i, &(ty, typmod))| {
260                            let attnum = i16::try_from(by + i + 1).unwrap_or(i16::MAX);
261                            let var = Var { relation: 0, attnum, levels_up: depth };
262                            let mut field = Expr::new(ExprKind::Var(var), ty);
263                            field.typmod = typmod;
264                            field
265                        })
266                        .collect();
267                    let row = Expr::new(ExprKind::Row(fields, CastForm::Implicit), expr.ty);
268                    let Some((insert, delete)) = &no_row else { return };
269                    let mut when = if by == width { insert.clone() } else { delete.clone() };
270                    rewrite(&mut when, 0, &mut |e: &mut Expr, _| {
271                        if let ExprKind::Var(v) = &mut e.kind {
272                            v.levels_up = depth;
273                        }
274                    });
275                    let null = Expr::new(ExprKind::Const(Value::Null), expr.ty);
276                    let case =
277                        Case { arg: None, whens: vec![(when, null)], default: Box::new(row) };
278                    expr.kind = ExprKind::Case(case);
279                }
280                Source::Shift(by) => {
281                    let attnum =
282                        i16::try_from(by).map_or(i16::MAX, |by| by.saturating_add(var.attnum));
283                    expr.kind = ExprKind::Var(Var { relation: 0, attnum, levels_up: depth });
284                }
285                Source::Carry(relation) => {
286                    let at = carry(&mut carried, relation, var.attnum, expr);
287                    let attnum = i16::try_from(base + 1 + at).unwrap_or(i16::MAX);
288                    expr.kind = ExprKind::Var(Var { relation: 0, attnum, levels_up: depth });
289                }
290            }
291        };
292        for target in &mut targets {
293            rewrite(&mut target.expr, 0, &mut change);
294        }
295        for expr in &mut shown {
296            rewrite(expr, 0, &mut |e: &mut Expr, depth| {
297                if let ExprKind::Var(var) = e.kind
298                    && var.levels_up == depth
299                {
300                    let named =
301                        [(old, RETURNING_OLD), (new, RETURNING_NEW), (action, MERGE_ACTION)];
302                    if let Some((_, relation)) =
303                        named.iter().find(|(r, _)| *r == Some(var.relation))
304                    {
305                        e.kind = ExprKind::Var(Var { relation: *relation, ..var });
306                        return;
307                    }
308                }
309                change(e, depth);
310            });
311        }
312        let [old, new] = names.map(|n| n.map(|(name, _)| name).unwrap_or_default());
313        let query = Query {
314            relations: vec![relation],
315            from: vec![FromItem::Relation(0)],
316            targets,
317            shown: Some(Box::new(Shown { targets: shown, old, new })),
318            ..Query::default()
319        };
320        Ok((query, carried))
321    }
322}
323
324/// As [`rewrite`], but `f` gives `true` for an expression whose expressions it does not see.
325fn rewrite_until(expr: &mut Expr, depth: usize, f: &mut dyn FnMut(&mut Expr, usize) -> bool) {
326    if f(expr, depth) {
327        return;
328    }
329    for child in expr.children_mut() {
330        rewrite_until(child, depth, f);
331    }
332    if let ExprKind::SubLink(sub) = &mut expr.kind {
333        sub.query.each_expr_mut(depth + 1, &mut |e, depth| rewrite_until(e, depth, &mut *f));
334    }
335}
336
337/// Calls `f` for the expression and for each expression in it, also in its subqueries, with the number of subqueries between the expression and this expression. `f` sees an expression before the expressions in it.
338pub(crate) fn rewrite(expr: &mut Expr, depth: usize, f: &mut dyn FnMut(&mut Expr, usize)) {
339    f(expr, depth);
340    for child in expr.children_mut() {
341        rewrite(child, depth, f);
342    }
343    if let ExprKind::SubLink(sub) = &mut expr.kind {
344        sub.query.each_expr_mut(depth + 1, &mut |e, depth| rewrite(e, depth, &mut *f));
345    }
346}