uni_query_functions/rewrite/
context.rs

1/// Rewrite context and configuration
2use crate::rewrite::error::RewriteError;
3use std::collections::HashMap;
4
5/// Contextual information available during query rewriting
6///
7/// The context carries rewrite statistics and configuration options for the
8/// current query.
9#[derive(Default)]
10pub struct RewriteContext {
11    /// Rewrite statistics (for observability)
12    pub stats: RewriteStats,
13
14    /// Configuration flags
15    pub config: RewriteConfig,
16}
17
18impl RewriteContext {
19    /// Create a new rewrite context with default configuration
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Create a new rewrite context with custom configuration
25    pub fn with_config(config: RewriteConfig) -> Self {
26        Self {
27            stats: RewriteStats::default(),
28            config,
29        }
30    }
31}
32
33/// Statistics collected during rewriting
34#[derive(Debug, Default, Clone)]
35pub struct RewriteStats {
36    /// Total number of function calls visited
37    pub functions_visited: usize,
38
39    /// Number of functions successfully rewritten
40    pub functions_rewritten: usize,
41
42    /// Number of functions that fell back to scalar execution
43    pub functions_skipped: usize,
44
45    /// Errors encountered during rewriting (non-fatal)
46    pub errors: Vec<RewriteError>,
47
48    /// Per-rule statistics
49    pub rule_stats: HashMap<String, RuleStats>,
50}
51
52impl RewriteStats {
53    /// Get or create rule stats entry
54    fn rule_stats_mut(&mut self, rule_name: &str) -> &mut RuleStats {
55        self.rule_stats.entry(rule_name.to_string()).or_default()
56    }
57
58    /// Record a successful rewrite for a rule
59    pub fn record_success(&mut self, rule_name: &str) {
60        self.functions_rewritten += 1;
61        self.rule_stats_mut(rule_name).record_success();
62    }
63
64    /// Record a failed rewrite for a rule
65    pub fn record_failure(&mut self, rule_name: &str, error: RewriteError) {
66        self.functions_skipped += 1;
67        self.rule_stats_mut(rule_name).record_failure(error.clone());
68        self.errors.push(error);
69    }
70
71    /// Record a visited function
72    pub fn record_visit(&mut self) {
73        self.functions_visited += 1;
74    }
75}
76
77/// Per-rule statistics
78#[derive(Debug, Default, Clone)]
79pub struct RuleStats {
80    /// Number of times this rule was attempted
81    pub attempts: usize,
82
83    /// Number of successful rewrites
84    pub successes: usize,
85
86    /// Failure counts by error type
87    pub failures: HashMap<String, usize>,
88}
89
90impl RuleStats {
91    fn record_success(&mut self) {
92        self.attempts += 1;
93        self.successes += 1;
94    }
95
96    fn record_failure(&mut self, error: RewriteError) {
97        self.attempts += 1;
98        let error_key = format!("{error:?}");
99        *self.failures.entry(error_key).or_default() += 1;
100    }
101}
102
103/// Configuration options for query rewriting
104#[derive(Debug, Clone)]
105pub struct RewriteConfig {
106    /// Enable temporal function rewrites
107    pub enable_temporal: bool,
108
109    /// Enable verbose logging of rewrite operations
110    pub verbose_logging: bool,
111}
112
113impl Default for RewriteConfig {
114    fn default() -> Self {
115        Self {
116            enable_temporal: true,
117            verbose_logging: false,
118        }
119    }
120}
121
122impl RewriteConfig {
123    /// Create a config with all rewrites enabled
124    pub fn all_enabled() -> Self {
125        Self {
126            enable_temporal: true,
127            verbose_logging: false,
128        }
129    }
130
131    /// Enable verbose logging
132    pub fn with_verbose_logging(mut self) -> Self {
133        self.verbose_logging = true;
134        self
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn test_context_default() {
144        let ctx = RewriteContext::default();
145        assert_eq!(ctx.stats.functions_visited, 0);
146        assert!(ctx.config.enable_temporal);
147    }
148
149    #[test]
150    fn test_stats_recording() {
151        let mut stats = RewriteStats::default();
152
153        stats.record_success("test.func");
154        assert_eq!(stats.functions_rewritten, 1);
155        assert_eq!(stats.rule_stats.get("test.func").unwrap().successes, 1);
156
157        stats.record_failure(
158            "test.func",
159            RewriteError::NotApplicable {
160                reason: "test".into(),
161            },
162        );
163        assert_eq!(stats.functions_skipped, 1);
164        assert_eq!(stats.errors.len(), 1);
165    }
166
167    #[test]
168    fn test_config_builders() {
169        let all_enabled = RewriteConfig::all_enabled();
170        assert!(all_enabled.enable_temporal);
171
172        let verbose = RewriteConfig::default().with_verbose_logging();
173        assert!(verbose.verbose_logging);
174    }
175}
uni_query_functions/rewrite/context.rs

uni_query_functions/rewrite/
context.rs
