uqa_sql/catalog/stored_ast/

mod.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Traversal and identity binding for durable SQL syntax trees.

use crate::{
    ast::{Expr, FrameBound, FromClause, SelectStmt, Statement},
    SQLError,
};
use std::collections::BTreeSet;
use uqa_core::RelationIdentity;
mod expressions;
mod merge;
mod routines;
mod sources;
mod types;
pub use expressions::*;
pub use merge::visit_stored_statement_merges;
pub use routines::*;
pub use sources::*;
pub use types::*;

pub type MergeCallback<'a> = &'a mut dyn FnMut(&mut crate::ast::MergeStmt) -> Result<(), SQLError>;

pub type ExpressionCallback<'a> = &'a mut dyn FnMut(&mut Expr) -> Result<(), SQLError>;
pub type SourceCallback<'a> = &'a mut dyn FnMut(&mut FromClause) -> Result<(), SQLError>;

pub struct StoredAstVisitor<'a, R, F> {
    pub source: Option<SourceCallback<'a>>,
    pub merge: Option<MergeCallback<'a>>,
    pub expression: Option<ExpressionCallback<'a>>,
    pub ty: Option<&'a mut dyn FnMut(&mut String)>,
    pub relation: &'a mut R,
    pub routine: &'a mut F,
}

