1use crate::{
2 ComparisonOp, CompileError, CompiledGroupedQuery, CompiledQuery, ExpansionSlot, Predicate,
3 QueryAst, QueryStep, ScalarValue, TraverseDirection, compile_grouped_query, compile_query,
4};
5
6#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct QueryBuilder {
13 ast: QueryAst,
14}
15
16impl QueryBuilder {
17 #[must_use]
19 pub fn nodes(kind: impl Into<String>) -> Self {
20 Self {
21 ast: QueryAst {
22 root_kind: kind.into(),
23 steps: Vec::new(),
24 expansions: Vec::new(),
25 final_limit: None,
26 },
27 }
28 }
29
30 #[must_use]
32 pub fn vector_search(mut self, query: impl Into<String>, limit: usize) -> Self {
33 self.ast.steps.push(QueryStep::VectorSearch {
34 query: query.into(),
35 limit,
36 });
37 self
38 }
39
40 #[must_use]
42 pub fn text_search(mut self, query: impl Into<String>, limit: usize) -> Self {
43 self.ast.steps.push(QueryStep::TextSearch {
44 query: query.into(),
45 limit,
46 });
47 self
48 }
49
50 #[must_use]
52 pub fn traverse(
53 mut self,
54 direction: TraverseDirection,
55 label: impl Into<String>,
56 max_depth: usize,
57 ) -> Self {
58 self.ast.steps.push(QueryStep::Traverse {
59 direction,
60 label: label.into(),
61 max_depth,
62 });
63 self
64 }
65
66 #[must_use]
68 pub fn filter_logical_id_eq(mut self, logical_id: impl Into<String>) -> Self {
69 self.ast
70 .steps
71 .push(QueryStep::Filter(Predicate::LogicalIdEq(logical_id.into())));
72 self
73 }
74
75 #[must_use]
77 pub fn filter_kind_eq(mut self, kind: impl Into<String>) -> Self {
78 self.ast
79 .steps
80 .push(QueryStep::Filter(Predicate::KindEq(kind.into())));
81 self
82 }
83
84 #[must_use]
86 pub fn filter_source_ref_eq(mut self, source_ref: impl Into<String>) -> Self {
87 self.ast
88 .steps
89 .push(QueryStep::Filter(Predicate::SourceRefEq(source_ref.into())));
90 self
91 }
92
93 #[must_use]
95 pub fn filter_json_text_eq(
96 mut self,
97 path: impl Into<String>,
98 value: impl Into<String>,
99 ) -> Self {
100 self.ast
101 .steps
102 .push(QueryStep::Filter(Predicate::JsonPathEq {
103 path: path.into(),
104 value: ScalarValue::Text(value.into()),
105 }));
106 self
107 }
108
109 #[must_use]
111 pub fn filter_json_bool_eq(mut self, path: impl Into<String>, value: bool) -> Self {
112 self.ast
113 .steps
114 .push(QueryStep::Filter(Predicate::JsonPathEq {
115 path: path.into(),
116 value: ScalarValue::Bool(value),
117 }));
118 self
119 }
120
121 #[must_use]
123 pub fn filter_json_integer_gt(mut self, path: impl Into<String>, value: i64) -> Self {
124 self.ast
125 .steps
126 .push(QueryStep::Filter(Predicate::JsonPathCompare {
127 path: path.into(),
128 op: ComparisonOp::Gt,
129 value: ScalarValue::Integer(value),
130 }));
131 self
132 }
133
134 #[must_use]
136 pub fn filter_json_integer_gte(mut self, path: impl Into<String>, value: i64) -> Self {
137 self.ast
138 .steps
139 .push(QueryStep::Filter(Predicate::JsonPathCompare {
140 path: path.into(),
141 op: ComparisonOp::Gte,
142 value: ScalarValue::Integer(value),
143 }));
144 self
145 }
146
147 #[must_use]
149 pub fn filter_json_integer_lt(mut self, path: impl Into<String>, value: i64) -> Self {
150 self.ast
151 .steps
152 .push(QueryStep::Filter(Predicate::JsonPathCompare {
153 path: path.into(),
154 op: ComparisonOp::Lt,
155 value: ScalarValue::Integer(value),
156 }));
157 self
158 }
159
160 #[must_use]
162 pub fn filter_json_integer_lte(mut self, path: impl Into<String>, value: i64) -> Self {
163 self.ast
164 .steps
165 .push(QueryStep::Filter(Predicate::JsonPathCompare {
166 path: path.into(),
167 op: ComparisonOp::Lte,
168 value: ScalarValue::Integer(value),
169 }));
170 self
171 }
172
173 #[must_use]
175 pub fn filter_json_timestamp_gt(self, path: impl Into<String>, value: i64) -> Self {
176 self.filter_json_integer_gt(path, value)
177 }
178
179 #[must_use]
181 pub fn filter_json_timestamp_gte(self, path: impl Into<String>, value: i64) -> Self {
182 self.filter_json_integer_gte(path, value)
183 }
184
185 #[must_use]
187 pub fn filter_json_timestamp_lt(self, path: impl Into<String>, value: i64) -> Self {
188 self.filter_json_integer_lt(path, value)
189 }
190
191 #[must_use]
193 pub fn filter_json_timestamp_lte(self, path: impl Into<String>, value: i64) -> Self {
194 self.filter_json_integer_lte(path, value)
195 }
196
197 #[must_use]
199 pub fn expand(
200 mut self,
201 slot: impl Into<String>,
202 direction: TraverseDirection,
203 label: impl Into<String>,
204 max_depth: usize,
205 ) -> Self {
206 self.ast.expansions.push(ExpansionSlot {
207 slot: slot.into(),
208 direction,
209 label: label.into(),
210 max_depth,
211 });
212 self
213 }
214
215 #[must_use]
217 pub fn limit(mut self, limit: usize) -> Self {
218 self.ast.final_limit = Some(limit);
219 self
220 }
221
222 #[must_use]
224 pub fn ast(&self) -> &QueryAst {
225 &self.ast
226 }
227
228 #[must_use]
230 pub fn into_ast(self) -> QueryAst {
231 self.ast
232 }
233
234 pub fn compile(&self) -> Result<CompiledQuery, CompileError> {
241 compile_query(&self.ast)
242 }
243
244 pub fn compile_grouped(&self) -> Result<CompiledGroupedQuery, CompileError> {
251 compile_grouped_query(&self.ast)
252 }
253}
254
255#[cfg(test)]
256#[allow(clippy::panic)]
257mod tests {
258 use crate::{Predicate, QueryBuilder, QueryStep, ScalarValue, TraverseDirection};
259
260 #[test]
261 fn builder_accumulates_expected_steps() {
262 let query = QueryBuilder::nodes("Meeting")
263 .text_search("budget", 5)
264 .traverse(TraverseDirection::Out, "HAS_TASK", 2)
265 .filter_json_text_eq("$.status", "active")
266 .limit(10);
267
268 assert_eq!(query.ast().steps.len(), 3);
269 assert_eq!(query.ast().final_limit, Some(10));
270 }
271
272 #[test]
273 fn builder_filter_json_bool_eq_produces_correct_predicate() {
274 let query = QueryBuilder::nodes("Feature").filter_json_bool_eq("$.enabled", true);
275
276 assert_eq!(query.ast().steps.len(), 1);
277 match &query.ast().steps[0] {
278 QueryStep::Filter(Predicate::JsonPathEq { path, value }) => {
279 assert_eq!(path, "$.enabled");
280 assert_eq!(*value, ScalarValue::Bool(true));
281 }
282 other => panic!("expected JsonPathEq/Bool, got {other:?}"),
283 }
284 }
285}