siderite_orm/queryset/
related.rs

1//! Loading related objects: [`Relation`] descriptors, `select_related`
2//! (one query with joins) and `prefetch_related` (one extra query per
3//! relation).
4//!
5//! A [`Relation`] names a foreign key *and* how to reach its slot in the
6//! model struct, because generic code cannot otherwise get at
7//! `book.author`:
8//!
9//! ```ignore
10//! let relation = Relation::new(Book::author, |book| &mut book.author);
11//! let books = Book::objects(&db).select_related(relation).all().await?;
12//! books[0].author.cached(); // Some(&Author)
13//! ```
14//!
15//! Chain relations with [`Relation::then`] for multi-hop `select_related`
16//! (`author` then `author.team`).
17
18use super::QuerySet;
19use crate::backend::Row;
20use crate::db::Db;
21use crate::error::{OrmError, QueryError};
22use crate::expr::{Expr, Field, FkSlot, Ident, RelHop, RelatedColumn, path_alias_of, pk_column};
23use crate::model::{FieldMeta, Model};
24use crate::persist::value_key;
25use crate::plan::SelectExpr;
26use crate::types::DbType;
27use crate::value::Value;
28use async_trait::async_trait;
29use std::any::Any;
30use std::collections::{HashMap, HashSet};
31use std::marker::PhantomData;
32use std::sync::Arc;
33
34type Loaded = Box<dyn Any + Send + Sync>;
35type Load = fn(&Row, &str) -> Result<Loaded, QueryError>;
36type Attach = Arc<dyn Fn(&mut dyn Any, Arc<dyn Any + Send + Sync>) + Send + Sync>;
37
38fn load<T: Model>(row: &Row, prefix: &str) -> Result<Loaded, QueryError> {
39    Ok(Box::new(T::from_row(row, prefix)?))
40}
41
42/// One foreign-key hop with the type-erased operations to load and attach it.
43#[derive(Clone)]
44struct Step {
45    hop: RelHop,
46    /// Columns of the target model (projected by `select_related`).
47    fields: &'static [FieldMeta],
48    load: Load,
49    attach: Attach,
50}
51
52/// A foreign key of model `M` pointing at `T`, together with access to its
53/// slot in `M`, so loaders can store the related object in the struct:
54///
55/// ```ignore
56/// let relation = Relation::new(Book::author, |book| &mut book.author);
57/// let books = Book::objects(&db).select_related(relation).all().await?;
58/// books[0].author.cached(); // Some(&Author)
59/// ```
60///
61/// Chain relations with [`then`](Self::then) for multi-hop `select_related`.
62pub struct Relation<M, T> {
63    steps: Vec<Step>,
64    _marker: PhantomData<fn() -> (M, T)>,
65}
66
67impl<M, T> Clone for Relation<M, T> {
68    fn clone(&self) -> Self {
69        Self {
70            steps: self.steps.clone(),
71            _marker: PhantomData,
72        }
73    }
74}
75
76impl<M: Model, T: Model> Relation<M, T> {
77    /// Relation through `field`; `slot` returns the struct field holding it:
78    /// `Relation::new(Book::author, |b| &mut b.author)`. Works for
79    /// `ForeignKey<T>` and `Option<ForeignKey<T>>` slots.
80    pub fn new<S: FkSlot<T> + 'static>(field: Field<M, S>, slot: fn(&mut M) -> &mut S) -> Self {
81        let attach: Attach = Arc::new(move |parent, child| {
82            if let (Some(parent), Ok(child)) = (parent.downcast_mut::<M>(), child.downcast::<T>())
83                && let Some(fk) = slot(parent).fk()
84            {
85                fk.set_cached(child);
86            }
87        });
88        Self {
89            steps: vec![Step {
90                hop: RelHop {
91                    fk_column: field.name().into(),
92                    table: T::META.table.into(),
93                    pk_column: pk_column::<T>(),
94                },
95                fields: T::META.fields,
96                load: load::<T>,
97                attach,
98            }],
99            _marker: PhantomData,
100        }
101    }
102
103    /// Continue through a relation of `T`: `author.then(team)` reaches
104    /// `book.author.team`.
105    pub fn then<U: Model>(mut self, next: Relation<T, U>) -> Relation<M, U> {
106        self.steps.extend(next.steps);
107        Relation {
108            steps: self.steps,
109            _marker: PhantomData,
110        }
111    }
112}
113
114/// A node of the `select_related` tree: one joined model and what hangs off it.
115#[derive(Clone)]
116pub(super) struct RelNode {
117    step: Step,
118    alias: String,
119    children: Vec<RelNode>,
120}
121
122impl RelNode {
123    /// Decode this node's object from `row` (with its own children attached),
124    /// or `None` when the `LEFT JOIN` found no row.
125    fn decode(&self, row: &Row) -> Result<Option<Loaded>, QueryError> {
126        let prefix = format!("{}__", self.alias);
127        let key = format!("{prefix}{}", self.step.hop.pk_column);
128        if row.get(&key).is_none_or(Value::is_null) {
129            return Ok(None);
130        }
131        let mut object = (self.step.load)(row, &prefix)?;
132        attach_all(&mut *object, &self.children, row)?;
133        Ok(Some(object))
134    }
135}
136
137/// Decode every node under `parent` and store the results in its slots.
138pub(super) fn attach_all(
139    parent: &mut dyn Any,
140    nodes: &[RelNode],
141    row: &Row,
142) -> Result<(), QueryError> {
143    for node in nodes {
144        if let Some(object) = node.decode(row)? {
145            (node.step.attach)(parent, Arc::from(object));
146        }
147    }
148    Ok(())
149}
150
151/// Insert `steps` into the tree, sharing nodes with earlier paths.
152fn insert_path(nodes: &mut Vec<RelNode>, steps: &[Step], path: &mut Vec<RelHop>) {
153    let Some((step, rest)) = steps.split_first() else {
154        return;
155    };
156    path.push(step.hop.clone());
157    let position = match nodes.iter().position(|n| n.step.hop == step.hop) {
158        Some(position) => position,
159        None => {
160            nodes.push(RelNode {
161                step: step.clone(),
162                alias: path_alias_of(path),
163                children: Vec::new(),
164            });
165            nodes.len() - 1
166        }
167    };
168    insert_path(&mut nodes[position].children, rest, path);
169    path.pop();
170}
171
172/// Projection of every column of the joined models in `nodes`, named
173/// `alias__column`, with the path each column is reached through.
174fn related_projection(nodes: &[RelNode], path: &mut Vec<RelHop>, out: &mut Vec<SelectExpr>) {
175    for node in nodes {
176        path.push(node.step.hop.clone());
177        for field in node.step.fields {
178            let column = RelatedColumn {
179                path: path.clone(),
180                column: field.column.into(),
181            };
182            let alias: Ident = column.output_name().into();
183            out.push(SelectExpr::new(Expr::Related(column), Some(alias)));
184        }
185        related_projection(&node.children, path, out);
186        path.pop();
187    }
188}
189
190impl<M: Model> QuerySet<M> {
191    /// Load `relation`'s target in the same query (`LEFT JOIN`), so
192    /// `ForeignKey::cached()` is filled without further queries. Repeated
193    /// calls share joins of a common path (`author`, then `author` → `team`).
194    #[must_use]
195    pub fn select_related<T: Model>(mut self, relation: Relation<M, T>) -> Self {
196        let mut earlier = Vec::new();
197        related_projection(&self.related, &mut Vec::new(), &mut earlier);
198        insert_path(&mut self.related, &relation.steps, &mut Vec::new());
199        let mut columns = Vec::new();
200        related_projection(&self.related, &mut Vec::new(), &mut columns);
201        self.with_plan(|mut plan| {
202            // Drop the columns of earlier calls; they are re-added below.
203            plan.projection
204                .retain(|s| !earlier.iter().any(|e| e.alias == s.alias));
205            if plan.projection.is_empty() {
206                plan.projection = super::model_projection::<M>();
207            }
208            plan.projection.extend(columns);
209            plan
210        })
211    }
212
213    /// Load `relation`'s targets with one extra query (`WHERE pk IN (..)`)
214    /// after the main query, instead of one query per row. Only single-hop
215    /// relations can be prefetched; use [`select_related`](Self::select_related)
216    /// for longer paths. Reverse foreign keys and many-to-many relations are
217    /// not prefetched (the model has no slot to hold them).
218    #[must_use]
219    pub fn prefetch_related<T: Model>(mut self, prefetch: impl Into<Prefetch<M, T>>) -> Self {
220        let prefetch = prefetch.into();
221        if prefetch.relation.steps.len() != 1 {
222            return self.fail("prefetch_related supports single-hop relations only");
223        }
224        self.prefetches.push(Arc::new(prefetch));
225        self
226    }
227
228    /// Run the prefetch queries for `models`.
229    pub(super) async fn run_prefetches(&self, models: &mut [M]) -> Result<(), OrmError> {
230        for prefetch in &self.prefetches {
231            prefetch.run(&self.db, models).await?;
232        }
233        Ok(())
234    }
235}
236
237/// `prefetch_related` request: a relation and, optionally, the queryset used
238/// to load the targets (to filter them, order them or use another handle).
239pub struct Prefetch<M: Model, T: Model> {
240    relation: Relation<M, T>,
241    queryset: Option<QuerySet<T>>,
242}
243
244impl<M: Model, T: Model> Prefetch<M, T> {
245    /// Prefetch `relation` with a plain queryset over the target.
246    pub fn new(relation: Relation<M, T>) -> Self {
247        Self {
248            relation,
249            queryset: None,
250        }
251    }
252
253    /// Load targets through `queryset`. Its filters apply; targets it does
254    /// not return stay unloaded. Its database capabilities and compiled
255    /// existing bind count govern prefetch batching. SQL backends must
256    /// implement Backend::read_parameter_count.
257    #[must_use]
258    pub fn queryset(mut self, queryset: QuerySet<T>) -> Self {
259        self.queryset = Some(queryset);
260        self
261    }
262}
263
264impl<M: Model, T: Model> From<Relation<M, T>> for Prefetch<M, T> {
265    fn from(relation: Relation<M, T>) -> Self {
266        Self::new(relation)
267    }
268}
269
270/// Type-erased prefetch, stored inside the queryset.
271#[async_trait]
272pub(super) trait Prefetcher<M: Model>: Send + Sync {
273    async fn run(&self, db: &Db, models: &mut [M]) -> Result<(), OrmError>;
274}
275
276#[async_trait]
277impl<M: Model, T: Model> Prefetcher<M> for Prefetch<M, T> {
278    async fn run(&self, db: &Db, models: &mut [M]) -> Result<(), OrmError> {
279        let Some(step) = self.relation.steps.first() else {
280            return Ok(());
281        };
282        let column = step.hop.fk_column.as_ref();
283        let mut seen = HashSet::new();
284        let mut wanted = Vec::new();
285        let keys: Vec<Option<String>> = models
286            .iter()
287            .map(|model| {
288                let value = model
289                    .to_values()
290                    .into_iter()
291                    .find(|(c, _)| *c == column)
292                    .map(|(_, v)| v)
293                    .filter(|v| !v.is_null())?;
294                let key = value_key(&value);
295                if seen.insert(key.clone()) {
296                    wanted.push(value);
297                }
298                Some(key)
299            })
300            .collect();
301        if wanted.is_empty() {
302            return Ok(());
303        }
304        let base = self.queryset.clone().unwrap_or_else(|| T::objects(db));
305        base.ready()?;
306        if base.empty {
307            return Ok(());
308        }
309        let caps = base.db.capabilities();
310        let used = match base.db.read_parameter_count(base.plan())? {
311            Some(count) => count,
312            None if matches!(caps.kind, crate::BackendKind::MongoDb) => 0,
313            None => {
314                return Err(QueryError::InvalidPlan(
315                    "prefetch requires compiled parameter counting".into(),
316                )
317                .into());
318            }
319        };
320        let chunk = caps
321            .max_params
322            .checked_sub(used)
323            .filter(|remaining| *remaining > 0)
324            .ok_or_else(|| QueryError::InvalidPlan("prefetch bind capacity exhausted".into()))?;
325        if wanted.len() > chunk && (base.plan.limit.is_some() || base.plan.offset.is_some()) {
326            return Err(
327                QueryError::InvalidPlan("sliced prefetch cannot span bind batches".into()).into(),
328            );
329        }
330        let mut loaded: HashMap<String, Arc<T>> = HashMap::new();
331        for wanted in wanted.chunks(chunk) {
332            let pk = Expr::col(step.hop.pk_column.clone());
333            let predicate = pk.is_in(wanted.iter().cloned());
334            for target in base.clone().filter(predicate).all().await? {
335                loaded.insert(value_key(&target.pk().to_value()), Arc::new(target));
336            }
337        }
338        for (model, key) in models.iter_mut().zip(&keys) {
339            let target = key.as_ref().and_then(|key| loaded.get(key));
340            if let Some(target) = target {
341                (step.attach)(model, Arc::<T>::clone(target));
342            }
343        }
344        Ok(())
345    }
346}
siderite_orm/queryset/related.rs

siderite_orm/queryset/
related.rs
