use crate::ast::FunctionBinding;
use crate::plan::{
AccessPathPlan, CommandPlan, ComputePlan, ConflictActionPlan, CtePlan, DeletePlan, InsertPlan,
JoinExecutionStrategy, MergePlan, ProjectionPlan, QueryBlockPlan, QueryPlan, RelationalPlan,
SourcePlan, UnifiedPlan, UpdatePlan,
};
use crate::SQLError;
use crate::{RowSchema, ScalarExpr, ScalarFrameBound};
use crate::{binding::context::BindingContext, routines::RoutineResolution};
use uqa_core::Value;
pub struct CatalogRoutineContext<'a, 'q> {
pub routines: &'a dyn RoutineResolution,
pub binding: &'a BindingContext<'q>,
}
pub struct BoundStatementRoutines {
pub query: Option<QueryPlan>,
pub references: Vec<BoundRoutineReference>,
}
#[derive(Debug, Clone)]
pub struct BoundRoutineReference {
pub name: String,
pub binding: Option<FunctionBinding>,
}
struct CommandRoutineInputs {
ctes: Vec<CtePlan>,
source: Option<SourcePlan>,
expressions: Vec<ScalarExpr>,
subqueries: Vec<QueryPlan>,
outer: RowSchema,
}
pub fn bind_catalog_statement_routines(
context: &CatalogRoutineContext<'_, '_>,
plan: &UnifiedPlan,
) -> Result<BoundStatementRoutines, SQLError> {
let query = match plan {
UnifiedPlan::Query(query) => {
let mut query = (**query).clone();
mark_query_relations_bound(&mut query);
crate::binding::bind_query_plan_routines_for_storage(
context.routines,
&mut query,
&[],
context.binding,
None,
)?;
Some(query)
}
UnifiedPlan::Command(command) => bind_command_statement_routines(context, command)?,
};
let mut references = Vec::new();
if let Some(query) = &query {
collect_query_routine_references(query, &mut references)?;
}
Ok(BoundStatementRoutines { query, references })
}
pub fn mark_catalog_statement_relations_bound(plan: &mut UnifiedPlan) -> Result<(), SQLError> {
match plan {
UnifiedPlan::Query(query) => mark_query_relations_bound(query),
UnifiedPlan::Command(command) => match command.as_mut() {
CommandPlan::Insert(plan) => {
plan.relations_bound = true;
for cte in &mut plan.ctes {
mark_cte_relations_bound(&mut cte.body);
}
if let Some(source) = &mut plan.source {
mark_query_relations_bound(source);
}
for subquery in &mut plan.subqueries {
mark_query_relations_bound(subquery);
}
}
CommandPlan::Update(plan) => {
plan.relations_bound = true;
for cte in &mut plan.ctes {
mark_cte_relations_bound(&mut cte.body);
}
if let Some(source) = &mut plan.source {
mark_source_relations_bound(source);
}
for subquery in &mut plan.subqueries {
mark_query_relations_bound(subquery);
}
}
CommandPlan::Delete(plan) => {
plan.relations_bound = true;
for cte in &mut plan.ctes {
mark_cte_relations_bound(&mut cte.body);
}
if let Some(source) = &mut plan.source {
mark_source_relations_bound(source);
}
for subquery in &mut plan.subqueries {
mark_query_relations_bound(subquery);
}
}
CommandPlan::Notify { .. } => {}
CommandPlan::Merge(plan) => {
for cte in &mut plan.ctes {
mark_cte_relations_bound(&mut cte.body);
}
mark_source_relations_bound(&mut plan.source);
for subquery in &mut plan.subqueries {
mark_query_relations_bound(subquery);
}
}
_ => {
return Err(SQLError::Internal(
"catalog-owned statement lowered to an unsupported command".into(),
));
}
},
}
Ok(())
}
fn bind_command_statement_routines(
context: &CatalogRoutineContext<'_, '_>,
command: &CommandPlan,
) -> Result<Option<QueryPlan>, SQLError> {
let Some(inputs) = command_statement_routine_inputs(context, command)? else {
return Ok(None);
};
let projections = inputs
.expressions
.into_iter()
.filter(|expression| expression_contains_routine(expression, &inputs.subqueries))
.filter(|expression| !matches!(expression, ScalarExpr::Default))
.map(|expr| ProjectionPlan { expr, alias: None })
.chain(std::iter::once(ProjectionPlan {
expr: ScalarExpr::Literal(Value::Int(1)),
alias: None,
}))
.collect();
let mut query = QueryPlan {
relations_bound: true,
ctes: inputs.ctes,
root: RelationalPlan::QueryBlock(Box::new(QueryBlockPlan {
projections,
from: inputs.source,
r#where: None,
compute: ComputePlan::Project,
group_by: Vec::new(),
grouping_sets: Vec::new(),
group_distinct: false,
having: None,
order_by: Vec::new(),
limit: None,
with_ties: false,
offset: None,
distinct: false,
distinct_on: Vec::new(),
subqueries: inputs.subqueries,
access: AccessPathPlan::Row,
locking: Vec::new(),
})),
};
mark_query_relations_bound(&mut query);
crate::binding::bind_query_plan_routines_for_storage(
context.routines,
&mut query,
&[],
context.binding,
Some(&inputs.outer),
)?;
Ok(Some(query))
}
fn command_statement_routine_inputs(
context: &CatalogRoutineContext<'_, '_>,
command: &CommandPlan,
) -> Result<Option<CommandRoutineInputs>, SQLError> {
match command {
CommandPlan::Insert(plan) => insert_statement_routine_inputs(context, plan).map(Some),
CommandPlan::Update(plan) => update_statement_routine_inputs(context, plan).map(Some),
CommandPlan::Delete(plan) => delete_statement_routine_inputs(context, plan).map(Some),
CommandPlan::Merge(plan) => Ok(Some(merge_statement_routine_inputs(plan))),
CommandPlan::Notify { .. } => Ok(None),
_ => Err(SQLError::Internal(
"catalog-owned statement lowered to an unsupported command".into(),
)),
}
}
fn merge_statement_routine_inputs(plan: &MergePlan) -> CommandRoutineInputs {
let target = SourcePlan::Table {
bound_columns: None,
name: plan.target.clone(),
qualifier: plan.target_qualifier.clone(),
alias: plan.target_alias.clone(),
column_aliases: Vec::new(),
include_descendants: plan.include_descendants,
};
let source = SourcePlan::Join {
left: Box::new(target),
right: plan.source.clone(),
kind: crate::ast::JoinKind::Cross,
on: None,
using: None,
natural: false,
alias: None,
column_aliases: Vec::new(),
lateral: false,
strategy: JoinExecutionStrategy::default(),
};
let mut expressions = vec![plan.join_condition.clone()];
for clause in &plan.when_clauses {
match clause {
crate::plan::MergeWhenPlan::UpdateMatched {
condition,
assignments,
}
| crate::plan::MergeWhenPlan::UpdateNotMatchedBySource {
condition,
assignments,
} => {
expressions.extend(condition.iter().cloned());
expressions.extend(
assignments
.iter()
.map(|assignment| assignment.value.clone()),
);
}
crate::plan::MergeWhenPlan::InsertNotMatched {
condition, values, ..
} => {
expressions.extend(condition.iter().cloned());
expressions.extend(values.iter().cloned());
}
crate::plan::MergeWhenPlan::DeleteMatched { condition }
| crate::plan::MergeWhenPlan::DeleteNotMatchedBySource { condition }
| crate::plan::MergeWhenPlan::NothingMatched { condition }
| crate::plan::MergeWhenPlan::NothingNotMatched { condition }
| crate::plan::MergeWhenPlan::NothingNotMatchedBySource { condition } => {
expressions.extend(condition.iter().cloned());
}
}
}
expressions.extend(
plan.returning
.iter()
.map(|projection| projection.expr.clone()),
);
CommandRoutineInputs {
ctes: plan.ctes.clone(),
source: Some(source),
expressions,
subqueries: plan.subqueries.clone(),
outer: RowSchema::default(),
}
}
fn insert_statement_routine_inputs(
context: &CatalogRoutineContext<'_, '_>,
plan: &InsertPlan,
) -> Result<CommandRoutineInputs, SQLError> {
let mut expressions = plan.rows.iter().flatten().cloned().collect::<Vec<_>>();
if let Some(conflict) = &plan.on_conflict {
expressions.extend(conflict.expressions.iter().cloned());
expressions.extend(conflict.predicate.iter().map(Box::as_ref).cloned());
if let ConflictActionPlan::Update {
assignments,
predicate,
} = &conflict.action
{
expressions.extend(
assignments
.iter()
.map(|assignment| assignment.value.clone()),
);
expressions.extend(predicate.iter().map(Box::as_ref).cloned());
}
}
expressions.extend(
plan.returning
.iter()
.map(|projection| projection.expr.clone()),
);
let source = plan.source.as_ref().map(|source| SourcePlan::Subquery {
body: Box::new((**source).clone()),
alias: Some("__uqa_catalog_statement_source".into()),
column_aliases: Vec::new(),
});
Ok(CommandRoutineInputs {
ctes: plan.ctes.clone(),
source,
expressions,
subqueries: plan.subqueries.clone(),
outer: statement_target_outer_schema(
context,
&plan.table,
&plan.target_qualifier,
&plan.returning_aliases,
)?,
})
}
fn update_statement_routine_inputs(
context: &CatalogRoutineContext<'_, '_>,
plan: &UpdatePlan,
) -> Result<CommandRoutineInputs, SQLError> {
let mut expressions = plan
.assignments
.iter()
.map(|assignment| assignment.value.clone())
.collect::<Vec<_>>();
expressions.extend(plan.predicate.iter().cloned());
expressions.extend(
plan.returning
.iter()
.map(|projection| projection.expr.clone()),
);
Ok(CommandRoutineInputs {
ctes: plan.ctes.clone(),
source: plan.source.as_deref().cloned(),
expressions,
subqueries: plan.subqueries.clone(),
outer: statement_target_outer_schema(
context,
&plan.table,
&plan.target_qualifier,
&plan.returning_aliases,
)?,
})
}
fn delete_statement_routine_inputs(
context: &CatalogRoutineContext<'_, '_>,
plan: &DeletePlan,
) -> Result<CommandRoutineInputs, SQLError> {
let mut expressions = plan.predicate.iter().cloned().collect::<Vec<_>>();
expressions.extend(
plan.returning
.iter()
.map(|projection| projection.expr.clone()),
);
Ok(CommandRoutineInputs {
ctes: plan.ctes.clone(),
source: plan.source.as_deref().cloned(),
expressions,
subqueries: plan.subqueries.clone(),
outer: statement_target_outer_schema(
context,
&plan.table,
&plan.target_qualifier,
&plan.returning_aliases,
)?,
})
}
fn expression_contains_routine(expression: &ScalarExpr, subqueries: &[QueryPlan]) -> bool {
let mut references = Vec::new();
collect_scalar_routine_references(expression, subqueries, &mut references).is_ok()
&& !references.is_empty()
}
fn statement_target_outer_schema(
context: &CatalogRoutineContext<'_, '_>,
table: &str,
target_qualifier: &str,
aliases: &crate::ast::ReturningAliases,
) -> Result<RowSchema, SQLError> {
let target = crate::binding::analyze_source_plan_schema(
context.routines,
&SourcePlan::Table {
bound_columns: None,
name: table.to_string(),
qualifier: target_qualifier.to_string(),
alias: None,
column_aliases: Vec::new(),
include_descendants: true,
},
&[],
context.binding,
None,
)?;
let target = RowSchema::with_types(target.columns().to_vec(), target.column_types().to_vec());
Ok(crate::semantics::returning_expression_schema(
&target,
target_qualifier,
aliases,
None,
))
}
pub fn collect_expression_routine_references(
expression: &crate::plan::ExpressionPlan,
) -> Result<Vec<BoundRoutineReference>, SQLError> {
let mut references = Vec::new();
collect_scalar_routine_references(&expression.scalar, &expression.subqueries, &mut references)?;
Ok(references)
}
fn collect_query_routine_references(
query: &QueryPlan,
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
for cte in &query.ctes {
collect_cte_routine_references(&cte.body, references)?;
if let Some(cycle) = &cte.cycle {
collect_scalar_routine_references(&cycle.mark_value, &[], references)?;
collect_scalar_routine_references(&cycle.mark_default, &[], references)?;
}
}
match &query.root {
RelationalPlan::QueryBlock(block) => {
if let Some(source) = &block.from {
collect_source_routine_references(source, &block.subqueries, references)?;
}
for projection in &block.projections {
collect_scalar_routine_references(&projection.expr, &block.subqueries, references)?;
}
if let Some(expression) = &block.r#where {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
for expression in &block.group_by {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
for expression in block.grouping_sets.iter().flatten() {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
if let Some(expression) = &block.having {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
for order in &block.order_by {
collect_scalar_routine_references(&order.expr, &block.subqueries, references)?;
}
if let Some(expression) = &block.limit {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
if let Some(expression) = &block.offset {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
for expression in &block.distinct_on {
collect_scalar_routine_references(expression, &block.subqueries, references)?;
}
}
RelationalPlan::SetOp {
left,
right,
order_by,
limit,
offset,
subqueries,
..
} => {
collect_query_routine_references(left, references)?;
collect_query_routine_references(right, references)?;
for order in order_by {
collect_scalar_routine_references(&order.expr, subqueries, references)?;
}
if let Some(expression) = limit {
collect_scalar_routine_references(expression, subqueries, references)?;
}
if let Some(expression) = offset {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
RelationalPlan::Values { rows, subqueries } => {
for expression in rows.iter().flatten() {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
}
Ok(())
}
fn collect_source_routine_references(
source: &SourcePlan,
subqueries: &[QueryPlan],
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
match source {
SourcePlan::Table { .. } => {}
SourcePlan::Join {
left, right, on, ..
} => {
collect_source_routine_references(left, subqueries, references)?;
collect_source_routine_references(right, subqueries, references)?;
if let Some(expression) = on {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
SourcePlan::Values { rows, .. } => {
for expression in rows.iter().flatten() {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
SourcePlan::Function {
name,
binding,
args,
..
} => {
references.push(BoundRoutineReference {
name: name.clone(),
binding: binding.clone(),
});
for expression in args {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
SourcePlan::FunctionGroup { functions, .. } => {
for function in functions {
references.push(BoundRoutineReference {
name: function.name.clone(),
binding: function.binding.clone(),
});
for expression in &function.args {
collect_scalar_routine_references(expression, subqueries, references)?;
}
}
}
SourcePlan::Subquery { body, .. } => {
collect_query_routine_references(body, references)?;
}
}
Ok(())
}
fn collect_scalar_routine_references(
expression: &ScalarExpr,
subqueries: &[QueryPlan],
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
match expression {
ScalarExpr::Func {
name,
binding,
args,
order_by,
filter,
..
} => {
for argument in args {
collect_scalar_routine_references(argument, subqueries, references)?;
}
for order in order_by {
collect_scalar_routine_references(&order.expr, subqueries, references)?;
}
if let Some(filter) = filter {
collect_scalar_routine_references(filter, subqueries, references)?;
}
references.push(BoundRoutineReference {
name: name.clone(),
binding: binding.clone(),
});
}
ScalarExpr::Array(items)
| ScalarExpr::Row(items)
| ScalarExpr::And(items)
| ScalarExpr::Or(items) => collect_many_routine_references(items, subqueries, references)?,
ScalarExpr::Binary { lhs, rhs, .. } => {
collect_scalar_routine_references(lhs, subqueries, references)?;
collect_scalar_routine_references(rhs, subqueries, references)?;
}
ScalarExpr::UnaryMinus(inner)
| ScalarExpr::Not(inner)
| ScalarExpr::IsNull { expr: inner, .. }
| ScalarExpr::Cast { expr: inner, .. } => {
collect_scalar_routine_references(inner, subqueries, references)?;
}
ScalarExpr::Between { expr, low, high } => {
collect_scalar_routine_references(expr, subqueries, references)?;
collect_scalar_routine_references(low, subqueries, references)?;
collect_scalar_routine_references(high, subqueries, references)?;
}
ScalarExpr::InList { expr, list, .. } => {
collect_scalar_routine_references(expr, subqueries, references)?;
for item in list {
collect_scalar_routine_references(item, subqueries, references)?;
}
}
ScalarExpr::WindowCall { name, args, spec } => {
collect_window_routine_references(name, args, spec, subqueries, references)?;
}
ScalarExpr::Case {
base,
when,
else_branch,
} => collect_case_routine_references(
base.as_deref(),
when,
else_branch.as_deref(),
subqueries,
references,
)?,
ScalarExpr::ScalarSubquery(index)
| ScalarExpr::Exists {
subquery: index, ..
} => {
let query = subqueries.get(*index).ok_or_else(|| {
SQLError::Internal(format!(
"stored catalog routine binding cannot resolve subquery slot {index}"
))
})?;
collect_query_routine_references(query, references)?;
}
ScalarExpr::InSubquery {
expr,
subquery: index,
..
} => {
collect_scalar_routine_references(expr, subqueries, references)?;
let query = subqueries.get(*index).ok_or_else(|| {
SQLError::Internal(format!(
"stored catalog routine binding cannot resolve subquery slot {index}"
))
})?;
collect_query_routine_references(query, references)?;
}
ScalarExpr::Star
| ScalarExpr::QualifiedStar(_)
| ScalarExpr::Default
| ScalarExpr::Column(_)
| ScalarExpr::Position(_)
| ScalarExpr::InternalColumn(_)
| ScalarExpr::QualifiedColumn { .. }
| ScalarExpr::Literal(_)
| ScalarExpr::TypedLiteral { .. }
| ScalarExpr::Param(_) => {}
}
Ok(())
}
fn collect_many_routine_references(
expressions: &[ScalarExpr],
subqueries: &[QueryPlan],
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
for expression in expressions {
collect_scalar_routine_references(expression, subqueries, references)?;
}
Ok(())
}
fn collect_case_routine_references(
base: Option<&ScalarExpr>,
when: &[(ScalarExpr, ScalarExpr)],
else_branch: Option<&ScalarExpr>,
subqueries: &[QueryPlan],
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
if let Some(base) = base {
collect_scalar_routine_references(base, subqueries, references)?;
}
for (condition, result) in when {
collect_scalar_routine_references(condition, subqueries, references)?;
collect_scalar_routine_references(result, subqueries, references)?;
}
if let Some(branch) = else_branch {
collect_scalar_routine_references(branch, subqueries, references)?;
}
Ok(())
}
fn collect_window_routine_references(
name: &str,
args: &[ScalarExpr],
spec: &crate::ScalarWindowSpec,
subqueries: &[QueryPlan],
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
for argument in args {
collect_scalar_routine_references(argument, subqueries, references)?;
}
for expression in &spec.partition_by {
collect_scalar_routine_references(expression, subqueries, references)?;
}
for order in &spec.order_by {
collect_scalar_routine_references(&order.expr, subqueries, references)?;
}
if let Some(frame) = &spec.frame {
for bound in [&frame.start, &frame.end] {
if let ScalarFrameBound::Preceding(inner) | ScalarFrameBound::Following(inner) = bound {
collect_scalar_routine_references(inner, subqueries, references)?;
}
}
}
references.push(BoundRoutineReference {
name: name.to_string(),
binding: None,
});
Ok(())
}
fn mark_cte_relations_bound(body: &mut crate::plan::CtePlanBody) {
match body {
crate::plan::CtePlanBody::Query(query) => mark_query_relations_bound(query),
crate::plan::CtePlanBody::Command(command) => {
match command.as_mut() {
CommandPlan::Insert(plan) => {
plan.relations_bound = true;
plan.target_relation_bound = true;
}
CommandPlan::Update(plan) => {
plan.relations_bound = true;
plan.target_relation_bound = true;
}
CommandPlan::Delete(plan) => {
plan.relations_bound = true;
plan.target_relation_bound = true;
}
_ => {}
}
if let Some(ctes) = command.ctes_mut() {
for cte in ctes {
mark_cte_relations_bound(&mut cte.body);
}
}
for query in command.query_inputs_mut() {
mark_query_relations_bound(query);
}
if let Some(source) = command.source_input_mut() {
mark_source_relations_bound(source);
}
}
}
}
fn collect_cte_routine_references(
body: &crate::plan::CtePlanBody,
references: &mut Vec<BoundRoutineReference>,
) -> Result<(), SQLError> {
match body {
crate::plan::CtePlanBody::Query(query) => {
collect_query_routine_references(query, references)
}
crate::plan::CtePlanBody::Command(command) => {
for cte in command.ctes() {
collect_cte_routine_references(&cte.body, references)?;
if let Some(cycle) = &cte.cycle {
collect_scalar_routine_references(&cycle.mark_value, &[], references)?;
collect_scalar_routine_references(&cycle.mark_default, &[], references)?;
}
}
for query in command.query_inputs() {
collect_query_routine_references(query, references)?;
}
if let Some(source) = command.source_input() {
collect_source_routine_references(source, command.scalar_subqueries(), references)?;
}
for expression in command.expressions() {
collect_scalar_routine_references(
expression,
command.scalar_subqueries(),
references,
)?;
}
Ok(())
}
}
}
fn mark_query_relations_bound(query: &mut QueryPlan) {
query.relations_bound = true;
for cte in &mut query.ctes {
mark_cte_relations_bound(&mut cte.body);
}
match &mut query.root {
RelationalPlan::QueryBlock(block) => {
if let Some(source) = &mut block.from {
mark_source_relations_bound(source);
}
for subquery in &mut block.subqueries {
mark_query_relations_bound(subquery);
}
}
RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
mark_query_relations_bound(left);
mark_query_relations_bound(right);
for subquery in subqueries {
mark_query_relations_bound(subquery);
}
}
RelationalPlan::Values { subqueries, .. } => {
for subquery in subqueries {
mark_query_relations_bound(subquery);
}
}
}
}
fn mark_source_relations_bound(source: &mut SourcePlan) {
match source {
SourcePlan::Join { left, right, .. } => {
mark_source_relations_bound(left);
mark_source_relations_bound(right);
}
SourcePlan::Subquery { body, .. } => mark_query_relations_bound(body),
SourcePlan::Table { .. }
| SourcePlan::Values { .. }
| SourcePlan::Function { .. }
| SourcePlan::FunctionGroup { .. } => {}
}
}
pub mod analysis;