impl<R, F> StoredAstVisitor<'_, R, F>
where
    R: FnMut(&mut String) -> Result<(), SQLError>,
    F: FnMut(&mut String, Option<&mut Option<crate::ast::FunctionBinding>>) -> Result<(), SQLError>,
{
    pub fn bind_statement(&mut self, statement: &mut Statement) -> Result<(), SQLError> {
        let ctes = BTreeSet::new();
        match statement {
            Statement::Select(query) => self.bind_select(query, &ctes),
            Statement::Insert(insert) => self.bind_insert(insert, &ctes),
            Statement::Update(update) => self.bind_update(update, &ctes),
            Statement::Delete(delete) => self.bind_delete(delete, &ctes),
            Statement::Notify { .. } => Ok(()),
            Statement::Values { rows } => {
                for expression in rows.iter_mut().flatten() {
                    self.bind_expr(expression, &ctes)?;
                }
                Ok(())
            }
            Statement::Merge(merge) => self.bind_merge(merge, &ctes),
            _ => Err(SQLError::Internal(
                "catalog-owned statement has an unsupported dependency shape".into(),
            )),
        }
    }

    fn bind_insert(
        &mut self,
        insert: &mut crate::ast::InsertStmt,
        inherited: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        (self.relation)(&mut insert.table)?;
        let visible = self.bind_ctes(&mut insert.with, inherited)?;
        if let Some(source) = insert.select_source.as_deref_mut() {
            self.bind_select(source, &visible)?;
        }
        for expression in insert
            .columns
            .iter_mut()
            .flat_map(crate::ast::AssignmentTarget::expressions_mut)
            .chain(insert.rows.iter_mut().flatten())
        {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(conflict) = &mut insert.on_conflict {
            for expression in &mut conflict.expressions {
                self.bind_expr(expression, &visible)?;
            }
            if let Some(predicate) = conflict.predicate.as_deref_mut() {
                self.bind_expr(predicate, &visible)?;
            }
            if let crate::ast::OnConflictAction::Update {
                assignments,
                r#where,
            } = &mut conflict.action
            {
                for expression in assignments.iter_mut().flat_map(|(target, value)| {
                    target.expressions_mut().chain(std::iter::once(value))
                }) {
                    self.bind_expr(expression, &visible)?;
                }
                if let Some(expression) = r#where {
                    self.bind_expr(expression, &visible)?;
                }
            }
        }
        for projection in &mut insert.returning {
            self.bind_expr(&mut projection.expr, &visible)?;
        }
        Ok(())
    }

    fn bind_update(
        &mut self,
        update: &mut crate::ast::UpdateStmt,
        inherited: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        (self.relation)(&mut update.table)?;
        let visible = self.bind_ctes(&mut update.with, inherited)?;
        if let Some(source) = &mut update.from {
            self.bind_from(source, &visible)?;
        }
        for expression in update
            .assignments
            .iter_mut()
            .flat_map(|(target, value)| target.expressions_mut().chain(std::iter::once(value)))
        {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(expression) = &mut update.r#where {
            self.bind_expr(expression, &visible)?;
        }
        for projection in &mut update.returning {
            self.bind_expr(&mut projection.expr, &visible)?;
        }
        Ok(())
    }

    fn bind_delete(
        &mut self,
        delete: &mut crate::ast::DeleteStmt,
        inherited: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        (self.relation)(&mut delete.table)?;
        let visible = self.bind_ctes(&mut delete.with, inherited)?;
        if let Some(source) = &mut delete.using {
            self.bind_from(source, &visible)?;
        }
        if let Some(expression) = &mut delete.r#where {
            self.bind_expr(expression, &visible)?;
        }
        for projection in &mut delete.returning {
            self.bind_expr(&mut projection.expr, &visible)?;
        }
        Ok(())
    }

    fn bind_ctes(
        &mut self,
        ctes: &mut [crate::ast::CTE],
        inherited: &BTreeSet<String>,
    ) -> Result<BTreeSet<String>, SQLError> {
        let mut visible = inherited.clone();
        let recursive = ctes.iter().any(|cte| cte.recursive).then(|| {
            ctes.iter()
                .map(|cte| cte.name.clone())
                .collect::<BTreeSet<_>>()
        });
        for cte in ctes {
            let body_scope = recursive.as_ref().map_or_else(
                || visible.clone(),
                |recursive| inherited.union(recursive).cloned().collect(),
            );
            match &mut cte.body {
                crate::ast::CteBody::Query(query) => self.bind_select(query, &body_scope)?,
                crate::ast::CteBody::Insert(plan) => self.bind_insert(plan, &body_scope)?,
                crate::ast::CteBody::Update(plan) => self.bind_update(plan, &body_scope)?,
                crate::ast::CteBody::Delete(plan) => self.bind_delete(plan, &body_scope)?,
                crate::ast::CteBody::Merge(plan) => self.bind_merge(plan, &body_scope)?,
            }
            if let Some(cycle) = &mut cte.cycle {
                self.bind_expr(&mut cycle.mark_value, &body_scope)?;
                self.bind_expr(&mut cycle.mark_default, &body_scope)?;
            }
            visible.insert(cte.name.clone());
        }
        Ok(visible)
    }

    fn bind_select(
        &mut self,
        select: &mut SelectStmt,
        inherited: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        let visible = self.bind_ctes(&mut select.with, inherited)?;
        if let Some(source) = &mut select.from {
            self.bind_from(source, &visible)?;
        }
        for projection in &mut select.projections {
            self.bind_expr(&mut projection.expr, &visible)?;
        }
        for expression in select.values.iter_mut().flatten() {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(expression) = &mut select.r#where {
            self.bind_expr(expression, &visible)?;
        }
        for expression in &mut select.group_by {
            self.bind_expr(expression, &visible)?;
        }
        for expression in select.grouping_sets.iter_mut().flatten() {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(expression) = &mut select.having {
            self.bind_expr(expression, &visible)?;
        }
        for order in &mut select.order_by {
            self.bind_expr(&mut order.expr, &visible)?;
        }
        if let Some(expression) = &mut select.limit {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(expression) = &mut select.offset {
            self.bind_expr(expression, &visible)?;
        }
        for expression in &mut select.distinct_on {
            self.bind_expr(expression, &visible)?;
        }
        if let Some(set) = &mut select.set_op {
            if let Some(left) = &mut set.left {
                self.bind_select(left, &visible)?;
            }
            self.bind_select(&mut set.right, &visible)?;
            for order in &mut set.combined_order_by {
                self.bind_expr(&mut order.expr, &visible)?;
            }
            if let Some(expression) = &mut set.combined_limit {
                self.bind_expr(expression, &visible)?;
            }
            if let Some(expression) = &mut set.combined_offset {
                self.bind_expr(expression, &visible)?;
            }
        }
        Ok(())
    }

    fn bind_from(
        &mut self,
        source: &mut FromClause,
        visible_ctes: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        if let FromClause::Table { name, .. } = source {
            let is_cte =
                RelationIdentity::parse_reference(name)
                    .ok()
                    .is_some_and(|(schema, relation)| {
                        schema.is_none() && visible_ctes.contains(&relation)
                    });
            if is_cte {
                return Ok(());
            }
        }
        if let Some(visit) = self.source.as_mut() {
            visit(source)?;
        }
        match source {
            FromClause::Table { name, .. } => {
                (self.relation)(name)?;
            }
            FromClause::Join {
                left, right, on, ..
            } => {
                self.bind_from(left, visible_ctes)?;
                self.bind_from(right, visible_ctes)?;
                if let Some(expression) = on {
                    self.bind_expr(expression, visible_ctes)?;
                }
            }
            FromClause::Values { rows, .. } => {
                for expression in rows.iter_mut().flatten() {
                    self.bind_expr(expression, visible_ctes)?;
                }
            }
            FromClause::Function {
                name,
                binding,
                relations,
                args,
                ..
            } => {
                (self.routine)(name, Some(binding))?;
                if let Some(relations) = relations {
                    (self.relation)(&mut relations.left)?;
                    (self.relation)(&mut relations.right)?;
                }
                for expression in args {
                    self.bind_expr(expression, visible_ctes)?;
                }
            }
            FromClause::FunctionGroup { functions, .. } => {
                for function in functions {
                    (self.routine)(&mut function.name, Some(&mut function.binding))?;
                    if let Some(relations) = &mut function.relations {
                        (self.relation)(&mut relations.left)?;
                        (self.relation)(&mut relations.right)?;
                    }
                    for expression in &mut function.args {
                        self.bind_expr(expression, visible_ctes)?;
                    }
                }
            }
            FromClause::Subquery { body, .. } => self.bind_select(body, visible_ctes)?,
        }
        Ok(())
    }

    fn bind_expression_type(&mut self, expression: &mut Expr) -> Result<(), SQLError> {
        if let Some(visit) = self.expression.as_mut() {
            visit(expression)?;
        }
        if let (Some(visit), Expr::Cast { ty, .. } | Expr::TypedLiteral { ty, .. }) =
            (self.ty.as_mut(), expression)
        {
            visit(ty);
        }
        Ok(())
    }

    pub fn bind_expr(
        &mut self,
        expression: &mut Expr,
        visible_ctes: &BTreeSet<String>,
    ) -> Result<(), SQLError> {
        self.bind_expression_type(expression)?;
        match expression {
            Expr::Func {
                name,
                binding,
                args,
                order_by,
                filter,
                ..
            } => {
                for argument in args {
                    self.bind_expr(argument, visible_ctes)?;
                }
                for order in order_by {
                    self.bind_expr(&mut order.expr, visible_ctes)?;
                }
                if let Some(filter) = filter {
                    self.bind_expr(filter, visible_ctes)?;
                }
                (self.routine)(name, Some(binding))?;
            }
            Expr::Array(items) | Expr::Row(items) | Expr::And(items) | Expr::Or(items) => {
                for item in items {
                    self.bind_expr(item, visible_ctes)?;
                }
            }
            Expr::Binary { lhs, rhs, .. } => {
                self.bind_expr(lhs, visible_ctes)?;
                self.bind_expr(rhs, visible_ctes)?;
            }
            Expr::UnaryMinus(inner)
            | Expr::Not(inner)
            | Expr::IsNull { expr: inner, .. }
            | Expr::Cast { expr: inner, .. } => self.bind_expr(inner, visible_ctes)?,
            Expr::Between { expr, low, high } => {
                self.bind_expr(expr, visible_ctes)?;
                self.bind_expr(low, visible_ctes)?;
                self.bind_expr(high, visible_ctes)?;
            }
            Expr::InList { expr, list, .. } => {
                self.bind_expr(expr, visible_ctes)?;
                for item in list {
                    self.bind_expr(item, visible_ctes)?;
                }
            }
            Expr::WindowCall { name, args, spec } => {
                for argument in args {
                    self.bind_expr(argument, visible_ctes)?;
                }
                for partition in &mut spec.partition_by {
                    self.bind_expr(partition, visible_ctes)?;
                }
                for order in &mut spec.order_by {
                    self.bind_expr(&mut order.expr, visible_ctes)?;
                }
                if let Some(frame) = &mut spec.frame {
                    for bound in [&mut frame.start, &mut frame.end] {
                        if let FrameBound::Preceding(inner) | FrameBound::Following(inner) = bound {
                            self.bind_expr(inner, visible_ctes)?;
                        }
                    }
                }
                (self.routine)(name, None)?;
            }
            Expr::Case {
                base,
                when,
                else_branch,
            } => {
                if let Some(base) = base {
                    self.bind_expr(base, visible_ctes)?;
                }
                for (condition, result) in when {
                    self.bind_expr(condition, visible_ctes)?;
                    self.bind_expr(result, visible_ctes)?;
                }
                if let Some(branch) = else_branch {
                    self.bind_expr(branch, visible_ctes)?;
                }
            }
            Expr::ScalarSubquery(body) | Expr::Exists { body, .. } => {
                self.bind_select(body, visible_ctes)?;
            }
            Expr::InSubquery { expr, body, .. } => {
                self.bind_expr(expr, visible_ctes)?;
                self.bind_select(body, visible_ctes)?;
            }
            Expr::Star
            | Expr::QualifiedStar(_)
            | Expr::Default
            | Expr::Column(_)
            | Expr::QualifiedColumn { .. }
            | Expr::InternalColumn(_)
            | Expr::Literal(_)
            | Expr::TypedLiteral { .. }
            | Expr::Param(_) => {}
        }
        Ok(())
    }
}