uqa_planner/statement_planning/
executable.rs1use 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
16pub 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;