Skip to main content

fathomdb_query/
builder.rs

1use crate::{
2    ComparisonOp, CompileError, CompiledGroupedQuery, CompiledQuery, ExpansionSlot, Predicate,
3    QueryAst, QueryStep, ScalarValue, TraverseDirection, compile_grouped_query, compile_query,
4};
5
6/// Fluent builder for constructing a [`QueryAst`].
7///
8/// Start with [`QueryBuilder::nodes`] and chain filtering, traversal, and
9/// expansion steps before calling [`compile`](QueryBuilder::compile) or
10/// [`compile_grouped`](QueryBuilder::compile_grouped).
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct QueryBuilder {
13    ast: QueryAst,
14}
15
16impl QueryBuilder {
17    /// Create a builder that queries nodes of the given kind.
18    #[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    /// Add a vector similarity search step.
31    #[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    /// Add a full-text search step.
41    #[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    /// Add a graph traversal step following edges of the given label.
51    #[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    /// Filter results to a single logical ID.
67    #[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    /// Filter results to nodes matching the given kind.
76    #[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    /// Filter results to nodes matching the given `source_ref`.
85    #[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    /// Filter results where a JSON property at `path` equals the given text value.
94    #[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    /// Filter results where a JSON property at `path` equals the given boolean value.
110    #[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    /// Filter results where a JSON integer at `path` is greater than `value`.
122    #[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    /// Filter results where a JSON integer at `path` is greater than or equal to `value`.
135    #[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    /// Filter results where a JSON integer at `path` is less than `value`.
148    #[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    /// Filter results where a JSON integer at `path` is less than or equal to `value`.
161    #[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    /// Filter results where a JSON timestamp at `path` is after `value` (epoch seconds).
174    #[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    /// Filter results where a JSON timestamp at `path` is at or after `value`.
180    #[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    /// Filter results where a JSON timestamp at `path` is before `value`.
186    #[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    /// Filter results where a JSON timestamp at `path` is at or before `value`.
192    #[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    /// Add an expansion slot that traverses edges of the given label for each root result.
198    #[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    /// Set the maximum number of result rows.
216    #[must_use]
217    pub fn limit(mut self, limit: usize) -> Self {
218        self.ast.final_limit = Some(limit);
219        self
220    }
221
222    /// Borrow the underlying [`QueryAst`].
223    #[must_use]
224    pub fn ast(&self) -> &QueryAst {
225        &self.ast
226    }
227
228    /// Consume the builder and return the underlying [`QueryAst`].
229    #[must_use]
230    pub fn into_ast(self) -> QueryAst {
231        self.ast
232    }
233
234    /// Compile this builder's AST into an executable [`CompiledQuery`].
235    ///
236    /// # Errors
237    ///
238    /// Returns [`CompileError`] if the query violates structural constraints
239    /// (e.g. too many traversal steps or too many bind parameters).
240    pub fn compile(&self) -> Result<CompiledQuery, CompileError> {
241        compile_query(&self.ast)
242    }
243
244    /// Compile this builder's AST into an executable grouped query.
245    ///
246    /// # Errors
247    ///
248    /// Returns [`CompileError`] if the query violates grouped-query structural
249    /// constraints such as duplicate slot names or excessive depth.
250    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}