Skip to main content

rango_orm_core/
related.rs

1//! Relation cache shared by the backends (`select_related` / `prefetch_related`).
2
3use std::any::{Any, TypeId, type_name};
4use std::collections::HashMap;
5use std::ops::{Deref, DerefMut};
6
7use crate::query::SqlValue;
8use crate::row::{FromRow, RangoRow, RowError};
9
10/// Type-erased relation storage, keyed by TypeId.
11///
12/// - Single FK relations (`related<R>`)      → key: `TypeId::of::<R>()`,       value: `Box<R>`
13/// - Reverse/M2M relations (`prefetched<R>`) → key: `TypeId::of::<Vec<R>>()`,  value: `Box<Vec<R>>`
14///   (different keys so both can coexist for the same R)
15pub type RelationCache = HashMap<TypeId, Box<dyn Any + Send + Sync>>;
16
17/// A model instance with an attached relation cache (level-1, in-process).
18///
19/// `Deref`s to the inner model so field access (`post.title`) works directly.
20/// Relations are populated by `QueryBuilder::all()` when `.select_related()`
21/// or `.prefetch_related()` were chained before execution.
22///
23/// # Example
24/// ```rust,ignore
25/// let posts = Post::filter(&pool)
26///     .select_related(Post::AUTHOR_ID)
27///     .prefetch_related(Comment::POST_ID)
28///     .all()
29///     .await?;
30///
31/// for post in &posts {
32///     println!("{} by {}", post.title, post.related::<Author>().name);
33///     for comment in post.prefetched::<Comment>() {
34///         println!("  - {}", comment.body);
35///     }
36/// }
37/// ```
38pub struct WithRelated<M> {
39    pub inner: M,
40    cache: RelationCache,
41}
42
43impl<M> WithRelated<M> {
44    /// The model, dropping the loaded relations.
45    pub fn into_inner(self) -> M {
46        self.inner
47    }
48
49    /// Used by backends; relations are attached by `QueryBuilder::all()`.
50    #[doc(hidden)]
51    pub fn new(inner: M) -> Self {
52        Self {
53            inner,
54            cache: HashMap::new(),
55        }
56    }
57
58    /// Returns the FK-related instance loaded by `select_related`.
59    ///
60    /// # Panics
61    /// Panics with a clear message if `.select_related(..)` was not called for `R`
62    /// on the query builder.
63    pub fn related<R: Any + Send + Sync + 'static>(&self) -> &R {
64        self.cache
65            .get(&TypeId::of::<R>())
66            .and_then(|b| b.downcast_ref::<R>())
67            .unwrap_or_else(|| {
68                panic!(
69                    "select_related<{ty}> was not loaded; \
70                 add .select_related(..) for {ty} to the QueryBuilder chain",
71                    ty = type_name::<R>()
72                )
73            })
74    }
75
76    /// Returns the prefetched slice loaded by `prefetch_related`.
77    ///
78    /// # Panics
79    /// Panics with a clear message if `.prefetch_related(..)` was not called for `R`
80    /// on the query builder.
81    pub fn prefetched<R: Any + Send + Sync + 'static>(&self) -> &[R] {
82        self.cache
83            .get(&TypeId::of::<Vec<R>>())
84            .and_then(|b| b.downcast_ref::<Vec<R>>())
85            .map(Vec::as_slice)
86            .unwrap_or_else(|| {
87                panic!(
88                    "prefetch_related<{ty}> was not loaded; \
89                 add .prefetch_related(..) for {ty} to the QueryBuilder chain",
90                    ty = type_name::<R>()
91                )
92            })
93    }
94
95    /// Low-level insert for type-erased callers in `QueryBuilder`.
96    /// `type_id` must match the actual concrete type stored in `boxed`.
97    #[doc(hidden)]
98    pub fn insert_raw(&mut self, type_id: TypeId, boxed: Box<dyn Any + Send + Sync>) {
99        self.cache.insert(type_id, boxed);
100    }
101}
102
103impl<M> Deref for WithRelated<M> {
104    type Target = M;
105    fn deref(&self) -> &M {
106        &self.inner
107    }
108}
109
110impl<M> DerefMut for WithRelated<M> {
111    fn deref_mut(&mut self) -> &mut M {
112        &mut self.inner
113    }
114}
115
116// ─── Type-erased helpers for the backends' relation specs ─────────────────────
117
118/// Type-erased `R::from_row`.
119#[doc(hidden)]
120pub type FromRowFn = fn(&dyn RangoRow) -> Result<Box<dyn Any + Send + Sync>, RowError>;
121
122/// Collects boxed `R`s into a boxed `Vec<R>`.
123#[doc(hidden)]
124pub type CollectVecFn = fn(Vec<Box<dyn Any + Send + Sync>>) -> Box<dyn Any + Send + Sync>;
125
126#[doc(hidden)]
127pub fn erased_from_row<R>(row: &dyn RangoRow) -> Result<Box<dyn Any + Send + Sync>, RowError>
128where
129    R: FromRow + Send + Sync + 'static,
130{
131    R::from_row(row).map(|v| Box::new(v) as Box<dyn Any + Send + Sync>)
132}
133
134#[doc(hidden)]
135pub fn collect_vec_erased<R>(items: Vec<Box<dyn Any + Send + Sync>>) -> Box<dyn Any + Send + Sync>
136where
137    R: Any + Send + Sync + 'static,
138{
139    let typed: Vec<R> = items
140        .into_iter()
141        .map(|b| {
142            *b.downcast::<R>()
143                .expect("prefetch_related type mismatch — this is a bug")
144        })
145        .collect();
146    Box::new(typed)
147}
148
149/// Hashable key for a primary / foreign key value (UUID, integer or text).
150/// M2M pivot naming shared by `M2M` ops, `prefetch_related` and `rango makemigrations`:
151/// `(pivot_table, from_col, to_col)` = `("{from}_{to}", "{from}_id", "{to}_id")`.
152/// A self-referencing relation (`from == to`) would get two columns with the same name,
153/// so its columns are `from_{x}_id` and `to_{x}_id`.
154pub fn m2m_pivot(from_table: &str, to_table: &str) -> (String, String, String) {
155    let (from_col, to_col) = if from_table == to_table {
156        (format!("from_{from_table}_id"), format!("to_{to_table}_id"))
157    } else {
158        (format!("{from_table}_id"), format!("{to_table}_id"))
159    };
160    (format!("{from_table}_{to_table}"), from_col, to_col)
161}
162
163#[doc(hidden)]
164pub fn sql_value_key(v: &SqlValue) -> Option<String> {
165    match v {
166        SqlValue::Uuid(u) => Some(u.to_string()),
167        SqlValue::BigInt(i) => Some(i.to_string()),
168        SqlValue::Int(i) => Some(i.to_string()),
169        SqlValue::SmallInt(i) => Some(i.to_string()),
170        SqlValue::Text(s) => Some(s.clone()),
171        _ => None,
172    }
173}
174
175/// Reads `col` with the same type as `like` (a key of the other side) and returns
176/// its [`sql_value_key`]. `None` when the column is NULL or does not decode.
177#[doc(hidden)]
178pub fn row_key(row: &dyn RangoRow, col: &str, like: &SqlValue) -> Option<String> {
179    match like {
180        SqlValue::Uuid(_) => row.get_uuid(col).ok().map(|v| v.to_string()),
181        SqlValue::BigInt(_) => row.get_i64(col).ok().map(|v| v.to_string()),
182        SqlValue::Int(_) => row.get_i32(col).ok().map(|v| v.to_string()),
183        SqlValue::SmallInt(_) => row.get_i16(col).ok().map(|v| v.to_string()),
184        SqlValue::Text(_) => row.get_string(col).ok(),
185        _ => None,
186    }
187}
188
189#[cfg(test)]
190mod pivot_tests {
191    use super::m2m_pivot;
192
193    #[test]
194    fn pivot_columns_are_distinct_for_self_reference() {
195        let s = |t: &str| t.to_string();
196        assert_eq!(m2m_pivot("a", "b"), (s("a_b"), s("a_id"), s("b_id")));
197        assert_eq!(
198            m2m_pivot("u", "u"),
199            (s("u_u"), s("from_u_id"), s("to_u_id"))
200        );
201    }
202}
203
204/// The SQL of the many-to-many pivot operations, for one dialect.
205///
206/// Every statement binds the from id first, then the to id.
207#[doc(hidden)]
208#[doc(hidden)]
209pub struct M2mSql {
210    dialect: crate::Dialect,
211    pivot: String,
212    from_col: String,
213    to_col: String,
214    from_table: String,
215    to_table: String,
216}
217
218impl M2mSql {
219    pub fn new(dialect: crate::Dialect, from_table: &str, to_table: &str) -> Self {
220        let (pivot, from_col, to_col) = m2m_pivot(from_table, to_table);
221        Self {
222            dialect,
223            pivot,
224            from_col,
225            to_col,
226            from_table: from_table.to_owned(),
227            to_table: to_table.to_owned(),
228        }
229    }
230
231    /// Unquoted pivot table name (for tracing).
232    pub fn pivot_table(&self) -> &str {
233        &self.pivot
234    }
235
236    fn q(&self, ident: &str) -> String {
237        self.dialect.quote_ident(ident)
238    }
239
240    fn ph(&self, n: usize) -> String {
241        match self.dialect {
242            crate::Dialect::Postgres => format!("${n}"),
243            _ => "?".to_owned(),
244        }
245    }
246
247    /// Idempotent insert of a `(from, to)` pair.
248    pub fn insert(&self) -> String {
249        let (pivot, from, to) = (
250            self.q(&self.pivot),
251            self.q(&self.from_col),
252            self.q(&self.to_col),
253        );
254        let (p1, p2) = (self.ph(1), self.ph(2));
255        match self.dialect {
256            // ON DUPLICATE KEY UPDATE, not INSERT IGNORE, to avoid silencing FK and other errors.
257            crate::Dialect::Mysql => format!(
258                "INSERT INTO {pivot} ({from}, {to}) VALUES ({p1}, {p2}) \
259                 ON DUPLICATE KEY UPDATE {from} = {from}"
260            ),
261            _ => format!(
262                "INSERT INTO {pivot} ({from}, {to}) VALUES ({p1}, {p2}) ON CONFLICT DO NOTHING"
263            ),
264        }
265    }
266
267    /// Delete every pair of one `from` id.
268    pub fn clear(&self) -> String {
269        format!(
270            "DELETE FROM {} WHERE {} = {}",
271            self.q(&self.pivot),
272            self.q(&self.from_col),
273            self.ph(1)
274        )
275    }
276
277    /// Delete one `(from, to)` pair.
278    pub fn remove(&self) -> String {
279        format!(
280            "DELETE FROM {} WHERE {} = {} AND {} = {}",
281            self.q(&self.pivot),
282            self.q(&self.from_col),
283            self.ph(1),
284            self.q(&self.to_col),
285            self.ph(2)
286        )
287    }
288
289    /// Does the `(from, to)` pair exist.
290    pub fn exists(&self) -> String {
291        format!(
292            "SELECT 1 FROM {} WHERE {} = {} AND {} = {} LIMIT 1",
293            self.q(&self.pivot),
294            self.q(&self.from_col),
295            self.ph(1),
296            self.q(&self.to_col),
297            self.ph(2)
298        )
299    }
300
301    /// The `to` rows of one `from` id (bind: from id).
302    pub fn all(&self, to_pk: &str) -> String {
303        format!(
304            r#"SELECT t.* FROM {to_table} t
305               INNER JOIN {pivot} p ON t.{to_pk} = p.{to_col}
306               WHERE p.{from_col} = {p1}"#,
307            to_table = self.q(&self.to_table),
308            pivot = self.q(&self.pivot),
309            to_pk = self.q(to_pk),
310            to_col = self.q(&self.to_col),
311            from_col = self.q(&self.from_col),
312            p1 = self.ph(1),
313        )
314    }
315
316    /// The `from` rows of one `to` id (bind: to id).
317    pub fn reverse(&self, from_pk: &str) -> String {
318        format!(
319            r#"SELECT f.* FROM {from_table} f
320               INNER JOIN {pivot} p ON f.{from_pk} = p.{from_col}
321               WHERE p.{to_col} = {p1}"#,
322            from_table = self.q(&self.from_table),
323            pivot = self.q(&self.pivot),
324            from_pk = self.q(from_pk),
325            from_col = self.q(&self.from_col),
326            to_col = self.q(&self.to_col),
327            p1 = self.ph(1),
328        )
329    }
330}
331
332#[cfg(test)]
333mod m2m_sql_tests {
334    use super::M2mSql;
335    use crate::Dialect;
336
337    #[test]
338    fn per_dialect_sql() {
339        let pg = M2mSql::new(Dialect::Postgres, "post", "tag");
340        assert_eq!(
341            pg.insert(),
342            "INSERT INTO \"post_tag\" (\"post_id\", \"tag_id\") VALUES ($1, $2) ON CONFLICT DO NOTHING"
343        );
344        assert_eq!(
345            pg.remove(),
346            "DELETE FROM \"post_tag\" WHERE \"post_id\" = $1 AND \"tag_id\" = $2"
347        );
348        let my = M2mSql::new(Dialect::Mysql, "post", "tag");
349        assert_eq!(
350            my.insert(),
351            "INSERT INTO `post_tag` (`post_id`, `tag_id`) VALUES (?, ?) \
352             ON DUPLICATE KEY UPDATE `post_id` = `post_id`"
353        );
354        let lite = M2mSql::new(Dialect::Sqlite, "post", "tag");
355        assert_eq!(
356            lite.clear(),
357            "DELETE FROM \"post_tag\" WHERE \"post_id\" = ?"
358        );
359    }
360}