1use std::collections::BTreeMap;
10
11use crate::ast::{
12 ColumnType, Expr, FromClause, InsertStmt, InternalColumnRef, InternalRelationId, JoinKind,
13 OnConflict, Projection, ReturningAliases, SelectStmt, Statement,
14};
15use crate::plpgsql::{bind_expr, bind_select, ResolvedVariable, VariableResolver};
16use crate::SQLError;
17use uqa_core::Value;
18
19use super::{RuleColumnMetadata, RuleRowValues, RuntimeRuleResolver};
20
21#[cfg(test)]
22mod tests;
23
24type BindRuleAction<'a> = dyn Fn(&mut dyn VariableResolver) -> Result<Statement, SQLError> + 'a;
25
26pub fn bind_insert_values_action(
27 matching_rows: &[usize],
28 rows: &[impl RuleRowValues],
29 columns: &BTreeMap<String, RuleColumnMetadata>,
30 bind_action: &BindRuleAction<'_>,
31) -> Result<Statement, SQLError> {
32 if matching_rows.is_empty() {
33 let mut bound = bind_action(&mut RuntimeRuleResolver {
34 old: None,
35 new: None,
36 old_doc_id: None,
37 new_doc_id: None,
38 columns,
39 })?;
40 let Statement::Insert(insert) = &mut bound else {
41 return Err(SQLError::Internal(
42 "rewrite rule INSERT VALUES action changed statement kind".into(),
43 ));
44 };
45 insert.rows.clear();
46 return Ok(bound);
47 }
48 let mut combined = None;
49 for row_index in matching_rows {
50 let row = rows.get(*row_index).ok_or_else(|| {
51 SQLError::Internal("rewrite rule lost its qualified row image".into())
52 })?;
53 let bound = bind_action(&mut runtime_rule_resolver(row, columns))?;
54 let Statement::Insert(mut insert) = bound else {
55 return Err(SQLError::Internal(
56 "rewrite rule INSERT VALUES action changed statement kind".into(),
57 ));
58 };
59 if let Some(Statement::Insert(existing)) = combined.as_mut() {
60 if BoundInsertContract::from(&*existing) != BoundInsertContract::from(&insert) {
61 return Err(SQLError::Internal(
62 "rewrite rule INSERT VALUES action produced row-dependent statement clauses"
63 .into(),
64 ));
65 }
66 existing.rows.append(&mut insert.rows);
67 } else {
68 combined = Some(Statement::Insert(insert));
69 }
70 }
71 combined.ok_or_else(|| SQLError::Internal("rewrite rule action lost its row source".into()))
72}
73
74#[derive(PartialEq)]
75struct BoundInsertContract<'a> {
76 table: &'a str,
77 target_relation_bound: bool,
78 target_qualifier: &'a str,
79 include_descendants: bool,
80 columns: &'a [crate::ast::AssignmentTarget],
81 with: &'a [crate::ast::CTE],
82 select_source: Option<&'a SelectStmt>,
83 on_conflict: Option<&'a OnConflict>,
84 returning: &'a [Projection],
85 returning_aliases: &'a ReturningAliases,
86}
87
88impl<'a> From<&'a InsertStmt> for BoundInsertContract<'a> {
89 fn from(insert: &'a InsertStmt) -> Self {
90 Self {
91 table: &insert.table,
92 target_relation_bound: insert.target_relation_bound,
93 target_qualifier: &insert.target_qualifier,
94 include_descendants: insert.include_descendants,
95 columns: &insert.columns,
96 with: &insert.with,
97 select_source: insert.select_source.as_deref(),
98 on_conflict: insert.on_conflict.as_ref(),
99 returning: &insert.returning,
100 returning_aliases: &insert.returning_aliases,
101 }
102 }
103}
104
105fn runtime_rule_resolver<'a>(
106 row: &'a dyn RuleRowValues,
107 columns: &'a BTreeMap<String, RuleColumnMetadata>,
108) -> RuntimeRuleResolver<'a> {
109 RuntimeRuleResolver {
110 old: row.old_row(),
111 new: row.new_row(),
112 old_doc_id: row.old_doc_id(),
113 new_doc_id: row.new_doc_id(),
114 columns,
115 }
116}
117
118struct RuleRowSource {
119 clause: FromClause,
120 relation: InternalRelationId,
121 old_columns: BTreeMap<String, InternalColumnRef>,
122 new_columns: BTreeMap<String, InternalColumnRef>,
123 old_row: InternalColumnRef,
124 new_row: InternalColumnRef,
125 source_index: InternalColumnRef,
126}
127
128pub struct BoundSetOrientedAction {
129 pub statement: Statement,
130 pub source_index: Expr,
131}
132
133pub fn bind_set_oriented_action(
134 matching_rows: &[usize],
135 rows: &[impl RuleRowValues],
136 columns: &BTreeMap<String, RuleColumnMetadata>,
137 bind_action: &BindRuleAction<'_>,
138) -> Result<BoundSetOrientedAction, SQLError> {
139 let source = rule_row_source(matching_rows, rows, columns)?;
140 let source_index = Expr::InternalColumn(source.source_index);
141 let mut bound = bind_action(&mut RuleSourceResolver {
142 old_columns: &source.old_columns,
143 new_columns: &source.new_columns,
144 old_row: source.old_row,
145 new_row: source.new_row,
146 })?;
147 attach_rule_row_source(&mut bound, source.clause, source.relation)?;
148 Ok(BoundSetOrientedAction {
149 statement: bound,
150 source_index,
151 })
152}
153
154fn rule_row_source(
155 matching_rows: &[usize],
156 rows: &[impl RuleRowValues],
157 columns: &BTreeMap<String, RuleColumnMetadata>,
158) -> Result<RuleRowSource, SQLError> {
159 let relation = InternalRelationId::allocate();
160 let mut internal_column_types = Vec::with_capacity(columns.len() * 2 + 3);
161 let mut internal_column_collations = Vec::with_capacity(columns.len() * 2 + 3);
162 let mut old_columns = BTreeMap::new();
163 let mut new_columns = BTreeMap::new();
164 for (index, (column, metadata)) in columns.iter().enumerate() {
165 let old = relation.column(index * 2);
166 let new = relation.column(index * 2 + 1);
167 old_columns.insert(column.clone(), old);
168 new_columns.insert(column.clone(), new);
169 internal_column_types.push(Some(metadata.ty.clone()));
170 internal_column_types.push(Some(metadata.ty.clone()));
171 internal_column_collations.extend([metadata.collation_oid; 2]);
172 }
173 let old_row = relation.column(internal_column_types.len());
174 internal_column_types.push(Some(ColumnType::Record));
175 let new_row = relation.column(internal_column_types.len());
176 internal_column_types.push(Some(ColumnType::Record));
177 let source_index = relation.column(internal_column_types.len());
178 internal_column_types.push(Some(ColumnType::BigInteger));
179 internal_column_collations.extend([None; 3]);
180 let values = matching_rows
181 .iter()
182 .map(|row_index| {
183 let row = rows.get(*row_index).ok_or_else(|| {
184 SQLError::Internal("rewrite rule lost its qualified row image".into())
185 })?;
186 let resolver = runtime_rule_resolver(row, columns);
187 let mut values = Vec::with_capacity(internal_column_types.len());
188 for column in columns.keys() {
189 values.push(resolver.field_expression(row.old_row(), row.old_doc_id(), column)?);
190 values.push(resolver.field_expression(row.new_row(), row.new_doc_id(), column)?);
191 }
192 values.push(ResolvedVariable::into_expression(
193 resolver.record(row.old_row(), row.old_doc_id())?,
194 ));
195 values.push(ResolvedVariable::into_expression(
196 resolver.record(row.new_row(), row.new_doc_id())?,
197 ));
198 let row_index = i64::try_from(*row_index).map_err(|_| {
199 SQLError::Internal("rewrite rule event row index exceeds BIGINT".into())
200 })?;
201 values.push(Expr::Literal(Value::Int(row_index)));
202 Ok(values)
203 })
204 .collect::<Result<Vec<_>, SQLError>>()?;
205 Ok(RuleRowSource {
206 clause: FromClause::Values {
207 rows: values,
208 alias: None,
209 column_aliases: Vec::new(),
210 internal_relation: Some(relation),
211 internal_column_types,
212 internal_column_collations,
213 },
214 relation,
215 old_columns,
216 new_columns,
217 old_row,
218 new_row,
219 source_index,
220 })
221}
222
223struct RuleSourceResolver<'a> {
224 old_columns: &'a BTreeMap<String, InternalColumnRef>,
225 new_columns: &'a BTreeMap<String, InternalColumnRef>,
226 old_row: InternalColumnRef,
227 new_row: InternalColumnRef,
228}
229
230impl VariableResolver for RuleSourceResolver<'_> {
231 fn resolve_name(&mut self, _name: &str) -> Result<Option<ResolvedVariable>, SQLError> {
232 Ok(None)
233 }
234
235 fn resolve_qualified(
236 &mut self,
237 _qualifier: &str,
238 _column: &str,
239 ) -> Result<Option<ResolvedVariable>, SQLError> {
240 Ok(None)
241 }
242
243 fn resolve_param(&mut self, _index: usize) -> Result<Option<ResolvedVariable>, SQLError> {
244 Ok(None)
245 }
246
247 fn rewrite_name(&mut self, name: &str) -> Result<Option<Expr>, SQLError> {
248 Ok(if name.eq_ignore_ascii_case("old") {
249 Some(Expr::InternalColumn(self.old_row))
250 } else if name.eq_ignore_ascii_case("new") {
251 Some(Expr::InternalColumn(self.new_row))
252 } else {
253 None
254 })
255 }
256
257 fn rewrite_qualified(
258 &mut self,
259 qualifier: &str,
260 column: &str,
261 ) -> Result<Option<Expr>, SQLError> {
262 let columns = if qualifier.eq_ignore_ascii_case("old") {
263 self.old_columns
264 } else if qualifier.eq_ignore_ascii_case("new") {
265 self.new_columns
266 } else {
267 return Ok(None);
268 };
269 let source_column = columns
270 .get(column)
271 .ok_or_else(|| SQLError::UnknownColumn(format!("{qualifier}.{column}")))?;
272 Ok(Some(Expr::InternalColumn(*source_column)))
273 }
274
275 fn rewrite_qualified_whole_row(&mut self, qualifier: &str) -> Result<Option<Expr>, SQLError> {
276 self.rewrite_name(qualifier)
277 }
278}
279
280fn attach_rule_row_source(
281 statement: &mut Statement,
282 source: FromClause,
283 relation: InternalRelationId,
284) -> Result<(), SQLError> {
285 match statement {
286 Statement::Select(select) => attach_select_rule_source(select, &source, relation),
287 Statement::Insert(insert) => {
288 let select = insert.select_source.as_mut().ok_or_else(|| {
289 SQLError::Internal("set-oriented rule INSERT action has no SELECT source".into())
290 })?;
291 attach_select_rule_source(select, &source, relation);
292 }
293 Statement::Update(update) => {
294 update.from = Some(prepend_rule_row_source(
295 update.from.take(),
296 source,
297 relation,
298 ));
299 }
300 Statement::Delete(delete) => {
301 delete.using = Some(prepend_rule_row_source(
302 delete.using.take(),
303 source,
304 relation,
305 ));
306 }
307 _ => {
308 return Err(SQLError::Internal(
309 "validated rewrite-rule action changed statement kind".into(),
310 ))
311 }
312 }
313 Ok(())
314}
315
316fn attach_select_rule_source(
317 select: &mut SelectStmt,
318 source: &FromClause,
319 relation: InternalRelationId,
320) {
321 if let Some(set_op) = select.set_op.as_mut() {
322 if let Some(left) = set_op.left.as_mut() {
323 attach_select_rule_source(left, source, relation);
324 } else {
325 select.from = Some(prepend_rule_row_source(
326 select.from.take(),
327 source.clone(),
328 relation,
329 ));
330 }
331 attach_select_rule_source(&mut set_op.right, source, relation);
332 } else {
333 select.from = Some(prepend_rule_row_source(
334 select.from.take(),
335 source.clone(),
336 relation,
337 ));
338 }
339}
340
341fn prepend_rule_row_source(
342 existing: Option<FromClause>,
343 source: FromClause,
344 relation: InternalRelationId,
345) -> FromClause {
346 let Some(existing) = existing else {
347 return source;
348 };
349 let lateral = from_references_internal_relation(&existing, relation);
350 FromClause::Join {
351 left: Box::new(source),
352 right: Box::new(existing),
353 kind: JoinKind::Cross,
354 on: None,
355 using: None,
356 natural: false,
357 alias: None,
358 column_aliases: Vec::new(),
359 lateral,
360 }
361}
362
363struct InternalReferenceResolver {
364 relation: InternalRelationId,
365 referenced: bool,
366}
367
368impl VariableResolver for InternalReferenceResolver {
369 fn resolve_name(&mut self, _name: &str) -> Result<Option<ResolvedVariable>, SQLError> {
370 Ok(None)
371 }
372
373 fn resolve_qualified(
374 &mut self,
375 _qualifier: &str,
376 _column: &str,
377 ) -> Result<Option<ResolvedVariable>, SQLError> {
378 Ok(None)
379 }
380
381 fn resolve_param(&mut self, _index: usize) -> Result<Option<ResolvedVariable>, SQLError> {
382 Ok(None)
383 }
384
385 fn rewrite_internal(&mut self, column: InternalColumnRef) -> Result<Option<Expr>, SQLError> {
386 if column.relation() == self.relation {
387 self.referenced = true;
388 }
389 Ok(None)
390 }
391}
392
393fn expr_references_internal_relation(expr: &Expr, relation: InternalRelationId) -> bool {
394 let mut resolver = InternalReferenceResolver {
395 relation,
396 referenced: false,
397 };
398 let _ = bind_expr(expr, &mut resolver);
399 resolver.referenced
400}
401
402fn select_references_internal_relation(select: &SelectStmt, relation: InternalRelationId) -> bool {
403 let mut resolver = InternalReferenceResolver {
404 relation,
405 referenced: false,
406 };
407 let _ = bind_select(select, &mut resolver);
408 resolver.referenced
409}
410
411fn from_references_internal_relation(from: &FromClause, relation: InternalRelationId) -> bool {
412 match from {
413 FromClause::Table { .. } => false,
414 FromClause::Join {
415 left, right, on, ..
416 } => {
417 from_references_internal_relation(left, relation)
418 || from_references_internal_relation(right, relation)
419 || on
420 .as_ref()
421 .is_some_and(|expr| expr_references_internal_relation(expr, relation))
422 }
423 FromClause::Values { rows, .. } => rows
424 .iter()
425 .flatten()
426 .any(|expr| expr_references_internal_relation(expr, relation)),
427 FromClause::Function { args, .. } => args
428 .iter()
429 .any(|expr| expr_references_internal_relation(expr, relation)),
430 FromClause::FunctionGroup { functions, .. } => functions.iter().any(|function| {
431 function
432 .args
433 .iter()
434 .any(|expr| expr_references_internal_relation(expr, relation))
435 }),
436 FromClause::Subquery { body, .. } => select_references_internal_relation(body, relation),
437 }
438}