1use super::ParameterizingTransformer;
2use crate::parser::ast::stmt::OrderDirection;
3use crate::parser::ast::stmt::{
4 DeleteStmt, FetchStmt, FromClause, GoStmt, InsertStmt, LookupStmt, MatchStmt, Pattern,
5 ReturnClause, ReturnItem, SetClause, Stmt, UpdateStmt, YieldClause,
6};
7use linkrs_core::types::expr::Expression;
8
9#[derive(Debug)]
10pub struct TemplateExtractor;
11
12impl TemplateExtractor {
13 pub fn extract(stmt: &Stmt) -> String {
15 match stmt {
16 Stmt::Match(m) => Self::extract_match_template(m),
17 Stmt::Go(g) => Self::extract_go_template(g),
18 Stmt::Lookup(l) => Self::extract_lookup_template(l),
19 Stmt::Fetch(f) => Self::extract_fetch_template(f),
20 Stmt::Insert(i) => Self::extract_insert_template(i),
21 Stmt::Delete(d) => Self::extract_delete_template(d),
22 Stmt::Update(u) => Self::extract_update_template(u),
23 _ => stmt.category().as_str().to_string(),
24 }
25 }
26
27 fn extract_match_template(stmt: &MatchStmt) -> String {
28 let mut transformer = ParameterizingTransformer::new();
29 let mut parts = Vec::new();
30
31 let pattern_template = Self::patterns_to_template(&stmt.patterns);
32 parts.push(format!("MATCH {}", pattern_template));
33
34 if let Some(ref where_expr) = stmt.where_clause {
35 let result = transformer.parameterize(where_expr);
36 let where_template = Self::expr_to_template_string(&result.expression);
37 parts.push(format!("WHERE {}", where_template));
38 }
39
40 if let Some(ref return_clause) = stmt.return_clause {
41 let return_template = Self::return_clause_to_template(return_clause, &mut transformer);
42 parts.push(return_template);
43 }
44
45 if let Some(ref order_by) = stmt.order_by {
46 let order_items: Vec<String> = order_by
47 .items
48 .iter()
49 .map(|item| {
50 let result = transformer.parameterize(&item.expression);
51 let expr_str = Self::expr_to_template_string(&result.expression);
52 let dir_str = match item.direction {
53 OrderDirection::Asc => "ASC",
54 OrderDirection::Desc => "DESC",
55 };
56 format!("{} {}", expr_str, dir_str)
57 })
58 .collect();
59 parts.push(format!("ORDER BY {}", order_items.join(", ")));
60 }
61
62 if let Some(skip) = &stmt.skip {
63 parts.push(format!("SKIP ${}", skip.count));
64 }
65
66 if let Some(limit) = &stmt.limit {
67 parts.push(format!("LIMIT ${}", limit.count));
68 }
69
70 if stmt.optional {
71 parts.insert(0, "OPTIONAL".to_string());
72 }
73
74 parts.join(" ")
75 }
76
77 fn extract_go_template(stmt: &GoStmt) -> String {
78 let mut transformer = ParameterizingTransformer::new();
79 let mut parts = Vec::new();
80
81 let step_template = match &stmt.steps {
82 crate::parser::ast::Steps::Fixed(n) => format!("{} STEPS", n),
83 crate::parser::ast::Steps::Range { min, max } => {
84 format!("{} TO {} STEPS", min, max)
85 }
86 crate::parser::ast::Steps::Variable(_) => "VARIABLE STEPS".to_string(),
87 };
88 parts.push(format!("GO {}", step_template));
89
90 let from_template = Self::from_clause_to_template(&stmt.from, &mut transformer);
91 parts.push(from_template);
92
93 if let Some(ref over) = stmt.over {
94 let edge_types = over.edge_types.join(", ");
95 let dir_str = match over.direction {
96 crate::parser::ast::EdgeDirection::Out => "",
97 crate::parser::ast::EdgeDirection::In => "REVERSELY ",
98 crate::parser::ast::EdgeDirection::Both => "BIDIRECT ",
99 };
100 parts.push(format!("OVER {}{}", dir_str, edge_types));
101 }
102
103 if let Some(ref where_expr) = stmt.where_clause {
104 let result = transformer.parameterize(where_expr);
105 let where_template = Self::expr_to_template_string(&result.expression);
106 parts.push(format!("WHERE {}", where_template));
107 }
108
109 if let Some(ref yield_clause) = stmt.yield_clause {
110 let yield_template = Self::yield_clause_to_template(yield_clause, &mut transformer);
111 parts.push(yield_template);
112 }
113
114 parts.join(" ")
115 }
116
117 fn extract_lookup_template(stmt: &LookupStmt) -> String {
118 let mut transformer = ParameterizingTransformer::new();
119 let mut parts = Vec::new();
120
121 let target_str = match &stmt.target {
122 crate::parser::ast::LookupTarget::Tag(name) => format!("ON {}", name),
123 crate::parser::ast::LookupTarget::Edge(name) => format!("ON {}", name),
124 crate::parser::ast::LookupTarget::Unspecified(name) => format!("ON {}", name),
125 };
126 parts.push(format!("LOOKUP {}", target_str));
127
128 if let Some(ref where_expr) = stmt.where_clause {
129 let result = transformer.parameterize(where_expr);
130 let where_template = Self::expr_to_template_string(&result.expression);
131 parts.push(format!("WHERE {}", where_template));
132 }
133
134 if let Some(ref yield_clause) = stmt.yield_clause {
135 let yield_template = Self::yield_clause_to_template(yield_clause, &mut transformer);
136 parts.push(yield_template);
137 }
138
139 parts.join(" ")
140 }
141
142 fn extract_fetch_template(stmt: &FetchStmt) -> String {
143 let mut transformer = ParameterizingTransformer::new();
144 let mut parts = Vec::new();
145
146 match &stmt.target {
147 crate::parser::ast::FetchTarget::Vertices {
148 ids, properties, ..
149 } => {
150 parts.push("FETCH VERTEX".to_string());
151
152 let id_templates: Vec<String> = ids
153 .iter()
154 .map(|id| {
155 let result = transformer.parameterize(id);
156 Self::expr_to_template_string(&result.expression)
157 })
158 .collect();
159 parts.push(id_templates.join(", "));
160
161 if let Some(props) = properties {
162 parts.push(format!("YIELD {}", props.join(", ")));
163 }
164 }
165 crate::parser::ast::FetchTarget::Edges {
166 src,
167 dst,
168 edge_type,
169 rank,
170 properties,
171 } => {
172 parts.push(format!("FETCH EDGE ON {}", edge_type));
173
174 let src_result = transformer.parameterize(src);
175 let dst_result = transformer.parameterize(dst);
176 parts.push(format!(
177 "{} -> {}",
178 Self::expr_to_template_string(&src_result.expression),
179 Self::expr_to_template_string(&dst_result.expression)
180 ));
181
182 if let Some(r) = rank {
183 let rank_result = transformer.parameterize(r);
184 parts.push(format!(
185 "@{}",
186 Self::expr_to_template_string(&rank_result.expression)
187 ));
188 }
189
190 if let Some(props) = properties {
191 parts.push(format!("YIELD {}", props.join(", ")));
192 }
193 }
194 }
195
196 parts.join(" ")
197 }
198
199 fn extract_insert_template(stmt: &InsertStmt) -> String {
200 let mut transformer = ParameterizingTransformer::new();
201 let mut parts = Vec::new();
202
203 match &stmt.target {
204 crate::parser::ast::InsertTarget::Vertices { tag, values } => {
205 parts.push("INSERT VERTEX".to_string());
206
207 parts.push(tag.tag_name.clone());
208
209 if !tag.prop_names.is_empty() {
210 parts.push(format!("({})", tag.prop_names.join(", ")));
211 }
212
213 parts.push("VALUES".to_string());
214
215 for row in values {
216 let vid_result = transformer.parameterize(&row.vid);
217 let vid_template = Self::expr_to_template_string(&vid_result.expression);
218
219 let value_templates: Vec<String> = row
220 .values
221 .iter()
222 .map(|v| {
223 let result = transformer.parameterize(v);
224 Self::expr_to_template_string(&result.expression)
225 })
226 .collect();
227 let tag_values_templates = [format!("({})", value_templates.join(", "))];
228
229 parts.push(format!(
230 "{}: {}",
231 vid_template,
232 tag_values_templates.join(", ")
233 ));
234 }
235 }
236 crate::parser::ast::InsertTarget::Edge {
237 edge_name,
238 prop_names,
239 edges,
240 } => {
241 parts.push(format!("INSERT EDGE {}", edge_name));
242
243 if !prop_names.is_empty() {
244 parts.push(format!("({})", prop_names.join(", ")));
245 }
246
247 parts.push("VALUES".to_string());
248
249 for (src, dst, rank, values) in edges {
250 let src_result = transformer.parameterize(src);
251 let dst_result = transformer.parameterize(dst);
252
253 let value_templates: Vec<String> = values
254 .iter()
255 .map(|v| {
256 let result = transformer.parameterize(v);
257 Self::expr_to_template_string(&result.expression)
258 })
259 .collect();
260
261 let mut edge_str = format!(
262 "{} -> {}",
263 Self::expr_to_template_string(&src_result.expression),
264 Self::expr_to_template_string(&dst_result.expression)
265 );
266
267 if let Some(r) = rank {
268 let rank_result = transformer.parameterize(r);
269 edge_str.push_str(&format!(
270 "@{}",
271 Self::expr_to_template_string(&rank_result.expression)
272 ));
273 }
274
275 if !value_templates.is_empty() {
276 edge_str.push_str(&format!(": ({})", value_templates.join(", ")));
277 }
278
279 parts.push(edge_str);
280 }
281 }
282 }
283
284 if stmt.if_not_exists {
285 parts.push("IF NOT EXISTS".to_string());
286 }
287
288 parts.join(" ")
289 }
290
291 fn extract_delete_template(stmt: &DeleteStmt) -> String {
292 let mut transformer = ParameterizingTransformer::new();
293 let mut parts = Vec::new();
294
295 match &stmt.target {
296 crate::parser::ast::DeleteTarget::Vertices { tag, vids } => {
297 parts.push(format!("DELETE VERTEX {} FROM", tag));
298
299 let id_templates: Vec<String> = vids
300 .iter()
301 .map(|expr| {
302 let result = transformer.parameterize(expr);
303 Self::expr_to_template_string(&result.expression)
304 })
305 .collect();
306 parts.push(id_templates.join(", "));
307 }
308 crate::parser::ast::DeleteTarget::Edges { edge_type, edges } => {
309 if let Some(et) = edge_type {
310 parts.push(format!("DELETE EDGE {}", et));
311 } else {
312 parts.push("DELETE EDGE".to_string());
313 }
314
315 for (src, dst, rank) in edges {
316 let src_result = transformer.parameterize(src);
317 let dst_result = transformer.parameterize(dst);
318
319 let mut edge_str = format!(
320 "{} -> {}",
321 Self::expr_to_template_string(&src_result.expression),
322 Self::expr_to_template_string(&dst_result.expression)
323 );
324
325 if let Some(r) = rank {
326 let rank_result = transformer.parameterize(r);
327 edge_str.push_str(&format!(
328 "@{}",
329 Self::expr_to_template_string(&rank_result.expression)
330 ));
331 }
332
333 parts.push(edge_str);
334 }
335 }
336 crate::parser::ast::DeleteTarget::Index(name) => {
337 parts.push(format!("DELETE INDEX {}", name));
338 }
339 }
340
341 if stmt.with_edge {
342 parts.push("WITH EDGE".to_string());
343 }
344
345 parts.join(" ")
346 }
347
348 fn extract_update_template(stmt: &UpdateStmt) -> String {
349 let mut transformer = ParameterizingTransformer::new();
350 let mut parts = Vec::new();
351
352 parts.push("UPDATE".to_string());
353
354 match &stmt.target {
355 crate::parser::ast::UpdateTarget::Vertex(expr) => {
356 let result = transformer.parameterize(expr);
357 parts.push(format!(
358 "VERTEX {}",
359 Self::expr_to_template_string(&result.expression)
360 ));
361 }
362 crate::parser::ast::UpdateTarget::Edge {
363 src,
364 dst,
365 edge_type,
366 rank,
367 } => {
368 let src_result = transformer.parameterize(src);
369 let dst_result = transformer.parameterize(dst);
370
371 let mut edge_str = format!(
372 "EDGE {} -> {}",
373 Self::expr_to_template_string(&src_result.expression),
374 Self::expr_to_template_string(&dst_result.expression)
375 );
376
377 if let Some(et) = edge_type {
378 edge_str.push_str(&format!(" OF {}", et));
379 }
380
381 if let Some(r) = rank {
382 let rank_result = transformer.parameterize(r);
383 edge_str.push_str(&format!(
384 "@{}",
385 Self::expr_to_template_string(&rank_result.expression)
386 ));
387 }
388
389 parts.push(edge_str);
390 }
391 crate::parser::ast::UpdateTarget::Tag(name) => {
392 parts.push(format!("TAG {}", name));
393 }
394 crate::parser::ast::UpdateTarget::TagOnVertex { vid, tag_name } => {
395 let vid_result = transformer.parameterize(vid);
396 parts.push(format!(
397 "VERTEX {} ON {}",
398 Self::expr_to_template_string(&vid_result.expression),
399 tag_name
400 ));
401 }
402 }
403
404 let set_template = Self::set_clause_to_template(&stmt.set_clause, &mut transformer);
405 parts.push(set_template);
406
407 if let Some(ref where_expr) = stmt.where_clause {
408 let result = transformer.parameterize(where_expr);
409 let where_template = Self::expr_to_template_string(&result.expression);
410 parts.push(format!("WHERE {}", where_template));
411 }
412
413 if stmt.is_upsert {
414 parts.push("UPSERT".to_string());
415 }
416
417 if let Some(ref yield_clause) = stmt.yield_clause {
418 let yield_template = Self::yield_clause_to_template(yield_clause, &mut transformer);
419 parts.push(yield_template);
420 }
421
422 parts.join(" ")
423 }
424
425 fn patterns_to_template(patterns: &[Pattern]) -> String {
426 patterns
427 .iter()
428 .map(Self::pattern_to_template)
429 .collect::<Vec<_>>()
430 .join(", ")
431 }
432
433 fn pattern_to_template(pattern: &Pattern) -> String {
434 match pattern {
435 Pattern::Node(node) => {
436 let mut parts = Vec::new();
437
438 if let Some(ref var) = node.variable {
439 parts.push(var.clone());
440 }
441
442 if !node.labels.is_empty() {
443 parts.push(format!(":{}", node.labels.join(":")));
444 }
445
446 if node.properties.is_some() {
447 parts.push("{...}".to_string());
448 }
449
450 if !node.predicates.is_empty() {
451 parts.push("WHERE ...".to_string());
452 }
453
454 format!("({})", parts.join(""))
455 }
456 Pattern::Edge(edge) => {
457 let mut parts = Vec::new();
458
459 if let Some(ref var) = edge.variable {
460 parts.push(var.clone());
461 }
462
463 if !edge.edge_types.is_empty() {
464 parts.push(format!(":{}", edge.edge_types.join("|")));
465 }
466
467 if edge.properties.is_some() {
468 parts.push("{...}".to_string());
469 }
470
471 let (prefix, suffix) = match edge.direction {
472 crate::parser::ast::EdgeDirection::Out => ("-[", "]->"),
473 crate::parser::ast::EdgeDirection::In => ("<-[", "]-"),
474 crate::parser::ast::EdgeDirection::Both => ("-[", "]-"),
475 };
476
477 format!("{}{}{}", prefix, parts.join(""), suffix)
478 }
479 Pattern::Path(path) => {
480 let elements: Vec<String> = path
481 .elements
482 .iter()
483 .map(|e| match e {
484 crate::parser::ast::PathElement::Node(n) => {
485 Self::pattern_to_template(&Pattern::Node(n.clone()))
486 }
487 crate::parser::ast::PathElement::Edge(e) => {
488 Self::pattern_to_template(&Pattern::Edge(e.clone()))
489 }
490 crate::parser::ast::PathElement::CreateEdge(_) => "(...)".to_string(),
491 crate::parser::ast::PathElement::Alternative(patterns) => {
492 let alts: Vec<String> =
493 patterns.iter().map(Self::pattern_to_template).collect();
494 format!("({})", alts.join(" | "))
495 }
496 crate::parser::ast::PathElement::Optional(elem) => {
497 format!("{}?", Self::path_element_to_template(elem))
498 }
499 crate::parser::ast::PathElement::Repeated(elem, rep) => {
500 let rep_str = match rep {
501 crate::parser::ast::RepetitionType::ZeroOrMore => "*",
502 crate::parser::ast::RepetitionType::OneOrMore => "+",
503 crate::parser::ast::RepetitionType::ZeroOrOne => "?",
504 crate::parser::ast::RepetitionType::Exactly(n) => {
505 &format!("{{{}}}", n)
506 }
507 crate::parser::ast::RepetitionType::Range(min, max) => {
508 &format!("{{{},{}}}", min, max)
509 }
510 };
511 format!("{}{}", Self::path_element_to_template(elem), rep_str)
512 }
513 crate::parser::ast::PathElement::Recursive(rc) => {
514 let mut parts = vec![rc.variable.clone()];
515 if let Some(ref edge_var) = rc.edge_variable {
516 parts.push(edge_var.clone());
517 }
518 if rc.filter_predicate.is_some() {
519 parts.push("WHERE".to_string());
520 }
521 if rc.node_projection.is_some() || rc.edge_projection.is_some() {
522 parts.push("|".to_string());
523 let mut projs = vec![];
524 if rc.node_projection.is_some() {
525 projs.push("{node}".to_string());
526 }
527 if rc.edge_projection.is_some() {
528 projs.push("{edge}".to_string());
529 }
530 parts.push(projs.join(", "));
531 }
532 format!("*({})", parts.join(", "))
533 }
534 })
535 .collect();
536 elements.join("")
537 }
538 Pattern::Variable(var) => format!("@{}", var.name),
539 }
540 }
541
542 fn path_element_to_template(elem: &crate::parser::ast::PathElement) -> String {
543 match elem {
544 crate::parser::ast::PathElement::Node(n) => {
545 Self::pattern_to_template(&Pattern::Node(n.clone()))
546 }
547 crate::parser::ast::PathElement::Edge(e) => {
548 Self::pattern_to_template(&Pattern::Edge(e.clone()))
549 }
550 _ => "(...)".to_string(),
551 }
552 }
553
554 pub(crate) fn expr_to_template_string(expr: &Expression) -> String {
555 match expr {
556 Expression::Variable(name) if name.starts_with('$') => name.clone(),
557 Expression::Variable(name) => name.clone(),
558 Expression::Literal(value) => format!("{:?}", value),
559 Expression::Property { object, property } => {
560 format!("{}.{}", Self::expr_to_template_string(object), property)
561 }
562 Expression::StructField { base, field } => {
563 format!("{}.{}", Self::expr_to_template_string(base), field)
564 }
565 Expression::Binary { left, op, right } => {
566 format!(
567 "({} {} {})",
568 Self::expr_to_template_string(left),
569 op,
570 Self::expr_to_template_string(right)
571 )
572 }
573 Expression::Unary { op, operand } => {
574 format!("({}{})", op, Self::expr_to_template_string(operand))
575 }
576 Expression::Function { name, args } => {
577 let arg_strs: Vec<String> = args
578 .iter()
579 .map(|a| Self::expr_to_template_string(a.as_expr()))
580 .collect();
581 format!("{}({})", name, arg_strs.join(", "))
582 }
583 Expression::Aggregate {
584 func,
585 args,
586 distinct,
587 filter,
588 } => {
589 let distinct_str = if *distinct { "DISTINCT " } else { "" };
590 let filter_str = filter
591 .as_ref()
592 .map(|f| format!(" FILTER (WHERE {})", Self::expr_to_template_string(f)))
593 .unwrap_or_default();
594 let arg_strs: Vec<String> =
595 args.iter().map(Self::expr_to_template_string).collect();
596 format!(
597 "{}({}{}{})",
598 func,
599 distinct_str,
600 arg_strs.join(", "),
601 filter_str
602 )
603 }
604 Expression::List(items) => {
605 let item_strs: Vec<String> =
606 items.iter().map(Self::expr_to_template_string).collect();
607 format!("[{}]", item_strs.join(", "))
608 }
609 Expression::Map(pairs) => {
610 let pair_strs: Vec<String> = pairs
611 .iter()
612 .map(|(k, v)| format!("{}: {}", k, Self::expr_to_template_string(v)))
613 .collect();
614 format!("{{{}}}", pair_strs.join(", "))
615 }
616 Expression::Case {
617 test_expr,
618 conditions,
619 default,
620 } => {
621 let mut parts = Vec::new();
622 parts.push("CASE".to_string());
623
624 if let Some(test) = test_expr {
625 parts.push(Self::expr_to_template_string(test));
626 }
627
628 for (cond, val) in conditions {
629 parts.push(format!(
630 "WHEN {} THEN {}",
631 Self::expr_to_template_string(cond),
632 Self::expr_to_template_string(val)
633 ));
634 }
635
636 if let Some(def) = default {
637 parts.push(format!("ELSE {}", Self::expr_to_template_string(def)));
638 }
639
640 parts.push("END".to_string());
641 parts.join(" ")
642 }
643 Expression::TypeCast {
644 expression,
645 target_type,
646 } => {
647 format!(
648 "CAST({} AS {:?})",
649 Self::expr_to_template_string(expression),
650 target_type
651 )
652 }
653 Expression::Subscript { collection, index } => {
654 format!(
655 "{}[{}]",
656 Self::expr_to_template_string(collection),
657 Self::expr_to_template_string(index)
658 )
659 }
660 Expression::Label(name) => name.clone(),
661 Expression::Parameter(name) => format!("@{}", name),
662 Expression::SessionVariable(name) => format!("${}", name),
663 _ => "...".to_string(),
664 }
665 }
666
667 fn return_clause_to_template(
668 clause: &ReturnClause,
669 transformer: &mut ParameterizingTransformer,
670 ) -> String {
671 let mut parts = Vec::new();
672
673 if clause.distinct {
674 parts.push("DISTINCT".to_string());
675 }
676
677 let item_strs: Vec<String> = clause
678 .items
679 .iter()
680 .map(|item| match item {
681 ReturnItem::Expression { expression, alias } => {
682 let result = transformer.parameterize(expression);
683 let mut expr_str = Self::expr_to_template_string(&result.expression);
684 if let Some(a) = alias {
685 expr_str.push_str(&format!(" AS {}", a));
686 }
687 expr_str
688 }
689 })
690 .collect();
691
692 parts.push(item_strs.join(", "));
693
694 format!("RETURN {}", parts.join(" "))
695 }
696
697 fn yield_clause_to_template(
698 clause: &YieldClause,
699 transformer: &mut ParameterizingTransformer,
700 ) -> String {
701 let mut parts = Vec::new();
702
703 let item_strs: Vec<String> = clause
704 .items
705 .iter()
706 .map(|item| {
707 let result = transformer.parameterize(&item.expression);
708 let mut expr_str = Self::expr_to_template_string(&result.expression);
709 if let Some(ref a) = item.alias {
710 expr_str.push_str(&format!(" AS {}", a));
711 }
712 expr_str
713 })
714 .collect();
715
716 parts.push(format!("YIELD {}", item_strs.join(", ")));
717
718 if let Some(ref where_expr) = clause.where_clause {
719 let result = transformer.parameterize(where_expr);
720 parts.push(format!(
721 "WHERE {}",
722 Self::expr_to_template_string(&result.expression)
723 ));
724 }
725
726 parts.join(" ").to_string()
727 }
728
729 fn from_clause_to_template(
730 clause: &FromClause,
731 transformer: &mut ParameterizingTransformer,
732 ) -> String {
733 let vertex_strs: Vec<String> = clause
734 .vertices
735 .iter()
736 .map(|expr| {
737 let result = transformer.parameterize(expr);
738 Self::expr_to_template_string(&result.expression)
739 })
740 .collect();
741
742 format!("FROM {}", vertex_strs.join(", "))
743 }
744
745 fn set_clause_to_template(
746 clause: &SetClause,
747 transformer: &mut ParameterizingTransformer,
748 ) -> String {
749 let assignment_strs: Vec<String> = clause
750 .assignments
751 .iter()
752 .map(|assign| {
753 let result = transformer.parameterize(&assign.value);
754 format!(
755 "{} = {}",
756 assign.property,
757 Self::expr_to_template_string(&result.expression)
758 )
759 })
760 .collect();
761
762 format!("SET {}", assignment_strs.join(", "))
763 }
764}
765
766impl Default for TemplateExtractor {
767 fn default() -> Self {
768 TemplateExtractor
769 }
770}