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

rws_server/model/
eager.rs
