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_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;