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 to nodes where `content_ref` is not NULL.
94    #[must_use]
95    pub fn filter_content_ref_not_null(mut self) -> Self {
96        self.ast
97            .steps
98            .push(QueryStep::Filter(Predicate::ContentRefNotNull));
99        self
100    }
101
102    /// Filter results to nodes matching the given `content_ref` URI.
103    #[must_use]
104    pub fn filter_content_ref_eq(mut self, content_ref: impl Into<String>) -> Self {
105        self.ast
106            .steps
107            .push(QueryStep::Filter(Predicate::ContentRefEq(
108                content_ref.into(),
109            )));
110        self
111    }
112
113    /// Filter results where a JSON property at `path` equals the given text value.
114    #[must_use]
115    pub fn filter_json_text_eq(
116        mut self,
117        path: impl Into<String>,
118        value: impl Into<String>,
119    ) -> Self {
120        self.ast
121            .steps
122            .push(QueryStep::Filter(Predicate::JsonPathEq {
123                path: path.into(),
124                value: ScalarValue::Text(value.into()),
125            }));
126        self
127    }
128
129    /// Filter results where a JSON property at `path` equals the given boolean value.
130    #[must_use]
131    pub fn filter_json_bool_eq(mut self, path: impl Into<String>, value: bool) -> Self {
132        self.ast
133            .steps
134            .push(QueryStep::Filter(Predicate::JsonPathEq {
135                path: path.into(),
136                value: ScalarValue::Bool(value),
137            }));
138        self
139    }
140
141    /// Filter results where a JSON integer at `path` is greater than `value`.
142    #[must_use]
143    pub fn filter_json_integer_gt(mut self, path: impl Into<String>, value: i64) -> Self {
144        self.ast
145            .steps
146            .push(QueryStep::Filter(Predicate::JsonPathCompare {
147                path: path.into(),
148                op: ComparisonOp::Gt,
149                value: ScalarValue::Integer(value),
150            }));
151        self
152    }
153
154    /// Filter results where a JSON integer at `path` is greater than or equal to `value`.
155    #[must_use]
156    pub fn filter_json_integer_gte(mut self, path: impl Into<String>, value: i64) -> Self {
157        self.ast
158            .steps
159            .push(QueryStep::Filter(Predicate::JsonPathCompare {
160                path: path.into(),
161                op: ComparisonOp::Gte,
162                value: ScalarValue::Integer(value),
163            }));
164        self
165    }
166
167    /// Filter results where a JSON integer at `path` is less than `value`.
168    #[must_use]
169    pub fn filter_json_integer_lt(mut self, path: impl Into<String>, value: i64) -> Self {
170        self.ast
171            .steps
172            .push(QueryStep::Filter(Predicate::JsonPathCompare {
173                path: path.into(),
174                op: ComparisonOp::Lt,
175                value: ScalarValue::Integer(value),
176            }));
177        self
178    }
179
180    /// Filter results where a JSON integer at `path` is less than or equal to `value`.
181    #[must_use]
182    pub fn filter_json_integer_lte(mut self, path: impl Into<String>, value: i64) -> Self {
183        self.ast
184            .steps
185            .push(QueryStep::Filter(Predicate::JsonPathCompare {
186                path: path.into(),
187                op: ComparisonOp::Lte,
188                value: ScalarValue::Integer(value),
189            }));
190        self
191    }
192
193    /// Filter results where a JSON timestamp at `path` is after `value` (epoch seconds).
194    #[must_use]
195    pub fn filter_json_timestamp_gt(self, path: impl Into<String>, value: i64) -> Self {
196        self.filter_json_integer_gt(path, value)
197    }
198
199    /// Filter results where a JSON timestamp at `path` is at or after `value`.
200    #[must_use]
201    pub fn filter_json_timestamp_gte(self, path: impl Into<String>, value: i64) -> Self {
202        self.filter_json_integer_gte(path, value)
203    }
204
205    /// Filter results where a JSON timestamp at `path` is before `value`.
206    #[must_use]
207    pub fn filter_json_timestamp_lt(self, path: impl Into<String>, value: i64) -> Self {
208        self.filter_json_integer_lt(path, value)
209    }
210
211    /// Filter results where a JSON timestamp at `path` is at or before `value`.
212    #[must_use]
213    pub fn filter_json_timestamp_lte(self, path: impl Into<String>, value: i64) -> Self {
214        self.filter_json_integer_lte(path, value)
215    }
216
217    /// Add an expansion slot that traverses edges of the given label for each root result.
218    #[must_use]
219    pub fn expand(
220        mut self,
221        slot: impl Into<String>,
222        direction: TraverseDirection,
223        label: impl Into<String>,
224        max_depth: usize,
225    ) -> Self {
226        self.ast.expansions.push(ExpansionSlot {
227            slot: slot.into(),
228            direction,
229            label: label.into(),
230            max_depth,
231        });
232        self
233    }
234
235    /// Set the maximum number of result rows.
236    #[must_use]
237    pub fn limit(mut self, limit: usize) -> Self {
238        self.ast.final_limit = Some(limit);
239        self
240    }
241
242    /// Borrow the underlying [`QueryAst`].
243    #[must_use]
244    pub fn ast(&self) -> &QueryAst {
245        &self.ast
246    }
247
248    /// Consume the builder and return the underlying [`QueryAst`].
249    #[must_use]
250    pub fn into_ast(self) -> QueryAst {
251        self.ast
252    }
253
254    /// Compile this builder's AST into an executable [`CompiledQuery`].
255    ///
256    /// # Errors
257    ///
258    /// Returns [`CompileError`] if the query violates structural constraints
259    /// (e.g. too many traversal steps or too many bind parameters).
260    pub fn compile(&self) -> Result<CompiledQuery, CompileError> {
261        compile_query(&self.ast)
262    }
263
264    /// Compile this builder's AST into an executable grouped query.
265    ///
266    /// # Errors
267    ///
268    /// Returns [`CompileError`] if the query violates grouped-query structural
269    /// constraints such as duplicate slot names or excessive depth.
270    pub fn compile_grouped(&self) -> Result<CompiledGroupedQuery, CompileError> {
271        compile_grouped_query(&self.ast)
272    }
273}
274
275#[cfg(test)]
276#[allow(clippy::panic)]
277mod tests {
278    use crate::{Predicate, QueryBuilder, QueryStep, ScalarValue, TraverseDirection};
279
280    #[test]
281    fn builder_accumulates_expected_steps() {
282        let query = QueryBuilder::nodes("Meeting")
283            .text_search("budget", 5)
284            .traverse(TraverseDirection::Out, "HAS_TASK", 2)
285            .filter_json_text_eq("$.status", "active")
286            .limit(10);
287
288        assert_eq!(query.ast().steps.len(), 3);
289        assert_eq!(query.ast().final_limit, Some(10));
290    }
291
292    #[test]
293    fn builder_filter_json_bool_eq_produces_correct_predicate() {
294        let query = QueryBuilder::nodes("Feature").filter_json_bool_eq("$.enabled", true);
295
296        assert_eq!(query.ast().steps.len(), 1);
297        match &query.ast().steps[0] {
298            QueryStep::Filter(Predicate::JsonPathEq { path, value }) => {
299                assert_eq!(path, "$.enabled");
300                assert_eq!(*value, ScalarValue::Bool(true));
301            }
302            other => panic!("expected JsonPathEq/Bool, got {other:?}"),
303        }
304    }
305}