Skip to main content

uqa_planner/statement_planning/
executable.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Analyze executable plans before optimizer inputs and constant evaluation.
8
9use super::{rule_inputs::RuleInputPlanningContext, StatementStatisticsContext};
10use crate::{AggregateClassifier, ConstantEvaluator, QueryPlan, UnifiedPlan};
11use uqa_sql::{
12    binding::statements::{analyze_executable_plan, AnalyzedResult, StatementAnalysisContext},
13    SQLError, SQLParam,
14};
15
16/// Borrow planner metadata after SQL analysis succeeds.
17pub trait StatementOptimizationContexts {
18    fn statistics(&self) -> StatementStatisticsContext<'_>;
19    fn rule_inputs(&self) -> RuleInputPlanningContext<'_>;
20    fn builtin_permissions(
21        &self,
22    ) -> Option<
23        std::sync::Arc<dyn uqa_sql::catalog::security::builtin_routines::BuiltinRoutineExecution>,
24    > {
25        None
26    }
27    fn routine_inlining(&self) -> Option<uqa_sql::routines::inlining::RoutineInliningContext<'_>> {
28        None
29    }
30}
31pub struct StatementPlanningContext<'a> {
32    pub analysis: StatementAnalysisContext<'a>,
33    pub aggregates: &'a dyn AggregateClassifier,
34    pub optimization: &'a dyn StatementOptimizationContexts,
35    pub constant_evaluator: ConstantEvaluator,
36}
37
38impl uqa_sql::plan::ExecutablePlanOptimizer for StatementPlanningContext<'_> {
39    fn plan_with_cached_analysis(
40        &self,
41        plan: UnifiedPlan,
42        params: &[SQLParam],
43        cached: Option<std::sync::Arc<uqa_sql::binding::statements::AnalyzedStatement>>,
44    ) -> Result<
45        (
46            UnifiedPlan,
47            Option<std::sync::Arc<uqa_sql::binding::statements::AnalyzedStatement>>,
48        ),
49        SQLError,
50    > {
51        let (plan, retained) = if let Some((plan, cached)) =
52            cached.and_then(|cached| cached.plan_for(params).cloned().map(|plan| (plan, cached)))
53        {
54            (plan, Some(cached))
55        } else {
56            let (plan, retained) = uqa_sql::binding::statements::analyze_for_statement_reuse(
57                &self.analysis,
58                plan,
59                params,
60            )?;
61            (plan, retained.map(std::sync::Arc::new))
62        };
63        Ok((optimize_plan(self, plan)?, retained))
64    }
65
66    fn plan_for_statement_cache(
67        &self,
68        mut plan: UnifiedPlan,
69        params: &[SQLParam],
70    ) -> Result<(UnifiedPlan, bool), SQLError> {
71        let (_, reusable) = uqa_sql::binding::statements::analyze_executable_plan_for_cache(
72            &self.analysis,
73            &mut plan,
74            params,
75        )?;
76        let reusable = reusable
77            && !self
78                .optimization
79                .routine_inlining()
80                .is_some_and(|routines| routines.requires_replanning(&plan));
81        Ok((optimize_plan(self, plan)?, reusable))
82    }
83
84    fn plan_for_execution(
85        &self,
86        plan: UnifiedPlan,
87        params: &[SQLParam],
88    ) -> Result<UnifiedPlan, SQLError> {
89        plan_for_execution(self, plan, params)
90    }
91
92    fn plan_with_result(
93        &self,
94        mut plan: UnifiedPlan,
95        params: &[SQLParam],
96    ) -> Result<(UnifiedPlan, AnalyzedResult), SQLError> {
97        let result = analyze_executable_plan(&self.analysis, &mut plan, params)?;
98        Ok((optimize_plan(self, plan)?, result))
99    }
100}
101
102pub fn plan_for_execution(
103    context: &StatementPlanningContext<'_>,
104    mut plan: UnifiedPlan,
105    params: &[SQLParam],
106) -> Result<UnifiedPlan, SQLError> {
107    analyze_executable_plan(&context.analysis, &mut plan, params)?;
108    optimize_plan(context, plan)
109}
110pub fn optimize_query(
111    context: &StatementPlanningContext<'_>,
112    query: &QueryPlan,
113) -> Result<QueryPlan, SQLError> {
114    match optimize_plan(context, UnifiedPlan::Query(Box::new(query.clone())))? {
115        UnifiedPlan::Query(query) => Ok(*query),
116        UnifiedPlan::Command(_) => Err(SQLError::Internal(
117            "query optimization produced a command".into(),
118        )),
119    }
120}
121pub fn optimize_plan(
122    context: &StatementPlanningContext<'_>,
123    plan: UnifiedPlan,
124) -> Result<UnifiedPlan, SQLError> {
125    super::optimize_plan_authorized(
126        context.optimization.statistics(),
127        &context.optimization.rule_inputs(),
128        context.aggregates,
129        context.constant_evaluator,
130        context.optimization.builtin_permissions(),
131        context.optimization.routine_inlining(),
132        plan,
133    )
134}
135
136#[cfg(test)]
137mod tests;