Skip to main content

linkrs_query/optimizer/
builder.rs

1//! Optimizer Engine Builder Module
2//!
3//! Provides a builder pattern for creating OptimizerEngine instances,
4//! improving flexibility and reducing coupling.
5
6use std::sync::Arc;
7
8use crate::optimizer::OptimizerEngine;
9use crate::optimizer::{CostModelConfig, CteCacheManager, StatisticsManager};
10use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
11
12/// Builder for creating OptimizerEngine instances
13pub struct OptimizerEngineBuilder {
14    cost_config: Option<CostModelConfig>,
15    expression_context: Option<Arc<ExpressionAnalysisContext>>,
16    stats_manager: Option<Arc<StatisticsManager>>,
17    cte_cache_manager: Option<Arc<CteCacheManager>>,
18    metrics_stats: Option<Arc<linkrs_metrics::StatsManager>>,
19}
20
21impl Default for OptimizerEngineBuilder {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl OptimizerEngineBuilder {
28    /// Create a new builder with default configuration.
29    pub fn new() -> Self {
30        Self {
31            cost_config: None,
32            expression_context: None,
33            stats_manager: None,
34            cte_cache_manager: None,
35            metrics_stats: None,
36        }
37    }
38
39    /// Set the cost model configuration.
40    pub fn with_cost_config(mut self, config: CostModelConfig) -> Self {
41        self.cost_config = Some(config);
42        self
43    }
44
45    /// Set the expression analysis context.
46    pub fn with_expression_context(mut self, ctx: Arc<ExpressionAnalysisContext>) -> Self {
47        self.expression_context = Some(ctx);
48        self
49    }
50
51    /// Set the statistics manager.
52    pub fn with_stats_manager(mut self, manager: Arc<StatisticsManager>) -> Self {
53        self.stats_manager = Some(manager);
54        self
55    }
56
57    /// Set the CTE cache manager.
58    pub fn with_cte_cache_manager(mut self, manager: Arc<CteCacheManager>) -> Self {
59        self.cte_cache_manager = Some(manager);
60        self
61    }
62
63    /// Set the observability sink for factorization fallback counters.
64    pub fn with_metrics_stats(mut self, stats: Arc<linkrs_metrics::StatsManager>) -> Self {
65        self.metrics_stats = Some(stats);
66        self
67    }
68
69    /// Build the OptimizerEngine with the configured settings.
70    pub fn build(self) -> OptimizerEngine {
71        let cost_config = self.cost_config.unwrap_or_default();
72        let expression_context = self
73            .expression_context
74            .unwrap_or_else(|| Arc::new(ExpressionAnalysisContext::new()));
75        let stats_manager = self
76            .stats_manager
77            .unwrap_or_else(|| Arc::new(StatisticsManager::new()));
78        let cte_cache_manager = self
79            .cte_cache_manager
80            .unwrap_or_else(|| Arc::new(CteCacheManager::new()));
81
82        OptimizerEngine::with_components(
83            expression_context,
84            stats_manager,
85            cte_cache_manager,
86            cost_config,
87            self.metrics_stats,
88        )
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_builder_default() {
98        let builder = OptimizerEngineBuilder::default();
99        let engine = builder.build();
100        assert_eq!(engine.cost_config(), &CostModelConfig::default());
101    }
102
103    #[test]
104    fn test_builder_with_cost_config() {
105        let config = CostModelConfig::for_ssd();
106        let builder = OptimizerEngineBuilder::new().with_cost_config(config);
107        let engine = builder.build();
108        assert_eq!(engine.cost_config(), &config);
109    }
110
111    #[test]
112    fn test_builder_with_expression_context() {
113        let ctx = Arc::new(ExpressionAnalysisContext::new());
114        let builder = OptimizerEngineBuilder::new().with_expression_context(ctx.clone());
115        let engine = builder.build();
116        assert!(Arc::ptr_eq(engine.expression_context(), &ctx));
117    }
118
119    #[test]
120    fn test_builder_with_stats_manager() {
121        let stats = Arc::new(StatisticsManager::new());
122        let builder = OptimizerEngineBuilder::new().with_stats_manager(stats.clone());
123        let engine = builder.build();
124        assert_eq!(
125            engine.stats_manager().as_ref() as *const _,
126            stats.as_ref() as *const _
127        );
128    }
129
130    #[test]
131    fn test_builder_chain() {
132        let config = CostModelConfig::for_ssd();
133        let ctx = Arc::new(ExpressionAnalysisContext::new());
134        let stats = Arc::new(StatisticsManager::new());
135
136        let engine = OptimizerEngineBuilder::new()
137            .with_cost_config(config)
138            .with_expression_context(ctx)
139            .with_stats_manager(stats)
140            .build();
141
142        assert_eq!(engine.cost_config(), &config);
143    }
144}