use crate::ast::FunctionBinding;
use crate::{
plan::{QueryPlan, RelationalPlan, SourcePlan},
ScalarExpr,
};
use uqa_core::{RelationIdentity, Value};
pub fn canonical_virtual_relation_reference(reference: &str) -> Option<String> {
let (schema, relation) = RelationIdentity::parse_reference(reference).ok()?;
let relation = relation.to_ascii_lowercase();
let schema = schema.map(|schema| schema.to_ascii_lowercase());
let information_schema = matches!(
relation.as_str(),
"schemata"
| "tables"
| "columns"
| "column_privileges"
| "role_column_grants"
| "views"
| "routines"
| "sequences"
| "table_constraints"
| "key_column_usage"
);
let pg_catalog = matches!(
relation.as_str(),
"pg_namespace"
| "pg_class"
| "pg_inherits"
| "pg_partitioned_table"
| "pg_attribute"
| "pg_attrdef"
| "pg_constraint"
| "pg_index"
| "pg_tables"
| "pg_views"
| "pg_indexes"
| "pg_type"
| "pg_proc"
| "pg_database"
| "pg_roles"
| "pg_user"
| "pg_settings"
| "pg_description"
| "pg_matviews"
| "pg_sequences"
);
match schema.as_deref() {
Some("information_schema") if information_schema => {
Some(format!("information_schema.{relation}"))
}
Some("pg_catalog") | None if pg_catalog => Some(format!("pg_catalog.{relation}")),
_ => None,
}
}
pub fn sequence_function_reference_mut(expression: &mut ScalarExpr) -> Option<&mut String> {
let ScalarExpr::Func { name, args, .. } = expression else {
return None;
};
let lower = name.to_ascii_lowercase();
let local = lower.strip_prefix("pg_catalog.").unwrap_or(&lower);
if !matches!(local, "nextval" | "currval" | "setval")
|| (lower.contains('.') && !lower.starts_with("pg_catalog."))
{
return None;
}
regclass_literal_mut(args.first_mut()?)
}
pub fn regclass_literal_mut(expression: &mut ScalarExpr) -> Option<&mut String> {
match expression {
ScalarExpr::Literal(Value::Str(reference)) => Some(reference),
ScalarExpr::Cast { expr, ty }
if ty.eq_ignore_ascii_case("regclass")
|| ty.eq_ignore_ascii_case("pg_catalog.regclass") =>
{
regclass_literal_mut(expr)
}
_ => None,
}
}
pub fn bind_query_plan_sequence_references<E>(
plan: &mut QueryPlan,
resolve: &mut impl FnMut(&str) -> Result<String, E>,
) -> Result<(), E> {
let mut error = None;
plan.rewrite_scalar_expressions(&mut |expression| {
if error.is_some() {
return;
}
let Some(reference) = sequence_function_reference_mut(expression) else {
return;
};
match resolve(reference) {
Ok(canonical) => *reference = canonical,
Err(binding_error) => error = Some(binding_error),
}
});
error.map_or(Ok(()), Err)
}
pub fn bind_query_plan_relations<E>(
plan: &mut QueryPlan,
inherited_ctes: &std::collections::BTreeSet<String>,
resolve: &mut impl FnMut(&str) -> Result<String, E>,
) -> Result<(), E> {
let mut visible_ctes = inherited_ctes.clone();
let recursive_ctes = plan.ctes.iter().any(|cte| cte.recursive).then(|| {
plan.ctes
.iter()
.map(|cte| cte.name.clone())
.collect::<std::collections::BTreeSet<_>>()
});
for cte in &mut plan.ctes {
let body_ctes = recursive_ctes.as_ref().map_or_else(
|| visible_ctes.clone(),
|ctes| inherited_ctes.union(ctes).cloned().collect(),
);
bind_cte_plan_relations(&mut cte.body, &body_ctes, resolve)?;
visible_ctes.insert(cte.name.clone());
}
bind_relational_plan_relations(&mut plan.root, &visible_ctes, resolve)?;
plan.relations_bound = true;
Ok(())
}
pub fn bind_relational_plan_relations<E>(
plan: &mut RelationalPlan,
visible_ctes: &std::collections::BTreeSet<String>,
resolve: &mut impl FnMut(&str) -> Result<String, E>,
) -> Result<(), E> {
match plan {
RelationalPlan::QueryBlock(block) => {
if let Some(source) = &mut block.from {
bind_source_plan_relations(source, visible_ctes, resolve)?;
}
for subquery in &mut block.subqueries {
bind_query_plan_relations(subquery, visible_ctes, resolve)?;
}
}
RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
bind_query_plan_relations(left, visible_ctes, resolve)?;
bind_query_plan_relations(right, visible_ctes, resolve)?;
for subquery in subqueries {
bind_query_plan_relations(subquery, visible_ctes, resolve)?;
}
}
RelationalPlan::Values { subqueries, .. } => {
for subquery in subqueries {
bind_query_plan_relations(subquery, visible_ctes, resolve)?;
}
}
}
Ok(())
}
pub fn bind_source_plan_relations<E>(
source: &mut SourcePlan,
visible_ctes: &std::collections::BTreeSet<String>,
resolve: &mut impl FnMut(&str) -> Result<String, E>,
) -> Result<(), E> {
match source {
SourcePlan::Table {
name, qualifier, ..
} => {
if qualifier.is_empty() {
*qualifier = RelationIdentity::parse_reference(name)
.map_or_else(|_| name.clone(), |(_, relation)| relation);
}
let is_cte =
RelationIdentity::parse_reference(name)
.ok()
.is_some_and(|(schema, relation)| {
schema.is_none() && visible_ctes.contains(&relation)
});
if !is_cte {
*name = resolve(name)?;
}
}
SourcePlan::Join { left, right, .. } => {
bind_source_plan_relations(left, visible_ctes, resolve)?;
bind_source_plan_relations(right, visible_ctes, resolve)?;
}
SourcePlan::Subquery { body, .. } => {
bind_query_plan_relations(body, visible_ctes, resolve)?;
}
SourcePlan::Function {
name,
output_name,
relations,
..
} => {
if output_name.is_empty() {
*output_name = RelationIdentity::parse_reference(name)
.map_or_else(|_| name.clone(), |(_, function)| function);
}
if let Some(relations) = relations {
relations.left = resolve(&relations.left)?;
relations.right = resolve(&relations.right)?;
}
}
SourcePlan::FunctionGroup { functions, .. } => {
for function in functions {
if function.output_name.is_empty() {
function.output_name = RelationIdentity::parse_reference(&function.name)
.map_or_else(|_| function.name.clone(), |(_, name)| name);
}
if let Some(relations) = &mut function.relations {
relations.left = resolve(&relations.left)?;
relations.right = resolve(&relations.right)?;
}
}
}
SourcePlan::Values { .. } => {}
}
Ok(())
}
pub fn relation_reference_matches(reference: &str, target: &RelationIdentity) -> bool {
match RelationIdentity::parse_reference(reference) {
Ok((Some(schema), name)) => schema == target.schema && name == target.name,
Ok((None, name)) => name == target.name,
Err(_) => true,
}
}
pub fn source_plan_references_relation(
source: &crate::plan::SourcePlan,
target: &RelationIdentity,
ctes: &std::collections::BTreeSet<String>,
) -> bool {
match source {
crate::plan::SourcePlan::Table { name, .. } => {
let is_cte = RelationIdentity::parse_reference(name)
.ok()
.is_some_and(|(schema, relation)| schema.is_none() && ctes.contains(&relation));
!is_cte && relation_reference_matches(name, target)
}
crate::plan::SourcePlan::Join { left, right, .. } => {
source_plan_references_relation(left, target, ctes)
|| source_plan_references_relation(right, target, ctes)
}
crate::plan::SourcePlan::Subquery { body, .. } => {
query_plan_references_relation(body, target, ctes)
}
crate::plan::SourcePlan::Function { relations, .. } => {
relations.as_ref().is_some_and(|relations| {
relation_reference_matches(&relations.left, target)
|| relation_reference_matches(&relations.right, target)
})
}
crate::plan::SourcePlan::FunctionGroup { functions, .. } => {
functions.iter().any(|function| {
function.relations.as_ref().is_some_and(|relations| {
relation_reference_matches(&relations.left, target)
|| relation_reference_matches(&relations.right, target)
})
})
}
crate::plan::SourcePlan::Values { .. } => false,
}
}
pub fn query_plan_references_relation(
query: &crate::plan::QueryPlan,
target: &RelationIdentity,
inherited_ctes: &std::collections::BTreeSet<String>,
) -> bool {
let mut ctes = inherited_ctes.clone();
ctes.extend(query.ctes.iter().map(|cte| cte.name.clone()));
if query
.ctes
.iter()
.any(|cte| cte_plan_references_relation(&cte.body, target, &ctes))
{
return true;
}
match &query.root {
crate::plan::RelationalPlan::QueryBlock(block) => {
block
.from
.as_ref()
.is_some_and(|source| source_plan_references_relation(source, target, &ctes))
|| block
.subqueries
.iter()
.any(|query| query_plan_references_relation(query, target, &ctes))
}
crate::plan::RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
query_plan_references_relation(left, target, &ctes)
|| query_plan_references_relation(right, target, &ctes)
|| subqueries
.iter()
.any(|query| query_plan_references_relation(query, target, &ctes))
}
crate::plan::RelationalPlan::Values { subqueries, .. } => subqueries
.iter()
.any(|query| query_plan_references_relation(query, target, &ctes)),
}
}
pub fn query_plan_references_sequence(plan: &QueryPlan, target: &RelationIdentity) -> bool {
let mut plan = plan.clone();
let mut referenced = false;
plan.rewrite_scalar_expressions(&mut |expression| {
if let Some(reference) = sequence_function_reference_mut(expression) {
referenced |= relation_reference_matches(reference, target);
}
});
referenced
}
pub use crate::routines::function_binding_matches;
fn function_binding_needs_object_identity(binding: &FunctionBinding) -> bool {
!binding.builtin
&& binding.dispatch.is_none()
&& binding.resolution_error.is_none()
&& binding.object_id.is_none()
}
fn source_plan_has_legacy_routine_identity(source: &SourcePlan) -> bool {
match source {
SourcePlan::Table { .. } | SourcePlan::Values { .. } => false,
SourcePlan::Join { left, right, .. } => {
source_plan_has_legacy_routine_identity(left)
|| source_plan_has_legacy_routine_identity(right)
}
SourcePlan::Subquery { body, .. } => query_plan_sources_have_legacy_routine_identity(body),
SourcePlan::Function { binding, .. } => binding
.as_ref()
.is_some_and(function_binding_needs_object_identity),
SourcePlan::FunctionGroup { functions, .. } => functions.iter().any(|function| {
function
.binding
.as_ref()
.is_some_and(function_binding_needs_object_identity)
}),
}
}
fn relational_plan_has_legacy_routine_identity(plan: &RelationalPlan) -> bool {
match plan {
RelationalPlan::QueryBlock(block) => {
block
.from
.as_ref()
.is_some_and(source_plan_has_legacy_routine_identity)
|| block
.subqueries
.iter()
.any(query_plan_sources_have_legacy_routine_identity)
}
RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
query_plan_sources_have_legacy_routine_identity(left)
|| query_plan_sources_have_legacy_routine_identity(right)
|| subqueries
.iter()
.any(query_plan_sources_have_legacy_routine_identity)
}
RelationalPlan::Values { subqueries, .. } => subqueries
.iter()
.any(query_plan_sources_have_legacy_routine_identity),
}
}
fn query_plan_sources_have_legacy_routine_identity(plan: &QueryPlan) -> bool {
plan.ctes.iter().any(|cte| {
cte_relational_inputs_any(
&cte.body,
&query_plan_sources_have_legacy_routine_identity,
&source_plan_has_legacy_routine_identity,
)
}) || relational_plan_has_legacy_routine_identity(&plan.root)
}
pub fn query_plan_has_legacy_routine_identity(plan: &QueryPlan) -> bool {
let mut scalar_plan = plan.clone();
let mut legacy = false;
scalar_plan.rewrite_scalar_expressions(&mut |expression| {
if let ScalarExpr::Func {
binding: Some(binding),
..
} = expression
{
legacy |= function_binding_needs_object_identity(binding);
}
});
legacy || query_plan_sources_have_legacy_routine_identity(plan)
}
fn source_plan_references_function(source: &SourcePlan, target: &FunctionBinding) -> bool {
match source {
SourcePlan::Table { .. } | SourcePlan::Values { .. } => false,
SourcePlan::Join { left, right, .. } => {
source_plan_references_function(left, target)
|| source_plan_references_function(right, target)
}
SourcePlan::Subquery { body, .. } => query_plan_sources_reference_function(body, target),
SourcePlan::Function { binding, .. } => binding
.as_ref()
.is_some_and(|binding| function_binding_matches(binding, target)),
SourcePlan::FunctionGroup { functions, .. } => functions.iter().any(|function| {
function
.binding
.as_ref()
.is_some_and(|binding| function_binding_matches(binding, target))
}),
}
}
fn relational_plan_references_function(plan: &RelationalPlan, target: &FunctionBinding) -> bool {
match plan {
RelationalPlan::QueryBlock(block) => {
block
.from
.as_ref()
.is_some_and(|source| source_plan_references_function(source, target))
|| block
.subqueries
.iter()
.any(|query| query_plan_sources_reference_function(query, target))
}
RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
query_plan_sources_reference_function(left, target)
|| query_plan_sources_reference_function(right, target)
|| subqueries
.iter()
.any(|query| query_plan_sources_reference_function(query, target))
}
RelationalPlan::Values { subqueries, .. } => subqueries
.iter()
.any(|query| query_plan_sources_reference_function(query, target)),
}
}
fn query_plan_sources_reference_function(plan: &QueryPlan, target: &FunctionBinding) -> bool {
plan.ctes.iter().any(|cte| {
cte_relational_inputs_any(
&cte.body,
&|query| query_plan_sources_reference_function(query, target),
&|source| source_plan_references_function(source, target),
)
}) || relational_plan_references_function(&plan.root, target)
}
pub fn query_plan_references_function(plan: &QueryPlan, target: &FunctionBinding) -> bool {
let mut scalar_plan = plan.clone();
let mut referenced = false;
scalar_plan.rewrite_scalar_expressions(&mut |expression| {
if let ScalarExpr::Func {
binding: Some(binding),
..
} = expression
{
referenced |= function_binding_matches(binding, target);
}
});
referenced || query_plan_sources_reference_function(plan, target)
}
fn rewrite_source_plan_routine_identity(
source: &mut SourcePlan,
target: &FunctionBinding,
new_name: &str,
) -> bool {
match source {
SourcePlan::Table { .. } | SourcePlan::Values { .. } => false,
SourcePlan::Join { left, right, .. } => {
rewrite_source_plan_routine_identity(left, target, new_name)
| rewrite_source_plan_routine_identity(right, target, new_name)
}
SourcePlan::Subquery { body, .. } => {
rewrite_query_source_routine_identity(body, target, new_name)
}
SourcePlan::Function { name, binding, .. } => {
let Some(binding) = binding.as_mut() else {
return false;
};
if !function_binding_matches(binding, target) {
return false;
}
*name = new_name.to_string();
binding.name = new_name.to_string();
true
}
SourcePlan::FunctionGroup { functions, .. } => {
let mut changed = false;
for function in functions {
let Some(binding) = function.binding.as_mut() else {
continue;
};
if function_binding_matches(binding, target) {
function.name = new_name.to_string();
binding.name = new_name.to_string();
changed = true;
}
}
changed
}
}
}
fn rewrite_relational_plan_source_routine_identity(
plan: &mut RelationalPlan,
target: &FunctionBinding,
new_name: &str,
) -> bool {
match plan {
RelationalPlan::QueryBlock(block) => {
let mut changed = block.from.as_mut().is_some_and(|source| {
rewrite_source_plan_routine_identity(source, target, new_name)
});
for subquery in &mut block.subqueries {
changed |= rewrite_query_source_routine_identity(subquery, target, new_name);
}
changed
}
RelationalPlan::SetOp {
left,
right,
subqueries,
..
} => {
let mut changed = rewrite_query_source_routine_identity(left, target, new_name)
| rewrite_query_source_routine_identity(right, target, new_name);
for subquery in subqueries {
changed |= rewrite_query_source_routine_identity(subquery, target, new_name);
}
changed
}
RelationalPlan::Values { subqueries, .. } => {
let mut changed = false;
for subquery in subqueries {
changed |= rewrite_query_source_routine_identity(subquery, target, new_name);
}
changed
}
}
}
fn rewrite_query_source_routine_identity(
plan: &mut QueryPlan,
target: &FunctionBinding,
new_name: &str,
) -> bool {
let mut changed = false;
for cte in &mut plan.ctes {
changed |= rewrite_cte_source_routine_identity(&mut cte.body, target, new_name);
}
changed | rewrite_relational_plan_source_routine_identity(&mut plan.root, target, new_name)
}
pub fn rewrite_query_plan_routine_identity(
plan: &mut QueryPlan,
target: &FunctionBinding,
new_name: &str,
) -> bool {
let mut changed = false;
plan.rewrite_scalar_expressions(&mut |expression| {
let ScalarExpr::Func {
name,
binding: Some(binding),
..
} = expression
else {
return;
};
if function_binding_matches(binding, target) {
*name = new_name.to_string();
binding.name = new_name.to_string();
changed = true;
}
});
changed | rewrite_query_source_routine_identity(plan, target, new_name)
}
pub fn bind_cte_plan_relations<E>(
body: &mut crate::plan::CtePlanBody,
inherited: &std::collections::BTreeSet<String>,
resolve: &mut impl FnMut(&str) -> Result<String, E>,
) -> Result<(), E> {
let crate::plan::CtePlanBody::Command(command) = body else {
let crate::plan::CtePlanBody::Query(query) = body else {
unreachable!()
};
return bind_query_plan_relations(query, inherited, resolve);
};
if let Some(target) = command.mutation_target_mut() {
*target = resolve(target)?;
}
match command.as_mut() {
crate::plan::CommandPlan::Insert(plan) => {
plan.target_relation_bound = true;
plan.relations_bound = true;
}
crate::plan::CommandPlan::Update(plan) => {
plan.target_relation_bound = true;
plan.relations_bound = true;
}
crate::plan::CommandPlan::Delete(plan) => {
plan.target_relation_bound = true;
plan.relations_bound = true;
}
_ => {}
}
let mut visible = inherited.clone();
if let Some(ctes) = command.ctes_mut() {
let recursive = ctes.iter().any(|cte| cte.recursive).then(|| {
ctes.iter()
.map(|cte| cte.name.clone())
.collect::<std::collections::BTreeSet<_>>()
});
for cte in ctes {
let scope = recursive.as_ref().map_or_else(
|| visible.clone(),
|names| inherited.union(names).cloned().collect(),
);
bind_cte_plan_relations(&mut cte.body, &scope, resolve)?;
visible.insert(cte.name.clone());
}
}
if let Some(source) = command.source_input_mut() {
bind_source_plan_relations(source, &visible, resolve)?;
}
for query in command.query_inputs_mut() {
bind_query_plan_relations(query, &visible, resolve)?;
}
Ok(())
}
fn cte_plan_references_relation(
body: &crate::plan::CtePlanBody,
target: &RelationIdentity,
inherited: &std::collections::BTreeSet<String>,
) -> bool {
match body {
crate::plan::CtePlanBody::Query(query) => {
query_plan_references_relation(query, target, inherited)
}
crate::plan::CtePlanBody::Command(command) => {
let mut visible = inherited.clone();
visible.extend(command.ctes().iter().map(|cte| cte.name.clone()));
command
.mutation_target()
.is_some_and(|name| relation_reference_matches(name, target))
|| command
.ctes()
.iter()
.any(|cte| cte_plan_references_relation(&cte.body, target, &visible))
|| command
.source_input()
.is_some_and(|source| source_plan_references_relation(source, target, &visible))
|| command
.query_inputs()
.iter()
.any(|query| query_plan_references_relation(query, target, &visible))
}
}
}
fn cte_relational_inputs_any(
body: &crate::plan::CtePlanBody,
query: &dyn Fn(&QueryPlan) -> bool,
source: &dyn Fn(&crate::plan::SourcePlan) -> bool,
) -> bool {
match body {
crate::plan::CtePlanBody::Query(plan) => query(plan),
crate::plan::CtePlanBody::Command(command) => {
command
.ctes()
.iter()
.any(|cte| cte_relational_inputs_any(&cte.body, query, source))
|| command.query_inputs().iter().any(|plan| query(plan))
|| command.source_input().is_some_and(source)
}
}
}
fn rewrite_cte_source_routine_identity(
body: &mut crate::plan::CtePlanBody,
target: &FunctionBinding,
new_name: &str,
) -> bool {
match body {
crate::plan::CtePlanBody::Query(query) => {
rewrite_query_source_routine_identity(query, target, new_name)
}
crate::plan::CtePlanBody::Command(command) => {
let mut changed = false;
if let Some(ctes) = command.ctes_mut() {
for cte in ctes {
changed |= rewrite_cte_source_routine_identity(&mut cte.body, target, new_name);
}
}
if let Some(source) = command.source_input_mut() {
changed |= rewrite_source_plan_routine_identity(source, target, new_name);
}
for query in command.query_inputs_mut() {
changed |= rewrite_query_source_routine_identity(query, target, new_name);
}
changed
}
}
}
pub mod restoration;