Skip to main content

uqa_sql/routines/lifecycle/
rewrites.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Surviving routine body candidates and exact registry identities for source rewrites.
8
9use super::RoutineRegistry;
10use crate::{
11    ast::{CreateFunction, FunctionBinding, FunctionBody},
12    catalog::events::RuleColumnDependency,
13    routines::{routine_signature_types, SQLUserFunction},
14    SQLError,
15};
16use std::{collections::BTreeSet, sync::Arc};
17use uqa_core::RelationIdentity;
18
19pub fn routine_column_drop_dependencies(
20    columns: BTreeSet<(String, String)>,
21) -> Result<BTreeSet<RuleColumnDependency>, SQLError> {
22    columns
23        .into_iter()
24        .map(|(table, column)| {
25            Ok(RuleColumnDependency {
26                relation: RelationIdentity::from_legacy_name(&table).map_err(SQLError::Internal)?,
27                column,
28            })
29        })
30        .collect()
31}
32
33pub fn rewritten_routine_definitions(
34    registry: &RoutineRegistry,
35    mut rewrite: impl FnMut(&mut crate::ast::Statement) -> Result<bool, SQLError>,
36) -> Result<Vec<CreateFunction>, SQLError> {
37    let mut definitions = Vec::new();
38    for overloads in registry.values() {
39        for function in overloads {
40            let mut definition = function.def.clone();
41            let mut changed = false;
42            if let FunctionBody::Statements(statements) = &mut definition.body {
43                for statement in statements {
44                    changed |= rewrite(statement)?;
45                }
46            }
47            if changed {
48                definitions.push(definition);
49            }
50        }
51    }
52    Ok(definitions)
53}
54
55pub fn routine_column_alias_drop_candidates(
56    columns: crate::binding::stored_columns::StoredColumnBindingContext<'_>,
57    registry: &RoutineRegistry,
58    dependencies: &BTreeSet<RuleColumnDependency>,
59    removed_routines: &[FunctionBinding],
60) -> Result<Vec<CreateFunction>, SQLError> {
61    let mut definitions = Vec::new();
62    for function in registry.values().flatten() {
63        if removed_routines.iter().any(|target| {
64            target.object_id == function.def.object_id
65                && target.name == function.def.name
66                && target.argument_types == routine_signature_types(&function.def)
67        }) {
68            continue;
69        }
70        let mut definition = function.def.clone();
71        let FunctionBody::Statements(statements) = &mut definition.body else {
72            continue;
73        };
74        let mut changed = false;
75        for statement in statements {
76            changed |=
77                crate::binding::stored_columns::remove_stored_statement_source_column_aliases(
78                    columns,
79                    statement,
80                    dependencies,
81                )?;
82        }
83        if changed {
84            definitions.push(definition);
85        }
86    }
87    Ok(definitions)
88}
89
90pub fn routine_body_rewrite_target<'a>(
91    rewritten: &'a mut RoutineRegistry,
92    definition: &CreateFunction,
93) -> Result<&'a mut Arc<SQLUserFunction>, SQLError> {
94    let signature = routine_signature_types(definition);
95    rewritten
96        .get_mut(&definition.name)
97        .and_then(|overloads| {
98            overloads.iter_mut().find(|function| {
99                function.def.object_id == definition.object_id
100                    && function.def.is_procedure == definition.is_procedure
101                    && routine_signature_types(&function.def) == signature
102            })
103        })
104        .ok_or_else(|| {
105            SQLError::Internal(format!(
106                "stored routine {} disappeared before its body rewrite",
107                definition.name
108            ))
109        })
110}