1use super::*;
2use crate::sql_lowering::{filter_to_expr, projection_to_select_item};
3
4pub struct TableQueryBuilder {
5 query: TableQuery,
6}
7
8impl TableQueryBuilder {
9 pub fn new(table: &str) -> Self {
11 Self {
12 query: TableQuery::new(table),
13 }
14 }
15
16 pub fn alias(mut self, alias: &str) -> Self {
18 self.query.alias = Some(alias.to_string());
19 self
20 }
21
22 pub fn select(mut self, column: &str) -> Self {
24 let field = FieldRef::column(
25 self.query.alias.as_deref().unwrap_or(&self.query.table),
26 column,
27 );
28 self.query.select_items.push(SelectItem::Expr {
29 expr: Expr::col(field.clone()),
30 alias: None,
31 });
32 self.query.columns.push(Projection::from_field(field));
33 self
34 }
35
36 pub fn select_all(mut self) -> Self {
38 self.query.select_items = vec![SelectItem::Wildcard];
39 self.query.columns.clear();
40 self
41 }
42
43 pub fn filter(mut self, f: Filter) -> Self {
45 let f_expr = filter_to_expr(&f);
46 self.query.where_expr = Some(match self.query.where_expr.take() {
47 Some(existing) => Expr::binop(BinOp::And, existing, f_expr),
48 None => f_expr,
49 });
50 self.query.filter = Some(match self.query.filter.take() {
51 Some(existing) => existing.and(f),
52 None => f,
53 });
54 self
55 }
56
57 pub fn order_by(mut self, clause: OrderByClause) -> Self {
59 self.query.order_by.push(clause);
60 self
61 }
62
63 pub fn limit(mut self, n: u64) -> Self {
65 self.query.limit = Some(n);
66 self
67 }
68
69 pub fn offset(mut self, n: u64) -> Self {
71 self.query.offset = Some(n);
72 self
73 }
74
75 pub fn join_graph(self, pattern: GraphPattern, on: JoinCondition) -> JoinQueryBuilder {
77 JoinQueryBuilder {
78 left: QueryExpr::Table(self.query),
79 right: QueryExpr::Graph(GraphQuery::new(pattern)),
80 on,
81 join_type: JoinType::Inner,
82 filter: None,
83 order_by: Vec::new(),
84 limit: None,
85 offset: None,
86 return_items: Vec::new(),
87 return_: Vec::new(),
88 }
89 }
90
91 pub fn join_table(self, table: &str, on: JoinCondition) -> JoinQueryBuilder {
93 JoinQueryBuilder {
94 left: QueryExpr::Table(self.query),
95 right: QueryExpr::Table(TableQuery::new(table)),
96 on,
97 join_type: JoinType::Inner,
98 filter: None,
99 order_by: Vec::new(),
100 limit: None,
101 offset: None,
102 return_items: Vec::new(),
103 return_: Vec::new(),
104 }
105 }
106
107 pub fn join_vector(self, query: VectorQuery, on: JoinCondition) -> JoinQueryBuilder {
109 JoinQueryBuilder {
110 left: QueryExpr::Table(self.query),
111 right: QueryExpr::Vector(query),
112 on,
113 join_type: JoinType::Inner,
114 filter: None,
115 order_by: Vec::new(),
116 limit: None,
117 offset: None,
118 return_items: Vec::new(),
119 return_: Vec::new(),
120 }
121 }
122
123 pub fn join_path(self, query: PathQuery, on: JoinCondition) -> JoinQueryBuilder {
125 JoinQueryBuilder {
126 left: QueryExpr::Table(self.query),
127 right: QueryExpr::Path(query),
128 on,
129 join_type: JoinType::Inner,
130 filter: None,
131 order_by: Vec::new(),
132 limit: None,
133 offset: None,
134 return_items: Vec::new(),
135 return_: Vec::new(),
136 }
137 }
138
139 pub fn join_hybrid(self, query: HybridQuery, on: JoinCondition) -> JoinQueryBuilder {
141 JoinQueryBuilder {
142 left: QueryExpr::Table(self.query),
143 right: QueryExpr::Hybrid(query),
144 on,
145 join_type: JoinType::Inner,
146 filter: None,
147 order_by: Vec::new(),
148 limit: None,
149 offset: None,
150 return_items: Vec::new(),
151 return_: Vec::new(),
152 }
153 }
154
155 pub fn build(self) -> QueryExpr {
157 QueryExpr::Table(self.query)
158 }
159}
160
161pub struct GraphQueryBuilder {
163 query: GraphQuery,
164}
165
166impl GraphQueryBuilder {
167 pub fn new() -> Self {
169 Self {
170 query: GraphQuery::new(GraphPattern::new()),
171 }
172 }
173
174 pub fn node(mut self, pattern: NodePattern) -> Self {
176 self.query.pattern.nodes.push(pattern);
177 self
178 }
179
180 pub fn edge(mut self, pattern: EdgePattern) -> Self {
182 self.query.pattern.edges.push(pattern);
183 self
184 }
185
186 pub fn filter(mut self, f: Filter) -> Self {
188 self.query.filter = Some(match self.query.filter.take() {
189 Some(existing) => existing.and(f),
190 None => f,
191 });
192 self
193 }
194
195 pub fn alias(mut self, alias: &str) -> Self {
197 self.query.alias = Some(alias.to_string());
198 self
199 }
200
201 pub fn limit(mut self, n: u64) -> Self {
203 self.query.limit = Some(n);
204 self
205 }
206
207 pub fn return_field(mut self, field: FieldRef) -> Self {
209 self.query.return_.push(Projection::from_field(field));
210 self
211 }
212
213 pub fn build(self) -> QueryExpr {
215 QueryExpr::Graph(self.query)
216 }
217}
218
219impl Default for GraphQueryBuilder {
220 fn default() -> Self {
221 Self::new()
222 }
223}
224
225pub struct JoinQueryBuilder {
227 left: QueryExpr,
228 right: QueryExpr,
229 on: JoinCondition,
230 join_type: JoinType,
231 filter: Option<Filter>,
232 order_by: Vec<OrderByClause>,
233 limit: Option<u64>,
234 offset: Option<u64>,
235 return_items: Vec<SelectItem>,
236 return_: Vec<Projection>,
237}
238
239impl JoinQueryBuilder {
240 pub fn join_type(mut self, jt: JoinType) -> Self {
242 self.join_type = jt;
243 self
244 }
245
246 pub fn right_alias(mut self, alias: &str) -> Self {
248 let alias = alias.to_string();
249 match &mut self.right {
250 QueryExpr::Table(table) => table.alias = Some(alias.clone()),
251 QueryExpr::Graph(graph) => graph.alias = Some(alias.clone()),
252 QueryExpr::Path(path) => path.alias = Some(alias.clone()),
253 QueryExpr::Vector(vector) => vector.alias = Some(alias.clone()),
254 QueryExpr::Hybrid(hybrid) => hybrid.alias = Some(alias.clone()),
255 QueryExpr::Join(_)
256 | QueryExpr::Insert(_)
257 | QueryExpr::Update(_)
258 | QueryExpr::Delete(_)
259 | QueryExpr::CreateTable(_)
260 | QueryExpr::CreateCollection(_)
261 | QueryExpr::CreateVector(_)
262 | QueryExpr::DropTable(_)
263 | QueryExpr::DropGraph(_)
264 | QueryExpr::DropVector(_)
265 | QueryExpr::DropDocument(_)
266 | QueryExpr::DropKv(_)
267 | QueryExpr::DropCollection(_)
268 | QueryExpr::Truncate(_)
269 | QueryExpr::AlterTable(_)
270 | QueryExpr::GraphCommand(_)
271 | QueryExpr::SearchCommand(_)
272 | QueryExpr::CreateIndex(_)
273 | QueryExpr::DropIndex(_)
274 | QueryExpr::ProbabilisticCommand(_)
275 | QueryExpr::Ask(_)
276 | QueryExpr::SetConfig { .. }
277 | QueryExpr::ShowConfig { .. }
278 | QueryExpr::SetSecret { .. }
279 | QueryExpr::DeleteSecret { .. }
280 | QueryExpr::ShowSecrets { .. }
281 | QueryExpr::SetTenant(_)
282 | QueryExpr::ShowTenant
283 | QueryExpr::CreateTimeSeries(_)
284 | QueryExpr::CreateMetric(_)
285 | QueryExpr::AlterMetric(_)
286 | QueryExpr::CreateSlo(_)
287 | QueryExpr::DropTimeSeries(_)
288 | QueryExpr::CreateQueue(_)
289 | QueryExpr::AlterQueue(_)
290 | QueryExpr::DropQueue(_)
291 | QueryExpr::QueueSelect(_)
292 | QueryExpr::QueueCommand(_)
293 | QueryExpr::KvCommand(_)
294 | QueryExpr::ConfigCommand(_)
295 | QueryExpr::CreateTree(_)
296 | QueryExpr::DropTree(_)
297 | QueryExpr::TreeCommand(_)
298 | QueryExpr::ExplainAlter(_)
299 | QueryExpr::TransactionControl(_)
300 | QueryExpr::MaintenanceCommand(_)
301 | QueryExpr::CreateSchema(_)
302 | QueryExpr::DropSchema(_)
303 | QueryExpr::CreateSequence(_)
304 | QueryExpr::DropSequence(_)
305 | QueryExpr::CopyFrom(_)
306 | QueryExpr::CreateView(_)
307 | QueryExpr::DropView(_)
308 | QueryExpr::RefreshMaterializedView(_)
309 | QueryExpr::CreatePolicy(_)
310 | QueryExpr::DropPolicy(_)
311 | QueryExpr::CreateServer(_)
312 | QueryExpr::DropServer(_)
313 | QueryExpr::CreateForeignTable(_)
314 | QueryExpr::DropForeignTable(_)
315 | QueryExpr::Grant(_)
316 | QueryExpr::Revoke(_)
317 | QueryExpr::AlterUser(_)
318 | QueryExpr::CreateUser(_)
319 | QueryExpr::CreateIamPolicy { .. }
320 | QueryExpr::DropIamPolicy { .. }
321 | QueryExpr::AttachPolicy { .. }
322 | QueryExpr::DetachPolicy { .. }
323 | QueryExpr::ShowPolicies { .. }
324 | QueryExpr::ShowEffectivePermissions { .. }
325 | QueryExpr::RankOf(_)
326 | QueryExpr::ApproxRankOf(_)
327 | QueryExpr::RankRange(_)
328 | QueryExpr::SimulatePolicy { .. }
329 | QueryExpr::LintPolicy { .. }
330 | QueryExpr::MigratePolicyMode { .. }
331 | QueryExpr::CreateMigration(_)
332 | QueryExpr::ApplyMigration(_)
333 | QueryExpr::RollbackMigration(_)
334 | QueryExpr::ExplainMigration(_)
335 | QueryExpr::EventsBackfill(_)
336 | QueryExpr::EventsBackfillStatus { .. } => {}
337 }
338 self
339 }
340
341 pub fn filter(mut self, f: Filter) -> Self {
343 self.filter = Some(match self.filter.take() {
344 Some(existing) => existing.and(f),
345 None => f,
346 });
347 self
348 }
349
350 pub fn order_by(mut self, clause: OrderByClause) -> Self {
352 self.order_by.push(clause);
353 self
354 }
355
356 pub fn limit(mut self, n: u64) -> Self {
358 self.limit = Some(n);
359 self
360 }
361
362 pub fn offset(mut self, n: u64) -> Self {
364 self.offset = Some(n);
365 self
366 }
367
368 pub fn return_field(mut self, field: FieldRef) -> Self {
370 let projection = Projection::from_field(field);
371 if let Some(item) = projection_to_select_item(&projection) {
372 self.return_items.push(item);
373 }
374 self.return_.push(projection);
375 self
376 }
377
378 pub fn select(mut self, column: &str) -> Self {
380 let projection = Projection::from_field(FieldRef::column("", column));
381 if let Some(item) = projection_to_select_item(&projection) {
382 self.return_items.push(item);
383 }
384 self.return_.push(projection);
385 self
386 }
387
388 pub fn build(self) -> QueryExpr {
390 QueryExpr::Join(JoinQuery {
391 left: Box::new(self.left),
392 right: Box::new(self.right),
393 join_type: self.join_type,
394 on: self.on,
395 filter: self.filter,
396 order_by: self.order_by,
397 limit: self.limit,
398 offset: self.offset,
399 return_items: self.return_items,
400 return_: self.return_,
401 })
402 }
403}
404
405pub struct PathQueryBuilder {
407 query: PathQuery,
408}
409
410impl PathQueryBuilder {
411 pub fn new(from: NodeSelector, to: NodeSelector) -> Self {
413 Self {
414 query: PathQuery::new(from, to),
415 }
416 }
417
418 pub fn via_label(mut self, label: impl Into<String>) -> Self {
420 self.query.via.push(label.into());
421 self
422 }
423
424 pub fn max_length(mut self, n: u32) -> Self {
426 self.query.max_length = n;
427 self
428 }
429
430 pub fn filter(mut self, f: Filter) -> Self {
432 self.query.filter = Some(f);
433 self
434 }
435
436 pub fn alias(mut self, alias: &str) -> Self {
438 self.query.alias = Some(alias.to_string());
439 self
440 }
441
442 pub fn build(self) -> QueryExpr {
444 QueryExpr::Path(self.query)
445 }
446}
447
448#[derive(Debug, Clone)]
481pub struct CteDefinition {
482 pub name: String,
484 pub columns: Vec<String>,
486 pub query: Box<QueryExpr>,
488 pub recursive: bool,
490}
491
492impl CteDefinition {
493 pub fn new(name: &str, query: QueryExpr) -> Self {
495 Self {
496 name: name.to_string(),
497 columns: Vec::new(),
498 query: Box::new(query),
499 recursive: false,
500 }
501 }
502
503 pub fn recursive(name: &str, query: QueryExpr) -> Self {
505 Self {
506 name: name.to_string(),
507 columns: Vec::new(),
508 query: Box::new(query),
509 recursive: true,
510 }
511 }
512
513 pub fn with_columns(mut self, columns: Vec<String>) -> Self {
515 self.columns = columns;
516 self
517 }
518}
519
520#[derive(Debug, Clone, Default)]
522pub struct WithClause {
523 pub ctes: Vec<CteDefinition>,
525 pub has_recursive: bool,
527}
528
529impl WithClause {
530 pub fn new() -> Self {
532 Self::default()
533 }
534
535 #[allow(clippy::should_implement_trait)]
539 pub fn add(mut self, cte: CteDefinition) -> Self {
540 if cte.recursive {
541 self.has_recursive = true;
542 }
543 self.ctes.push(cte);
544 self
545 }
546
547 pub fn is_empty(&self) -> bool {
549 self.ctes.is_empty()
550 }
551
552 pub fn get(&self, name: &str) -> Option<&CteDefinition> {
554 self.ctes.iter().find(|c| c.name == name)
555 }
556}
557
558#[derive(Debug, Clone)]
560pub struct QueryWithCte {
561 pub with_clause: Option<WithClause>,
563 pub query: QueryExpr,
565}
566
567impl QueryWithCte {
568 pub fn simple(query: QueryExpr) -> Self {
570 Self {
571 with_clause: None,
572 query,
573 }
574 }
575
576 pub fn with_ctes(with_clause: WithClause, query: QueryExpr) -> Self {
578 Self {
579 with_clause: Some(with_clause),
580 query,
581 }
582 }
583}
584
585pub struct CteQueryBuilder {
587 with_clause: WithClause,
588}
589
590impl CteQueryBuilder {
591 pub fn new() -> Self {
593 Self {
594 with_clause: WithClause::new(),
595 }
596 }
597
598 pub fn cte(mut self, name: &str, query: QueryExpr) -> Self {
600 self.with_clause = self.with_clause.add(CteDefinition::new(name, query));
601 self
602 }
603
604 pub fn recursive_cte(mut self, name: &str, query: QueryExpr) -> Self {
606 self.with_clause = self.with_clause.add(CteDefinition::recursive(name, query));
607 self
608 }
609
610 pub fn cte_with_columns(mut self, name: &str, columns: Vec<String>, query: QueryExpr) -> Self {
612 let cte = CteDefinition::new(name, query).with_columns(columns);
613 self.with_clause = self.with_clause.add(cte);
614 self
615 }
616
617 pub fn build(self, main_query: QueryExpr) -> QueryWithCte {
619 QueryWithCte::with_ctes(self.with_clause, main_query)
620 }
621}
622
623impl Default for CteQueryBuilder {
624 fn default() -> Self {
625 Self::new()
626 }
627}
628
629#[cfg(test)]
630mod tests {
631 use super::*;
632
633 fn eq_filter(column: &str, value: Value) -> Filter {
634 Filter::compare(FieldRef::column("", column), CompareOp::Eq, value)
635 }
636
637 fn join_condition() -> JoinCondition {
638 JoinCondition::new(
639 FieldRef::column("left", "id"),
640 FieldRef::column("right", "id"),
641 )
642 }
643
644 #[test]
645 fn table_builder_covers_selection_filters_order_limit_and_offset() {
646 let query = TableQueryBuilder::new("hosts")
647 .alias("h")
648 .select("ip")
649 .filter(eq_filter("os", Value::text("linux")))
650 .filter(eq_filter("active", Value::Boolean(true)))
651 .order_by(OrderByClause::desc(FieldRef::column("h", "last_seen")))
652 .limit(10)
653 .offset(5)
654 .build();
655
656 let QueryExpr::Table(table) = query else {
657 panic!("expected table query");
658 };
659 assert_eq!(table.table, "hosts");
660 assert_eq!(table.alias.as_deref(), Some("h"));
661 assert_eq!(table.select_items.len(), 1);
662 assert_eq!(table.columns.len(), 1);
663 assert!(matches!(table.filter, Some(Filter::And(_, _))));
664 assert!(matches!(
665 table.where_expr,
666 Some(Expr::BinaryOp { op: BinOp::And, .. })
667 ));
668 assert_eq!(table.order_by.len(), 1);
669 assert_eq!(table.limit, Some(10));
670 assert_eq!(table.offset, Some(5));
671
672 let QueryExpr::Table(table) = TableQueryBuilder::new("hosts").select_all().build() else {
673 panic!("expected table query");
674 };
675 assert_eq!(table.select_items, vec![SelectItem::Wildcard]);
676 assert!(table.columns.is_empty());
677 }
678
679 #[test]
680 fn graph_builder_combines_filters_alias_limit_and_returns() {
681 let query = GraphQueryBuilder::new()
682 .node(NodePattern::new("h").of_label("Host"))
683 .edge(EdgePattern::new("h", "s").of_label("HAS_SERVICE"))
684 .filter(eq_filter("critical", Value::Boolean(true)))
685 .filter(eq_filter("active", Value::Boolean(true)))
686 .alias("g")
687 .limit(3)
688 .return_field(FieldRef::node_prop("h", "ip"))
689 .build();
690
691 let QueryExpr::Graph(graph) = query else {
692 panic!("expected graph query");
693 };
694 assert_eq!(graph.alias.as_deref(), Some("g"));
695 assert_eq!(graph.pattern.nodes.len(), 1);
696 assert_eq!(graph.pattern.edges.len(), 1);
697 assert!(matches!(graph.filter, Some(Filter::And(_, _))));
698 assert_eq!(graph.limit, Some(3));
699 assert_eq!(graph.return_.len(), 1);
700 }
701
702 #[test]
703 fn join_builder_aliases_supported_right_sources_and_builds_options() {
704 let condition = join_condition();
705 let cases = vec![
706 TableQueryBuilder::new("hosts").join_table("services", condition.clone()),
707 TableQueryBuilder::new("hosts").join_graph(GraphPattern::new(), condition.clone()),
708 TableQueryBuilder::new("hosts").join_path(
709 PathQuery::new(NodeSelector::by_id("a"), NodeSelector::by_id("b")),
710 condition.clone(),
711 ),
712 TableQueryBuilder::new("hosts").join_vector(
713 VectorQuery::new("embeddings", VectorSource::text("ssh")),
714 condition.clone(),
715 ),
716 TableQueryBuilder::new("hosts").join_hybrid(
717 HybridQuery::new(
718 QueryExpr::Table(TableQuery::new("hosts")),
719 VectorQuery::new("embeddings", VectorSource::text("ssh")),
720 ),
721 condition.clone(),
722 ),
723 ];
724
725 for builder in cases {
726 let query = builder
727 .right_alias("rhs")
728 .join_type(JoinType::FullOuter)
729 .filter(eq_filter("ok", Value::Boolean(true)))
730 .order_by(OrderByClause::asc(FieldRef::column("", "id")))
731 .limit(4)
732 .offset(2)
733 .return_field(FieldRef::column("hosts", "id"))
734 .select("name")
735 .build();
736
737 let QueryExpr::Join(join) = query else {
738 panic!("expected join query");
739 };
740 assert_eq!(join.join_type, JoinType::FullOuter);
741 assert!(join.filter.is_some());
742 assert_eq!(join.order_by.len(), 1);
743 assert_eq!(join.limit, Some(4));
744 assert_eq!(join.offset, Some(2));
745 assert_eq!(join.return_.len(), 2);
746 assert_eq!(join.return_items.len(), 2);
747 match *join.right {
748 QueryExpr::Table(table) => assert_eq!(table.alias.as_deref(), Some("rhs")),
749 QueryExpr::Graph(graph) => assert_eq!(graph.alias.as_deref(), Some("rhs")),
750 QueryExpr::Path(path) => assert_eq!(path.alias.as_deref(), Some("rhs")),
751 QueryExpr::Vector(vector) => assert_eq!(vector.alias.as_deref(), Some("rhs")),
752 QueryExpr::Hybrid(hybrid) => assert_eq!(hybrid.alias.as_deref(), Some("rhs")),
753 other => panic!("unexpected right source: {other:?}"),
754 }
755 }
756 }
757
758 #[test]
759 fn join_builder_right_alias_ignores_non_source_variants() {
760 let builder = JoinQueryBuilder {
761 left: QueryExpr::Table(TableQuery::new("left")),
762 right: QueryExpr::SetTenant(Some("acme".to_string())),
763 on: join_condition(),
764 join_type: JoinType::Inner,
765 filter: None,
766 order_by: Vec::new(),
767 limit: None,
768 offset: None,
769 return_items: Vec::new(),
770 return_: Vec::new(),
771 };
772 let query = builder.right_alias("ignored").build();
773 let QueryExpr::Join(join) = query else {
774 panic!("expected join query");
775 };
776 assert!(matches!(*join.right, QueryExpr::SetTenant(Some(ref tenant)) if tenant == "acme"));
777 }
778
779 #[test]
780 fn path_builder_sets_alias_via_filter_and_length() {
781 let query = PathQueryBuilder::new(NodeSelector::by_id("a"), NodeSelector::by_id("b"))
782 .via_label("CONNECTS_TO")
783 .max_length(7)
784 .filter(eq_filter("kind", Value::text("vpn")))
785 .alias("p")
786 .build();
787
788 let QueryExpr::Path(path) = query else {
789 panic!("expected path query");
790 };
791 assert_eq!(path.alias.as_deref(), Some("p"));
792 assert_eq!(path.via, vec!["CONNECTS_TO"]);
793 assert_eq!(path.max_length, 7);
794 assert!(path.filter.is_some());
795 }
796
797 #[test]
798 fn cte_helpers_track_recursive_state_and_lookup() {
799 let base = QueryExpr::Table(TableQuery::new("hosts"));
800 let cte = CteDefinition::new("active", base.clone())
801 .with_columns(vec!["id".to_string(), "ip".to_string()]);
802 assert_eq!(cte.name, "active");
803 assert_eq!(cte.columns, vec!["id", "ip"]);
804 assert!(!cte.recursive);
805
806 let recursive = CteDefinition::recursive("walk", base.clone());
807 assert!(recursive.recursive);
808
809 let clause = WithClause::new().add(cte).add(recursive);
810 assert!(!clause.is_empty());
811 assert!(clause.has_recursive);
812 assert!(clause.get("active").is_some());
813 assert!(clause.get("missing").is_none());
814
815 let simple = QueryWithCte::simple(base.clone());
816 assert!(simple.with_clause.is_none());
817 assert!(matches!(simple.query, QueryExpr::Table(_)));
818
819 let with_ctes = QueryWithCte::with_ctes(clause.clone(), base.clone());
820 assert!(with_ctes.with_clause.is_some());
821
822 let built = CteQueryBuilder::new()
823 .cte("one", base.clone())
824 .recursive_cte("two", base.clone())
825 .cte_with_columns("three", vec!["id".to_string()], base.clone())
826 .build(base);
827 let clause = built.with_clause.expect("with clause");
828 assert_eq!(clause.ctes.len(), 3);
829 assert!(clause.has_recursive);
830 assert_eq!(clause.get("three").expect("cte").columns, vec!["id"]);
831 }
832}