1use crate::{
9 ast::{BinaryOp, ColumnDef, TableConstraintSet},
10 binding::snapshot::BindingSnapshot,
11 catalog::index::EnforcedKey,
12 plan::AggregateClassifier,
13 plan::{ConflictPlan, ExpressionPlan, InsertPlan},
14 routines::RoutineResolution,
15 RowSchema, SQLError, SQLParam, ScalarExpr as Expr,
16};
17use std::collections::BTreeSet;
18use uqa_core::Value;
19
20pub trait ConflictCatalog {
21 fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
22 fn enforced_keys(&self, table: &str) -> Result<Vec<EnforcedKey>, String>;
23 fn try_declared_table_constraints(&self, table: &str) -> Result<TableConstraintSet, String>;
24}
25pub trait InferenceBindingScope {
27 fn binding_scope(&self) -> Result<BindingSnapshot, SQLError>;
28}
29#[derive(Clone, Copy)]
30pub struct InferenceContext<'a> {
31 pub catalog: &'a dyn ConflictCatalog,
32 pub aggregates: &'a dyn AggregateClassifier,
33 pub routines: &'a dyn RoutineResolution,
34 pub binding: &'a dyn InferenceBindingScope,
35}
36
37pub fn prepare_inference_predicate<'a>(
39 context: InferenceContext<'_>,
40 statement: &'a InsertPlan,
41 params: &[SQLParam],
42) -> Result<std::borrow::Cow<'a, InsertPlan>, SQLError> {
43 if statement
44 .on_conflict
45 .as_ref()
46 .is_none_or(|conflict| conflict.predicate.is_none() && conflict.expressions.is_empty())
47 {
48 return Ok(std::borrow::Cow::Borrowed(statement));
49 }
50 let columns = context
51 .catalog
52 .try_describe_table(&statement.table)
53 .map_err(SQLError::Internal)?
54 .ok_or_else(|| SQLError::UnknownTable(statement.table.clone()))?;
55 let schema = RowSchema::with_column_definitions(Some(&statement.target_qualifier), &columns);
56 let mut statement = statement.clone();
57 if let Some(conflict) = &mut statement.on_conflict {
58 for expression in conflict
59 .expressions
60 .iter_mut()
61 .chain(conflict.predicate.iter_mut().map(Box::as_mut))
62 {
63 prepare_inference_expression(
64 context,
65 expression,
66 &statement.target_qualifier,
67 &schema,
68 &columns,
69 params,
70 )?;
71 }
72 let expressions = std::mem::take(&mut conflict.expressions);
74 for expression in expressions {
75 if let Expr::Column(column) = expression {
76 conflict.conflict_columns.push(column);
77 } else {
78 conflict.expressions.push(expression);
79 }
80 }
81 }
82 Ok(std::borrow::Cow::Owned(statement))
83}
84
85fn prepare_inference_expression(
86 context: InferenceContext<'_>,
87 expression: &mut Expr,
88 qualifier: &str,
89 schema: &RowSchema,
90 columns: &[crate::ast::ColumnDef],
91 params: &[SQLParam],
92) -> Result<(), SQLError> {
93 let mut has_subquery = false;
94 expression.visit(&mut |part| {
95 has_subquery |= matches!(
96 part,
97 Expr::ScalarSubquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. }
98 );
99 });
100 if has_subquery {
101 return Err(predicate_error(
102 "0A000",
103 "cannot use subquery in index inference",
104 ));
105 }
106 if crate::semantics::aggregates::contains_aggregate(context.aggregates, expression) {
107 return Err(predicate_error(
108 "42803",
109 "aggregate functions are not allowed in index inference",
110 ));
111 }
112 if crate::semantics::windows::expr_has_window(expression) {
113 return Err(predicate_error(
114 "42P20",
115 "window functions are not allowed in index inference",
116 ));
117 }
118 let mut plan = ExpressionPlan {
119 scalar: expression.clone(),
120 subqueries: Vec::new(),
121 };
122 let binding = context.binding.binding_scope()?;
123 crate::binding::bind_expression_plan_routines_for_storage(
124 context.routines,
125 &mut plan,
126 params,
127 &binding.context(),
128 schema,
129 )?;
130 crate::plan::rewrite_scalar_expression(&mut plan.scalar, &mut |expression| {
131 if let Expr::QualifiedColumn {
132 qualifier: source,
133 column,
134 } = expression
135 {
136 if source == qualifier {
137 *expression = Expr::Column(column.clone());
138 }
139 }
140 });
141 crate::plan::rewrite_scalar_expression(&mut plan.scalar, &mut |expression| {
142 if let Expr::Cast { expr, ty, .. } = expression {
143 if let Expr::Column(name) = expr.as_ref() {
144 if columns.iter().any(|column| {
145 column.name == *name
146 && crate::ast::ColumnType::from_sql_name(ty).ok().as_ref()
147 == Some(&column.ty)
148 }) {
149 *expression = Expr::Column(name.clone());
150 }
151 }
152 }
153 });
154 *expression = inference_identity(plan.scalar);
155 Ok(())
156}
157
158fn inference_identity(mut expression: Expr) -> Expr {
159 crate::plan::rewrite_scalar_expression(&mut expression, &mut |node| {
161 while matches!(node, Expr::Collated { .. }) {
162 let Expr::Collated { expr, .. } = std::mem::replace(node, Expr::Literal(Value::Null))
163 else {
164 unreachable!("derived result envelope");
165 };
166 *node = *expr;
167 }
168 if let Expr::Cast { implicit, .. } = node {
170 *implicit = false;
171 }
172 });
173 expression
174}
175
176fn predicate_error(sqlstate: &str, message: &str) -> SQLError {
177 SQLError::Routine {
178 sqlstate: sqlstate.into(),
179 message: message.into(),
180 }
181}
182
183pub fn conflict_key_indices(
185 catalog: &dyn ConflictCatalog,
186 table: &str,
187 keys: &[EnforcedKey],
188 conflict: &ConflictPlan,
189) -> Result<Vec<usize>, SQLError> {
190 if let Some(name) = &conflict.constraint {
191 return constraint_target_index(catalog, table, keys, name);
192 }
193 if conflict.conflict_columns.is_empty() && conflict.expressions.is_empty() {
194 return Ok((0..keys.len()).collect());
195 }
196 validate_conflict_columns(catalog, table, &conflict.conflict_columns)?;
197 let target = conflict
198 .conflict_columns
199 .iter()
200 .map(String::as_str)
201 .collect::<BTreeSet<_>>();
202 let indexes = keys
203 .iter()
204 .enumerate()
205 .filter_map(|(index, key)| {
206 (key.columns
207 .iter()
208 .map(String::as_str)
209 .collect::<BTreeSet<_>>()
210 == target
211 && {
212 let expressions = key
213 .keys
214 .iter()
215 .filter_map(|key| match key {
216 crate::ast::IndexKey::Expression(expr) => Some(inference_identity(
217 ExpressionPlan::lower((**expr).clone()).scalar,
218 )),
219 crate::ast::IndexKey::Column(_) => None,
220 })
221 .collect::<Vec<_>>();
222 expressions
223 .iter()
224 .all(|expr| conflict.expressions.contains(expr))
225 && conflict
226 .expressions
227 .iter()
228 .all(|expr| expressions.contains(expr))
229 }
230 && key.predicate.as_deref().is_none_or(|required| {
231 conflict.predicate.as_deref().is_some_and(|given| {
232 implies(
233 given,
234 &inference_identity(ExpressionPlan::lower(required.clone()).scalar),
235 )
236 })
237 }))
238 .then_some(index)
239 })
240 .collect::<Vec<_>>();
241 if indexes.is_empty() {
242 return Err(SQLError::Routine {
243 sqlstate: "42P10".into(),
244 message:
245 "there is no unique or exclusion constraint matching the ON CONFLICT specification"
246 .into(),
247 });
248 }
249 Ok(indexes)
250}
251
252fn implies(given: &Expr, required: &Expr) -> bool {
254 if given == required || matches!(required, Expr::Literal(Value::Bool(true))) {
255 return true;
256 }
257 if matches!(given, Expr::Literal(Value::Bool(false) | Value::Null)) {
258 return true;
259 }
260 if let Expr::And(parts) = required {
261 return parts.iter().all(|part| implies(given, part));
262 }
263 if let Expr::Or(parts) = given {
264 return parts.iter().all(|part| implies(part, required));
265 }
266 if let Expr::And(parts) = given {
267 if parts.iter().any(|part| implies(part, required)) {
268 return true;
269 }
270 }
271 if let Expr::Or(parts) = required {
272 return parts.iter().any(|part| implies(given, part));
273 }
274 if let Some((left, given_op, given_value)) = comparison(given) {
275 if let Some((right, required_op, required_value)) = comparison(required) {
276 return left == right
277 && comparison_implies(given_op, given_value, required_op, required_value);
278 }
279 if let Expr::IsNull {
280 expr,
281 negated: true,
282 } = required
283 {
284 return left == expr.as_ref() && !matches!(given_value, Value::Null);
285 }
286 }
287 false
288}
289
290fn comparison(expr: &Expr) -> Option<(&Expr, BinaryOp, &Value)> {
291 let Expr::Binary { op, lhs, rhs } = expr else {
292 return None;
293 };
294 if !matches!(
295 op,
296 BinaryOp::Equal
297 | BinaryOp::NotEqual
298 | BinaryOp::Less
299 | BinaryOp::LessEqual
300 | BinaryOp::Greater
301 | BinaryOp::GreaterEqual
302 ) {
303 return None;
304 }
305 if let Expr::Literal(value) = rhs.as_ref() {
306 return Some((lhs, *op, value));
307 }
308 if let Expr::Literal(value) = lhs.as_ref() {
309 let reversed = match op {
310 BinaryOp::Less => BinaryOp::Greater,
311 BinaryOp::LessEqual => BinaryOp::GreaterEqual,
312 BinaryOp::Greater => BinaryOp::Less,
313 BinaryOp::GreaterEqual => BinaryOp::LessEqual,
314 other => *other,
315 };
316 return Some((rhs, reversed, value));
317 }
318 None
319}
320
321fn comparison_implies(given: BinaryOp, left: &Value, required: BinaryOp, right: &Value) -> bool {
322 use BinaryOp::{Equal, Greater, GreaterEqual, Less, LessEqual, NotEqual};
323 if matches!(left, Value::Null) || matches!(right, Value::Null) {
324 return false;
325 }
326 if std::mem::discriminant(left) != std::mem::discriminant(right) {
327 return false;
328 }
329 let order = left.cmp(right);
330 match (given, required) {
331 (Equal, Equal) | (NotEqual, NotEqual) => order.is_eq(),
332 (Equal, NotEqual) => !order.is_eq(),
333 (Equal | GreaterEqual, Greater) | (GreaterEqual, NotEqual) => order.is_gt(),
334 (Equal | LessEqual, Less) | (LessEqual, NotEqual) => order.is_lt(),
335 (Equal | Greater | GreaterEqual, GreaterEqual) | (Greater, Greater | NotEqual) => {
336 order.is_ge()
337 }
338 (Equal | Less | LessEqual, LessEqual) | (Less, Less | NotEqual) => order.is_le(),
339 _ => false,
340 }
341}
342
343pub fn validate_conflict_target(
344 catalog: &dyn ConflictCatalog,
345 table: &str,
346 conflict: &ConflictPlan,
347) -> Result<(), SQLError> {
348 let keys = catalog
349 .enforced_keys(table)
350 .map_err(|error| SQLError::Internal(format!("conflict target keys: {error}")))?;
351 conflict_key_indices(catalog, table, &keys, conflict).map(|_| ())
352}
353
354fn constraint_target_index(
355 catalog: &dyn ConflictCatalog,
356 table: &str,
357 keys: &[EnforcedKey],
358 name: &str,
359) -> Result<Vec<usize>, SQLError> {
360 if let Some(index) = keys
361 .iter()
362 .position(|key| key.constraint_owned && key.name.as_deref() == Some(name))
363 {
364 return Ok(vec![index]);
365 }
366 let snapshot = catalog
367 .try_declared_table_constraints(table)
368 .map_err(SQLError::Internal)?;
369 let columns = catalog
370 .try_describe_table(table)
371 .map_err(SQLError::Internal)?
372 .ok_or_else(|| SQLError::UnknownTable(table.into()))?;
373 let exists = snapshot
374 .checks
375 .iter()
376 .any(|check| check.name.as_deref() == Some(name))
377 || snapshot
378 .foreign_keys
379 .iter()
380 .any(|key| key.name.as_deref() == Some(name))
381 || columns.iter().any(|column| {
382 column.not_null_name.as_deref() == Some(name)
383 || column.check_name.as_deref() == Some(name)
384 || column
385 .references
386 .as_ref()
387 .is_some_and(|reference| reference.name.as_deref() == Some(name))
388 });
389 if exists {
390 return Err(SQLError::Routine {
391 sqlstate: "42809".into(),
392 message: "constraint in ON CONFLICT clause has no associated index".into(),
393 });
394 }
395 Err(SQLError::Routine {
396 sqlstate: "42704".into(),
397 message: format!("constraint \"{name}\" for table \"{table}\" does not exist"),
398 })
399}
400
401fn validate_conflict_columns(
402 catalog: &dyn ConflictCatalog,
403 table: &str,
404 names: &[String],
405) -> Result<(), SQLError> {
406 let columns = catalog
407 .try_describe_table(table)
408 .map_err(SQLError::Internal)?
409 .ok_or_else(|| SQLError::UnknownTable(table.into()))?;
410 for name in names {
411 if !columns.iter().any(|column| column.name == *name)
412 && !matches!(
413 name.as_str(),
414 "ctid" | "tableoid" | "xmin" | "xmax" | "cmin" | "cmax"
415 )
416 {
417 return Err(SQLError::UnknownColumn(name.clone()));
418 }
419 }
420 Ok(())
421}
422
423#[cfg(test)]
424mod inference;
425
426#[cfg(test)]
427mod tests {
428 use super::*;
429
430 #[test]
431 fn inference_implication_ignores_cast_origin_but_preserves_cast_types() {
432 let predicate = |implicit, ty: &str, bound| {
433 inference_identity(Expr::Binary {
434 op: BinaryOp::Greater,
435 lhs: Box::new(Expr::Cast {
436 implicit,
437 expr: Box::new(Expr::Column("value".into())),
438 ty: ty.into(),
439 }),
440 rhs: Box::new(Expr::Literal(Value::Int(bound))),
441 })
442 };
443 for implicit in [false, true] {
444 let required = predicate(implicit, "bigint", 0);
445 let equivalent = predicate(!implicit, "bigint", 0);
446 assert_eq!(required, equivalent);
447 assert!(implies(&equivalent, &required));
448 assert!(implies(&predicate(!implicit, "bigint", 1), &required));
449 assert!(!implies(&predicate(!implicit, "bigint", -1), &required));
450 assert!(!implies(&predicate(!implicit, "integer", 0), &required));
451 }
452 }
453}