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_for_statement_cache(
40        &self,
41        mut plan: UnifiedPlan,
42        params: &[SQLParam],
43    ) -> Result<(UnifiedPlan, bool), SQLError> {
44        let (_, reusable) = uqa_sql::binding::statements::analyze_executable_plan_for_cache(
45            &self.analysis,
46            &mut plan,
47            params,
48        )?;
49        let reusable = reusable
50            && !self
51                .optimization
52                .routine_inlining()
53                .is_some_and(|routines| routines.requires_replanning(&plan));
54        Ok((optimize_plan(self, plan)?, reusable))
55    }
56
57    fn plan_for_execution(
58        &self,
59        plan: UnifiedPlan,
60        params: &[SQLParam],
61    ) -> Result<UnifiedPlan, SQLError> {
62        plan_for_execution(self, plan, params)
63    }
64
65    fn plan_with_result(
66        &self,
67        mut plan: UnifiedPlan,
68        params: &[SQLParam],
69    ) -> Result<(UnifiedPlan, AnalyzedResult), SQLError> {
70        let result = analyze_executable_plan(&self.analysis, &mut plan, params)?;
71        Ok((optimize_plan(self, plan)?, result))
72    }
73}
74
75pub fn plan_for_execution(
76    context: &StatementPlanningContext<'_>,
77    mut plan: UnifiedPlan,
78    params: &[SQLParam],
79) -> Result<UnifiedPlan, SQLError> {
80    analyze_executable_plan(&context.analysis, &mut plan, params)?;
81    optimize_plan(context, plan)
82}
83pub fn optimize_query(
84    context: &StatementPlanningContext<'_>,
85    query: &QueryPlan,
86) -> Result<QueryPlan, SQLError> {
87    match optimize_plan(context, UnifiedPlan::Query(Box::new(query.clone())))? {
88        UnifiedPlan::Query(query) => Ok(*query),
89        UnifiedPlan::Command(_) => Err(SQLError::Internal(
90            "query optimization produced a command".into(),
91        )),
92    }
93}
94pub fn optimize_plan(
95    context: &StatementPlanningContext<'_>,
96    plan: UnifiedPlan,
97) -> Result<UnifiedPlan, SQLError> {
98    super::optimize_plan_authorized(
99        context.optimization.statistics(),
100        &context.optimization.rule_inputs(),
101        context.aggregates,
102        context.constant_evaluator,
103        context.optimization.builtin_permissions(),
104        context.optimization.routine_inlining(),
105        plan,
106    )
107}
108
109#[cfg(test)]
110mod tests;