Skip to main content

dol/
planner.rs

1//! Logical query plan generation.
2//!
3//! Converts a parsed [`Query`] into a [`LogicalPlan`] for display and
4//! optimisation. Plans can be shown via `dol --explain "observe containers"`.
5//! Includes a simple filter push-down optimisation.
6//!
7//! # Example
8//!
9//! ```ignore
10//! let plan = planner::plan(&query);
11//! println!("{plan}"); // Shows execution plan
12//! ```
13
14use std::fmt;
15
16use crate::ast::{
17    AlertRule, AnalyzeQuery, CollectionTarget, EventsQuery, InspectQuery, LogsQuery, ObserveQuery,
18    PipelineNode, Query, SortDirection,
19};
20
21#[derive(Debug, Clone, PartialEq)]
22pub enum LogicalPlan {
23    Observe(ObservePlan),
24    Events(EventsPlan),
25    Inspect(InspectPlan),
26    Logs(LogsPlan),
27    Compose(ComposePlan),
28    Ping,
29    Analyze(AnalyzePlan),
30    Alert(AlertPlan),
31    Fields(crate::ast::CollectionTarget),
32}
33
34#[derive(Debug, Clone, PartialEq)]
35pub struct ObservePlan {
36    pub target: CollectionTarget,
37    pub join: Option<crate::ast::JoinClause>,
38    pub steps: Vec<PlanStep>,
39}
40
41#[derive(Debug, Clone, PartialEq)]
42pub struct EventsPlan {
43    pub target: CollectionTarget,
44    pub steps: Vec<PlanStep>,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct InspectPlan {
49    pub target: crate::ast::SingularTarget,
50    pub at: Option<String>,
51}
52
53#[derive(Debug, Clone, PartialEq)]
54pub struct AnalyzePlan {
55    pub target: crate::ast::AnalysisTarget,
56    pub verb: crate::ast::AnalysisVerb,
57    pub subject: Option<String>,
58    pub steps: Vec<PlanStep>,
59}
60
61#[derive(Debug, Clone, PartialEq)]
62pub struct AlertPlan {
63    pub rule: AlertRule,
64}
65
66#[derive(Debug, Clone, PartialEq)]
67pub struct ComposePlan {
68    pub project: String,
69    pub target: crate::ast::ComposeTarget,
70    pub steps: Vec<PlanStep>,
71}
72
73#[derive(Debug, Clone, PartialEq)]
74pub struct LogsPlan {
75    pub container: String,
76    pub tail: Option<u64>,
77    pub steps: Vec<PlanStep>,
78}
79
80#[derive(Debug, Clone, PartialEq)]
81pub enum PlanStep {
82    Fetch(CollectionTarget),
83    Filter(crate::ast::Expression),
84    In {
85        field: String,
86        values: Vec<crate::ast::Value>,
87    },
88    Select(Vec<String>),
89    GroupBy(Vec<String>),
90    SortBy {
91        field: String,
92        direction: SortDirection,
93    },
94    Limit(u64),
95    Alert(String),
96    If {
97        condition: crate::ast::Expression,
98        then_branch: Vec<Self>,
99        else_branch: Option<Vec<Self>>,
100    },
101    Set {
102        field: String,
103        value: crate::ast::SetValue,
104    },
105    Debug,
106    Assert(crate::ast::Expression),
107    RowNumber(String),
108    Rank(String),
109    Lag {
110        field: String,
111        alias: String,
112        offset: u64,
113    },
114    Lead {
115        field: String,
116        alias: String,
117        offset: u64,
118    },
119}
120
121impl fmt::Display for LogicalPlan {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        match self {
124            Self::Observe(p) => {
125                writeln!(f, "ObservePlan {{ target: {:?} }}", p.target)?;
126                if let Some(join) = &p.join {
127                    writeln!(f, "  JOIN {:?} ON ...", join.right)?;
128                }
129                for step in &p.steps {
130                    writeln!(f, "  {step}")?;
131                }
132                Ok(())
133            }
134            Self::Events(p) => {
135                writeln!(f, "EventsPlan {{ target: {:?} }}", p.target)?;
136                for step in &p.steps {
137                    writeln!(f, "  {step}")?;
138                }
139                Ok(())
140            }
141            Self::Inspect(p) => {
142                write!(f, "InspectPlan {{ target: {:?}, at: ", p.target.kind)?;
143                match &p.at {
144                    Some(t) => write!(f, "Some(\"{t}\")")?,
145                    None => write!(f, "None")?,
146                }
147                writeln!(f, " }}")
148            }
149            Self::Analyze(p) => {
150                writeln!(
151                    f,
152                    "AnalyzePlan {{ verb: {:?}, subject: {:?} }}",
153                    p.verb, p.subject
154                )?;
155                for step in &p.steps {
156                    writeln!(f, "  {step}")?;
157                }
158                Ok(())
159            }
160            Self::Alert(p) => {
161                writeln!(f, "AlertPlan {{ condition: {:?} }}", p.rule.condition)
162            }
163            Self::Logs(p) => {
164                writeln!(
165                    f,
166                    "LogsPlan {{ container: {}, tail: {:?} }}",
167                    p.container, p.tail
168                )?;
169                for step in &p.steps {
170                    writeln!(f, "  {step}")?;
171                }
172                Ok(())
173            }
174            Self::Compose(p) => {
175                writeln!(
176                    f,
177                    "ComposePlan {{ project: {}, target: {:?} }}",
178                    p.project, p.target
179                )?;
180                for step in &p.steps {
181                    writeln!(f, "  {step}")?;
182                }
183                Ok(())
184            }
185            Self::Ping => {
186                writeln!(f, "Ping {{ test Docker connectivity }}")
187            }
188            Self::Fields(target) => {
189                writeln!(f, "FieldsPlan {{ target: {target:?} }}")
190            }
191        }
192    }
193}
194
195impl fmt::Display for PlanStep {
196    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197        match self {
198            Self::Debug => write!(f, "Debug"),
199            Self::Fetch(target) => write!(f, "Fetch({target:?})"),
200            Self::Filter(_) => write!(f, "Filter(<expression>)"),
201            Self::In { field, values } => {
202                write!(f, "In({field}, [")?;
203                let vals: Vec<String> = values.iter().map(|v| format!("{v:?}")).collect();
204                write!(f, "{}", vals.join(", "))?;
205                write!(f, "])")
206            }
207            Self::Select(fields) => write!(f, "Select({})", fields.join(", ")),
208            Self::GroupBy(fields) => write!(f, "GroupBy({})", fields.join(", ")),
209            Self::SortBy { field, direction } => write!(f, "SortBy({field}, {direction:?})"),
210            Self::Assert(_) => write!(f, "Assert(<condition>)"),
211            Self::Limit(n) => write!(f, "Limit({n})"),
212            Self::Alert(msg) => write!(f, "Alert(\"{msg}\")"),
213            Self::If {
214                condition,
215                then_branch,
216                else_branch,
217            } => {
218                write!(f, "If({condition:?}, then=[")?;
219                for (i, step) in then_branch.iter().enumerate() {
220                    if i > 0 {
221                        write!(f, ", ")?;
222                    }
223                    write!(f, "{step}")?;
224                }
225                write!(f, "]")?;
226                if let Some(else_b) = else_branch {
227                    write!(f, ", else=[")?;
228                    for (i, step) in else_b.iter().enumerate() {
229                        if i > 0 {
230                            write!(f, ", ")?;
231                        }
232                        write!(f, "{step}")?;
233                    }
234                    write!(f, "]")?;
235                }
236                write!(f, ")")
237            }
238            Self::RowNumber(alias) => write!(f, "RowNumber({alias})"),
239            Self::Rank(field) => write!(f, "Rank({field})"),
240            Self::Lag {
241                field,
242                alias,
243                offset,
244            } => {
245                write!(f, "Lag({field}, {alias}, offset={offset})")
246            }
247            Self::Lead {
248                field,
249                alias,
250                offset,
251            } => {
252                write!(f, "Lead({field}, {alias}, offset={offset})")
253            }
254            Self::Set { field, value: _ } => write!(f, "Set({field}, <value>)"),
255        }
256    }
257}
258
259#[must_use]
260pub fn plan(query: &Query) -> LogicalPlan {
261    match query {
262        Query::Observe(q) => plan_observe(q),
263        Query::Events(q) => plan_events(q),
264        Query::Inspect(q) => plan_inspect(q),
265        Query::Analyze(q) => plan_analyze(q),
266        Query::Alert(rule) => LogicalPlan::Alert(AlertPlan { rule: rule.clone() }),
267        Query::Compose(q) => plan_compose(q),
268        Query::Logs(q) => plan_logs(q),
269        Query::Ping => LogicalPlan::Ping,
270        Query::Fields(target) => LogicalPlan::Fields(*target),
271    }
272}
273
274fn plan_observe(query: &ObserveQuery) -> LogicalPlan {
275    let mut steps = vec![PlanStep::Fetch(query.target)];
276
277    if let Some(filter) = &query.filter {
278        steps.push(PlanStep::Filter(filter.clone()));
279    }
280
281    for node in &query.pipeline {
282        steps.push(node_to_step(node));
283    }
284
285    optimize_steps(&mut steps);
286
287    LogicalPlan::Observe(ObservePlan {
288        target: query.target,
289        join: query.join.clone(),
290        steps,
291    })
292}
293
294fn plan_events(query: &EventsQuery) -> LogicalPlan {
295    let mut steps = vec![PlanStep::Fetch(query.target)];
296
297    if let Some(filter) = &query.filter {
298        steps.push(PlanStep::Filter(filter.clone()));
299    }
300
301    for node in &query.pipeline {
302        steps.push(node_to_step(node));
303    }
304
305    LogicalPlan::Events(EventsPlan {
306        target: query.target,
307        steps,
308    })
309}
310
311fn plan_inspect(query: &InspectQuery) -> LogicalPlan {
312    LogicalPlan::Inspect(InspectPlan {
313        target: query.target.clone(),
314        at: query.at.clone(),
315    })
316}
317
318fn plan_logs(query: &LogsQuery) -> LogicalPlan {
319    let mut steps = Vec::new();
320
321    if let Some(filter) = &query.filter {
322        steps.push(PlanStep::Filter(filter.clone()));
323    }
324
325    for node in &query.pipeline {
326        steps.push(node_to_step(node));
327    }
328
329    LogicalPlan::Logs(LogsPlan {
330        container: query.container.clone(),
331        tail: query.tail,
332        steps,
333    })
334}
335
336fn plan_compose(query: &crate::ast::ComposeQuery) -> LogicalPlan {
337    let mut steps = Vec::new();
338
339    for node in &query.pipeline {
340        steps.push(node_to_step(node));
341    }
342
343    LogicalPlan::Compose(ComposePlan {
344        project: query.project.clone(),
345        target: query.target,
346        steps,
347    })
348}
349
350fn plan_analyze(query: &AnalyzeQuery) -> LogicalPlan {
351    let mut steps = Vec::new();
352    for node in &query.pipeline {
353        steps.push(node_to_step(node));
354    }
355
356    LogicalPlan::Analyze(AnalyzePlan {
357        target: query.target.clone(),
358        verb: query.verb,
359        subject: query.subject.clone(),
360        steps,
361    })
362}
363
364fn node_to_step(node: &PipelineNode) -> PlanStep {
365    match node {
366        PipelineNode::Where(expr) => match expr {
367            crate::ast::Expression::In {
368                expr: inner_expr,
369                values,
370            } => {
371                let field = match inner_expr.as_ref() {
372                    crate::ast::Expression::Field(f) => f.clone(),
373                    _ => return PlanStep::Filter(expr.clone()),
374                };
375                PlanStep::In {
376                    field,
377                    values: values.clone(),
378                }
379            }
380            other => PlanStep::Filter(other.clone()),
381        },
382        PipelineNode::Select(fields) => PlanStep::Select(fields.clone()),
383        PipelineNode::GroupBy { fields, .. } => PlanStep::GroupBy(fields.clone()),
384        PipelineNode::SortBy { fields } => {
385            // Take the first sort field for the plan step (primary sort)
386            if let Some((field, direction)) = fields.first() {
387                PlanStep::SortBy {
388                    field: field.clone(),
389                    direction: *direction,
390                }
391            } else {
392                PlanStep::SortBy {
393                    field: String::new(),
394                    direction: SortDirection::Asc,
395                }
396            }
397        }
398        PipelineNode::Limit(n) | PipelineNode::Offset(n) => PlanStep::Limit(*n),
399        PipelineNode::Alert(msg) => PlanStep::Alert(msg.clone()),
400        PipelineNode::If {
401            condition,
402            then_branch,
403            else_branch,
404        } => PlanStep::If {
405            condition: condition.clone(),
406            then_branch: then_branch.iter().map(node_to_step).collect(),
407            else_branch: else_branch
408                .as_ref()
409                .map(|nodes| nodes.iter().map(node_to_step).collect()),
410        },
411        PipelineNode::Set { field, value } => PlanStep::Set {
412            field: field.clone(),
413            value: value.clone(),
414        },
415        PipelineNode::Fill {
416            field,
417            default,
418            condition: _condition,
419        } => PlanStep::Set {
420            field: field.clone(),
421            value: default.clone(),
422        },
423        PipelineNode::RowNumber { alias } => PlanStep::RowNumber(alias.clone()),
424        PipelineNode::Rank { field, alias: _ } => PlanStep::Rank(field.clone()),
425        PipelineNode::Lag {
426            field,
427            alias,
428            offset,
429        } => PlanStep::Lag {
430            field: field.clone(),
431            alias: alias.clone(),
432            offset: *offset,
433        },
434        PipelineNode::Lead {
435            field,
436            alias,
437            offset,
438        } => PlanStep::Lead {
439            field: field.clone(),
440            alias: alias.clone(),
441            offset: *offset,
442        },
443        PipelineNode::Let { name, value } => PlanStep::Set {
444            field: name.clone(),
445            value: crate::ast::SetValue::Expr(value.clone()),
446        },
447        PipelineNode::Having(expr) => PlanStep::Filter(expr.clone()),
448        PipelineNode::Debug => PlanStep::Debug,
449        PipelineNode::Assert(expr) => PlanStep::Assert(expr.clone()),
450        PipelineNode::Distinct => PlanStep::Select(vec!["*".to_owned()]),
451    }
452}
453
454fn optimize_steps(steps: &mut Vec<PlanStep>) {
455    push_filters_early(steps);
456}
457fn push_filters_early(steps: &mut Vec<PlanStep>) {
458    let first_sort = steps
459        .iter()
460        .position(|s| matches!(s, PlanStep::SortBy { .. }));
461    let first_limit = steps.iter().position(|s| matches!(s, PlanStep::Limit(_)));
462
463    let barrier = match (first_sort, first_limit) {
464        (Some(s), Some(l)) => s.min(l),
465        (Some(s), None) => s,
466        (None, Some(l)) => l,
467        (None, None) => return,
468    };
469
470    let filter_indices: Vec<usize> = (barrier + 1..steps.len())
471        .filter(|&i| matches!(steps[i], PlanStep::Filter(_)))
472        .collect();
473
474    if filter_indices.is_empty() {
475        return;
476    }
477
478    let filters: Vec<PlanStep> = filter_indices
479        .iter()
480        .rev()
481        .map(|&i| steps.remove(i))
482        .collect();
483    for (offset, filter) in filters.into_iter().enumerate() {
484        steps.insert(1 + offset, filter);
485    }
486}
487
488#[cfg(test)]
489mod tests {
490    use crate::parser;
491
492    use super::*;
493
494    #[test]
495    fn plans_observe_query() {
496        let Query::Observe(q) = parser::parse("observe containers where state = running")
497            .unwrap()
498            .query
499        else {
500            panic!("expected observe");
501        };
502        let p = plan_observe(&q);
503        assert!(matches!(p, LogicalPlan::Observe(_)));
504        if let LogicalPlan::Observe(op) = p {
505            assert_eq!(op.target, CollectionTarget::Containers);
506            assert!(op.steps.iter().any(|s| matches!(s, PlanStep::Fetch(_))));
507            assert!(op.steps.iter().any(|s| matches!(s, PlanStep::Filter(_))));
508        }
509    }
510
511    #[test]
512    fn plans_events_query() {
513        let Query::Events(q) = parser::parse("events containers where action = \"die\"")
514            .unwrap()
515            .query
516        else {
517            panic!("expected events");
518        };
519        let p = plan_events(&q);
520        if let LogicalPlan::Events(ep) = p {
521            assert_eq!(ep.target, CollectionTarget::Containers);
522        }
523    }
524
525    #[test]
526    fn plans_inspect_query() {
527        let Query::Inspect(q) = parser::parse("inspect container api at \"2026-05-31T02:00:00Z\"")
528            .unwrap()
529            .query
530        else {
531            panic!("expected inspect");
532        };
533        let p = plan_inspect(&q);
534        if let LogicalPlan::Inspect(ip) = p {
535            assert!(ip.at.is_some());
536        }
537    }
538
539    #[test]
540    fn plans_alert_rule() {
541        let Query::Alert(rule) =
542            parser::parse("alert when cpu > 85% for 2m then print \"High CPU\"")
543                .unwrap()
544                .query
545        else {
546            panic!("expected alert");
547        };
548        let p = plan(&Query::Alert(rule));
549        assert!(matches!(p, LogicalPlan::Alert(_)));
550    }
551
552    #[test]
553    fn optimize_pushes_filters_before_sort() {
554        let mut steps = vec![
555            PlanStep::Fetch(CollectionTarget::Containers),
556            PlanStep::Select(vec!["name".into()]),
557            PlanStep::SortBy {
558                field: "cpu".into(),
559                direction: SortDirection::Desc,
560            },
561            PlanStep::Filter(crate::ast::Expression::Comparison {
562                left: Box::new(crate::ast::Expression::Field("state".into())),
563                operator: crate::ast::Operator::Eq,
564                right: Box::new(crate::ast::Expression::Literal(crate::ast::Value::String(
565                    "running".into(),
566                ))),
567            }),
568        ];
569        push_filters_early(&mut steps);
570
571        assert!(matches!(steps[0], PlanStep::Fetch(_)));
572        assert!(matches!(steps[1], PlanStep::Filter(_)));
573        assert!(matches!(steps[2], PlanStep::Select(_)));
574    }
575
576    #[test]
577    fn plans_compose_query() {
578        let q = parser::parse("compose myapp").unwrap();
579        let p = plan(&q.query);
580        if let LogicalPlan::Compose(cp) = p {
581            assert_eq!(cp.project, "myapp");
582            assert_eq!(cp.target, crate::ast::ComposeTarget::Containers);
583            assert!(cp.steps.is_empty());
584        } else {
585            panic!("expected Compose plan, got {:?}", p);
586        }
587    }
588
589    #[test]
590    fn plans_compose_services_with_pipeline() {
591        let q =
592            parser::parse("compose myapp services | where cpu > 80% | select name, cpu").unwrap();
593        let p = plan(&q.query);
594        if let LogicalPlan::Compose(cp) = p {
595            assert_eq!(cp.project, "myapp");
596            assert_eq!(cp.target, crate::ast::ComposeTarget::Services);
597            assert_eq!(cp.steps.len(), 2);
598        } else {
599            panic!("expected Compose plan, got {:?}", p);
600        }
601    }
602
603    #[test]
604    fn plans_compose_with_sort_limit() {
605        let q = parser::parse("compose myapp | sort by name asc | limit 5").unwrap();
606        let p = plan(&q.query);
607        if let LogicalPlan::Compose(cp) = p {
608            assert_eq!(cp.steps.len(), 2);
609            assert!(matches!(cp.steps[0], PlanStep::SortBy { .. }));
610            assert!(matches!(cp.steps[1], PlanStep::Limit(5)));
611        } else {
612            panic!("expected Compose plan, got {:?}", p);
613        }
614    }
615
616    #[test]
617    fn compose_plan_display() {
618        let q = parser::parse("compose myapp services").unwrap();
619        let p = plan(&q.query);
620        let s = format!("{p}");
621        assert!(s.contains("myapp"));
622        assert!(s.contains("Services"));
623    }
624
625    #[test]
626    fn plans_observe_join_images() {
627        let q = parser::parse("observe containers join images on id = id").unwrap();
628        let p = plan(&q.query);
629        if let LogicalPlan::Observe(op) = p {
630            assert!(op.join.is_some());
631            let join = op.join.unwrap();
632            assert_eq!(join.right, CollectionTarget::Images);
633        } else {
634            panic!("expected Observe plan, got {:?}", p);
635        }
636    }
637
638    #[test]
639    fn plans_observe_join_without_pipeline() {
640        let q = parser::parse("observe containers join networks on name = name").unwrap();
641        let p = plan(&q.query);
642        if let LogicalPlan::Observe(op) = p {
643            assert!(op.join.is_some());
644            assert_eq!(op.steps.len(), 1);
645            assert!(matches!(
646                op.steps[0],
647                PlanStep::Fetch(CollectionTarget::Containers)
648            ));
649        } else {
650            panic!("expected Observe plan, got {:?}", p);
651        }
652    }
653
654    #[test]
655    fn observe_join_plan_display() {
656        let q = parser::parse("observe containers join images on id = id").unwrap();
657        let p = plan(&q.query);
658        let s = format!("{p}");
659        assert!(s.contains("JOIN"));
660        assert!(s.contains("Images"));
661    }
662
663    #[test]
664    fn plans_logs_query() {
665        let Query::Logs(q) =
666            parser::parse("logs container my-api tail 200 | where message contains \"error\" | select line, message | limit 10")
667                .unwrap()
668                .query
669        else {
670            panic!("expected logs");
671        };
672        let p = plan_logs(&q);
673        if let LogicalPlan::Logs(lp) = p {
674            assert_eq!(lp.container, "my-api");
675            assert_eq!(lp.tail, Some(200));
676            assert_eq!(lp.steps.len(), 3);
677            // Steps: Filter (where), Select, Limit
678            assert!(matches!(lp.steps[0], PlanStep::Filter(_)));
679            assert!(matches!(lp.steps[1], PlanStep::Select(_)));
680            assert!(matches!(lp.steps[2], PlanStep::Limit(10)));
681        } else {
682            panic!("expected LogsPlan, got {p:?}");
683        }
684    }
685
686    #[test]
687    fn logs_plan_display() {
688        let q = parser::parse("logs container my-api tail 50").unwrap();
689        let p = plan(&q.query);
690        let s = format!("{p}");
691        assert!(s.contains("my-api"));
692        assert!(s.contains("tail: Some(50)"));
693    }
694
695    #[test]
696    fn plans_logs_no_tail() {
697        let Query::Logs(q) = parser::parse("logs container my-api").unwrap().query else {
698            panic!("expected logs");
699        };
700        let p = plan_logs(&q);
701        if let LogicalPlan::Logs(lp) = p {
702            assert_eq!(lp.container, "my-api");
703            assert_eq!(lp.tail, None);
704            assert!(lp.steps.is_empty());
705        } else {
706            panic!("expected LogsPlan, got {p:?}");
707        }
708    }
709}