use crate::{
ast::{BinaryOp, ColumnDef, TableConstraintSet},
binding::snapshot::BindingSnapshot,
catalog::index::EnforcedKey,
plan::AggregateClassifier,
plan::{ConflictPlan, ExpressionPlan, InsertPlan},
routines::RoutineResolution,
RowSchema, SQLError, SQLParam, ScalarExpr as Expr,
};
use std::collections::BTreeSet;
use uqa_core::Value;
pub trait ConflictCatalog {
fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
fn enforced_keys(&self, table: &str) -> Result<Vec<EnforcedKey>, String>;
fn try_declared_table_constraints(&self, table: &str) -> Result<TableConstraintSet, String>;
}
pub trait InferenceBindingScope {
fn binding_scope(&self) -> Result<BindingSnapshot, SQLError>;
}
#[derive(Clone, Copy)]
pub struct InferenceContext<'a> {
pub catalog: &'a dyn ConflictCatalog,
pub aggregates: &'a dyn AggregateClassifier,
pub routines: &'a dyn RoutineResolution,
pub binding: &'a dyn InferenceBindingScope,
}
pub fn prepare_inference_predicate<'a>(
context: InferenceContext<'_>,
statement: &'a InsertPlan,
params: &[SQLParam],
) -> Result<std::borrow::Cow<'a, InsertPlan>, SQLError> {
if statement
.on_conflict
.as_ref()
.is_none_or(|conflict| conflict.predicate.is_none() && conflict.expressions.is_empty())
{
return Ok(std::borrow::Cow::Borrowed(statement));
}
let columns = context
.catalog
.try_describe_table(&statement.table)
.map_err(|error| SQLError::Internal(error.to_string()))?
.ok_or_else(|| SQLError::UnknownTable(statement.table.clone()))?;
let schema = RowSchema::with_qualified_types(
&statement.target_qualifier,
columns.iter().map(|column| column.name.clone()).collect(),
columns
.iter()
.map(|column| Some(column.ty.clone()))
.collect(),
);
let mut statement = statement.clone();
if let Some(conflict) = &mut statement.on_conflict {
for expression in conflict
.expressions
.iter_mut()
.chain(conflict.predicate.iter_mut().map(Box::as_mut))
{
prepare_inference_expression(
context,
expression,
&statement.target_qualifier,
&schema,
&columns,
params,
)?;
}
let expressions = std::mem::take(&mut conflict.expressions);
for expression in expressions {
if let Expr::Column(column) = expression {
conflict.conflict_columns.push(column);
} else {
conflict.expressions.push(expression);
}
}
}
Ok(std::borrow::Cow::Owned(statement))
}
fn prepare_inference_expression(
context: InferenceContext<'_>,
expression: &mut Expr,
qualifier: &str,
schema: &RowSchema,
columns: &[crate::ast::ColumnDef],
params: &[SQLParam],
) -> Result<(), SQLError> {
let mut has_subquery = false;
expression.visit(&mut |part| {
has_subquery |= matches!(
part,
Expr::ScalarSubquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. }
);
});
if has_subquery {
return Err(predicate_error(
"0A000",
"cannot use subquery in index inference",
));
}
if crate::semantics::aggregates::contains_aggregate(context.aggregates, expression) {
return Err(predicate_error(
"42803",
"aggregate functions are not allowed in index inference",
));
}
if crate::semantics::windows::expr_has_window(expression) {
return Err(predicate_error(
"42P20",
"window functions are not allowed in index inference",
));
}
let mut plan = ExpressionPlan {
scalar: expression.clone(),
subqueries: Vec::new(),
};
let binding = context.binding.binding_scope()?;
crate::binding::bind_expression_plan_routines_for_storage(
context.routines,
&mut plan,
params,
&binding.context(),
schema,
)?;
crate::plan::rewrite_scalar_expression(&mut plan.scalar, &mut |expression| {
if let Expr::QualifiedColumn {
qualifier: source,
column,
} = expression
{
if source == qualifier {
*expression = Expr::Column(column.clone());
}
}
});
crate::plan::rewrite_scalar_expression(&mut plan.scalar, &mut |expression| {
if let Expr::Cast { expr, ty } = expression {
if let Expr::Column(name) = expr.as_ref() {
if columns.iter().any(|column| {
column.name == *name
&& crate::ast::ColumnType::from_sql_name(ty).ok().as_ref()
== Some(&column.ty)
}) {
*expression = Expr::Column(name.clone());
}
}
}
});
*expression = plan.scalar;
Ok(())
}
fn predicate_error(sqlstate: &str, message: &str) -> SQLError {
SQLError::Routine {
sqlstate: sqlstate.into(),
message: message.into(),
}
}
pub fn conflict_key_indices(
catalog: &dyn ConflictCatalog,
table: &str,
keys: &[EnforcedKey],
conflict: &ConflictPlan,
) -> Result<Vec<usize>, SQLError> {
if let Some(name) = &conflict.constraint {
return constraint_target_index(catalog, table, keys, name);
}
if conflict.conflict_columns.is_empty() && conflict.expressions.is_empty() {
return Ok((0..keys.len()).collect());
}
validate_conflict_columns(catalog, table, &conflict.conflict_columns)?;
let target = conflict
.conflict_columns
.iter()
.map(String::as_str)
.collect::<BTreeSet<_>>();
let indexes = keys
.iter()
.enumerate()
.filter_map(|(index, key)| {
(key.columns
.iter()
.map(String::as_str)
.collect::<BTreeSet<_>>()
== target
&& {
let expressions = key
.keys
.iter()
.filter_map(|key| match key {
crate::ast::IndexKey::Expression(expr) => {
Some(ExpressionPlan::lower((**expr).clone()).scalar)
}
crate::ast::IndexKey::Column(_) => None,
})
.collect::<Vec<_>>();
expressions
.iter()
.all(|expr| conflict.expressions.contains(expr))
&& conflict
.expressions
.iter()
.all(|expr| expressions.contains(expr))
}
&& key.predicate.as_deref().is_none_or(|required| {
conflict.predicate.as_deref().is_some_and(|given| {
implies(given, &ExpressionPlan::lower(required.clone()).scalar)
})
}))
.then_some(index)
})
.collect::<Vec<_>>();
if indexes.is_empty() {
return Err(SQLError::Routine {
sqlstate: "42P10".into(),
message:
"there is no unique or exclusion constraint matching the ON CONFLICT specification"
.into(),
});
}
Ok(indexes)
}
fn implies(given: &Expr, required: &Expr) -> bool {
if given == required || matches!(required, Expr::Literal(Value::Bool(true))) {
return true;
}
if matches!(given, Expr::Literal(Value::Bool(false) | Value::Null)) {
return true;
}
if let Expr::And(parts) = required {
return parts.iter().all(|part| implies(given, part));
}
if let Expr::Or(parts) = given {
return parts.iter().all(|part| implies(part, required));
}
if let Expr::And(parts) = given {
if parts.iter().any(|part| implies(part, required)) {
return true;
}
}
if let Expr::Or(parts) = required {
return parts.iter().any(|part| implies(given, part));
}
if let Some((left, given_op, given_value)) = comparison(given) {
if let Some((right, required_op, required_value)) = comparison(required) {
return left == right
&& comparison_implies(given_op, given_value, required_op, required_value);
}
if let Expr::IsNull {
expr,
negated: true,
} = required
{
return left == expr.as_ref() && !matches!(given_value, Value::Null);
}
}
false
}
fn comparison(expr: &Expr) -> Option<(&Expr, BinaryOp, &Value)> {
let Expr::Binary { op, lhs, rhs } = expr else {
return None;
};
if !matches!(
op,
BinaryOp::Equal
| BinaryOp::NotEqual
| BinaryOp::Less
| BinaryOp::LessEqual
| BinaryOp::Greater
| BinaryOp::GreaterEqual
) {
return None;
}
if let Expr::Literal(value) = rhs.as_ref() {
return Some((lhs, *op, value));
}
if let Expr::Literal(value) = lhs.as_ref() {
let reversed = match op {
BinaryOp::Less => BinaryOp::Greater,
BinaryOp::LessEqual => BinaryOp::GreaterEqual,
BinaryOp::Greater => BinaryOp::Less,
BinaryOp::GreaterEqual => BinaryOp::LessEqual,
other => *other,
};
return Some((rhs, reversed, value));
}
None
}
fn comparison_implies(given: BinaryOp, left: &Value, required: BinaryOp, right: &Value) -> bool {
use BinaryOp::{Equal, Greater, GreaterEqual, Less, LessEqual, NotEqual};
if matches!(left, Value::Null) || matches!(right, Value::Null) {
return false;
}
if std::mem::discriminant(left) != std::mem::discriminant(right) {
return false;
}
let order = left.cmp(right);
match (given, required) {
(Equal, Equal) | (NotEqual, NotEqual) => order.is_eq(),
(Equal, NotEqual) => !order.is_eq(),
(Equal | GreaterEqual, Greater) | (GreaterEqual, NotEqual) => order.is_gt(),
(Equal | LessEqual, Less) | (LessEqual, NotEqual) => order.is_lt(),
(Equal | Greater | GreaterEqual, GreaterEqual) | (Greater, Greater | NotEqual) => {
order.is_ge()
}
(Equal | Less | LessEqual, LessEqual) | (Less, Less | NotEqual) => order.is_le(),
_ => false,
}
}
pub fn validate_conflict_target(
catalog: &dyn ConflictCatalog,
table: &str,
conflict: &ConflictPlan,
) -> Result<(), SQLError> {
let keys = catalog
.enforced_keys(table)
.map_err(|error| SQLError::Internal(format!("conflict target keys: {error}")))?;
conflict_key_indices(catalog, table, &keys, conflict).map(|_| ())
}
fn constraint_target_index(
catalog: &dyn ConflictCatalog,
table: &str,
keys: &[EnforcedKey],
name: &str,
) -> Result<Vec<usize>, SQLError> {
if let Some(index) = keys
.iter()
.position(|key| key.constraint_owned && key.name.as_deref() == Some(name))
{
return Ok(vec![index]);
}
let snapshot = catalog
.try_declared_table_constraints(table)
.map_err(|error| SQLError::Internal(error.to_string()))?;
let columns = catalog
.try_describe_table(table)
.map_err(|error| SQLError::Internal(error.to_string()))?
.ok_or_else(|| SQLError::UnknownTable(table.into()))?;
let exists = snapshot
.checks
.iter()
.any(|check| check.name.as_deref() == Some(name))
|| snapshot
.foreign_keys
.iter()
.any(|key| key.name.as_deref() == Some(name))
|| columns.iter().any(|column| {
column.not_null_name.as_deref() == Some(name)
|| column.check_name.as_deref() == Some(name)
|| column
.references
.as_ref()
.is_some_and(|reference| reference.name.as_deref() == Some(name))
});
if exists {
return Err(SQLError::Routine {
sqlstate: "42809".into(),
message: "constraint in ON CONFLICT clause has no associated index".into(),
});
}
Err(SQLError::Routine {
sqlstate: "42704".into(),
message: format!("constraint \"{name}\" for table \"{table}\" does not exist"),
})
}
fn validate_conflict_columns(
catalog: &dyn ConflictCatalog,
table: &str,
names: &[String],
) -> Result<(), SQLError> {
let columns = catalog
.try_describe_table(table)
.map_err(|error| SQLError::Internal(error.to_string()))?
.ok_or_else(|| SQLError::UnknownTable(table.into()))?;
for name in names {
if !columns.iter().any(|column| column.name == *name)
&& !matches!(
name.as_str(),
"ctid" | "tableoid" | "xmin" | "xmax" | "cmin" | "cmax"
)
{
return Err(SQLError::UnknownColumn(name.clone()));
}
}
Ok(())
}