uqa-sql 0.3.5

PostgreSQL-compatible SQL compiler built on libpg_query
Documentation
//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Relation and routine dependencies of stored query plans, including mutation CTEs.

mod ctes;
use super::{RuleDependencies, RuleRoutineDependency};
use crate::{
    plan::{QueryPlan, RelationalPlan, SourcePlan},
    SQLError,
};
use ctes::{collect_cte_relation_dependencies, collect_cte_source_routine_dependencies};
use std::collections::BTreeSet;
use uqa_core::RelationIdentity;

pub fn collect_query_relation_dependencies(
    query: &QueryPlan,
    dependencies: &mut RuleDependencies,
    inherited_ctes: &BTreeSet<String>,
) -> Result<(), SQLError> {
    let mut visible_ctes = inherited_ctes.clone();
    let recursive = query.ctes.iter().any(|cte| cte.recursive).then(|| {
        query
            .ctes
            .iter()
            .map(|cte| cte.name.clone())
            .collect::<BTreeSet<_>>()
    });
    for cte in &query.ctes {
        let body_scope = recursive.as_ref().map_or_else(
            || visible_ctes.clone(),
            |recursive| inherited_ctes.union(recursive).cloned().collect(),
        );
        collect_cte_relation_dependencies(&cte.body, dependencies, &body_scope)?;
        visible_ctes.insert(cte.name.clone());
    }
    match &query.root {
        RelationalPlan::QueryBlock(block) => {
            if let Some(source) = &block.from {
                collect_source_relation_dependencies(source, dependencies, &visible_ctes)?;
            }
            for subquery in &block.subqueries {
                collect_query_relation_dependencies(subquery, dependencies, &visible_ctes)?;
            }
        }
        RelationalPlan::SetOp {
            left,
            right,
            subqueries,
            ..
        } => {
            collect_query_relation_dependencies(left, dependencies, &visible_ctes)?;
            collect_query_relation_dependencies(right, dependencies, &visible_ctes)?;
            for subquery in subqueries {
                collect_query_relation_dependencies(subquery, dependencies, &visible_ctes)?;
            }
        }
        RelationalPlan::Values { subqueries, .. } => {
            for subquery in subqueries {
                collect_query_relation_dependencies(subquery, dependencies, &visible_ctes)?;
            }
        }
    }
    Ok(())
}

pub fn collect_expression_routine_dependencies(
    expression: &crate::plan::ExpressionPlan,
    dependencies: &mut RuleDependencies,
) {
    let mut scalar = expression.scalar.clone();
    crate::plan::rewrite_scalar_expression(&mut scalar, &mut |expression| {
        if let crate::ScalarExpr::Func {
            binding: Some(binding),
            ..
        } = expression
        {
            insert_routine_dependency(binding, dependencies);
        }
    });
    for query in &expression.subqueries {
        collect_query_routine_dependencies(query, dependencies);
    }
}

pub fn collect_query_routine_dependencies(query: &QueryPlan, dependencies: &mut RuleDependencies) {
    let mut scalar_plan = query.clone();
    scalar_plan.rewrite_scalar_expressions(&mut |expression| {
        if let crate::ScalarExpr::Func {
            binding: Some(binding),
            ..
        } = expression
        {
            insert_routine_dependency(binding, dependencies);
        }
    });
    for cte in &query.ctes {
        collect_cte_source_routine_dependencies(&cte.body, dependencies);
    }
    collect_relational_source_routine_dependencies(&query.root, dependencies);
}

fn collect_query_source_routine_dependencies(
    query: &QueryPlan,
    dependencies: &mut RuleDependencies,
) {
    for cte in &query.ctes {
        collect_cte_source_routine_dependencies(&cte.body, dependencies);
    }
    collect_relational_source_routine_dependencies(&query.root, dependencies);
}

