siderite_orm/queryset/
mod.rs

1//! [`QuerySet`]: a lazy, immutable query over one model.
2//!
3//! Builder methods consume and return the queryset and never perform I/O;
4//! terminal methods (`all`, `first`, `get`, `count`, `exists`, `update`, ...)
5//! run the compiled [`QueryPlan`] on the queryset's [`Db`]. Clone a queryset
6//! to branch from a common base.
7//!
8//! Mistakes that can only be detected while building (an unknown annotation
9//! name, a filter on a window function) are remembered and returned by the
10//! next terminal method, so builder chains stay infallible.
11//!
12//! # Annotations
13//!
14//! [`annotate`](QuerySet::annotate) and [`alias`](QuerySet::alias) name an
15//! expression; later `filter` / `order_by` / `project` calls refer to it with
16//! `Expr::col("name")` and the expression is inlined (PostgreSQL does not
17//! allow output aliases in `WHERE` / `HAVING`). A filter on an aggregate
18//! annotation becomes `HAVING`; when the projection contains aggregates the
19//! remaining projected columns become the `GROUP BY` and the model's default
20//! ordering is dropped, as in Django.
21
22mod combine;
23mod cursor;
24mod fetch;
25mod related;
26mod rows;
27mod write;
28
29pub use cursor::CursorPage;
30pub use fetch::Page;
31pub use related::{Prefetch, Relation};
32
33use crate::db::Db;
34use crate::error::QueryError;
35use crate::expr::{Column, Expr, Ident};
36use crate::model::Model;
37use crate::plan::{DistinctMode, LockMode, OrderDirection, OrderExpr, QueryPlan, SelectExpr};
38use related::{Prefetcher, RelNode};
39use std::marker::PhantomData;
40use std::sync::Arc;
41
42/// Lazy query over model `M`.
43pub struct QuerySet<M: Model> {
44    db: Db,
45    plan: QueryPlan,
46    /// Named expressions from `annotate` / `alias`, already fully resolved.
47    annotations: Vec<(Ident, Expr)>,
48    /// The ordering still comes from `ModelMeta::ordering` (dropped for
49    /// grouped queries).
50    default_ordering: bool,
51    /// `none()`: terminals return empty results without touching the database.
52    empty: bool,
53    /// Joined models loaded by `select_related`.
54    related: Vec<RelNode>,
55    /// Extra queries run by `prefetch_related`.
56    prefetches: Vec<Arc<dyn Prefetcher<M>>>,
57    /// The token set by `after` / `before`, read by `cursor_page`.
58    cursor: Option<(cursor::Way, String)>,
59    /// First problem found while building; returned by terminals.
60    error: Option<QueryError>,
61    _model: PhantomData<fn() -> M>,
62}
63
64impl<M: Model> Clone for QuerySet<M> {
65    fn clone(&self) -> Self {
66        Self {
67            db: self.db.clone(),
68            plan: self.plan.clone(),
69            annotations: self.annotations.clone(),
70            default_ordering: self.default_ordering,
71            empty: self.empty,
72            related: self.related.clone(),
73            prefetches: self.prefetches.clone(),
74            cursor: self.cursor.clone(),
75            error: self.error.clone(),
76            _model: PhantomData,
77        }
78    }
79}
80
81impl<M: Model> std::fmt::Debug for QuerySet<M> {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("QuerySet")
84            .field("model", &M::META.name)
85            .field("plan", &self.plan)
86            .finish_non_exhaustive()
87    }
88}
89
90impl<M: Model> QuerySet<M> {
91    /// Every row of `M`'s table, in the model's default ordering.
92    pub fn new(db: &Db) -> Self {
93        let mut plan = QueryPlan::from_table(M::META.table);
94        for (column, direction) in M::META.ordering {
95            plan = plan.order_by(Expr::col(*column), *direction);
96        }
97        Self {
98            db: db.clone(),
99            plan,
100            annotations: Vec::new(),
101            default_ordering: !M::META.ordering.is_empty(),
102            empty: false,
103            related: Vec::new(),
104            prefetches: Vec::new(),
105            cursor: None,
106            error: None,
107            _model: PhantomData,
108        }
109    }
110
111    /// The database handle this queryset runs on.
112    pub fn db(&self) -> &Db {
113        &self.db
114    }
115
116    /// The plan built so far.
117    pub fn plan(&self) -> &QueryPlan {
118        &self.plan
119    }
120
121    /// Run on `db` instead (Django `using`), e.g. a replica or the current
122    /// transaction handle.
123    #[must_use]
124    pub fn using(mut self, db: &Db) -> Self {
125        self.db = db.clone();
126        self
127    }
128
129    /// Replace the plan (for extension code building on this queryset).
130    #[must_use]
131    pub fn with_plan(mut self, f: impl FnOnce(QueryPlan) -> QueryPlan) -> Self {
132        self.plan = f(self.plan);
133        self.settle()
134    }
135
136    /// Record the first building problem.
137    fn fail(mut self, message: impl Into<String>) -> Self {
138        self.error
139            .get_or_insert_with(|| QueryError::InvalidPlan(message.into()));
140        self
141    }
142
143    /// The building problem, if any; terminals call this first.
144    fn ready(&self) -> Result<(), QueryError> {
145        self.error.clone().map_or(Ok(()), Err)
146    }
147
148    /// Normalise the plan after an edit: turn related columns into joins and
149    /// derive `GROUP BY` from the projection.
150    fn settle(mut self) -> Self {
151        let plan = std::mem::replace(&mut self.plan, QueryPlan::from_table(""));
152        let mut plan = plan.resolve_relations();
153        let aggregated = plan.having.is_some()
154            || plan
155                .projection
156                .iter()
157                .map(|s| &s.expr)
158                .chain(plan.ordering.iter().map(|o| &o.expr))
159                .any(Expr::contains_aggregate);
160        if aggregated && plan.projection.is_empty() {
161            plan.projection = model_projection::<M>();
162        }
163        plan.grouping = if aggregated {
164            plan.projection
165                .iter()
166                .map(|s| &s.expr)
167                .filter(|e| !e.contains_aggregate() && !e.contains_window())
168                .cloned()
169                .collect()
170        } else {
171            Vec::new()
172        };
173        if !plan.grouping.is_empty() && self.default_ordering {
174            plan.ordering.clear();
175        }
176        self.plan = plan;
177        self
178    }
179
180    /// Inline annotation references in `expr`.
181    fn resolve(&self, mut expr: Expr) -> Expr {
182        inline_annotations(&self.annotations, &mut expr);
183        expr
184    }
185
186    /// Keep rows matching `predicate` (ANDed with earlier filters).
187    ///
188    /// A predicate over an aggregate annotation is applied as `HAVING`.
189    #[must_use]
190    pub fn filter(self, predicate: Expr) -> Self {
191        let predicate = self.resolve(predicate);
192        if predicate.contains_window() {
193            return self.fail("cannot filter on a window function; wrap the query as a subquery");
194        }
195        self.with_plan(|mut plan| {
196            if predicate.contains_aggregate() {
197                plan.having = Some(match plan.having.take() {
198                    Some(existing) => existing.and(predicate),
199                    None => predicate,
200                });
201                plan
202            } else {
203                plan.filter(predicate)
204            }
205        })
206    }
207
208    /// Drop rows matching `predicate`.
209    #[must_use]
210    pub fn exclude(self, predicate: Expr) -> Self {
211        self.filter(!predicate)
212    }
213
214    /// Replace the ordering (`order_by([User::name.asc(), User::id.desc()])`).
215    #[must_use]
216    pub fn order_by(mut self, terms: impl IntoIterator<Item = OrderExpr>) -> Self {
217        let terms: Vec<OrderExpr> = terms
218            .into_iter()
219            .map(|term| OrderExpr {
220                expr: self.resolve(term.expr),
221                direction: term.direction,
222            })
223            .collect();
224        self.default_ordering = false;
225        self.with_plan(|mut plan| {
226            plan.ordering = terms;
227            plan
228        })
229    }
230
231    /// Flip every ordering direction. An unordered queryset is ordered by
232    /// primary key descending, so `reverse().first()` is the last row.
233    #[must_use]
234    pub fn reverse(mut self) -> Self {
235        self.default_ordering = false;
236        let pk = M::META.pk().map(|f| f.column);
237        self.with_plan(|mut plan| {
238            if plan.ordering.is_empty() {
239                plan.ordering.extend(pk.map(|c| Expr::col(c).desc()));
240                return plan;
241            }
242            for term in &mut plan.ordering {
243                term.direction = match term.direction {
244                    OrderDirection::Asc => OrderDirection::Desc,
245                    OrderDirection::Desc => OrderDirection::Asc,
246                };
247            }
248            plan
249        })
250    }
251
252    /// Return at most `n` rows.
253    #[must_use]
254    pub fn limit(self, n: u64) -> Self {
255        self.with_plan(|p| p.limit(n))
256    }
257
258    /// Skip the first `n` rows.
259    #[must_use]
260    pub fn offset(self, n: u64) -> Self {
261        self.with_plan(|p| p.offset(n))
262    }
263
264    /// Remove duplicate rows (`SELECT DISTINCT`).
265    #[must_use]
266    pub fn distinct(self) -> Self {
267        self.with_plan(|p| p.distinct(DistinctMode::All))
268    }
269
270    /// Keep the first row per distinct value of `exprs` (`DISTINCT ON`,
271    /// PostgreSQL only; combine with a matching `order_by`).
272    #[must_use]
273    pub fn distinct_on<E: Into<Expr>>(self, exprs: impl IntoIterator<Item = E>) -> Self {
274        let exprs = exprs.into_iter().map(|e| self.resolve(e.into())).collect();
275        self.with_plan(|p| p.distinct(DistinctMode::On(exprs)))
276    }
277
278    /// A queryset that yields nothing and never queries the database.
279    #[must_use]
280    pub fn none(mut self) -> Self {
281        self.empty = true;
282        self
283    }
284
285    /// Name the root table in this query, so a subquery over the same table
286    /// can refer to the enclosing row with [`Expr::outer`].
287    #[must_use]
288    pub fn aliased(self, alias: impl Into<Ident>) -> Self {
289        let alias = alias.into();
290        self.with_plan(|mut p| {
291            p.source.alias = Some(alias);
292            p
293        })
294    }
295
296    /// Lock the selected rows (`SELECT .. FOR UPDATE`). Backends without row
297    /// locking (SQLite) reject the query before running it. Locks are held
298    /// until the surrounding transaction ends, so use this inside
299    /// [`Db::transaction`].
300    #[must_use]
301    pub fn select_for_update(self) -> Self {
302        self.with_plan(|p| p.lock(LockMode::ForUpdate))
303    }
304
305    /// `FOR UPDATE NOWAIT`: fail instead of waiting for locked rows.
306    #[must_use]
307    pub fn nowait(self) -> Self {
308        self.with_plan(|p| p.lock(LockMode::ForUpdateNoWait))
309    }
310
311    /// `FOR UPDATE SKIP LOCKED`: skip rows locked by other transactions.
312    #[must_use]
313    pub fn skip_locked(self) -> Self {
314        self.with_plan(|p| p.lock(LockMode::ForUpdateSkipLocked))
315    }
316
317    /// Add `expr` to the result as column `name` (Django `annotate`).
318    ///
319    /// Aggregates group by the other projected columns (see the module
320    /// docs). Read the value with [`all_annotated`](Self::all_annotated) or
321    /// [`rows`](Self::rows).
322    #[must_use]
323    pub fn annotate(self, name: impl Into<Ident>, expr: impl Into<Expr>) -> Self {
324        self.name_expression(name.into(), expr.into(), true)
325    }
326
327    /// Name `expr` for use in `filter` / `order_by` without selecting it
328    /// (Django `alias`).
329    #[must_use]
330    pub fn alias(self, name: impl Into<Ident>, expr: impl Into<Expr>) -> Self {
331        self.name_expression(name.into(), expr.into(), false)
332    }
333
334    fn name_expression(mut self, name: Ident, expr: Expr, selected: bool) -> Self {
335        if M::META.column(&name).is_some() || self.annotations.iter().any(|(n, _)| *n == name) {
336            return self.fail(format!("annotation `{name}` clashes with an existing name"));
337        }
338        let expr = self.resolve(expr);
339        self.annotations.push((name.clone(), expr.clone()));
340        self.with_plan(|mut plan| {
341            if selected {
342                if plan.projection.is_empty() {
343                    plan.projection = model_projection::<M>();
344                }
345                plan.projection.push(SelectExpr::new(expr, Some(name)));
346            }
347            plan
348        })
349    }
350}
351
352/// Every column of `M` as an explicit projection.
353fn model_projection<M: Model>() -> Vec<SelectExpr> {
354    M::META
355        .fields
356        .iter()
357        .map(|f| SelectExpr::new(Expr::col(f.column), None))
358        .collect()
359}
360
361/// Replace unqualified column references named like an annotation by the
362/// annotated expression. Nested subqueries keep their own names.
363fn inline_annotations(annotations: &[(Ident, Expr)], expr: &mut Expr) {
364    if let Expr::Column(Column { source: None, name }) = expr
365        && let Some((_, annotated)) = annotations.iter().find(|(n, _)| n == name)
366    {
367        *expr = annotated.clone();
368        return;
369    }
370    for child in expr.children_mut() {
371        inline_annotations(annotations, child);
372    }
373}
374
375impl Db {
376    /// Queryset over `M` on this handle (same as `M::objects(&db)`).
377    pub fn objects<M: Model>(&self) -> QuerySet<M> {
378        QuerySet::new(self)
379    }
380}
siderite_orm/queryset/mod.rs

siderite_orm/queryset/
mod.rs
