1use 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
19pub const RETURNING_OLD: usize = usize::MAX - 1;
21pub const RETURNING_NEW: usize = usize::MAX - 2;
23pub const MERGE_ACTION: usize = usize::MAX - 3;
25
26#[derive(Clone, Debug, PartialEq)]
28pub struct Shown {
29 pub targets: Vec<Expr>,
30 pub old: String,
32 pub new: String,
33}
34
35enum Source {
37 Row,
39 Carry(usize),
41 Null,
43 Shift(usize),
45}
46
47impl Analyzer<'_> {
48 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 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 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 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 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 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 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 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 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
324fn 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
337pub(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}