hawdb_cascades/
context.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct OptimizerConfig {
17    pub max_groups: usize,
18}
19
20impl Default for OptimizerConfig {
21    fn default() -> Self {
22        Self { max_groups: 128 }
23    }
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum QueryFamily {
28    GraphRead,
29    GraphWrite,
30    VectorSearch,
31    Schema,
32    System,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum StatementClass {
37    Read,
38    Write,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum ExplainMode {
43    None,
44    Explain,
45    Analyze,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum TraceSink {
50    Disabled,
51    Collect,
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct ResourceHints {
56    pub priority: u8,
57    pub max_memory_bytes: Option<u64>,
58    pub max_parallelism: usize,
59}
60
61impl Default for ResourceHints {
62    fn default() -> Self {
63        Self {
64            priority: 128,
65            max_memory_bytes: None,
66            max_parallelism: 1,
67        }
68    }
69}
70
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub struct OptimizerContext {
73    normalized_query: Option<String>,
74    query_digest: Option<String>,
75    query_family: QueryFamily,
76    statement_class: StatementClass,
77    resource_hints: ResourceHints,
78    explain_mode: ExplainMode,
79    trace_sink: TraceSink,
80    optimizer_config: OptimizerConfig,
81}
82
83impl Default for OptimizerContext {
84    fn default() -> Self {
85        Self {
86            normalized_query: None,
87            query_digest: None,
88            query_family: QueryFamily::GraphRead,
89            statement_class: StatementClass::Read,
90            resource_hints: ResourceHints::default(),
91            explain_mode: ExplainMode::None,
92            trace_sink: TraceSink::Collect,
93            optimizer_config: OptimizerConfig::default(),
94        }
95    }
96}
97
98impl OptimizerContext {
99    pub fn from_config(optimizer_config: OptimizerConfig) -> Self {
100        Self {
101            optimizer_config,
102            ..Self::default()
103        }
104    }
105
106    pub fn with_normalized_query(mut self, normalized_query: impl Into<String>) -> Self {
107        self.normalized_query = Some(normalized_query.into());
108        self
109    }
110
111    pub fn with_query_identity(
112        mut self,
113        normalized_query: impl Into<String>,
114        query_digest: impl Into<String>,
115    ) -> Self {
116        self.normalized_query = Some(normalized_query.into());
117        self.query_digest = Some(query_digest.into());
118        self
119    }
120
121    pub fn with_query_family(mut self, query_family: QueryFamily) -> Self {
122        self.query_family = query_family;
123        self.statement_class = match query_family {
124            QueryFamily::GraphWrite | QueryFamily::Schema => StatementClass::Write,
125            QueryFamily::GraphRead | QueryFamily::VectorSearch | QueryFamily::System => {
126                StatementClass::Read
127            }
128        };
129        self
130    }
131
132    pub fn with_resource_hints(mut self, resource_hints: ResourceHints) -> Self {
133        self.resource_hints = resource_hints;
134        self
135    }
136
137    pub fn with_explain_mode(mut self, explain_mode: ExplainMode) -> Self {
138        self.explain_mode = explain_mode;
139        self
140    }
141
142    pub fn with_trace_sink(mut self, trace_sink: TraceSink) -> Self {
143        self.trace_sink = trace_sink;
144        self
145    }
146
147    pub fn normalized_query(&self) -> Option<&str> {
148        self.normalized_query.as_deref()
149    }
150
151    pub fn query_digest(&self) -> Option<&str> {
152        self.query_digest.as_deref()
153    }
154
155    pub fn query_family(&self) -> QueryFamily {
156        self.query_family
157    }
158
159    pub fn statement_class(&self) -> StatementClass {
160        self.statement_class
161    }
162
163    pub fn resource_hints(&self) -> &ResourceHints {
164        &self.resource_hints
165    }
166
167    pub fn explain_mode(&self) -> ExplainMode {
168        self.explain_mode
169    }
170
171    pub fn trace_sink(&self) -> TraceSink {
172        self.trace_sink
173    }
174
175    pub fn optimizer_config(&self) -> &OptimizerConfig {
176        &self.optimizer_config
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    #[test]
185    fn query_family_derives_statement_class_with_query_identity() {
186        let context = OptimizerContext::default()
187            .with_query_family(QueryFamily::VectorSearch)
188            .with_query_identity("match (m) return m", "q1:test");
189
190        assert_eq!(context.statement_class(), StatementClass::Read);
191        assert_eq!(context.query_digest(), Some("q1:test"));
192        assert_eq!(context.normalized_query(), Some("match (m) return m"));
193    }
194
195    #[test]
196    fn resource_hints_are_bounded_by_explicit_parallelism() {
197        let context = OptimizerContext::default().with_resource_hints(ResourceHints {
198            priority: 200,
199            max_memory_bytes: Some(16 * 1024 * 1024),
200            max_parallelism: 2,
201        });
202
203        assert_eq!(context.resource_hints().max_parallelism, 2);
204        assert_eq!(
205            context.resource_hints().max_memory_bytes,
206            Some(16 * 1024 * 1024)
207        );
208    }
209}
hawdb_cascades/context.rs

hawdb_cascades/
context.rs