fn collect_relational_source_routine_dependencies(
    plan: &RelationalPlan,
    dependencies: &mut RuleDependencies,
) {
    match plan {
        RelationalPlan::QueryBlock(block) => {
            if let Some(source) = &block.from {
                collect_source_routine_dependencies(source, dependencies);
            }
            for subquery in &block.subqueries {
                collect_query_source_routine_dependencies(subquery, dependencies);
            }
        }
        RelationalPlan::SetOp {
            left,
            right,
            subqueries,
            ..
        } => {
            collect_query_source_routine_dependencies(left, dependencies);
            collect_query_source_routine_dependencies(right, dependencies);
            for subquery in subqueries {
                collect_query_source_routine_dependencies(subquery, dependencies);
            }
        }
        RelationalPlan::Values { subqueries, .. } => {
            for subquery in subqueries {
                collect_query_source_routine_dependencies(subquery, dependencies);
            }
        }
    }
}

fn collect_source_routine_dependencies(source: &SourcePlan, dependencies: &mut RuleDependencies) {
    match source {
        SourcePlan::Table { .. } | SourcePlan::Values { .. } => {}
        SourcePlan::Join { left, right, .. } => {
            collect_source_routine_dependencies(left, dependencies);
            collect_source_routine_dependencies(right, dependencies);
        }
        SourcePlan::Subquery { body, .. } => {
            collect_query_source_routine_dependencies(body, dependencies);
        }
        SourcePlan::Function { binding, .. } => {
            if let Some(binding) = binding {
                insert_routine_dependency(binding, dependencies);
            }
        }
        SourcePlan::FunctionGroup { functions, .. } => {
            for function in functions {
                if let Some(binding) = &function.binding {
                    insert_routine_dependency(binding, dependencies);
                }
            }
        }
    }
}

fn insert_routine_dependency(
    binding: &crate::ast::FunctionBinding,
    dependencies: &mut RuleDependencies,
) {
    if !binding.builtin {
        dependencies.routines.insert(RuleRoutineDependency {
            object_id: binding.object_id,
            name: binding.name.clone(),
            argument_types: binding.argument_types.clone(),
        });
    }
}

fn collect_source_relation_dependencies(
    source: &SourcePlan,
    dependencies: &mut RuleDependencies,
    visible_ctes: &BTreeSet<String>,
) -> Result<(), SQLError> {
    match source {
        SourcePlan::Table { name, .. } => {
            if crate::binding::view_dependencies::canonical_virtual_relation_reference(name)
                .is_some()
            {
                return Ok(());
            }
            let (schema, relation) = RelationIdentity::parse_reference(name).map_err(|error| {
                SQLError::Internal(format!("decode stored rule dependency `{name}`: {error}"))
            })?;
            if schema.is_none() && visible_ctes.contains(&relation) {
                return Ok(());
            }
            let schema = schema.ok_or_else(|| {
                SQLError::Internal(format!(
                    "stored rule relation dependency `{name}` is not catalog-bound"
                ))
            })?;
            dependencies
                .relations
                .insert(RelationIdentity::new(schema, relation));
        }
        SourcePlan::Join { left, right, .. } => {
            collect_source_relation_dependencies(left, dependencies, visible_ctes)?;
            collect_source_relation_dependencies(right, dependencies, visible_ctes)?;
        }
        SourcePlan::Subquery { body, .. } => {
            collect_query_relation_dependencies(body, dependencies, visible_ctes)?;
        }
        SourcePlan::Function { relations, .. } => {
            if let Some(relations) = relations {
                collect_canonical_relation(&relations.left, dependencies)?;
                collect_canonical_relation(&relations.right, dependencies)?;
            }
        }
        SourcePlan::FunctionGroup { functions, .. } => {
            for function in functions {
                if let Some(relations) = &function.relations {
                    collect_canonical_relation(&relations.left, dependencies)?;
                    collect_canonical_relation(&relations.right, dependencies)?;
                }
            }
        }
        SourcePlan::Values { .. } => {}
    }
    Ok(())
}

fn collect_canonical_relation(
    reference: &str,
    dependencies: &mut RuleDependencies,
) -> Result<(), SQLError> {
    let relation = RelationIdentity::from_legacy_name(reference).map_err(|error| {
        SQLError::Internal(format!(
            "decode stored rule dependency `{reference}`: {error}"
        ))
    })?;
    dependencies.relations.insert(relation);
    Ok(())
}