Skip to main content

rws_server/model/
eager.rs

1//! Eager loading (18.30.0): a [`Relation`] names how one model reaches
2//! another; [`QueryBuilder::with`](super::QueryBuilder::with) loads it for a
3//! whole result set in one more query per relation (`WHERE fk IN (…)`), not
4//! one per row.
5//!
6//! ```rust,ignore
7//! use rust_web_server::model::{Order, Relation};
8//!
9//! impl Post {
10//!     pub fn comments() -> Relation<Post, Comment> { Relation::has_many("post_id") }
11//!     pub fn author() -> Relation<Post, User> { Relation::belongs_to("user_id") }
12//!     pub fn tags() -> Relation<Post, Tag> { Relation::belongs_to_many("post_tags", "post_id", "tag_id") }
13//! }
14//!
15//! // Three queries: the posts, their comments, their authors.
16//! let posts = Post::query(&pool)
17//!     .order_by("id", Order::Desc)
18//!     .with(Post::comments())
19//!     .with(Post::author())
20//!     .fetch_all()
21//!     .await?;
22//! for post in &posts {
23//!     let comments: &[Comment] = post.related(&Post::comments());
24//!     let author: Option<&User> = post.one(&Post::author());
25//!     println!("{} by {:?}: {} comments", post.title, author.map(|a| &a.name), comments.len());
26//! }
27//! ```
28//!
29//! A relation's own relations load with [`EagerQuery::with_nested`]
30//! (18.35.0): one more query per nested relation, and
31//! [`Loaded::related_loaded`] gives the related rows with theirs:
32//!
33//! ```rust,ignore
34//! // Four queries: authors, their posts, the posts' comments, the posts' labels.
35//! let authors = Author::query(&pool)
36//!     .with_nested(Author::posts(), |posts| posts.with(Post::comments()).with(Post::labels()))
37//!     .fetch_all()
38//!     .await?;
39//! for post in authors[0].related_loaded(&Author::posts()) {
40//!     println!("{}: {} comments", post.title, post.related(&Post::comments()).len());
41//! }
42//! ```
43//!
44//! Related rows that are soft-deleted (`#[soft_delete]`) are left out, as
45//! everywhere else. Keys are sent in batches of [`IN_BATCH`], so a result
46//! set larger than that costs one query per batch.
47
48use std::any::Any;
49use std::collections::HashMap;
50use std::future::Future;
51use std::marker::PhantomData;
52use std::pin::Pin;
53
54use super::pool::DbPool;
55use super::query::QueryBuilder;
56use super::repository::placeholder;
57use super::{DbError, Model, ModelRow, Value};
58use crate::pagination::Page;
59
60/// How many keys one `IN (…)` list carries; more are sent in more queries.
61pub const IN_BATCH: usize = 500;
62
63/// The column eager loading of a many-to-many relation adds to each related
64/// row: the owner's key, from the join table.
65const OWNER_KEY: &str = "rws_owner_key";
66
67// ── RelationKind ──────────────────────────────────────────────────────────────
68
69/// What a [`Relation`] is.
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
71#[non_exhaustive]
72pub enum RelationKind {
73    /// The related table's foreign key holds the owner's primary key; any
74    /// number of rows.
75    HasMany,
76    /// As `HasMany`, at most one row (the one with the lowest key).
77    HasOne,
78    /// The owner's foreign key holds the related row's primary key.
79    BelongsTo,
80    /// Through a join table holding both keys.
81    BelongsToMany,
82}
83
84// ── Relation ──────────────────────────────────────────────────────────────────
85
86/// How a model `O` reaches related rows of `R`: for
87/// [`QueryBuilder::with`](super::QueryBuilder::with), and to load the rows
88/// of one owner ([`load`](Relation::load)) or of many
89/// ([`load_for`](Relation::load_for)). Usually returned by an associated
90/// function of the owner, so call sites name it:
91///
92/// ```rust,ignore
93/// impl User {
94///     pub fn posts() -> Relation<User, Post> { Relation::has_many("user_id") }
95/// }
96/// let posts: Vec<Post> = User::posts().load(&pool, &user).await?;
97/// ```
98pub struct Relation<O, R> {
99    kind: RelationKind,
100    /// `HasMany`/`HasOne`: the related table's column; `BelongsTo`: the
101    /// owner's; `BelongsToMany`: the join table's column holding the
102    /// owner's key.
103    foreign_key: &'static str,
104    /// `BelongsToMany`: the join table and its column holding the related
105    /// row's key.
106    join: Option<(&'static str, &'static str)>,
107    _models: PhantomData<fn() -> (O, R)>,
108}
109
110impl<O, R> Clone for Relation<O, R> {
111    fn clone(&self) -> Self {
112        *self
113    }
114}
115
116impl<O, R> Copy for Relation<O, R> {}
117
118impl<O, R> std::fmt::Debug for Relation<O, R> {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct("Relation")
121            .field("kind", &self.kind)
122            .field("foreign_key", &self.foreign_key)
123            .field("join", &self.join)
124            .finish()
125    }
126}
127
128impl<O: Model, R: Model> Relation<O, R> {
129    /// One owner, many `R`: `R`'s column `foreign_key` holds the owner's
130    /// primary key (`Relation::<User, Post>::has_many("user_id")`).
131    pub fn has_many(foreign_key: &'static str) -> Self {
132        Relation { kind: RelationKind::HasMany, foreign_key, join: None, _models: PhantomData }
133    }
134
135    /// As [`has_many`](Self::has_many), at most one `R` per owner (the one
136    /// with the lowest primary key when there are more).
137    pub fn has_one(foreign_key: &'static str) -> Self {
138        Relation { kind: RelationKind::HasOne, foreign_key, join: None, _models: PhantomData }
139    }
140
141    /// The owner's column `foreign_key` holds `R`'s primary key
142    /// (`Relation::<Post, User>::belongs_to("user_id")`); `NULL` means
143    /// none.
144    pub fn belongs_to(foreign_key: &'static str) -> Self {
145        Relation { kind: RelationKind::BelongsTo, foreign_key, join: None, _models: PhantomData }
146    }
147
148    /// Many-to-many through `join_table`, whose column `owner_key` holds
149    /// the owner's primary key and `related_key` `R`'s
150    /// (`Relation::<Post, Tag>::belongs_to_many("post_tags", "post_id",
151    /// "tag_id")`). [`BelongsToMany`](super::BelongsToMany) adds and
152    /// removes the links.
153    pub fn belongs_to_many(join_table: &'static str, owner_key: &'static str, related_key: &'static str) -> Self {
154        Relation { kind: RelationKind::BelongsToMany, foreign_key: owner_key, join: Some((join_table, related_key)), _models: PhantomData }
155    }
156
157    /// What the relation is.
158    pub fn kind(&self) -> RelationKind {
159        self.kind
160    }
161
162    /// The column holding the key that links the two: `R`'s for
163    /// `HasMany`/`HasOne`, `O`'s for `BelongsTo`, the join table's owner
164    /// column for `BelongsToMany`.
165    pub fn foreign_key(&self) -> &'static str {
166        self.foreign_key
167    }
168
169    /// `BelongsToMany`: the join table and its column holding `R`'s key.
170    pub fn join_table(&self) -> Option<(&'static str, &'static str)> {
171        self.join
172    }
173
174    /// Identifies the relation among an owner's loaded ones.
175    fn key(&self) -> String {
176        format!("{:?}:{}:{}:{:?}", self.kind, R::table_name(), self.foreign_key, self.join)
177    }
178
179    /// The related rows of one owner (one query).
180    pub async fn load(&self, pool: &DbPool, owner: &O) -> Result<Vec<R>, DbError>
181    where
182        R: Clone,
183    {
184        Ok(self.load_for(pool, std::slice::from_ref(owner)).await?.pop().unwrap_or_default())
185    }
186
187    /// The related rows of every owner, in the owners' order — one query
188    /// per [`IN_BATCH`] distinct keys, the related rows by primary key. A
189    /// row related to several owners (a shared author, a tag) is cloned
190    /// into each.
191    pub async fn load_for(&self, pool: &DbPool, owners: &[O]) -> Result<Vec<Vec<R>>, DbError>
192    where
193        R: Clone,
194    {
195        // The key each owner is matched by.
196        let owner_keys: Vec<Value> = match self.kind {
197            RelationKind::BelongsTo => {
198                let mut keys = Vec::with_capacity(owners.len());
199                for owner in owners {
200                    let value = owner.to_values().into_iter().find(|(c, _)| c.eq_ignore_ascii_case(self.foreign_key)).map(|(_, v)| v);
201                    keys.push(value.ok_or_else(|| {
202                        DbError::new(format!("belongs_to: {} has no column {}", O::table_name(), self.foreign_key))
203                    })?);
204                }
205                keys
206            }
207            _ => owners.iter().map(|o| o.primary_key_value()).collect(),
208        };
209        let distinct = distinct_keys(&owner_keys);
210        let pk = R::primary_key_name();
211        let (rows, match_column) = match self.kind {
212            RelationKind::HasMany | RelationKind::HasOne => {
213                let live = live_and::<R>("");
214                let rows = fetch_in(pool, &distinct, |list| {
215                    format!("SELECT * FROM {} WHERE {} IN ({list}){live} ORDER BY {pk}", R::table_name(), self.foreign_key)
216                })
217                .await?;
218                (rows, self.foreign_key)
219            }
220            RelationKind::BelongsTo => {
221                let live = live_and::<R>("");
222                let rows = fetch_in(pool, &distinct, |list| {
223                    format!("SELECT * FROM {} WHERE {pk} IN ({list}){live} ORDER BY {pk}", R::table_name())
224                })
225                .await?;
226                (rows, pk)
227            }
228            RelationKind::BelongsToMany => {
229                let (join, related_key) = self.join.unwrap_or(("", ""));
230                let live = live_and::<R>("r.");
231                let rows = fetch_in(pool, &distinct, |list| {
232                    format!(
233                        "SELECT r.*, j.{owner} AS {OWNER_KEY} FROM {} r INNER JOIN {join} j ON j.{related_key} = r.{pk} \
234                         WHERE j.{owner} IN ({list}){live} ORDER BY r.{pk}",
235                        R::table_name(),
236                        owner = self.foreign_key,
237                    )
238                })
239                .await?;
240                (rows, OWNER_KEY)
241            }
242        };
243        let mut by_key: HashMap<String, Vec<R>> = HashMap::new();
244        for row in &rows {
245            let Some(key) = column(row, match_column).and_then(key_of) else { continue };
246            by_key.entry(key).or_default().push(R::from_row(row)?);
247        }
248        let one = matches!(self.kind, RelationKind::HasOne | RelationKind::BelongsTo);
249        Ok(owner_keys
250            .iter()
251            .map(|k| {
252                let mut related = key_of(k.clone()).and_then(|k| by_key.get(&k)).cloned().unwrap_or_default();
253                if one {
254                    related.truncate(1);
255                }
256                related
257            })
258            .collect())
259    }
260}
261
262/// `" AND <prefix>deleted_at IS NULL"` for a soft-delete model.
263fn live_and<T: Model>(prefix: &str) -> String {
264    T::soft_delete_column().map(|c| format!(" AND {prefix}{c} IS NULL")).unwrap_or_default()
265}
266
267/// A row's value of `name`.
268fn column(row: &ModelRow, name: &str) -> Option<Value> {
269    row.columns.iter().find(|(c, _)| c.eq_ignore_ascii_case(name)).map(|(_, v)| v.clone())
270}
271
272/// A value as a map key: `None` for `NULL` (it matches nothing). Integers
273/// and their text spelling are different keys — the drivers decode a column
274/// as one or the other, consistently.
275pub(crate) fn key_of(value: Value) -> Option<String> {
276    match value {
277        Value::Null => None,
278        Value::Bool(b) => Some(format!("i{}", b as i64)),
279        Value::Int(n) => Some(format!("i{n}")),
280        Value::Float(f) => Some(format!("f{f}")),
281        Value::Text(s) => Some(format!("t{s}")),
282        Value::Bytes(b) => Some(format!("b{}", b.iter().map(|x| format!("{x:02x}")).collect::<String>())),
283    }
284}
285
286/// `values` without `NULL`s and repeats, in first-seen order.
287pub(crate) fn distinct_keys(values: &[Value]) -> Vec<Value> {
288    let mut seen = std::collections::HashSet::new();
289    values.iter().filter(|v| key_of((*v).clone()).is_some_and(|k| seen.insert(k))).cloned().collect()
290}
291
292/// Rows of `sql(list)` for every batch of `keys`, `list` being the batch's
293/// placeholders.
294async fn fetch_in(pool: &DbPool, keys: &[Value], sql: impl Fn(String) -> String) -> Result<Vec<ModelRow>, DbError> {
295    let backend = pool.backend();
296    let mut rows = Vec::new();
297    for batch in keys.chunks(IN_BATCH) {
298        let list = (1..=batch.len()).map(|i| placeholder(backend, i)).collect::<Vec<_>>().join(", ");
299        rows.extend(pool.query_rows(&sql(list), batch).await?);
300    }
301    Ok(rows)
302}
303
304// ── Loaded ────────────────────────────────────────────────────────────────────
305
306/// A model with the relations [`QueryBuilder::with`](super::QueryBuilder::with)
307/// loaded for it. Derefs to the model.
308pub struct Loaded<T> {
309    model: T,
310    relations: Vec<(String, Box<dyn Any + Send + Sync>)>,
311}
312
313impl<T: Model> Loaded<T> {
314    /// The model.
315    pub fn model(&self) -> &T {
316        &self.model
317    }
318
319    /// The model, relations dropped.
320    pub fn into_inner(self) -> T {
321        self.model
322    }
323
324    /// The rows of `relation` (empty when there are none, or when the query
325    /// didn't load it — see [`is_loaded`](Self::is_loaded)).
326    pub fn related<R: Model + 'static>(&self, relation: &Relation<T, R>) -> &[R] {
327        self.get(relation).map(Vec::as_slice).unwrap_or(&[])
328    }
329
330    /// The first row of `relation`: the one of a `belongs_to` / `has_one`.
331    pub fn one<R: Model + 'static>(&self, relation: &Relation<T, R>) -> Option<&R> {
332        self.related(relation).first()
333    }
334
335    /// The rows of `relation` with the relations
336    /// [`with_nested`](EagerQuery::with_nested) loaded for them (empty when
337    /// it was loaded without, or not at all).
338    pub fn related_loaded<R: Model + 'static>(&self, relation: &Relation<T, R>) -> &[Loaded<R>] {
339        let key = nested_key(&relation.key());
340        self.relations.iter().find(|(k, _)| *k == key).and_then(|(_, v)| v.downcast_ref::<Vec<Loaded<R>>>()).map(Vec::as_slice).unwrap_or(&[])
341    }
342
343    /// Whether the query loaded `relation`.
344    pub fn is_loaded<R: Model + 'static>(&self, relation: &Relation<T, R>) -> bool {
345        self.get(relation).is_some()
346    }
347
348    /// The rows of `relation`, taken out (empty when not loaded).
349    pub fn take<R: Model + 'static>(&mut self, relation: &Relation<T, R>) -> Vec<R> {
350        let key = relation.key();
351        match self.relations.iter().position(|(k, _)| *k == key) {
352            Some(i) => self.relations.remove(i).1.downcast::<Vec<R>>().map(|b| *b).unwrap_or_default(),
353            None => Vec::new(),
354        }
355    }
356
357    fn get<R: Model + 'static>(&self, relation: &Relation<T, R>) -> Option<&Vec<R>> {
358        let key = relation.key();
359        self.relations.iter().find(|(k, _)| *k == key).and_then(|(_, v)| v.downcast_ref::<Vec<R>>())
360    }
361}
362
363impl<T> std::ops::Deref for Loaded<T> {
364    type Target = T;
365    fn deref(&self) -> &T {
366        &self.model
367    }
368}
369
370impl<T: std::fmt::Debug> std::fmt::Debug for Loaded<T> {
371    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
372        f.debug_struct("Loaded")
373            .field("model", &self.model)
374            .field("relations", &self.relations.iter().map(|(k, _)| k.as_str()).collect::<Vec<_>>())
375            .finish()
376    }
377}
378
379// ── EagerQuery ────────────────────────────────────────────────────────────────
380
381type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
382type Erased = Box<dyn Any + Send + Sync>;
383/// What a loader answers: for each owner, its entries by key.
384type PerOwner = Vec<Vec<(String, Erased)>>;
385
386/// A relation of `O` with its related type erased: what it loads for each
387/// owner, under one or more keys.
388trait Loader<O>: Send + Sync {
389    fn key(&self) -> String;
390    fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>>;
391}
392
393impl<O: Model + 'static, R: Model + Clone + 'static> Loader<O> for Relation<O, R> {
394    fn key(&self) -> String {
395        Relation::key(self)
396    }
397
398    fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>> {
399        Box::pin(async move {
400            let key = Relation::key(self);
401            let related = self.load_for(pool, owners).await?;
402            Ok(related.into_iter().map(|r| vec![(key.clone(), Box::new(r) as Erased)]).collect())
403        })
404    }
405}
406
407/// The key a relation's rows *with their own relations* are kept under.
408fn nested_key(key: &str) -> String {
409    format!("nested:{key}")
410}
411
412/// A relation loaded with relations of its own ([`EagerQuery::with_nested`]).
413struct NestedLoader<O, R: Model> {
414    relation: Relation<O, R>,
415    loaders: Vec<Box<dyn Loader<R>>>,
416}
417
418impl<O: Model + 'static, R: Model + Clone + 'static> Loader<O> for NestedLoader<O, R> {
419    fn key(&self) -> String {
420        nested_key(&Relation::key(&self.relation))
421    }
422
423    fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>> {
424        Box::pin(async move {
425            let key = Relation::key(&self.relation);
426            let per_owner = self.relation.load_for(pool, owners).await?;
427            // Every owner's rows in one list, their relations loaded at once,
428            // then back to their owners.
429            let counts: Vec<usize> = per_owner.iter().map(Vec::len).collect();
430            let flat: Vec<R> = per_owner.iter().flatten().cloned().collect();
431            let mut loaded = attach(pool, flat, &self.loaders).await?.into_iter();
432            let mut out = Vec::with_capacity(owners.len());
433            for (plain, n) in per_owner.into_iter().zip(counts) {
434                let nested: Vec<Loaded<R>> = loaded.by_ref().take(n).collect();
435                out.push(vec![(key.clone(), Box::new(plain) as Erased), (nested_key(&key), Box::new(nested) as Erased)]);
436            }
437            Ok(out)
438        })
439    }
440}
441
442/// The relations to load for a relation's rows: built in
443/// [`EagerQuery::with_nested`]'s closure.
444pub struct Nested<R: Model> {
445    loaders: Vec<Box<dyn Loader<R>>>,
446}
447
448impl<R: Model + 'static> Nested<R> {
449    /// Load `relation` for the related rows too.
450    pub fn with<S: Model + Clone + 'static>(mut self, relation: Relation<R, S>) -> Self {
451        let key = relation.key();
452        if !self.loaders.iter().any(|l| l.key() == key) {
453            self.loaders.push(Box::new(relation));
454        }
455        self
456    }
457
458    /// Load `relation` for the related rows, with relations of its own.
459    pub fn with_nested<S: Model + Clone + 'static>(mut self, relation: Relation<R, S>, nested: impl FnOnce(Nested<S>) -> Nested<S>) -> Self {
460        let loaders = nested(Nested { loaders: Vec::new() }).loaders;
461        let loader = NestedLoader { relation, loaders };
462        let key = Loader::<R>::key(&loader);
463        if !self.loaders.iter().any(|l| l.key() == key) {
464            self.loaders.push(Box::new(loader));
465        }
466        self
467    }
468}
469
470/// A [`QueryBuilder`] with relations to load: from
471/// [`QueryBuilder::with`](super::QueryBuilder::with). Its `fetch_*` answer
472/// [`Loaded`] models.
473pub struct EagerQuery<'a, T: Model> {
474    query: QueryBuilder<'a, T>,
475    loaders: Vec<Box<dyn Loader<T>>>,
476}
477
478impl<'a, T: Model + 'static> EagerQuery<'a, T> {
479    pub(crate) fn new(query: QueryBuilder<'a, T>) -> Self {
480        EagerQuery { query, loaders: Vec::new() }
481    }
482
483    /// Load `relation` too (once, however often it is named).
484    pub fn with<R: Model + Clone + 'static>(mut self, relation: Relation<T, R>) -> Self {
485        let key = relation.key();
486        if !self.loaders.iter().any(|l| l.key() == key) {
487            self.loaders.push(Box::new(relation));
488        }
489        self
490    }
491
492    /// Load `relation`, and for its rows the relations `nested` names
493    /// (18.35.0): `.with_nested(Author::posts(), |p| p.with(Post::comments()))`.
494    /// Read them with [`Loaded::related_loaded`]; [`Loaded::related`] still
495    /// gives the plain rows.
496    pub fn with_nested<R: Model + Clone + 'static>(mut self, relation: Relation<T, R>, nested: impl FnOnce(Nested<R>) -> Nested<R>) -> Self {
497        let loaders = nested(Nested { loaders: Vec::new() }).loaders;
498        let loader = NestedLoader { relation, loaders };
499        let key = Loader::<T>::key(&loader);
500        if !self.loaders.iter().any(|l| l.key() == key) {
501            self.loaders.push(Box::new(loader));
502        }
503        self
504    }
505
506    /// The query's rows, each with its relations: one query for the rows,
507    /// then one per relation.
508    pub async fn fetch_all(self) -> Result<Vec<Loaded<T>>, DbError> {
509        let pool = self.query.pool();
510        let models = self.query.fetch_all().await?;
511        attach(pool, models, &self.loaders).await
512    }
513
514    /// The query's first row, with its relations.
515    pub async fn fetch_one(self) -> Result<Option<Loaded<T>>, DbError> {
516        let pool = self.query.pool();
517        let model = self.query.fetch_one().await?;
518        Ok(attach(pool, model.into_iter().collect(), &self.loaders).await?.pop())
519    }
520
521    /// [`QueryBuilder::paginate`], each row with its relations.
522    pub async fn paginate(self, page: u64, per_page: u64) -> Result<Page<Loaded<T>>, DbError> {
523        let pool = self.query.pool();
524        let found = self.query.paginate(page, per_page).await?;
525        let (page, per_page, total) = (found.page, found.per_page, found.total_items);
526        let items = attach(pool, found.items, &self.loaders).await?;
527        Ok(Page::new(items, page, per_page, total))
528    }
529}
530
531async fn attach<T: Model>(pool: &DbPool, models: Vec<T>, loaders: &[Box<dyn Loader<T>>]) -> Result<Vec<Loaded<T>>, DbError> {
532    let mut relations: Vec<Vec<(String, Erased)>> = models.iter().map(|_| Vec::with_capacity(loaders.len())).collect();
533    if !models.is_empty() {
534        for loader in loaders {
535            for (slot, entries) in relations.iter_mut().zip(loader.load(pool, &models).await?) {
536                slot.extend(entries);
537            }
538        }
539    }
540    Ok(models.into_iter().zip(relations).map(|(model, relations)| Loaded { model, relations }).collect())
541}