1use crate::ir::{IrGroup, IrRowSource, IrSelect, IrShapePointer, IrStmt};
37
38#[derive(Debug, Clone, PartialEq)]
45pub struct ShapePathAlias {
46 pub sql_alias: String,
47 pub path: String,
48 pub marker_offset: Option<usize>,
49}
50
51const ROOT_PATH: &str = "root";
52
53pub fn collect_shape_path_aliases(stmt: &IrStmt) -> Vec<ShapePathAlias> {
61 let mut out = Vec::new();
62 match stmt {
63 IrStmt::Select(sel) => collect_select(sel, ROOT_PATH, None, &mut out),
64 IrStmt::Insert(ins) => {
65 out.push(ShapePathAlias {
66 sql_alias: ins.target.alias.clone(),
67 path: ROOT_PATH.to_string(),
68 marker_offset: None,
69 });
70 collect_shape(&ins.returning, ROOT_PATH, &mut out);
71 }
72 IrStmt::Update(upd) => {
73 out.push(ShapePathAlias {
74 sql_alias: upd.target.alias.clone(),
75 path: ROOT_PATH.to_string(),
76 marker_offset: None,
77 });
78 collect_shape(&upd.returning, ROOT_PATH, &mut out);
79 }
80 IrStmt::Delete(del) => {
81 out.push(ShapePathAlias {
82 sql_alias: del.target.alias.clone(),
83 path: ROOT_PATH.to_string(),
84 marker_offset: None,
85 });
86 collect_shape(&del.returning, ROOT_PATH, &mut out);
87 }
88 IrStmt::Group(g) => {
89 let IrGroup { source, shape, .. } = g;
90 out.push(ShapePathAlias {
91 sql_alias: source.alias.clone(),
92 path: ROOT_PATH.to_string(),
93 marker_offset: None,
94 });
95 collect_shape(shape, ROOT_PATH, &mut out);
96 }
97 IrStmt::PathSelect(_)
98 | IrStmt::For(_)
99 | IrStmt::FunctionSelect(_)
100 | IrStmt::VectorSearch(_)
101 | IrStmt::FtsSearch(_)
102 | IrStmt::ScalarUnion(_) => {}
103 }
104 out
105}
106
107pub fn root_marker_offset(stmt: &crate::parse::Stmt) -> Option<usize> {
117 use crate::parse::{Expr, Stmt};
118 match stmt {
119 Stmt::Analyze(inner) => root_marker_offset(inner),
120 Stmt::Select(sel) => match &sel.result {
121 Expr::Shape(shape) => shape.marker_offset,
122 _ => None,
123 },
124 _ => None,
128 }
129}
130
131fn collect_select(sel: &IrSelect, path: &str, marker_offset: Option<usize>, out: &mut Vec<ShapePathAlias>) {
132 for row in &sel.rows {
133 if let IrRowSource::Bound { source, shape } = row {
134 out.push(ShapePathAlias {
135 sql_alias: source.alias.clone(),
136 path: path.to_string(),
137 marker_offset,
138 });
139 collect_shape(shape, path, out);
140 }
141 }
144}
145
146fn collect_shape(shape: &[IrShapePointer], parent_path: &str, out: &mut Vec<ShapePathAlias>) {
147 for ptr in shape {
148 match ptr {
149 IrShapePointer::SingleLink(p) => {
150 let path = format!("{parent_path}.{}", p.alias);
151 collect_select(&p.subquery, &path, p.marker_offset, out);
152 }
153 IrShapePointer::MultiLink(p) => {
154 let path = format!("{parent_path}.{}", p.alias);
155 collect_select(&p.subquery, &path, p.marker_offset, out);
156 }
157 IrShapePointer::Scalar(_) | IrShapePointer::Computed(_) | IrShapePointer::ScalarSet(_) => {}
161 IrShapePointer::Asserted(a) => collect_shape(std::slice::from_ref(&a.inner), parent_path, out),
164 }
165 }
166}
167
168use std::collections::HashMap;
178
179#[derive(Debug, Clone, serde::Deserialize)]
180struct RawExplainRoot {
181 #[serde(rename = "Plan")]
182 plan: RawPlanNode,
183}
184
185#[derive(Debug, Clone, serde::Deserialize)]
186struct RawPlanNode {
187 #[serde(rename = "Alias")]
188 alias: Option<String>,
189 #[serde(rename = "Relation Name")]
190 relation_name: Option<String>,
191 #[serde(rename = "Startup Cost", default)]
192 startup_cost: f64,
193 #[serde(rename = "Total Cost", default)]
194 total_cost: f64,
195 #[serde(rename = "Plan Rows", default)]
196 plan_rows: f64,
197 #[serde(rename = "Plan Width", default)]
198 plan_width: i64,
199 #[serde(rename = "Actual Startup Time")]
200 actual_startup_time: Option<f64>,
201 #[serde(rename = "Actual Total Time")]
202 actual_total_time: Option<f64>,
203 #[serde(rename = "Actual Rows")]
204 actual_rows: Option<f64>,
205 #[serde(rename = "Actual Loops")]
206 actual_loops: Option<f64>,
207 #[serde(rename = "Plans", default)]
208 plans: Vec<RawPlanNode>,
209}
210
211#[derive(Debug, Clone, PartialEq, serde::Serialize)]
217pub struct PlanCost {
218 pub startup_cost: f64,
219 pub total_cost: f64,
220 pub plan_rows: f64,
221 pub plan_width: i64,
222 pub actual_startup_time: Option<f64>,
223 pub actual_total_time: Option<f64>,
224 pub actual_rows: Option<f64>,
225 pub actual_loops: Option<f64>,
226}
227
228impl From<&RawPlanNode> for PlanCost {
229 fn from(raw: &RawPlanNode) -> Self {
230 PlanCost {
231 startup_cost: raw.startup_cost,
232 total_cost: raw.total_cost,
233 plan_rows: raw.plan_rows,
234 plan_width: raw.plan_width,
235 actual_startup_time: raw.actual_startup_time,
236 actual_total_time: raw.actual_total_time,
237 actual_rows: raw.actual_rows,
238 actual_loops: raw.actual_loops,
239 }
240 }
241}
242
243#[derive(Debug, Clone, PartialEq, serde::Serialize)]
247pub struct CoarseGrainedNode {
248 pub path: String,
249 pub marker_offset: Option<usize>,
254 pub relations: Vec<String>,
257 pub cost: PlanCost,
258 pub children: Vec<ChildEntry>,
259}
260
261#[derive(Debug, Clone, PartialEq, serde::Serialize)]
267pub struct ChildEntry {
268 pub name: String,
269 pub node: CoarseGrainedNode,
270}
271
272pub fn build_coarse_grained(raw_json: &str, path_aliases: &[ShapePathAlias]) -> Result<CoarseGrainedNode, String> {
278 let roots: Vec<RawExplainRoot> =
279 serde_json::from_str(raw_json).map_err(|e| format!("malformed EXPLAIN (FORMAT JSON) output: {e}"))?;
280 let root = roots
281 .into_iter()
282 .next()
283 .ok_or_else(|| "EXPLAIN produced no plan".to_string())?;
284
285 let alias_to_path: HashMap<&str, &str> = path_aliases
286 .iter()
287 .map(|p| (p.sql_alias.as_str(), p.path.as_str()))
288 .collect();
289 let path_to_marker: HashMap<&str, Option<usize>> = path_aliases
290 .iter()
291 .map(|p| (p.path.as_str(), p.marker_offset))
292 .collect();
293
294 Ok(build_node(&root.plan, ROOT_PATH, &alias_to_path, &path_to_marker))
295}
296
297fn build_node(
298 raw: &RawPlanNode,
299 path: &str,
300 alias_to_path: &HashMap<&str, &str>,
301 path_to_marker: &HashMap<&str, Option<usize>>,
302) -> CoarseGrainedNode {
303 let mut relations = Vec::new();
304 let mut seen_relations = std::collections::HashSet::new();
305 let mut children = Vec::new();
306 collect_plan_nodes(
307 raw,
308 path,
309 alias_to_path,
310 &mut relations,
311 &mut seen_relations,
312 &mut children,
313 path_to_marker,
314 );
315 let marker_offset = path_to_marker.get(path).copied().flatten();
316 CoarseGrainedNode {
317 path: path.to_string(),
318 marker_offset,
319 relations,
320 cost: PlanCost::from(raw),
321 children,
322 }
323}
324
325fn collect_plan_nodes(
342 raw: &RawPlanNode,
343 path: &str,
344 alias_to_path: &HashMap<&str, &str>,
345 relations: &mut Vec<String>,
346 seen_relations: &mut std::collections::HashSet<String>,
347 children: &mut Vec<ChildEntry>,
348 path_to_marker: &HashMap<&str, Option<usize>>,
349) {
350 if let Some(rel) = &raw.relation_name
351 && seen_relations.insert(rel.clone())
352 {
353 relations.push(rel.clone());
354 }
355 for child in &raw.plans {
356 match resolve_subtree_path(child, alias_to_path) {
357 Some(child_path) if child_path != path => {
358 let name = child_path.rsplit('.').next().unwrap_or(child_path).to_string();
359 children.push(ChildEntry {
360 name,
361 node: build_node(child, child_path, alias_to_path, path_to_marker),
362 });
363 }
364 _ => collect_plan_nodes(
365 child,
366 path,
367 alias_to_path,
368 relations,
369 seen_relations,
370 children,
371 path_to_marker,
372 ),
373 }
374 }
375}
376
377fn resolve_subtree_path<'a>(node: &RawPlanNode, alias_to_path: &HashMap<&str, &'a str>) -> Option<&'a str> {
385 if let Some(path) = node.alias.as_deref().and_then(|a| alias_to_path.get(a).copied()) {
386 return Some(path);
387 }
388 node.plans
389 .iter()
390 .find_map(|child| resolve_subtree_path(child, alias_to_path))
391}
392
393#[cfg(test)]
394mod tests {
395 use super::*;
396 use crate::schema::{LinkDescriptor, MultiLinkDescriptor, PropertyDescriptor, SchemaDescriptor, TypeDescriptor};
397 use crate::{ir, parse};
398
399 fn id_prop() -> PropertyDescriptor {
400 PropertyDescriptor {
401 name: "id".into(),
402 pg_type: "uuid".into(),
403 nullable: false,
404 default_sql: Some("uuidv7()".into()),
405 default_pyql: None,
406 description: None,
407 check_constraints: vec![],
408 is_exclusive: true,
409 is_pk: true,
410 is_readonly: true,
411 rewrites: vec![],
412 tuple_members: None,
413 column_type: None,
414 }
415 }
416
417 fn name_prop() -> PropertyDescriptor {
418 PropertyDescriptor {
419 name: "name".into(),
420 pg_type: "text".into(),
421 nullable: false,
422 default_sql: None,
423 default_pyql: None,
424 description: None,
425 check_constraints: vec![],
426 is_exclusive: false,
427 is_pk: false,
428 is_readonly: false,
429 rewrites: vec![],
430 tuple_members: None,
431 column_type: None,
432 }
433 }
434
435 fn make_schema() -> SchemaDescriptor {
436 SchemaDescriptor {
437 types: vec![
438 TypeDescriptor {
439 name: "Hero".into(),
440 module: "default".into(),
441 table: "hero".into(),
442 abstract_: false,
443 materialized: false,
444 description: None,
445 parents: vec![],
446 interfaces: vec![],
447 bases: vec![],
448 properties: vec![id_prop(), name_prop()],
449 links: vec![],
450 multilinks: vec![MultiLinkDescriptor {
451 name: "villains".into(),
452 target: "default::Villain".into(),
453 through: None,
454 nullable: false,
455 description: None,
456 default_pyql: None,
457 on_delete: vec![],
458 is_exclusive: false,
459 }],
460 computed: vec![],
461 constraints: vec![],
462 indexes: vec![],
463 partition: None,
464 vector_indexes: vec![],
465 search_indexes: vec![],
466 triggers: vec![],
467 junction: false,
468 signals: vec![],
469 },
470 TypeDescriptor {
471 name: "Villain".into(),
472 module: "default".into(),
473 table: "villain".into(),
474 abstract_: false,
475 materialized: false,
476 description: None,
477 parents: vec![],
478 interfaces: vec![],
479 bases: vec![],
480 properties: vec![id_prop(), name_prop()],
481 links: vec![LinkDescriptor {
482 name: "nemesis".into(),
483 target: "default::Hero".into(),
484 nullable: true,
485 through: None,
486 description: None,
487 default_pyql: None,
488 is_exclusive: false,
489 is_readonly: false,
490 rewrites: vec![],
491 on_delete: vec![],
492 }],
493 multilinks: vec![],
494 computed: vec![],
495 constraints: vec![],
496 indexes: vec![],
497 partition: None,
498 vector_indexes: vec![],
499 search_indexes: vec![],
500 triggers: vec![],
501 junction: false,
502 signals: vec![],
503 },
504 ],
505 scalars: vec![],
506 enums: vec![],
507 named_tuples: vec![],
508 globals: vec![],
509 functions: vec![],
510 aliases: vec![],
511 channels: vec![],
512 ..Default::default()
513 }
514 }
515
516 fn compile(query: &str) -> ir::IrOutput {
517 let schema = make_schema();
518 let ast = parse::parse(query).unwrap();
519 let ast = match ast {
524 parse::Stmt::Analyze(inner) => *inner,
525 other => other,
526 };
527 ir::compile(&ast, &schema).unwrap()
528 }
529
530 #[test]
531 fn test_root_only_query_produces_a_single_root_alias() {
532 let ir = compile("select Hero { name }");
533 let paths = collect_shape_path_aliases(&ir.stmt);
534 assert_eq!(paths.len(), 1);
535 assert_eq!(paths[0].path, "root");
536 }
537
538 #[test]
539 fn test_nested_multilink_produces_root_and_nested_paths() {
540 let ir = compile("select Hero { name, villains: { name, nemesis: { name } } }");
541 let paths = collect_shape_path_aliases(&ir.stmt);
542 let by_path: Vec<&str> = paths.iter().map(|p| p.path.as_str()).collect();
543 assert_eq!(by_path, vec!["root", "root.villains", "root.villains.nemesis"]);
544 assert_ne!(paths[0].sql_alias, paths[1].sql_alias);
546 assert_ne!(paths[1].sql_alias, paths[2].sql_alias);
547 }
548
549 #[test]
550 fn test_marker_offsets_survive_from_ast_to_the_alias_map() {
551 let query = "select Hero { name, villains: { name } }";
552 let ir = compile(query);
553 let paths = collect_shape_path_aliases(&ir.stmt);
554 let villains = paths.iter().find(|p| p.path == "root.villains").unwrap();
555 let offset = villains
556 .marker_offset
557 .expect("villains path should carry a marker offset");
558 assert_eq!(&query[offset..offset + "villains".len()], "villains");
559 }
560
561 #[test]
562 fn test_analyze_wrapped_query_compiles_to_the_same_ir_as_the_bare_inner_query() {
563 let wrapped = compile("analyze select Hero { name, villains: { name } }");
568 let bare = compile("select Hero { name, villains: { name } }");
569 assert_eq!(
570 collect_shape_path_aliases(&wrapped.stmt).len(),
571 collect_shape_path_aliases(&bare.stmt).len()
572 );
573 }
574
575 fn explain_json_fixture(root_alias: &str, villain_alias: &str) -> String {
586 format!(
587 r#"[
588 {{
589 "Plan": {{
590 "Node Type": "Seq Scan",
591 "Alias": "{root_alias}",
592 "Relation Name": "hero",
593 "Startup Cost": 0.0,
594 "Total Cost": 12.5,
595 "Plan Rows": 100,
596 "Plan Width": 40,
597 "Actual Startup Time": 0.01,
598 "Actual Total Time": 0.05,
599 "Actual Rows": 100,
600 "Actual Loops": 1,
601 "Plans": [
602 {{
603 "Node Type": "Nested Loop",
604 "Startup Cost": 0.0,
605 "Total Cost": 8.2,
606 "Plan Rows": 3,
607 "Plan Width": 32,
608 "Actual Startup Time": 0.0,
609 "Actual Total Time": 0.01,
610 "Actual Rows": 3,
611 "Actual Loops": 100,
612 "Plans": [
613 {{
614 "Node Type": "Seq Scan",
615 "Alias": "hv",
616 "Relation Name": "hero.villains",
617 "Startup Cost": 0.0,
618 "Total Cost": 2.0,
619 "Plan Rows": 3,
620 "Plan Width": 16,
621 "Actual Startup Time": 0.0,
622 "Actual Total Time": 0.0,
623 "Actual Rows": 3,
624 "Actual Loops": 100,
625 "Plans": []
626 }},
627 {{
628 "Node Type": "Seq Scan",
629 "Alias": "{villain_alias}",
630 "Relation Name": "villain",
631 "Startup Cost": 0.0,
632 "Total Cost": 2.0,
633 "Plan Rows": 1,
634 "Plan Width": 32,
635 "Actual Startup Time": 0.0,
636 "Actual Total Time": 0.0,
637 "Actual Rows": 1,
638 "Actual Loops": 3,
639 "Plans": []
640 }}
641 ]
642 }}
643 ]
644 }},
645 "Planning Time": 0.1,
646 "Execution Time": 0.2
647 }}
648 ]"#
649 )
650 }
651
652 #[test]
653 fn test_build_coarse_grained_groups_nodes_by_shape_path() {
654 let ir = compile("select Hero { name, villains: { name } }");
655 let paths = collect_shape_path_aliases(&ir.stmt);
656 let root_alias = paths.iter().find(|p| p.path == "root").unwrap().sql_alias.clone();
657 let villains_alias = paths
658 .iter()
659 .find(|p| p.path == "root.villains")
660 .unwrap()
661 .sql_alias
662 .clone();
663
664 let raw_json = explain_json_fixture(&root_alias, &villains_alias);
665 let tree = build_coarse_grained(&raw_json, &paths).unwrap();
666
667 assert_eq!(tree.path, "root");
668 assert_eq!(tree.relations, vec!["hero".to_string()]);
669 assert_eq!(tree.cost.total_cost, 12.5);
670 assert_eq!(tree.children.len(), 1);
671
672 let ChildEntry {
673 name,
674 node: villains_node,
675 } = &tree.children[0];
676 assert_eq!(name, "villains");
677 assert_eq!(villains_node.path, "root.villains");
678 assert_eq!(
683 villains_node.relations,
684 vec!["hero.villains".to_string(), "villain".to_string()]
685 );
686 assert_eq!(villains_node.cost.total_cost, 8.2);
689 assert!(villains_node.children.is_empty());
690 }
691
692 #[test]
693 fn test_build_coarse_grained_carries_marker_offsets_for_repl_rendering() {
694 let query = "select Hero { name, villains: { name } }";
695 let ir = compile(query);
696 let mut paths = collect_shape_path_aliases(&ir.stmt);
697 let ast = crate::parse::parse(query).unwrap();
701 if let Some(root) = paths.iter_mut().find(|p| p.path == "root") {
702 root.marker_offset = root_marker_offset(&ast);
703 }
704 let root_alias = paths.iter().find(|p| p.path == "root").unwrap().sql_alias.clone();
705 let villains_alias = paths
706 .iter()
707 .find(|p| p.path == "root.villains")
708 .unwrap()
709 .sql_alias
710 .clone();
711
712 let raw_json = explain_json_fixture(&root_alias, &villains_alias);
713 let tree = build_coarse_grained(&raw_json, &paths).unwrap();
714
715 let root_offset = tree.marker_offset.expect("root should carry a marker offset");
716 assert_eq!(&query[root_offset..root_offset + "Hero".len()], "Hero");
717
718 let villains_offset = tree.children[0]
719 .node
720 .marker_offset
721 .expect("villains should carry a marker offset");
722 assert_eq!(&query[villains_offset..villains_offset + "villains".len()], "villains");
723 }
724
725 #[test]
726 fn test_build_coarse_grained_with_no_path_aliases_still_produces_a_root_node() {
727 let raw_json = explain_json_fixture("t0", "t1");
731 let tree = build_coarse_grained(&raw_json, &[]).unwrap();
732 assert_eq!(tree.path, "root");
733 assert!(tree.children.is_empty());
734 }
735
736 #[test]
737 fn test_build_coarse_grained_rejects_malformed_json() {
738 assert!(build_coarse_grained("not json", &[]).is_err());
739 }
740}