hawdb_optimizer/
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
15use crate::OptimizerConfig;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum QueryFamily {
19    GraphRead,
20    GraphWrite,
21    VectorSearch,
22    Schema,
23    System,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum StatementClass {
28    Read,
29    Write,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum ExplainMode {
34    None,
35    Explain,
36    Analyze,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum TraceSink {
41    Disabled,
42    Collect,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct ResourceHints {
47    pub priority: u8,
48    pub max_memory_bytes: Option<u64>,
49    pub max_parallelism: usize,
50}
51
52impl Default for ResourceHints {
53    fn default() -> Self {
54        Self {
55            priority: 128,
56            max_memory_bytes: None,
57            max_parallelism: 1,
58        }
59    }
60}
61
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct OptimizerContext {
64    normalized_query: Option<String>,
65    query_digest: Option<String>,
66    query_family: QueryFamily,
67    statement_class: StatementClass,
68    resource_hints: ResourceHints,
69    explain_mode: ExplainMode,
70    trace_sink: TraceSink,
71    optimizer_config: OptimizerConfig,
72}
73
74impl Default for OptimizerContext {
75    fn default() -> Self {
76        Self {
77            normalized_query: None,
78            query_digest: None,
79            query_family: QueryFamily::GraphRead,
80            statement_class: StatementClass::Read,
81            resource_hints: ResourceHints::default(),
82            explain_mode: ExplainMode::None,
83            trace_sink: TraceSink::Collect,
84            optimizer_config: OptimizerConfig::default(),
85        }
86    }
87}
88
89impl OptimizerContext {
90    pub fn from_config(optimizer_config: OptimizerConfig) -> Self {
91        Self {
92            optimizer_config,
93            ..Self::default()
94        }
95    }
96
97    pub fn with_normalized_query(mut self, normalized_query: impl Into<String>) -> Self {
98        self.normalized_query = Some(normalized_query.into());
99        self
100    }
101
102    pub fn with_query_identity(
103        mut self,
104        normalized_query: impl Into<String>,
105        query_digest: impl Into<String>,
106    ) -> Self {
107        self.normalized_query = Some(normalized_query.into());
108        self.query_digest = Some(query_digest.into());
109        self
110    }
111
112    pub fn with_query_family(mut self, query_family: QueryFamily) -> Self {
113        self.query_family = query_family;
114        self.statement_class = match query_family {
115            QueryFamily::GraphWrite | QueryFamily::Schema => StatementClass::Write,
116            QueryFamily::GraphRead | QueryFamily::VectorSearch | QueryFamily::System => {
117                StatementClass::Read
118            }
119        };
120        self
121    }
122
123    pub fn with_resource_hints(mut self, resource_hints: ResourceHints) -> Self {
124        self.resource_hints = resource_hints;
125        self
126    }
127
128    pub fn with_explain_mode(mut self, explain_mode: ExplainMode) -> Self {
129        self.explain_mode = explain_mode;
130        self
131    }
132
133    pub fn with_trace_sink(mut self, trace_sink: TraceSink) -> Self {
134        self.trace_sink = trace_sink;
135        self
136    }
137
138    pub fn normalized_query(&self) -> Option<&str> {
139        self.normalized_query.as_deref()
140    }
141
142    pub fn query_digest(&self) -> Option<&str> {
143        self.query_digest.as_deref()
144    }
145
146    pub fn query_family(&self) -> QueryFamily {
147        self.query_family
148    }
149
150    pub fn statement_class(&self) -> StatementClass {
151        self.statement_class
152    }
153
154    pub fn resource_hints(&self) -> &ResourceHints {
155        &self.resource_hints
156    }
157
158    pub fn explain_mode(&self) -> ExplainMode {
159        self.explain_mode
160    }
161
162    pub fn trace_sink(&self) -> TraceSink {
163        self.trace_sink
164    }
165
166    pub fn optimizer_config(&self) -> &OptimizerConfig {
167        &self.optimizer_config
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn query_family_derives_statement_class_with_query_identity() {
177        let context = OptimizerContext::default()
178            .with_query_family(QueryFamily::VectorSearch)
179            .with_query_identity("match (m) return m", "q1:test");
180
181        assert_eq!(context.statement_class(), StatementClass::Read);
182        assert_eq!(context.query_digest(), Some("q1:test"));
183        assert_eq!(context.normalized_query(), Some("match (m) return m"));
184    }
185
186    #[test]
187    fn resource_hints_are_bounded_by_explicit_parallelism() {
188        let context = OptimizerContext::default().with_resource_hints(ResourceHints {
189            priority: 200,
190            max_memory_bytes: Some(16 * 1024 * 1024),
191            max_parallelism: 2,
192        });
193
194        assert_eq!(context.resource_hints().max_parallelism, 2);
195        assert_eq!(
196            context.resource_hints().max_memory_bytes,
197            Some(16 * 1024 * 1024)
198        );
199    }
200}
hawdb_optimizer/context.rs

hawdb_optimizer/
context.rs
