1use crate::{
2 ComparisonOp, CompileError, CompiledGroupedQuery, CompiledQuery, ExpansionSlot, Predicate,
3 QueryAst, QueryStep, ScalarValue, TextQuery, TraverseDirection, compile_grouped_query,
4 compile_query,
5};
6
7#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct QueryBuilder {
14 ast: QueryAst,
15}
16
17impl QueryBuilder {
18 #[must_use]
20 pub fn nodes(kind: impl Into<String>) -> Self {
21 Self {
22 ast: QueryAst {
23 root_kind: kind.into(),
24 steps: Vec::new(),
25 expansions: Vec::new(),
26 final_limit: None,
27 },
28 }
29 }
30
31 #[must_use]
33 pub fn vector_search(mut self, query: impl Into<String>, limit: usize) -> Self {
34 self.ast.steps.push(QueryStep::VectorSearch {
35 query: query.into(),
36 limit,
37 });
38 self
39 }
40
41 #[must_use]
48 pub fn text_search(mut self, query: impl Into<String>, limit: usize) -> Self {
49 let query = TextQuery::parse(&query.into());
50 self.ast.steps.push(QueryStep::TextSearch { query, limit });
51 self
52 }
53
54 #[must_use]
56 pub fn traverse(
57 mut self,
58 direction: TraverseDirection,
59 label: impl Into<String>,
60 max_depth: usize,
61 ) -> Self {
62 self.ast.steps.push(QueryStep::Traverse {
63 direction,
64 label: label.into(),
65 max_depth,
66 });
67 self
68 }
69
70 #[must_use]
72 pub fn filter_logical_id_eq(mut self, logical_id: impl Into<String>) -> Self {
73 self.ast
74 .steps
75 .push(QueryStep::Filter(Predicate::LogicalIdEq(logical_id.into())));
76 self
77 }
78
79 #[must_use]
81 pub fn filter_kind_eq(mut self, kind: impl Into<String>) -> Self {
82 self.ast
83 .steps
84 .push(QueryStep::Filter(Predicate::KindEq(kind.into())));
85 self
86 }
87
88 #[must_use]
90 pub fn filter_source_ref_eq(mut self, source_ref: impl Into<String>) -> Self {
91 self.ast
92 .steps
93 .push(QueryStep::Filter(Predicate::SourceRefEq(source_ref.into())));
94 self
95 }
96
97 #[must_use]
99 pub fn filter_content_ref_not_null(mut self) -> Self {
100 self.ast
101 .steps
102 .push(QueryStep::Filter(Predicate::ContentRefNotNull));
103 self
104 }
105
106 #[must_use]
108 pub fn filter_content_ref_eq(mut self, content_ref: impl Into<String>) -> Self {
109 self.ast
110 .steps
111 .push(QueryStep::Filter(Predicate::ContentRefEq(
112 content_ref.into(),
113 )));
114 self
115 }
116
117 #[must_use]
119 pub fn filter_json_text_eq(
120 mut self,
121 path: impl Into<String>,
122 value: impl Into<String>,
123 ) -> Self {
124 self.ast
125 .steps
126 .push(QueryStep::Filter(Predicate::JsonPathEq {
127 path: path.into(),
128 value: ScalarValue::Text(value.into()),
129 }));
130 self
131 }
132
133 #[must_use]
135 pub fn filter_json_bool_eq(mut self, path: impl Into<String>, value: bool) -> Self {
136 self.ast
137 .steps
138 .push(QueryStep::Filter(Predicate::JsonPathEq {
139 path: path.into(),
140 value: ScalarValue::Bool(value),
141 }));
142 self
143 }
144
145 #[must_use]
147 pub fn filter_json_integer_gt(mut self, path: impl Into<String>, value: i64) -> Self {
148 self.ast
149 .steps
150 .push(QueryStep::Filter(Predicate::JsonPathCompare {
151 path: path.into(),
152 op: ComparisonOp::Gt,
153 value: ScalarValue::Integer(value),
154 }));
155 self
156 }
157
158 #[must_use]
160 pub fn filter_json_integer_gte(mut self, path: impl Into<String>, value: i64) -> Self {
161 self.ast
162 .steps
163 .push(QueryStep::Filter(Predicate::JsonPathCompare {
164 path: path.into(),
165 op: ComparisonOp::Gte,
166 value: ScalarValue::Integer(value),
167 }));
168 self
169 }
170
171 #[must_use]
173 pub fn filter_json_integer_lt(mut self, path: impl Into<String>, value: i64) -> Self {
174 self.ast
175 .steps
176 .push(QueryStep::Filter(Predicate::JsonPathCompare {
177 path: path.into(),
178 op: ComparisonOp::Lt,
179 value: ScalarValue::Integer(value),
180 }));
181 self
182 }
183
184 #[must_use]
186 pub fn filter_json_integer_lte(mut self, path: impl Into<String>, value: i64) -> Self {
187 self.ast
188 .steps
189 .push(QueryStep::Filter(Predicate::JsonPathCompare {
190 path: path.into(),
191 op: ComparisonOp::Lte,
192 value: ScalarValue::Integer(value),
193 }));
194 self
195 }
196
197 #[must_use]
199 pub fn filter_json_timestamp_gt(self, path: impl Into<String>, value: i64) -> Self {
200 self.filter_json_integer_gt(path, value)
201 }
202
203 #[must_use]
205 pub fn filter_json_timestamp_gte(self, path: impl Into<String>, value: i64) -> Self {
206 self.filter_json_integer_gte(path, value)
207 }
208
209 #[must_use]
211 pub fn filter_json_timestamp_lt(self, path: impl Into<String>, value: i64) -> Self {
212 self.filter_json_integer_lt(path, value)
213 }
214
215 #[must_use]
217 pub fn filter_json_timestamp_lte(self, path: impl Into<String>, value: i64) -> Self {
218 self.filter_json_integer_lte(path, value)
219 }
220
221 #[must_use]
223 pub fn expand(
224 mut self,
225 slot: impl Into<String>,
226 direction: TraverseDirection,
227 label: impl Into<String>,
228 max_depth: usize,
229 ) -> Self {
230 self.ast.expansions.push(ExpansionSlot {
231 slot: slot.into(),
232 direction,
233 label: label.into(),
234 max_depth,
235 });
236 self
237 }
238
239 #[must_use]
241 pub fn limit(mut self, limit: usize) -> Self {
242 self.ast.final_limit = Some(limit);
243 self
244 }
245
246 #[must_use]
248 pub fn ast(&self) -> &QueryAst {
249 &self.ast
250 }
251
252 #[must_use]
254 pub fn into_ast(self) -> QueryAst {
255 self.ast
256 }
257
258 pub fn compile(&self) -> Result<CompiledQuery, CompileError> {
265 compile_query(&self.ast)
266 }
267
268 pub fn compile_grouped(&self) -> Result<CompiledGroupedQuery, CompileError> {
275 compile_grouped_query(&self.ast)
276 }
277}
278
279#[cfg(test)]
280#[allow(clippy::panic)]
281mod tests {
282 use crate::{Predicate, QueryBuilder, QueryStep, ScalarValue, TextQuery, TraverseDirection};
283
284 #[test]
285 fn builder_accumulates_expected_steps() {
286 let query = QueryBuilder::nodes("Meeting")
287 .text_search("budget", 5)
288 .traverse(TraverseDirection::Out, "HAS_TASK", 2)
289 .filter_json_text_eq("$.status", "active")
290 .limit(10);
291
292 assert_eq!(query.ast().steps.len(), 3);
293 assert_eq!(query.ast().final_limit, Some(10));
294 }
295
296 #[test]
297 fn builder_filter_json_bool_eq_produces_correct_predicate() {
298 let query = QueryBuilder::nodes("Feature").filter_json_bool_eq("$.enabled", true);
299
300 assert_eq!(query.ast().steps.len(), 1);
301 match &query.ast().steps[0] {
302 QueryStep::Filter(Predicate::JsonPathEq { path, value }) => {
303 assert_eq!(path, "$.enabled");
304 assert_eq!(*value, ScalarValue::Bool(true));
305 }
306 other => panic!("expected JsonPathEq/Bool, got {other:?}"),
307 }
308 }
309
310 #[test]
311 fn builder_text_search_parses_into_typed_query() {
312 let query = QueryBuilder::nodes("Meeting").text_search("ship NOT blocked", 10);
313
314 match &query.ast().steps[0] {
315 QueryStep::TextSearch { query, limit } => {
316 assert_eq!(*limit, 10);
317 assert_eq!(
318 *query,
319 TextQuery::And(vec![
320 TextQuery::Term("ship".into()),
321 TextQuery::Not(Box::new(TextQuery::Term("blocked".into()))),
322 ])
323 );
324 }
325 other => panic!("expected TextSearch, got {other:?}"),
326 }
327 }
328}