Skip to main content

geode_client/
query_builder.rs

1//! Query builder for constructing GQL queries programmatically.
2
3use std::collections::HashMap;
4
5/// Edge direction in patterns
6#[derive(Debug, Clone, Copy)]
7pub enum EdgeDirection {
8    Outgoing,
9    Incoming,
10    Undirected,
11}
12
13/// Query builder for GQL
14pub struct QueryBuilder {
15    clauses: Vec<(String, String)>,
16    params: HashMap<String, serde_json::Value>,
17}
18
19impl QueryBuilder {
20    pub fn new() -> Self {
21        Self {
22            clauses: Vec::new(),
23            params: HashMap::new(),
24        }
25    }
26
27    pub fn match_pattern(mut self, pattern: &str) -> Self {
28        self.clauses
29            .push(("MATCH".to_string(), pattern.to_string()));
30        self
31    }
32
33    pub fn optional_match(mut self, pattern: &str) -> Self {
34        self.clauses
35            .push(("OPTIONAL MATCH".to_string(), pattern.to_string()));
36        self
37    }
38
39    pub fn where_clause(mut self, condition: &str) -> Self {
40        self.clauses
41            .push(("WHERE".to_string(), condition.to_string()));
42        self
43    }
44
45    pub fn with(mut self, expressions: &[&str]) -> Self {
46        self.clauses
47            .push(("WITH".to_string(), expressions.join(", ")));
48        self
49    }
50
51    pub fn return_(mut self, expressions: &[&str]) -> Self {
52        self.clauses
53            .push(("RETURN".to_string(), expressions.join(", ")));
54        self
55    }
56
57    pub fn order_by(mut self, expressions: &[&str]) -> Self {
58        self.clauses
59            .push(("ORDER BY".to_string(), expressions.join(", ")));
60        self
61    }
62
63    pub fn limit(mut self, n: usize) -> Self {
64        self.clauses.push(("LIMIT".to_string(), n.to_string()));
65        self
66    }
67
68    pub fn with_param<V: serde::Serialize>(mut self, name: &str, value: V) -> Self {
69        if let Ok(v) = serde_json::to_value(value) {
70            self.params.insert(name.to_string(), v);
71        }
72        self
73    }
74
75    pub fn build(self) -> (String, HashMap<String, serde_json::Value>) {
76        let query = self
77            .clauses
78            .iter()
79            .map(|(clause_type, content)| {
80                if content.is_empty() {
81                    clause_type.clone()
82                } else {
83                    format!("{} {}", clause_type, content)
84                }
85            })
86            .collect::<Vec<_>>()
87            .join("\n");
88        (query, self.params)
89    }
90}
91
92impl Default for QueryBuilder {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98/// Pattern builder for graph patterns
99pub struct PatternBuilder {
100    elements: Vec<String>,
101}
102
103impl PatternBuilder {
104    pub fn new() -> Self {
105        Self {
106            elements: Vec::new(),
107        }
108    }
109
110    pub fn node(mut self, variable: &str, label: &str) -> Self {
111        let pattern = if label.is_empty() {
112            format!("({})", variable)
113        } else {
114            format!("({}:{})", variable, label)
115        };
116        self.elements.push(pattern);
117        self
118    }
119
120    pub fn edge(mut self, variable: &str, edge_type: &str, direction: EdgeDirection) -> Self {
121        let pattern = match direction {
122            EdgeDirection::Outgoing => {
123                if edge_type.is_empty() {
124                    "->".to_string()
125                } else {
126                    format!("-[{}:{}]->", variable, edge_type)
127                }
128            }
129            EdgeDirection::Incoming => {
130                if edge_type.is_empty() {
131                    "<-".to_string()
132                } else {
133                    format!("<-[{}:{}]-", variable, edge_type)
134                }
135            }
136            EdgeDirection::Undirected => {
137                if edge_type.is_empty() {
138                    "-".to_string()
139                } else {
140                    format!("-[{}:{}]-", variable, edge_type)
141                }
142            }
143        };
144        self.elements.push(pattern);
145        self
146    }
147
148    pub fn build(self) -> String {
149        self.elements.join("")
150    }
151}
152
153impl Default for PatternBuilder {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159/// Predicate builder for WHERE conditions
160pub struct PredicateBuilder {
161    conditions: Vec<String>,
162}
163
164impl PredicateBuilder {
165    pub fn new() -> Self {
166        Self {
167            conditions: Vec::new(),
168        }
169    }
170
171    pub fn greater_than(mut self, left: &str, right: &str) -> Self {
172        self.conditions.push(format!("{} > {}", left, right));
173        self
174    }
175
176    pub fn is_not_null(mut self, expr: &str) -> Self {
177        self.conditions.push(format!("{} IS NOT NULL", expr));
178        self
179    }
180
181    pub fn build_and(self) -> String {
182        self.conditions.join(" AND ")
183    }
184}
185
186impl Default for PredicateBuilder {
187    fn default() -> Self {
188        Self::new()
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    // ==================== QueryBuilder Tests ====================
197
198    #[test]
199    fn test_query_builder_new() {
200        let builder = QueryBuilder::new();
201        let (query, params) = builder.build();
202        assert!(query.is_empty());
203        assert!(params.is_empty());
204    }
205
206    #[test]
207    fn test_query_builder_default() {
208        let builder = QueryBuilder::default();
209        let (query, _) = builder.build();
210        assert!(query.is_empty());
211    }
212
213    #[test]
214    fn test_query_builder_match_pattern() {
215        let (query, _) = QueryBuilder::new().match_pattern("(n:Person)").build();
216        assert_eq!(query, "MATCH (n:Person)");
217    }
218
219    #[test]
220    fn test_query_builder_optional_match() {
221        let (query, _) = QueryBuilder::new()
222            .optional_match("(n:Person)-[:KNOWS]->(m)")
223            .build();
224        assert_eq!(query, "OPTIONAL MATCH (n:Person)-[:KNOWS]->(m)");
225    }
226
227    #[test]
228    fn test_query_builder_where_clause() {
229        let (query, _) = QueryBuilder::new()
230            .match_pattern("(n:Person)")
231            .where_clause("n.age > 25")
232            .build();
233        assert_eq!(query, "MATCH (n:Person)\nWHERE n.age > 25");
234    }
235
236    #[test]
237    fn test_query_builder_with() {
238        let (query, _) = QueryBuilder::new()
239            .match_pattern("(n:Person)")
240            .with(&["n.name AS name", "n.age AS age"])
241            .build();
242        assert_eq!(query, "MATCH (n:Person)\nWITH n.name AS name, n.age AS age");
243    }
244
245    #[test]
246    fn test_query_builder_with_single() {
247        let (query, _) = QueryBuilder::new()
248            .match_pattern("(n)")
249            .with(&["n"])
250            .build();
251        assert_eq!(query, "MATCH (n)\nWITH n");
252    }
253
254    #[test]
255    fn test_query_builder_return() {
256        let (query, _) = QueryBuilder::new()
257            .match_pattern("(n:Person)")
258            .return_(&["n.name", "n.age"])
259            .build();
260        assert_eq!(query, "MATCH (n:Person)\nRETURN n.name, n.age");
261    }
262
263    #[test]
264    fn test_query_builder_return_single() {
265        let (query, _) = QueryBuilder::new()
266            .match_pattern("(n)")
267            .return_(&["n"])
268            .build();
269        assert_eq!(query, "MATCH (n)\nRETURN n");
270    }
271
272    #[test]
273    fn test_query_builder_order_by() {
274        let (query, _) = QueryBuilder::new()
275            .match_pattern("(n:Person)")
276            .return_(&["n.name", "n.age"])
277            .order_by(&["n.age DESC", "n.name ASC"])
278            .build();
279        assert!(query.contains("ORDER BY n.age DESC, n.name ASC"));
280    }
281
282    #[test]
283    fn test_query_builder_limit() {
284        let (query, _) = QueryBuilder::new()
285            .match_pattern("(n:Person)")
286            .return_(&["n"])
287            .limit(10)
288            .build();
289        assert!(query.contains("LIMIT 10"));
290    }
291
292    #[test]
293    fn test_query_builder_limit_zero() {
294        let (query, _) = QueryBuilder::new()
295            .match_pattern("(n)")
296            .return_(&["n"])
297            .limit(0)
298            .build();
299        assert!(query.contains("LIMIT 0"));
300    }
301
302    #[test]
303    fn test_query_builder_with_param_int() {
304        let (_, params) = QueryBuilder::new()
305            .match_pattern("(n:Person)")
306            .where_clause("n.age > $min_age")
307            .with_param("min_age", 25)
308            .build();
309        assert_eq!(params.get("min_age").unwrap(), &serde_json::json!(25));
310    }
311
312    #[test]
313    fn test_query_builder_with_param_string() {
314        let (_, params) = QueryBuilder::new()
315            .match_pattern("(n:Person)")
316            .where_clause("n.name = $name")
317            .with_param("name", "Alice")
318            .build();
319        assert_eq!(params.get("name").unwrap(), &serde_json::json!("Alice"));
320    }
321
322    #[test]
323    fn test_query_builder_with_param_bool() {
324        let (_, params) = QueryBuilder::new().with_param("active", true).build();
325        assert_eq!(params.get("active").unwrap(), &serde_json::json!(true));
326    }
327
328    #[test]
329    fn test_query_builder_multiple_params() {
330        let (_, params) = QueryBuilder::new()
331            .match_pattern("(n:Person)")
332            .where_clause("n.age > $min_age AND n.name = $name")
333            .with_param("min_age", 25)
334            .with_param("name", "Alice")
335            .build();
336        assert_eq!(params.len(), 2);
337        assert_eq!(params.get("min_age").unwrap(), &serde_json::json!(25));
338        assert_eq!(params.get("name").unwrap(), &serde_json::json!("Alice"));
339    }
340
341    #[test]
342    fn test_query_builder_full_query() {
343        let (query, params) = QueryBuilder::new()
344            .match_pattern("(p:Person)")
345            .where_clause("p.age > $min_age")
346            .return_(&["p.name", "p.age"])
347            .order_by(&["p.age DESC"])
348            .limit(100)
349            .with_param("min_age", 25)
350            .build();
351
352        assert!(query.contains("MATCH (p:Person)"));
353        assert!(query.contains("WHERE p.age > $min_age"));
354        assert!(query.contains("RETURN p.name, p.age"));
355        assert!(query.contains("ORDER BY p.age DESC"));
356        assert!(query.contains("LIMIT 100"));
357        assert_eq!(params.get("min_age").unwrap(), &serde_json::json!(25));
358    }
359
360    #[test]
361    fn test_query_builder_complex_query() {
362        let (query, _) = QueryBuilder::new()
363            .match_pattern("(a:Person)-[r:KNOWS]->(b:Person)")
364            .where_clause("a.name = 'Alice'")
365            .optional_match("(b)-[:WORKS_AT]->(c:Company)")
366            .return_(&["a.name", "b.name", "c.name"])
367            .build();
368
369        let lines: Vec<&str> = query.lines().collect();
370        assert_eq!(lines[0], "MATCH (a:Person)-[r:KNOWS]->(b:Person)");
371        assert_eq!(lines[1], "WHERE a.name = 'Alice'");
372        assert_eq!(lines[2], "OPTIONAL MATCH (b)-[:WORKS_AT]->(c:Company)");
373        assert_eq!(lines[3], "RETURN a.name, b.name, c.name");
374    }
375
376    #[test]
377    fn test_query_builder_empty_content_clause() {
378        // Test that empty content produces just the clause type
379        let mut builder = QueryBuilder::new();
380        builder.clauses.push(("RETURN".to_string(), "".to_string()));
381        let (query, _) = builder.build();
382        assert_eq!(query, "RETURN");
383    }
384
385    #[test]
386    fn test_query_builder_chaining() {
387        // Test that methods return Self for chaining
388        let builder = QueryBuilder::new()
389            .match_pattern("(n)")
390            .where_clause("n.x > 0")
391            .return_(&["n"]);
392        let (query, _) = builder.build();
393        assert!(query.contains("MATCH"));
394        assert!(query.contains("WHERE"));
395        assert!(query.contains("RETURN"));
396    }
397
398    #[test]
399    fn test_query_builder_with_array_param() {
400        let (_, params) = QueryBuilder::new().with_param("ids", vec![1, 2, 3]).build();
401        assert_eq!(params.get("ids").unwrap(), &serde_json::json!([1, 2, 3]));
402    }
403
404    #[test]
405    fn test_query_builder_with_null_param() {
406        let (_, params) = QueryBuilder::new()
407            .with_param("value", serde_json::Value::Null)
408            .build();
409        assert_eq!(params.get("value").unwrap(), &serde_json::json!(null));
410    }
411
412    // ==================== PatternBuilder Tests ====================
413
414    #[test]
415    fn test_pattern_builder_new() {
416        let builder = PatternBuilder::new();
417        let pattern = builder.build();
418        assert!(pattern.is_empty());
419    }
420
421    #[test]
422    fn test_pattern_builder_default() {
423        let builder = PatternBuilder::default();
424        let pattern = builder.build();
425        assert!(pattern.is_empty());
426    }
427
428    #[test]
429    fn test_pattern_builder_node_with_label() {
430        let pattern = PatternBuilder::new().node("n", "Person").build();
431        assert_eq!(pattern, "(n:Person)");
432    }
433
434    #[test]
435    fn test_pattern_builder_node_without_label() {
436        let pattern = PatternBuilder::new().node("n", "").build();
437        assert_eq!(pattern, "(n)");
438    }
439
440    #[test]
441    fn test_pattern_builder_edge_outgoing_with_type() {
442        let pattern = PatternBuilder::new()
443            .node("a", "Person")
444            .edge("r", "KNOWS", EdgeDirection::Outgoing)
445            .node("b", "Person")
446            .build();
447        assert_eq!(pattern, "(a:Person)-[r:KNOWS]->(b:Person)");
448    }
449
450    #[test]
451    fn test_pattern_builder_edge_outgoing_without_type() {
452        let pattern = PatternBuilder::new()
453            .node("a", "")
454            .edge("", "", EdgeDirection::Outgoing)
455            .node("b", "")
456            .build();
457        assert_eq!(pattern, "(a)->(b)");
458    }
459
460    #[test]
461    fn test_pattern_builder_edge_incoming_with_type() {
462        let pattern = PatternBuilder::new()
463            .node("a", "Person")
464            .edge("r", "KNOWS", EdgeDirection::Incoming)
465            .node("b", "Person")
466            .build();
467        assert_eq!(pattern, "(a:Person)<-[r:KNOWS]-(b:Person)");
468    }
469
470    #[test]
471    fn test_pattern_builder_edge_incoming_without_type() {
472        let pattern = PatternBuilder::new()
473            .node("a", "")
474            .edge("", "", EdgeDirection::Incoming)
475            .node("b", "")
476            .build();
477        assert_eq!(pattern, "(a)<-(b)");
478    }
479
480    #[test]
481    fn test_pattern_builder_edge_undirected_with_type() {
482        let pattern = PatternBuilder::new()
483            .node("a", "Person")
484            .edge("r", "KNOWS", EdgeDirection::Undirected)
485            .node("b", "Person")
486            .build();
487        assert_eq!(pattern, "(a:Person)-[r:KNOWS]-(b:Person)");
488    }
489
490    #[test]
491    fn test_pattern_builder_edge_undirected_without_type() {
492        let pattern = PatternBuilder::new()
493            .node("a", "")
494            .edge("", "", EdgeDirection::Undirected)
495            .node("b", "")
496            .build();
497        assert_eq!(pattern, "(a)-(b)");
498    }
499
500    #[test]
501    fn test_pattern_builder_chain() {
502        let pattern = PatternBuilder::new()
503            .node("a", "Person")
504            .edge("r1", "KNOWS", EdgeDirection::Outgoing)
505            .node("b", "Person")
506            .edge("r2", "WORKS_AT", EdgeDirection::Outgoing)
507            .node("c", "Company")
508            .build();
509        assert_eq!(
510            pattern,
511            "(a:Person)-[r1:KNOWS]->(b:Person)-[r2:WORKS_AT]->(c:Company)"
512        );
513    }
514
515    #[test]
516    fn test_pattern_builder_mixed_directions() {
517        let pattern = PatternBuilder::new()
518            .node("a", "Person")
519            .edge("", "", EdgeDirection::Outgoing)
520            .node("b", "Person")
521            .edge("", "", EdgeDirection::Incoming)
522            .node("c", "Person")
523            .build();
524        assert_eq!(pattern, "(a:Person)->(b:Person)<-(c:Person)");
525    }
526
527    #[test]
528    fn test_edge_direction_debug() {
529        assert_eq!(format!("{:?}", EdgeDirection::Outgoing), "Outgoing");
530        assert_eq!(format!("{:?}", EdgeDirection::Incoming), "Incoming");
531        assert_eq!(format!("{:?}", EdgeDirection::Undirected), "Undirected");
532    }
533
534    #[test]
535    fn test_edge_direction_copy() {
536        let dir = EdgeDirection::Outgoing;
537        let dir_copy = dir;
538        assert!(matches!(dir, EdgeDirection::Outgoing));
539        assert!(matches!(dir_copy, EdgeDirection::Outgoing));
540    }
541
542    // ==================== PredicateBuilder Tests ====================
543
544    #[test]
545    fn test_predicate_builder_new() {
546        let builder = PredicateBuilder::new();
547        let predicate = builder.build_and();
548        assert!(predicate.is_empty());
549    }
550
551    #[test]
552    fn test_predicate_builder_default() {
553        let builder = PredicateBuilder::default();
554        let predicate = builder.build_and();
555        assert!(predicate.is_empty());
556    }
557
558    #[test]
559    fn test_predicate_builder_greater_than() {
560        let predicate = PredicateBuilder::new()
561            .greater_than("n.age", "25")
562            .build_and();
563        assert_eq!(predicate, "n.age > 25");
564    }
565
566    #[test]
567    fn test_predicate_builder_is_not_null() {
568        let predicate = PredicateBuilder::new().is_not_null("n.email").build_and();
569        assert_eq!(predicate, "n.email IS NOT NULL");
570    }
571
572    #[test]
573    fn test_predicate_builder_multiple_conditions() {
574        let predicate = PredicateBuilder::new()
575            .greater_than("n.age", "25")
576            .is_not_null("n.email")
577            .build_and();
578        assert_eq!(predicate, "n.age > 25 AND n.email IS NOT NULL");
579    }
580
581    #[test]
582    fn test_predicate_builder_chain() {
583        let predicate = PredicateBuilder::new()
584            .greater_than("n.age", "18")
585            .greater_than("n.salary", "50000")
586            .is_not_null("n.department")
587            .build_and();
588        assert_eq!(
589            predicate,
590            "n.age > 18 AND n.salary > 50000 AND n.department IS NOT NULL"
591        );
592    }
593
594    #[test]
595    fn test_predicate_builder_with_param_placeholders() {
596        let predicate = PredicateBuilder::new()
597            .greater_than("n.age", "$min_age")
598            .build_and();
599        assert_eq!(predicate, "n.age > $min_age");
600    }
601
602    // ==================== Integration Tests ====================
603
604    #[test]
605    fn test_query_with_pattern_builder() {
606        let pattern = PatternBuilder::new()
607            .node("p", "Person")
608            .edge("r", "KNOWS", EdgeDirection::Outgoing)
609            .node("f", "Person")
610            .build();
611
612        let (query, _) = QueryBuilder::new()
613            .match_pattern(&pattern)
614            .return_(&["p.name", "f.name"])
615            .build();
616
617        assert!(query.contains("(p:Person)-[r:KNOWS]->(f:Person)"));
618    }
619
620    #[test]
621    fn test_query_with_predicate_builder() {
622        let predicate = PredicateBuilder::new()
623            .greater_than("p.age", "$min_age")
624            .is_not_null("p.email")
625            .build_and();
626
627        let (query, params) = QueryBuilder::new()
628            .match_pattern("(p:Person)")
629            .where_clause(&predicate)
630            .return_(&["p"])
631            .with_param("min_age", 25)
632            .build();
633
634        assert!(query.contains("p.age > $min_age AND p.email IS NOT NULL"));
635        assert_eq!(params.get("min_age").unwrap(), &serde_json::json!(25));
636    }
637
638    #[test]
639    fn test_full_integration() {
640        let pattern = PatternBuilder::new()
641            .node("p", "Person")
642            .edge("", "", EdgeDirection::Outgoing)
643            .node("c", "City")
644            .build();
645
646        let predicate = PredicateBuilder::new()
647            .greater_than("p.age", "$min_age")
648            .is_not_null("c.name")
649            .build_and();
650
651        let (query, params) = QueryBuilder::new()
652            .match_pattern(&pattern)
653            .where_clause(&predicate)
654            .return_(&["p.name", "c.name"])
655            .order_by(&["p.age DESC"])
656            .limit(50)
657            .with_param("min_age", 30)
658            .build();
659
660        assert!(query.contains("MATCH (p:Person)->(c:City)"));
661        assert!(query.contains("WHERE p.age > $min_age AND c.name IS NOT NULL"));
662        assert!(query.contains("RETURN p.name, c.name"));
663        assert!(query.contains("ORDER BY p.age DESC"));
664        assert!(query.contains("LIMIT 50"));
665        assert_eq!(params.get("min_age").unwrap(), &serde_json::json!(30));
666    }
667}