siderite_orm/
plan.rs

1//! [`QueryPlan`]: the backend-neutral intermediate representation.
2//!
3//! Every high-level query API compiles into a `QueryPlan`; backends compile a
4//! plan into SQL, a MongoDB pipeline, etc. Relational-only parts (joins,
5//! locking, `DISTINCT ON`) are reported through
6//! [`QueryPlan::required_features`] so they can be checked against
7//! [`BackendCapabilities`] before execution.
8
9use crate::capabilities::{BackendCapabilities, Feature};
10use crate::error::BackendCapabilityError;
11use crate::expr::path_alias_of;
12use crate::expr::{AggFunc, Aggregate, Column, Expr, Ident, Lookup, RelatedColumn, WindowFunc};
13
14/// What the query reads from.
15#[derive(Debug, Clone, PartialEq)]
16pub struct QuerySource {
17    /// Table or collection name (for a subquery source, its alias).
18    pub name: Ident,
19    /// Optional alias.
20    pub alias: Option<Ident>,
21    /// Derived table: the source is the result of this query.
22    pub subquery: Option<Box<QueryPlan>>,
23}
24
25impl QuerySource {
26    /// Unaliased source.
27    pub fn table(name: impl Into<Ident>) -> Self {
28        Self {
29            name: name.into(),
30            alias: None,
31            subquery: None,
32        }
33    }
34
35    /// Derived table `(plan) AS alias`.
36    pub fn subquery(plan: QueryPlan, alias: impl Into<Ident>) -> Self {
37        Self {
38            name: alias.into(),
39            alias: None,
40            subquery: Some(Box::new(plan)),
41        }
42    }
43
44    /// Name other parts of the query use to refer to this source.
45    pub fn reference(&self) -> &Ident {
46        self.alias.as_ref().unwrap_or(&self.name)
47    }
48}
49
50/// A projected expression.
51#[derive(Debug, Clone, PartialEq)]
52pub struct SelectExpr {
53    /// Expression.
54    pub expr: Expr,
55    /// Output name.
56    pub alias: Option<Ident>,
57}
58
59impl SelectExpr {
60    /// Projection of `expr` under an optional output name.
61    pub fn new(expr: Expr, alias: Option<Ident>) -> Self {
62        Self { expr, alias }
63    }
64}
65
66/// A column of the root model, named after itself.
67impl From<&'static str> for SelectExpr {
68    fn from(column: &'static str) -> Self {
69        Self::new(Expr::col(column), Some(column.into()))
70    }
71}
72
73/// An expression under an explicit output name: `("total", Sum::of(..))`.
74impl<N: Into<Ident>, E: Into<Expr>> From<(N, E)> for SelectExpr {
75    fn from((name, expr): (N, E)) -> Self {
76        Self::new(expr.into(), Some(name.into()))
77    }
78}
79
80/// Join type.
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum JoinKind {
83    /// `INNER JOIN`.
84    Inner,
85    /// `LEFT OUTER JOIN`.
86    Left,
87}
88
89/// A join to another source (relational capability).
90#[derive(Debug, Clone, PartialEq)]
91pub struct JoinExpr {
92    /// Join kind.
93    pub kind: JoinKind,
94    /// Joined source.
95    pub source: QuerySource,
96    /// Join condition.
97    pub on: Expr,
98}
99
100/// Sort direction.
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum OrderDirection {
103    /// Ascending.
104    Asc,
105    /// Descending.
106    Desc,
107}
108
109/// One ordering term.
110#[derive(Debug, Clone, PartialEq)]
111pub struct OrderExpr {
112    /// Expression.
113    pub expr: Expr,
114    /// Direction.
115    pub direction: OrderDirection,
116}
117
118/// Distinct mode.
119#[derive(Debug, Clone, PartialEq, Default)]
120pub enum DistinctMode {
121    /// No de-duplication.
122    #[default]
123    None,
124    /// `DISTINCT`.
125    All,
126    /// `DISTINCT ON (...)` (PostgreSQL only).
127    On(Vec<Expr>),
128}
129
130/// Row-lock request.
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub enum LockMode {
133    /// `FOR UPDATE`.
134    ForUpdate,
135    /// `FOR UPDATE NOWAIT`.
136    ForUpdateNoWait,
137    /// `FOR UPDATE SKIP LOCKED`.
138    ForUpdateSkipLocked,
139}
140
141/// How a query is combined with another one.
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum SetOp {
144    /// `UNION` (distinct rows).
145    Union,
146    /// `UNION ALL`.
147    UnionAll,
148    /// `INTERSECT`.
149    Intersect,
150    /// `EXCEPT`.
151    Except,
152}
153
154/// Another query combined with the plan by a [`SetOp`].
155#[derive(Debug, Clone, PartialEq)]
156pub struct Compound {
157    /// Operator.
158    pub op: SetOp,
159    /// Right-hand query.
160    pub plan: QueryPlan,
161}
162
163/// Backend-neutral read query.
164///
165/// Built immutably: every builder method consumes and returns `self`, so a
166/// `QuerySet` can clone a base plan and branch safely.
167#[derive(Debug, Clone, PartialEq)]
168pub struct QueryPlan {
169    /// Root source.
170    pub source: QuerySource,
171    /// Projection; empty means "all columns of the root source".
172    pub projection: Vec<SelectExpr>,
173    /// Joins.
174    pub joins: Vec<JoinExpr>,
175    /// `WHERE`.
176    pub filter: Option<Expr>,
177    /// `GROUP BY`.
178    pub grouping: Vec<Expr>,
179    /// `HAVING`.
180    pub having: Option<Expr>,
181    /// `ORDER BY`.
182    pub ordering: Vec<OrderExpr>,
183    /// `LIMIT`.
184    pub limit: Option<u64>,
185    /// `OFFSET`.
186    pub offset: Option<u64>,
187    /// Distinct mode.
188    pub distinct: DistinctMode,
189    /// Row lock.
190    pub lock: Option<LockMode>,
191    /// Queries combined with this one (`UNION` ...). `ordering`, `limit` and
192    /// `offset` of this plan then apply to the combined result, and every
193    /// member must project the same number of columns.
194    pub compound: Vec<Compound>,
195    /// Database the plan was built against, recorded when a `QuerySet`
196    /// becomes a subquery. Never compiled; `Db` uses it to reject a subquery
197    /// that belongs to another database.
198    pub origin: Option<PlanOrigin>,
199}
200
201/// Identity of the database a subquery plan was built against.
202#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
203pub struct PlanOrigin(usize);
204
205impl PlanOrigin {
206    pub(crate) fn new(id: usize) -> Self {
207        Self(id)
208    }
209}
210
211impl QueryPlan {
212    /// Plan selecting every row of `table`.
213    pub fn from_table(table: impl Into<Ident>) -> Self {
214        Self {
215            source: QuerySource::table(table),
216            projection: Vec::new(),
217            joins: Vec::new(),
218            filter: None,
219            grouping: Vec::new(),
220            having: None,
221            ordering: Vec::new(),
222            limit: None,
223            offset: None,
224            distinct: DistinctMode::None,
225            lock: None,
226            compound: Vec::new(),
227            origin: None,
228        }
229    }
230
231    /// Plan selecting every row of the derived table `(inner) AS alias`.
232    pub fn from_subquery(inner: QueryPlan, alias: impl Into<Ident>) -> Self {
233        let mut plan = Self::from_table("");
234        plan.source = QuerySource::subquery(inner, alias);
235        plan
236    }
237
238    /// Add a projected expression.
239    #[must_use]
240    pub fn select(mut self, expr: impl Into<Expr>, alias: Option<&'static str>) -> Self {
241        self.projection.push(SelectExpr {
242            expr: expr.into(),
243            alias: alias.map(Into::into),
244        });
245        self
246    }
247
248    /// AND a predicate into the filter.
249    #[must_use]
250    pub fn filter(mut self, predicate: Expr) -> Self {
251        self.filter = Some(match self.filter.take() {
252            Some(existing) => existing.and(predicate),
253            None => predicate,
254        });
255        self
256    }
257
258    /// AND the negation of a predicate into the filter (Django `exclude`).
259    #[must_use]
260    pub fn exclude(self, predicate: Expr) -> Self {
261        self.filter(!predicate)
262    }
263
264    /// Append an ordering term.
265    #[must_use]
266    pub fn order_by(mut self, expr: impl Into<Expr>, direction: OrderDirection) -> Self {
267        self.ordering.push(OrderExpr {
268            expr: expr.into(),
269            direction,
270        });
271        self
272    }
273
274    /// Set `LIMIT`.
275    #[must_use]
276    pub fn limit(mut self, n: u64) -> Self {
277        self.limit = Some(n);
278        self
279    }
280
281    /// Set `OFFSET`.
282    #[must_use]
283    pub fn offset(mut self, n: u64) -> Self {
284        self.offset = Some(n);
285        self
286    }
287
288    /// Add a join.
289    #[must_use]
290    pub fn join(mut self, kind: JoinKind, source: QuerySource, on: Expr) -> Self {
291        self.joins.push(JoinExpr { kind, source, on });
292        self
293    }
294
295    /// Request a row lock.
296    #[must_use]
297    pub fn lock(mut self, mode: LockMode) -> Self {
298        self.lock = Some(mode);
299        self
300    }
301
302    /// Set distinct mode.
303    #[must_use]
304    pub fn distinct(mut self, mode: DistinctMode) -> Self {
305        self.distinct = mode;
306        self
307    }
308
309    /// Every expression of this plan (not of nested plans).
310    fn exprs(&self) -> impl Iterator<Item = &Expr> {
311        let distinct_on = match &self.distinct {
312            DistinctMode::On(exprs) => exprs.as_slice(),
313            _ => &[],
314        };
315        self.projection
316            .iter()
317            .map(|s| &s.expr)
318            .chain(self.joins.iter().map(|j| &j.on))
319            .chain(self.filter.iter())
320            .chain(self.having.iter())
321            .chain(self.grouping.iter())
322            .chain(self.ordering.iter().map(|o| &o.expr))
323            .chain(distinct_on)
324    }
325
326    /// Features this plan needs from a backend (deduplicated).
327    pub fn required_features(&self) -> Vec<Feature> {
328        let mut out = Vec::new();
329        let mut need = |f| {
330            if !out.contains(&f) {
331                out.push(f);
332            }
333        };
334        if !self.joins.is_empty() {
335            need(Feature::Joins);
336        }
337        match self.lock {
338            Some(LockMode::ForUpdate) => need(Feature::RowLocking),
339            Some(_) => {
340                need(Feature::RowLocking);
341                need(Feature::LockModifiers);
342            }
343            None => {}
344        }
345        if matches!(self.distinct, DistinctMode::On(_)) {
346            need(Feature::DistinctOn);
347        }
348        if !self.compound.is_empty() {
349            need(Feature::SetOperations);
350        }
351        let nested = self
352            .compound
353            .iter()
354            .map(|c| &c.plan)
355            .chain(self.source.subquery.as_deref());
356        for plan in nested {
357            plan.required_features().into_iter().for_each(&mut need);
358        }
359        for e in self.exprs() {
360            visit(e, &mut need);
361        }
362        out
363    }
364
365    /// Replace [`Expr::Related`] columns by joined-table columns, adding the
366    /// needed `LEFT JOIN`s (deduplicated by alias) to this plan and to nested
367    /// plans. Idempotent; querysets call it after every builder step.
368    #[must_use]
369    pub fn resolve_relations(mut self) -> Self {
370        let root = self.source.reference().clone();
371        let mut joins = std::mem::take(&mut self.joins);
372        // Resolving a join's `ON` may append joins; index so the list can grow.
373        for i in 0..joins.len() {
374            let mut on = std::mem::replace(&mut joins[i].on, Expr::Or(Vec::new()));
375            resolve_expr(&mut on, &root, &mut joins);
376            joins[i].on = on;
377        }
378        let mut exprs: Vec<&mut Expr> = Vec::new();
379        exprs.extend(self.projection.iter_mut().map(|s| &mut s.expr));
380        exprs.extend(self.filter.iter_mut());
381        exprs.extend(self.having.iter_mut());
382        exprs.extend(self.grouping.iter_mut());
383        exprs.extend(self.ordering.iter_mut().map(|o| &mut o.expr));
384        if let DistinctMode::On(on) = &mut self.distinct {
385            exprs.extend(on.iter_mut());
386        }
387        for e in exprs {
388            resolve_expr(e, &root, &mut joins);
389        }
390        self.joins = joins;
391        self.compound = std::mem::take(&mut self.compound)
392            .into_iter()
393            .map(|c| Compound {
394                op: c.op,
395                plan: c.plan.resolve_relations(),
396            })
397            .collect();
398        if let Some(inner) = self.source.subquery.take() {
399            self.source.subquery = Some(Box::new(inner.resolve_relations()));
400        }
401        self
402    }
403
404    /// Whether this plan, or any plan nested in it, was built against a
405    /// database other than `origin`.
406    pub fn has_foreign_origin(&self, origin: PlanOrigin) -> bool {
407        if self.origin.is_some_and(|o| o != origin) {
408            return true;
409        }
410        let exprs = self
411            .projection
412            .iter()
413            .map(|s| &s.expr)
414            .chain(self.filter.iter())
415            .chain(self.having.iter())
416            .chain(self.grouping.iter())
417            .chain(self.ordering.iter().map(|o| &o.expr))
418            .chain(self.joins.iter().map(|j| &j.on));
419        exprs.into_iter().any(|e| e.has_foreign_origin(origin))
420            || self
421                .compound
422                .iter()
423                .any(|c| c.plan.has_foreign_origin(origin))
424            || self
425                .source
426                .subquery
427                .as_deref()
428                .is_some_and(|p| p.has_foreign_origin(origin))
429            || self
430                .joins
431                .iter()
432                .filter_map(|j| j.source.subquery.as_deref())
433                .any(|p| p.has_foreign_origin(origin))
434    }
435
436    /// Check the plan against backend capabilities.
437    pub fn check(&self, caps: &BackendCapabilities) -> Result<(), BackendCapabilityError> {
438        self.required_features()
439            .into_iter()
440            .try_for_each(|f| caps.require(f))
441    }
442}
443
444impl Expr {
445    /// Whether a subquery inside this expression was built against a
446    /// database other than `origin`.
447    pub fn has_foreign_origin(&self, origin: PlanOrigin) -> bool {
448        self.subplan().is_some_and(|p| p.has_foreign_origin(origin))
449            || self
450                .children()
451                .into_iter()
452                .any(|c| c.has_foreign_origin(origin))
453    }
454
455    /// Features this expression (including nested subqueries) needs.
456    pub fn required_features(&self) -> Vec<Feature> {
457        let mut out = Vec::new();
458        visit(self, &mut |f| {
459            if !out.contains(&f) {
460                out.push(f);
461            }
462        });
463        out
464    }
465}
466
467fn visit(expr: &Expr, need: &mut impl FnMut(Feature)) {
468    match expr {
469        Expr::Lookup {
470            lookup: Lookup::Regex(_),
471            ..
472        } => need(Feature::Regex),
473        Expr::Related(_) => need(Feature::Joins),
474        Expr::Aggregate(agg) => aggregate_features(agg, need),
475        Expr::Window(window) => {
476            need(Feature::WindowFunctions);
477            if let WindowFunc::Aggregate(agg) = &window.func {
478                aggregate_features(agg, need);
479            }
480        }
481        _ => {}
482    }
483    if let Some(plan) = expr.subplan() {
484        need(Feature::Subqueries);
485        plan.required_features().into_iter().for_each(&mut *need);
486    }
487    for child in expr.children() {
488        visit(child, need);
489    }
490}
491
492fn aggregate_features(agg: &Aggregate, need: &mut impl FnMut(Feature)) {
493    match agg.func {
494        AggFunc::StdDev { .. } | AggFunc::Variance { .. } => need(Feature::StatisticalAggregates),
495        AggFunc::ArrayAgg => need(Feature::Arrays),
496        _ => {}
497    }
498}
499
500fn resolve_expr(expr: &mut Expr, root: &Ident, joins: &mut Vec<JoinExpr>) {
501    if let Expr::Related(related) = expr {
502        *expr = Expr::Column(join_path(related, root, joins));
503        return;
504    }
505    if let Some(sub) = expr.subplan_mut() {
506        let taken = std::mem::replace(sub, QueryPlan::from_table(""));
507        *sub = taken.resolve_relations();
508    }
509    for child in expr.children_mut() {
510        resolve_expr(child, root, joins);
511    }
512}
513
514/// Add the joins for `related.path` (skipping existing aliases) and return
515/// the qualified column.
516fn join_path(related: &RelatedColumn, root: &Ident, joins: &mut Vec<JoinExpr>) -> Column {
517    let mut parent = root.to_string();
518    for depth in 1..=related.path.len() {
519        let hop = &related.path[depth - 1];
520        let alias = path_alias_of(&related.path[..depth]);
521        if !joins
522            .iter()
523            .any(|j| j.source.alias.as_deref() == Some(alias.as_str()))
524        {
525            joins.push(JoinExpr {
526                kind: JoinKind::Left,
527                source: QuerySource {
528                    name: hop.table.clone(),
529                    alias: Some(alias.clone().into()),
530                    subquery: None,
531                },
532                on: Expr::Column(Column::qualified(parent.clone(), hop.fk_column.clone())).eq(
533                    Expr::Column(Column::qualified(alias.clone(), hop.pk_column.clone())),
534                ),
535            });
536        }
537        parent = alias;
538    }
539    Column::qualified(parent, related.column.clone())
540}
541
542#[cfg(test)]
543mod tests {
544    use super::*;
545    use crate::capabilities::BackendKind;
546
547    #[test]
548    fn filters_accumulate_with_and() {
549        let p = QueryPlan::from_table("users")
550            .filter(Expr::col("a").eq(1))
551            .filter(Expr::col("b").eq(2));
552        assert!(matches!(p.filter, Some(Expr::And(ref v)) if v.len() == 2));
553    }
554
555    #[test]
556    fn branching_is_independent() {
557        let base = QueryPlan::from_table("users").filter(Expr::col("active").eq(true));
558        let a = base.clone().limit(10);
559        let b = base.clone().order_by(Expr::col("id"), OrderDirection::Desc);
560        assert_eq!(base.limit, None);
561        assert_eq!(a.limit, Some(10));
562        assert!(b.limit.is_none() && b.ordering.len() == 1);
563    }
564
565    #[test]
566    fn capability_check_detects_nested_regex_and_locks() {
567        let sub = QueryPlan::from_table("t").filter(Expr::col("x").lookup_regex("^a"));
568        let p = QueryPlan::from_table("users")
569            .filter(Expr::Exists(Box::new(sub)))
570            .lock(LockMode::ForUpdateSkipLocked);
571        let feats = p.required_features();
572        assert!(feats.contains(&Feature::Regex));
573        assert!(feats.contains(&Feature::LockModifiers));
574        let err = p.check(&BackendCapabilities::sqlite());
575        assert_eq!(
576            err,
577            Err(BackendCapabilityError::RowLockingUnsupported {
578                backend: BackendKind::Sqlite
579            })
580        );
581        assert!(p.check(&BackendCapabilities::postgres()).is_ok());
582    }
583
584    trait RegexExt {
585        fn lookup_regex(self, p: &str) -> Expr;
586    }
587    impl RegexExt for Expr {
588        fn lookup_regex(self, p: &str) -> Expr {
589            Expr::Lookup {
590                expr: Box::new(self),
591                lookup: Lookup::Regex(p.into()),
592            }
593        }
594    }
595}
siderite_orm/plan.rs

siderite_orm/
plan.rs
