sea_orm/query/
loader.rs

1use super::get_key_from_model;
2use crate::{
3    ColumnTrait, Condition, ConnectionTrait, DbBackend, DbErr, EntityTrait, Identity, JoinType,
4    ModelTrait, QueryFilter, QuerySelect, Related, RelatedSelfVia, RelationDef, RelationTrait,
5    RelationType, Select, dynamic, query::column_tuple_in_condition, query_err,
6};
7use sea_query::{ColumnRef, DynIden, Expr, ExprTrait, IntoColumnRef, TableRef, ValueTuple};
8use std::{collections::HashMap, str::FromStr};
9
10// TODO: Replace DynIden::inner with a better API that without clone
11
12/// Either an `Entity` or a `Select<Entity>`; accepted by [`LoaderTrait`] so
13/// you can pass the entity directly when no filtering is needed, or a
14/// pre-filtered [`Select`] when it is.
15pub trait EntityOrSelect<E: EntityTrait> {
16    /// Lift `self` into a [`Select<E>`]. For an entity, this is `E::find()`;
17    /// for an existing `Select<E>`, the value is returned as-is.
18    fn select(self) -> Select<E>;
19}
20
21type LoaderEntity<T> = <<T as LoaderTrait>::Model as ModelTrait>::Entity;
22type LoaderModel<T> = <<<T as LoaderTrait>::Model as ModelTrait>::Entity as EntityTrait>::Model;
23type LoaderRelation<T> =
24    <<<T as LoaderTrait>::Model as ModelTrait>::Entity as EntityTrait>::Relation;
25
26/// Batch-load related entities for a slice of parent models, avoiding the
27/// N+1 query problem.
28///
29/// Given a `Vec<Parent>`, [`load_one`](Self::load_one) /
30/// [`load_many`](Self::load_many) / [`load_many_to_many`](Self::load_many_to_many)
31/// issue a single `WHERE … IN (…)` query for each relation hop and
32/// reassemble the results, grouped per parent.
33pub trait LoaderTrait {
34    /// Source model
35    type Model: ModelTrait;
36
37    /// Used to eager load self_ref relations
38    fn load_self<S, C>(
39        &self,
40        stmt: S,
41        relation_enum: LoaderRelation<Self>,
42        db: &C,
43    ) -> Result<Vec<Option<LoaderModel<Self>>>, DbErr>
44    where
45        C: ConnectionTrait,
46        LoaderModel<Self>: Send,
47        LoaderRelation<Self>: Send,
48        S: EntityOrSelect<LoaderEntity<Self>>;
49
50    /// Used to eager load self_ref relations in reverse, output is `Vec<Model>` instead of
51    /// `Option<Model>` of `load_self`.
52    fn load_self_many<S, C>(
53        &self,
54        stmt: S,
55        relation_enum: LoaderRelation<Self>,
56        db: &C,
57    ) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
58    where
59        C: ConnectionTrait,
60        LoaderModel<Self>: Send,
61        LoaderRelation<Self>: Send,
62        S: EntityOrSelect<LoaderEntity<Self>>;
63
64    /// Used to eager load self_ref + via relations
65    fn load_self_via<V, C>(&self, via: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
66    where
67        C: ConnectionTrait,
68        V: EntityTrait,
69        LoaderModel<Self>: Send,
70        LoaderEntity<Self>: RelatedSelfVia<V>;
71
72    /// Used to eager load self_ref + via relations, but in reverse
73    fn load_self_via_rev<V, C>(&self, via: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
74    where
75        C: ConnectionTrait,
76        V: EntityTrait,
77        LoaderModel<Self>: Send,
78        LoaderEntity<Self>: RelatedSelfVia<V>;
79
80    /// Used to eager load has_one relations
81    fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
82    where
83        C: ConnectionTrait,
84        R: EntityTrait,
85        S: EntityOrSelect<R>,
86        <Self::Model as ModelTrait>::Entity: Related<R>;
87
88    /// Used to eager load has_many relations
89    fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
90    where
91        C: ConnectionTrait,
92        R: EntityTrait,
93        S: EntityOrSelect<R>,
94        <Self::Model as ModelTrait>::Entity: Related<R>;
95
96    /// Used to eager load many_to_many relations. In SeaORM 2.0 `load_many` already support M-N
97    /// relations so this method is not needed, only kept as legacy.
98    fn load_many_to_many<R, S, V, C>(
99        &self,
100        stmt: S,
101        via: V,
102        db: &C,
103    ) -> Result<Vec<Vec<R::Model>>, DbErr>
104    where
105        C: ConnectionTrait,
106        R: EntityTrait,
107        S: EntityOrSelect<R>,
108        V: EntityTrait,
109        V::Model: Send,
110        <Self::Model as ModelTrait>::Entity: Related<R>;
111}
112
113type LoaderExEntity<T> = <<T as LoaderTraitEx>::Model as ModelTrait>::Entity;
114type LoaderExModel<T> = <<<T as LoaderTraitEx>::Model as ModelTrait>::Entity as EntityTrait>::Model;
115type LoaderExModelEx<T> =
116    <<<T as LoaderTraitEx>::Model as ModelTrait>::Entity as EntityTrait>::ModelEx;
117type LoaderExRelation<T> =
118    <<<T as LoaderTraitEx>::Model as ModelTrait>::Entity as EntityTrait>::Relation;
119
120#[doc(hidden)]
121pub trait LoaderTraitEx {
122    type Model: ModelTrait;
123
124    fn load_self_ex<S, C>(
125        &self,
126        stmt: S,
127        relation_enum: LoaderExRelation<Self>,
128        db: &C,
129    ) -> Result<Vec<Option<LoaderExModelEx<Self>>>, DbErr>
130    where
131        C: ConnectionTrait,
132        LoaderExModel<Self>: Send,
133        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
134        LoaderExRelation<Self>: Send,
135        S: EntityOrSelect<LoaderExEntity<Self>>;
136
137    fn load_self_many_ex<S, C>(
138        &self,
139        stmt: S,
140        relation_enum: LoaderExRelation<Self>,
141        db: &C,
142    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
143    where
144        C: ConnectionTrait,
145        LoaderExModel<Self>: Send,
146        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
147        LoaderExRelation<Self>: Send,
148        S: EntityOrSelect<LoaderExEntity<Self>>;
149
150    fn load_self_via_ex<V, C>(
151        &self,
152        via: V,
153        is_reverse: bool,
154        db: &C,
155    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
156    where
157        C: ConnectionTrait,
158        V: EntityTrait,
159        LoaderExModel<Self>: Send,
160        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
161        LoaderExEntity<Self>: RelatedSelfVia<V>;
162
163    fn load_one_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::ModelEx>>, DbErr>
164    where
165        C: ConnectionTrait,
166        R: EntityTrait,
167        S: EntityOrSelect<R>,
168        R::ModelEx: From<R::Model>,
169        <Self::Model as ModelTrait>::Entity: Related<R>;
170
171    fn load_one_ex_with_rel<R>(
172        &self,
173        stmt: impl EntityOrSelect<R>,
174        rel_def: RelationDef,
175        db: &impl ConnectionTrait,
176    ) -> Result<Vec<Option<R::ModelEx>>, DbErr>
177    where
178        R: EntityTrait,
179        R::ModelEx: From<R::Model>;
180
181    fn load_many_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
182    where
183        C: ConnectionTrait,
184        R: EntityTrait,
185        S: EntityOrSelect<R>,
186        R::ModelEx: From<R::Model>,
187        <Self::Model as ModelTrait>::Entity: Related<R>;
188
189    fn load_many_ex_with_rel<R>(
190        &self,
191        stmt: impl EntityOrSelect<R>,
192        rel_def: RelationDef,
193        db: &impl ConnectionTrait,
194    ) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
195    where
196        R: EntityTrait,
197        R::ModelEx: From<R::Model>;
198}
199
200type NestedEntity<T> = <<T as NestedLoaderTrait>::Model as ModelTrait>::Entity;
201type NestedModel<T> =
202    <<<T as NestedLoaderTrait>::Model as ModelTrait>::Entity as EntityTrait>::Model;
203type NestedModelEx<T> =
204    <<<T as NestedLoaderTrait>::Model as ModelTrait>::Entity as EntityTrait>::ModelEx;
205type NestedLoaderRelation<T> =
206    <<<T as NestedLoaderTrait>::Model as ModelTrait>::Entity as EntityTrait>::Relation;
207
208#[doc(hidden)]
209pub trait NestedLoaderTrait {
210    type Model: ModelTrait;
211
212    fn load_self_ex<S, C>(
213        &self,
214        stmt: S,
215        relation_enum: NestedLoaderRelation<Self>,
216        db: &C,
217    ) -> Result<Vec<Vec<Option<NestedModelEx<Self>>>>, DbErr>
218    where
219        C: ConnectionTrait,
220        NestedModel<Self>: Send,
221        NestedModelEx<Self>: From<NestedModel<Self>>,
222        NestedLoaderRelation<Self>: Send,
223        S: EntityOrSelect<<<Self as NestedLoaderTrait>::Model as ModelTrait>::Entity>;
224
225    fn load_self_many_ex<S, C>(
226        &self,
227        stmt: S,
228        relation_enum: NestedLoaderRelation<Self>,
229        db: &C,
230    ) -> Result<Vec<Vec<Vec<NestedModelEx<Self>>>>, DbErr>
231    where
232        C: ConnectionTrait,
233        NestedModel<Self>: Send,
234        NestedModelEx<Self>: From<NestedModel<Self>>,
235        NestedLoaderRelation<Self>: Send,
236        S: EntityOrSelect<<<Self as NestedLoaderTrait>::Model as ModelTrait>::Entity>;
237
238    fn load_self_via_ex<V, C>(
239        &self,
240        via: V,
241        is_reverse: bool,
242        db: &C,
243    ) -> Result<Vec<Vec<Vec<NestedModelEx<Self>>>>, DbErr>
244    where
245        C: ConnectionTrait,
246        V: EntityTrait,
247        NestedModel<Self>: Send,
248        NestedModelEx<Self>: From<NestedModel<Self>>,
249        NestedEntity<Self>: RelatedSelfVia<V>;
250
251    fn load_one_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<Option<R::ModelEx>>>, DbErr>
252    where
253        C: ConnectionTrait,
254        R: EntityTrait,
255        S: EntityOrSelect<R>,
256        R::ModelEx: From<R::Model>,
257        <Self::Model as ModelTrait>::Entity: Related<R>;
258
259    fn load_one_ex_with_rel<R>(
260        &self,
261        stmt: impl EntityOrSelect<R>,
262        rel_def: RelationDef,
263        db: &impl ConnectionTrait,
264    ) -> Result<Vec<Vec<Option<R::ModelEx>>>, DbErr>
265    where
266        R: EntityTrait,
267        R::ModelEx: From<R::Model>;
268
269    fn load_many_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<Vec<R::ModelEx>>>, DbErr>
270    where
271        C: ConnectionTrait,
272        R: EntityTrait,
273        S: EntityOrSelect<R>,
274        R::ModelEx: From<R::Model>,
275        <Self::Model as ModelTrait>::Entity: Related<R>;
276
277    fn load_many_ex_with_rel<R>(
278        &self,
279        stmt: impl EntityOrSelect<R>,
280        rel_def: RelationDef,
281        db: &impl ConnectionTrait,
282    ) -> Result<Vec<Vec<Vec<R::ModelEx>>>, DbErr>
283    where
284        R: EntityTrait,
285        R::ModelEx: From<R::Model>;
286}
287
288impl<E> EntityOrSelect<E> for E
289where
290    E: EntityTrait,
291{
292    fn select(self) -> Select<E> {
293        E::find().order_by_id_asc()
294    }
295}
296
297impl<E> EntityOrSelect<E> for Select<E>
298where
299    E: EntityTrait,
300{
301    fn select(self) -> Select<E> {
302        self
303    }
304}
305
306impl<M> LoaderTrait for Vec<M>
307where
308    M: ModelTrait,
309{
310    type Model = M;
311
312    fn load_self<S, C>(
313        &self,
314        stmt: S,
315        relation_enum: LoaderRelation<Self>,
316        db: &C,
317    ) -> Result<Vec<Option<LoaderModel<Self>>>, DbErr>
318    where
319        C: ConnectionTrait,
320        LoaderModel<Self>: Send,
321        LoaderRelation<Self>: Send,
322        S: EntityOrSelect<LoaderEntity<Self>>,
323    {
324        LoaderTrait::load_self(&self.as_slice(), stmt, relation_enum, db)
325    }
326
327    fn load_self_many<S, C>(
328        &self,
329        stmt: S,
330        relation_enum: LoaderRelation<Self>,
331        db: &C,
332    ) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
333    where
334        C: ConnectionTrait,
335        LoaderModel<Self>: Send,
336        LoaderRelation<Self>: Send,
337        S: EntityOrSelect<LoaderEntity<Self>>,
338    {
339        LoaderTrait::load_self_many(&self.as_slice(), stmt, relation_enum, db)
340    }
341
342    fn load_self_via<V, C>(&self, via: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
343    where
344        C: ConnectionTrait,
345        V: EntityTrait,
346        LoaderModel<Self>: Send,
347        LoaderEntity<Self>: RelatedSelfVia<V>,
348    {
349        LoaderTrait::load_self_via(&self.as_slice(), via, db)
350    }
351
352    fn load_self_via_rev<V, C>(&self, via: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
353    where
354        C: ConnectionTrait,
355        V: EntityTrait,
356        LoaderModel<Self>: Send,
357        LoaderEntity<Self>: RelatedSelfVia<V>,
358    {
359        LoaderTrait::load_self_via_rev(&self.as_slice(), via, db)
360    }
361
362    fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
363    where
364        C: ConnectionTrait,
365        R: EntityTrait,
366        S: EntityOrSelect<R>,
367        <Self::Model as ModelTrait>::Entity: Related<R>,
368    {
369        LoaderTrait::load_one(&self.as_slice(), stmt, db)
370    }
371
372    fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
373    where
374        C: ConnectionTrait,
375        R: EntityTrait,
376        S: EntityOrSelect<R>,
377        <Self::Model as ModelTrait>::Entity: Related<R>,
378    {
379        LoaderTrait::load_many(&self.as_slice(), stmt, db)
380    }
381
382    fn load_many_to_many<R, S, V, C>(
383        &self,
384        stmt: S,
385        via: V,
386        db: &C,
387    ) -> Result<Vec<Vec<R::Model>>, DbErr>
388    where
389        C: ConnectionTrait,
390        R: EntityTrait,
391        S: EntityOrSelect<R>,
392        V: EntityTrait,
393        V::Model: Send,
394        <Self::Model as ModelTrait>::Entity: Related<R>,
395    {
396        LoaderTrait::load_many_to_many(&self.as_slice(), stmt, via, db)
397    }
398}
399
400impl<M> LoaderTrait for &[M]
401where
402    M: ModelTrait,
403{
404    type Model = M;
405
406    fn load_self<S, C>(
407        &self,
408        stmt: S,
409        relation_enum: LoaderRelation<Self>,
410        db: &C,
411    ) -> Result<Vec<Option<LoaderModel<Self>>>, DbErr>
412    where
413        C: ConnectionTrait,
414        LoaderModel<Self>: Send,
415        LoaderRelation<Self>: Send,
416        S: EntityOrSelect<LoaderEntity<Self>>,
417    {
418        let rel_def = relation_enum.def();
419        check_self_ref(&rel_def)?;
420        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
421    }
422
423    fn load_self_many<S, C>(
424        &self,
425        stmt: S,
426        relation_enum: LoaderRelation<Self>,
427        db: &C,
428    ) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
429    where
430        C: ConnectionTrait,
431        LoaderModel<Self>: Send,
432        LoaderRelation<Self>: Send,
433        S: EntityOrSelect<LoaderEntity<Self>>,
434    {
435        let mut rel_def = relation_enum.def();
436        if rel_def.from_tbl != rel_def.to_tbl {
437            return Err(query_err("Relation must be self referencing"));
438        }
439        if !rel_def.is_owner {
440            // flip belongs_to
441            rel_def = rel_def.rev();
442        }
443        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
444    }
445
446    fn load_self_via<V, C>(&self, _: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
447    where
448        C: ConnectionTrait,
449        V: EntityTrait,
450        LoaderModel<Self>: Send,
451        LoaderEntity<Self>: RelatedSelfVia<V>,
452    {
453        let rel_def = <LoaderEntity<Self> as RelatedSelfVia<V>>::to();
454        let rel_via = <LoaderEntity<Self> as RelatedSelfVia<V>>::via();
455        loader_impl_impl(
456            self.iter(),
457            EntityOrSelect::select(LoaderEntity::<Self>::default()),
458            rel_def,
459            Some(rel_via),
460            db,
461        )
462    }
463
464    fn load_self_via_rev<V, C>(&self, _: V, db: &C) -> Result<Vec<Vec<LoaderModel<Self>>>, DbErr>
465    where
466        C: ConnectionTrait,
467        V: EntityTrait,
468        LoaderModel<Self>: Send,
469        LoaderEntity<Self>: RelatedSelfVia<V>,
470    {
471        let rel_def = <LoaderEntity<Self> as RelatedSelfVia<V>>::via().rev();
472        let rel_via = <LoaderEntity<Self> as RelatedSelfVia<V>>::to().rev();
473        loader_impl_impl(
474            self.iter(),
475            EntityOrSelect::select(LoaderEntity::<Self>::default()),
476            rel_def,
477            Some(rel_via),
478            db,
479        )
480    }
481
482    fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
483    where
484        C: ConnectionTrait,
485        R: EntityTrait,
486        S: EntityOrSelect<R>,
487        <Self::Model as ModelTrait>::Entity: Related<R>,
488    {
489        let rel_def = <<Self::Model as ModelTrait>::Entity as Related<R>>::to();
490        if rel_def.rel_type != RelationType::HasOne {
491            return Err(query_err("Relation is HasMany instead of HasOne"));
492        }
493        loader_impl(self.iter(), stmt.select(), db)
494    }
495
496    fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
497    where
498        C: ConnectionTrait,
499        R: EntityTrait,
500        S: EntityOrSelect<R>,
501        <Self::Model as ModelTrait>::Entity: Related<R>,
502    {
503        loader_impl(self.iter(), stmt.select(), db)
504    }
505
506    fn load_many_to_many<R, S, V, C>(
507        &self,
508        stmt: S,
509        via: V,
510        db: &C,
511    ) -> Result<Vec<Vec<R::Model>>, DbErr>
512    where
513        C: ConnectionTrait,
514        R: EntityTrait,
515        S: EntityOrSelect<R>,
516        V: EntityTrait,
517        V::Model: Send,
518        <Self::Model as ModelTrait>::Entity: Related<R>,
519    {
520        if let Some(via_rel) = <<Self::Model as ModelTrait>::Entity as Related<R>>::via() {
521            let rel_def = <<Self::Model as ModelTrait>::Entity as Related<R>>::to();
522            if rel_def.rel_type != RelationType::HasOne {
523                return Err(query_err("Relation to is not HasOne"));
524            }
525
526            if !cmp_table_ref(&via_rel.to_tbl, &via.table_ref()) {
527                return Err(query_err(format!(
528                    "The given via Entity is incorrect: expected: {:?}, given: {:?}",
529                    via_rel.to_tbl,
530                    via.table_ref()
531                )));
532            }
533
534            if self.is_empty() {
535                return Ok(Vec::new());
536            }
537
538            let pkeys = self
539                .iter()
540                .map(|model| get_key_from_model(&via_rel.from_col, model))
541                .collect::<Result<Vec<_>, _>>()?;
542
543            // Map of M::PK -> Vec<R::PK>
544            let mut keymap: HashMap<ValueTuple, Vec<ValueTuple>> = Default::default();
545
546            let keys: Vec<ValueTuple> = {
547                let condition = prepare_condition::<M>(
548                    &via_rel.to_tbl,
549                    &via_rel.from_col,
550                    &via_rel.to_col,
551                    &pkeys,
552                    db.get_database_backend(),
553                )?;
554                let stmt = V::find().filter(condition);
555                let data = stmt.all(db)?;
556                for model in data {
557                    let pk = get_key_from_model(&via_rel.to_col, &model)?;
558                    let entry = keymap.entry(pk).or_default();
559
560                    let fk = get_key_from_model(&rel_def.from_col, &model)?;
561                    entry.push(fk);
562                }
563
564                keymap.values().flatten().cloned().collect()
565            };
566
567            let condition = prepare_condition::<V::Model>(
568                &rel_def.to_tbl,
569                &rel_def.from_col,
570                &rel_def.to_col,
571                &keys,
572                db.get_database_backend(),
573            )?;
574
575            let stmt = QueryFilter::filter(stmt.select(), condition);
576
577            let models = stmt.all(db)?;
578
579            // Map of R::PK -> R::Model
580            let data = models.into_iter().try_fold(
581                HashMap::<ValueTuple, <R as EntityTrait>::Model>::new(),
582                |mut acc, model| {
583                    get_key_from_model(&rel_def.to_col, &model).map(|key| {
584                        acc.insert(key, model);
585
586                        acc
587                    })
588                },
589            )?;
590
591            let result: Vec<Vec<R::Model>> = pkeys
592                .into_iter()
593                .map(|pkey| {
594                    let fkeys = keymap.get(&pkey).cloned().unwrap_or_default();
595
596                    let models: Vec<_> = fkeys
597                        .into_iter()
598                        .filter_map(|fkey| data.get(&fkey).cloned())
599                        .collect();
600
601                    models
602                })
603                .collect();
604
605            Ok(result)
606        } else {
607            return Err(query_err("Relation is not ManyToMany"));
608        }
609    }
610}
611
612impl<M> LoaderTraitEx for &[M]
613where
614    M: ModelTrait,
615{
616    type Model = M;
617
618    fn load_self_ex<S, C>(
619        &self,
620        stmt: S,
621        relation_enum: LoaderExRelation<Self>,
622        db: &C,
623    ) -> Result<Vec<Option<LoaderExModelEx<Self>>>, DbErr>
624    where
625        C: ConnectionTrait,
626        LoaderExModel<Self>: Send,
627        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
628        LoaderExRelation<Self>: Send,
629        S: EntityOrSelect<LoaderExEntity<Self>>,
630    {
631        let rel_def = relation_enum.def();
632        check_self_ref(&rel_def)?;
633        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
634    }
635
636    fn load_self_many_ex<S, C>(
637        &self,
638        stmt: S,
639        relation_enum: LoaderExRelation<Self>,
640        db: &C,
641    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
642    where
643        C: ConnectionTrait,
644        LoaderExModel<Self>: Send,
645        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
646        LoaderExRelation<Self>: Send,
647        S: EntityOrSelect<LoaderExEntity<Self>>,
648    {
649        let rel_def = relation_enum.def();
650        check_self_ref_many(&rel_def)?;
651        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
652    }
653
654    fn load_self_via_ex<V, C>(
655        &self,
656        _: V,
657        is_reverse: bool,
658        db: &C,
659    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
660    where
661        C: ConnectionTrait,
662        V: EntityTrait,
663        LoaderExModel<Self>: Send,
664        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
665        LoaderExEntity<Self>: RelatedSelfVia<V>,
666    {
667        let (rel_def, rel_via) = if !is_reverse {
668            (
669                <LoaderExEntity<Self> as RelatedSelfVia<V>>::to(),
670                <LoaderExEntity<Self> as RelatedSelfVia<V>>::via(),
671            )
672        } else {
673            (
674                <LoaderEntity<Self> as RelatedSelfVia<V>>::via().rev(),
675                <LoaderEntity<Self> as RelatedSelfVia<V>>::to().rev(),
676            )
677        };
678        loader_impl_impl(
679            self.iter(),
680            EntityOrSelect::select(LoaderExEntity::<Self>::default()),
681            rel_def,
682            Some(rel_via),
683            db,
684        )
685    }
686
687    fn load_one_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::ModelEx>>, DbErr>
688    where
689        C: ConnectionTrait,
690        R: EntityTrait,
691        S: EntityOrSelect<R>,
692        R::ModelEx: From<R::Model>,
693        <Self::Model as ModelTrait>::Entity: Related<R>,
694    {
695        let rel_def = <<Self::Model as ModelTrait>::Entity as Related<R>>::to();
696        if rel_def.rel_type != RelationType::HasOne {
697            return Err(query_err("Relation is HasMany instead of HasOne"));
698        }
699        loader_impl(self.iter(), stmt.select(), db)
700    }
701
702    fn load_one_ex_with_rel<R>(
703        &self,
704        stmt: impl EntityOrSelect<R>,
705        rel_def: RelationDef,
706        db: &impl ConnectionTrait,
707    ) -> Result<Vec<Option<R::ModelEx>>, DbErr>
708    where
709        R: EntityTrait,
710        R::ModelEx: From<R::Model>,
711    {
712        if rel_def.rel_type != RelationType::HasOne {
713            return Err(query_err("Relation is HasMany instead of HasOne"));
714        }
715        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
716    }
717
718    fn load_many_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
719    where
720        C: ConnectionTrait,
721        R: EntityTrait,
722        S: EntityOrSelect<R>,
723        R::ModelEx: From<R::Model>,
724        <Self::Model as ModelTrait>::Entity: Related<R>,
725    {
726        loader_impl(self.iter(), stmt.select(), db)
727    }
728
729    fn load_many_ex_with_rel<R>(
730        &self,
731        stmt: impl EntityOrSelect<R>,
732        rel_def: RelationDef,
733        db: &impl ConnectionTrait,
734    ) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
735    where
736        R: EntityTrait,
737        R::ModelEx: From<R::Model>,
738    {
739        loader_impl_impl(self.iter(), stmt.select(), rel_def, None, db)
740    }
741}
742
743impl<M> LoaderTraitEx for &[Option<M>]
744where
745    M: ModelTrait,
746{
747    type Model = M;
748
749    fn load_self_ex<S, C>(
750        &self,
751        stmt: S,
752        relation_enum: LoaderExRelation<Self>,
753        db: &C,
754    ) -> Result<Vec<Option<LoaderExModelEx<Self>>>, DbErr>
755    where
756        C: ConnectionTrait,
757        LoaderExModel<Self>: Send,
758        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
759        LoaderExRelation<Self>: Send,
760        S: EntityOrSelect<LoaderExEntity<Self>>,
761    {
762        let rel_def = relation_enum.def();
763        check_self_ref(&rel_def)?;
764        let items: Vec<Option<_>> = loader_impl_impl(
765            self.iter().filter_map(|o| o.as_ref()),
766            stmt.select(),
767            rel_def,
768            None,
769            db,
770        )?;
771        Ok(assemble_options(self, items))
772    }
773
774    fn load_self_many_ex<S, C>(
775        &self,
776        stmt: S,
777        relation_enum: LoaderExRelation<Self>,
778        db: &C,
779    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
780    where
781        C: ConnectionTrait,
782        LoaderExModel<Self>: Send,
783        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
784        LoaderExRelation<Self>: Send,
785        S: EntityOrSelect<LoaderExEntity<Self>>,
786    {
787        let rel_def = relation_enum.def();
788        check_self_ref_many(&rel_def)?;
789        let items: Vec<Vec<_>> = loader_impl_impl(
790            self.iter().filter_map(|o| o.as_ref()),
791            stmt.select(),
792            rel_def,
793            None,
794            db,
795        )?;
796        Ok(assemble_options(self, items))
797    }
798
799    fn load_self_via_ex<V, C>(
800        &self,
801        _: V,
802        is_reverse: bool,
803        db: &C,
804    ) -> Result<Vec<Vec<LoaderExModelEx<Self>>>, DbErr>
805    where
806        C: ConnectionTrait,
807        V: EntityTrait,
808        LoaderExModel<Self>: Send,
809        LoaderExModelEx<Self>: From<LoaderExModel<Self>>,
810        LoaderExEntity<Self>: RelatedSelfVia<V>,
811    {
812        let (rel_def, rel_via) = if !is_reverse {
813            (
814                <LoaderExEntity<Self> as RelatedSelfVia<V>>::to(),
815                <LoaderExEntity<Self> as RelatedSelfVia<V>>::via(),
816            )
817        } else {
818            (
819                <LoaderExEntity<Self> as RelatedSelfVia<V>>::via().rev(),
820                <LoaderExEntity<Self> as RelatedSelfVia<V>>::to().rev(),
821            )
822        };
823        let items: Vec<Vec<_>> = loader_impl_impl(
824            self.iter().filter_map(|o| o.as_ref()),
825            EntityOrSelect::select(LoaderExEntity::<Self>::default()),
826            rel_def,
827            Some(rel_via),
828            db,
829        )?;
830        Ok(assemble_options(self, items))
831    }
832
833    fn load_one_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::ModelEx>>, DbErr>
834    where
835        C: ConnectionTrait,
836        R: EntityTrait,
837        S: EntityOrSelect<R>,
838        R::ModelEx: From<R::Model>,
839        <Self::Model as ModelTrait>::Entity: Related<R>,
840    {
841        let rel_def = <<Self::Model as ModelTrait>::Entity as Related<R>>::to();
842        if rel_def.rel_type != RelationType::HasOne {
843            return Err(query_err("Relation is HasMany instead of HasOne"));
844        }
845        let items: Vec<Option<R::ModelEx>> =
846            loader_impl(self.iter().filter_map(|o| o.as_ref()), stmt.select(), db)?;
847        Ok(assemble_options(self, items))
848    }
849
850    fn load_one_ex_with_rel<R>(
851        &self,
852        stmt: impl EntityOrSelect<R>,
853        rel_def: RelationDef,
854        db: &impl ConnectionTrait,
855    ) -> Result<Vec<Option<R::ModelEx>>, DbErr>
856    where
857        R: EntityTrait,
858        R::ModelEx: From<R::Model>,
859    {
860        if rel_def.rel_type != RelationType::HasOne {
861            return Err(query_err("Relation is HasMany instead of HasOne"));
862        }
863        let items: Vec<Option<R::ModelEx>> = loader_impl_impl(
864            self.iter().filter_map(|o| o.as_ref()),
865            stmt.select(),
866            rel_def,
867            None,
868            db,
869        )?;
870        Ok(assemble_options(self, items))
871    }
872
873    fn load_many_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
874    where
875        C: ConnectionTrait,
876        R: EntityTrait,
877        S: EntityOrSelect<R>,
878        R::ModelEx: From<R::Model>,
879        <Self::Model as ModelTrait>::Entity: Related<R>,
880    {
881        let items: Vec<Vec<R::ModelEx>> =
882            loader_impl(self.iter().filter_map(|o| o.as_ref()), stmt.select(), db)?;
883        Ok(assemble_options(self, items))
884    }
885
886    fn load_many_ex_with_rel<R>(
887        &self,
888        stmt: impl EntityOrSelect<R>,
889        rel_def: RelationDef,
890        db: &impl ConnectionTrait,
891    ) -> Result<Vec<Vec<R::ModelEx>>, DbErr>
892    where
893        R: EntityTrait,
894        R::ModelEx: From<R::Model>,
895    {
896        let items: Vec<Vec<R::ModelEx>> = loader_impl_impl(
897            self.iter().filter_map(|o| o.as_ref()),
898            stmt.select(),
899            rel_def,
900            None,
901            db,
902        )?;
903        Ok(assemble_options(self, items))
904    }
905}
906
907impl<M> NestedLoaderTrait for &[Vec<M>]
908where
909    M: ModelTrait,
910{
911    type Model = M;
912
913    fn load_self_ex<S, C>(
914        &self,
915        stmt: S,
916        relation_enum: NestedLoaderRelation<Self>,
917        db: &C,
918    ) -> Result<Vec<Vec<Option<NestedModelEx<Self>>>>, DbErr>
919    where
920        C: ConnectionTrait,
921        NestedModel<Self>: Send,
922        NestedModelEx<Self>: From<NestedModel<Self>>,
923        NestedLoaderRelation<Self>: Send,
924        S: EntityOrSelect<<<Self as NestedLoaderTrait>::Model as ModelTrait>::Entity>,
925    {
926        let rel_def = relation_enum.def();
927        check_self_ref(&rel_def)?;
928        let items: Vec<Option<_>> =
929            loader_impl_impl(self.iter().flatten(), stmt.select(), rel_def, None, db)?;
930        Ok(assemble_vectors(self, items))
931    }
932
933    fn load_self_many_ex<S, C>(
934        &self,
935        stmt: S,
936        relation_enum: NestedLoaderRelation<Self>,
937        db: &C,
938    ) -> Result<Vec<Vec<Vec<NestedModelEx<Self>>>>, DbErr>
939    where
940        C: ConnectionTrait,
941        NestedModel<Self>: Send,
942        NestedModelEx<Self>: From<NestedModel<Self>>,
943        NestedLoaderRelation<Self>: Send,
944        S: EntityOrSelect<<<Self as NestedLoaderTrait>::Model as ModelTrait>::Entity>,
945    {
946        let rel_def = relation_enum.def();
947        check_self_ref_many(&rel_def)?;
948        let items: Vec<Vec<_>> =
949            loader_impl_impl(self.iter().flatten(), stmt.select(), rel_def, None, db)?;
950        Ok(assemble_vectors(self, items))
951    }
952
953    fn load_self_via_ex<V, C>(
954        &self,
955        _: V,
956        is_reverse: bool,
957        db: &C,
958    ) -> Result<Vec<Vec<Vec<NestedModelEx<Self>>>>, DbErr>
959    where
960        C: ConnectionTrait,
961        V: EntityTrait,
962        NestedModel<Self>: Send,
963        NestedModelEx<Self>: From<NestedModel<Self>>,
964        NestedEntity<Self>: RelatedSelfVia<V>,
965    {
966        let (rel_def, rel_via) = if !is_reverse {
967            (
968                <NestedEntity<Self> as RelatedSelfVia<V>>::to(),
969                <NestedEntity<Self> as RelatedSelfVia<V>>::via(),
970            )
971        } else {
972            (
973                <NestedEntity<Self> as RelatedSelfVia<V>>::via().rev(),
974                <NestedEntity<Self> as RelatedSelfVia<V>>::to().rev(),
975            )
976        };
977        let items: Vec<Vec<_>> = loader_impl_impl(
978            self.iter().flatten(),
979            EntityOrSelect::select(NestedEntity::<Self>::default()),
980            rel_def,
981            Some(rel_via),
982            db,
983        )?;
984        Ok(assemble_vectors(self, items))
985    }
986
987    fn load_one_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<Option<R::ModelEx>>>, DbErr>
988    where
989        C: ConnectionTrait,
990        R: EntityTrait,
991        S: EntityOrSelect<R>,
992        R::ModelEx: From<R::Model>,
993        <Self::Model as ModelTrait>::Entity: Related<R>,
994    {
995        let rel_def = <<Self::Model as ModelTrait>::Entity as Related<R>>::to();
996        if rel_def.rel_type != RelationType::HasOne {
997            return Err(query_err("Relation is HasMany instead of HasOne"));
998        }
999        let items: Vec<Option<R::ModelEx>> = loader_impl(self.iter().flatten(), stmt.select(), db)?;
1000        Ok(assemble_vectors(self, items))
1001    }
1002
1003    fn load_one_ex_with_rel<R>(
1004        &self,
1005        stmt: impl EntityOrSelect<R>,
1006        rel_def: RelationDef,
1007        db: &impl ConnectionTrait,
1008    ) -> Result<Vec<Vec<Option<R::ModelEx>>>, DbErr>
1009    where
1010        R: EntityTrait,
1011        R::ModelEx: From<R::Model>,
1012    {
1013        if rel_def.rel_type != RelationType::HasOne {
1014            return Err(query_err("Relation is HasMany instead of HasOne"));
1015        }
1016        let items: Vec<Option<R::ModelEx>> =
1017            loader_impl_impl(self.iter().flatten(), stmt.select(), rel_def, None, db)?;
1018        Ok(assemble_vectors(self, items))
1019    }
1020
1021    fn load_many_ex<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<Vec<R::ModelEx>>>, DbErr>
1022    where
1023        C: ConnectionTrait,
1024        R: EntityTrait,
1025        S: EntityOrSelect<R>,
1026        R::ModelEx: From<R::Model>,
1027        <Self::Model as ModelTrait>::Entity: Related<R>,
1028    {
1029        let items: Vec<Vec<R::ModelEx>> = loader_impl(self.iter().flatten(), stmt.select(), db)?;
1030        Ok(assemble_vectors(self, items))
1031    }
1032
1033    fn load_many_ex_with_rel<R>(
1034        &self,
1035        stmt: impl EntityOrSelect<R>,
1036        rel_def: RelationDef,
1037        db: &impl ConnectionTrait,
1038    ) -> Result<Vec<Vec<Vec<R::ModelEx>>>, DbErr>
1039    where
1040        R: EntityTrait,
1041        R::ModelEx: From<R::Model>,
1042    {
1043        let items: Vec<Vec<R::ModelEx>> =
1044            loader_impl_impl(self.iter().flatten(), stmt.select(), rel_def, None, db)?;
1045        Ok(assemble_vectors(self, items))
1046    }
1047}
1048
1049fn assemble_options<I, T: Default>(input: &[Option<I>], items: Vec<T>) -> Vec<T> {
1050    let mut items = items.into_iter();
1051    let mut output = Vec::new();
1052    for input in input.iter() {
1053        if input.is_some() {
1054            output.push(items.next().unwrap_or_default());
1055        } else {
1056            output.push(T::default());
1057        }
1058    }
1059    output
1060}
1061
1062fn assemble_vectors<I, T: Default>(input: &[Vec<I>], items: Vec<T>) -> Vec<Vec<T>> {
1063    let mut items = items.into_iter();
1064
1065    let mut output = Vec::new();
1066
1067    for input in input.iter() {
1068        output.push(Vec::new());
1069
1070        for _inner in input.iter() {
1071            output
1072                .last_mut()
1073                .expect("Pushed above")
1074                .push(items.next().unwrap_or_default());
1075        }
1076    }
1077
1078    output
1079}
1080
1081trait Container: Default + Clone {
1082    type Item;
1083    fn add(&mut self, item: Self::Item);
1084}
1085
1086impl<T: Clone> Container for Vec<T> {
1087    type Item = T;
1088    fn add(&mut self, item: Self::Item) {
1089        self.push(item);
1090    }
1091}
1092
1093impl<T: Clone> Container for Option<T> {
1094    type Item = T;
1095    fn add(&mut self, item: Self::Item) {
1096        self.replace(item);
1097    }
1098}
1099
1100fn loader_impl<'a, Model, Iter, R, C, T, Output>(
1101    items: Iter,
1102    stmt: Select<R>,
1103    db: &C,
1104) -> Result<Vec<T>, DbErr>
1105where
1106    Model: ModelTrait + 'a,
1107    Iter: Iterator<Item = &'a Model> + 'a,
1108    C: ConnectionTrait,
1109    R: EntityTrait,
1110    Model::Entity: Related<R>,
1111    Output: From<R::Model>,
1112    T: Container<Item = Output>,
1113{
1114    loader_impl_impl(
1115        items,
1116        stmt,
1117        <Model::Entity as Related<R>>::to(),
1118        <Model::Entity as Related<R>>::via(),
1119        db,
1120    )
1121}
1122
1123// All variants monomorphizes to this implementation, which is a constant number of variants
1124// per Entity, permutating on different shapes, e.g. Vec<Model>, Option<ModelEx>
1125fn loader_impl_impl<'a, Model, Iter, R, C, T, Output>(
1126    items: Iter,
1127    stmt: Select<R>,
1128    rel_def: RelationDef,
1129    via_def: Option<RelationDef>,
1130    db: &C,
1131) -> Result<Vec<T>, DbErr>
1132where
1133    Model: ModelTrait + 'a,
1134    Iter: Iterator<Item = &'a Model> + 'a,
1135    C: ConnectionTrait,
1136    R: EntityTrait,
1137    Output: From<R::Model>,
1138    T: Container<Item = Output>,
1139{
1140    let (keys, hashmap) = if let Some(via_def) = via_def {
1141        let keys = items
1142            .map(|model| get_key_from_model(&via_def.from_col, model))
1143            .collect::<Result<Vec<_>, _>>()?;
1144
1145        if keys.is_empty() {
1146            return Ok(Vec::new());
1147        }
1148
1149        let condition = prepare_condition::<Model>(
1150            &via_def.to_tbl,
1151            &via_def.from_col,
1152            &via_def.to_col,
1153            &keys,
1154            db.get_database_backend(),
1155        )?;
1156
1157        let stmt = QueryFilter::filter(stmt.join_rev(JoinType::InnerJoin, rel_def), condition);
1158
1159        // The idea is to do a SelectTwo with join, then extract key via a dynamic model
1160        // i.e. select (baker + cake_baker) and extract cake_id from result rows
1161        // SELECT "baker"."id", "baker"."name", "baker"."contact_details", "baker"."bakery_id",
1162        //     "cakes_bakers"."cake_id" <- extra select
1163        // FROM "baker" <- target
1164        // INNER JOIN "cakes_bakers" <- junction
1165        //     ON "cakes_bakers"."baker_id" = "baker"."id" <- relation
1166        // WHERE "cakes_bakers"."cake_id" IN (..)
1167
1168        let data = stmt
1169            .select_also_dyn_model(
1170                via_def.to_tbl.sea_orm_table().clone(),
1171                dynamic::ModelType {
1172                    // we uses the left Model's type but the right Model's field
1173                    fields: extract_col_type::<Model>(&via_def.from_col, &via_def.to_col)?,
1174                },
1175            )
1176            .all(db)?;
1177
1178        let mut hashmap: HashMap<ValueTuple, T> =
1179            keys.iter()
1180                .fold(HashMap::new(), |mut acc, key: &ValueTuple| {
1181                    acc.insert(key.clone(), T::default());
1182                    acc
1183                });
1184
1185        for (item, key) in data {
1186            let key = dyn_model_to_key(key)?;
1187
1188            let vec = hashmap.get_mut(&key).ok_or_else(|| {
1189                DbErr::RecordNotFound(format!("Loader: failed to find model for {key:?}"))
1190            })?;
1191
1192            vec.add(item.into());
1193        }
1194
1195        (keys, hashmap)
1196    } else {
1197        let keys = items
1198            .map(|model| get_key_from_model(&rel_def.from_col, model))
1199            .collect::<Result<Vec<_>, _>>()?;
1200
1201        if keys.is_empty() {
1202            return Ok(Vec::new());
1203        }
1204
1205        let condition = prepare_condition::<Model>(
1206            &rel_def.to_tbl,
1207            &rel_def.from_col,
1208            &rel_def.to_col,
1209            &keys,
1210            db.get_database_backend(),
1211        )?;
1212
1213        let stmt = QueryFilter::filter(stmt, condition);
1214
1215        let data = stmt.all(db)?;
1216
1217        let mut hashmap: HashMap<ValueTuple, T> = Default::default();
1218
1219        for item in data {
1220            let key = get_key_from_model(&rel_def.to_col, &item)?;
1221            let holder = hashmap.entry(key).or_default();
1222            holder.add(item.into());
1223        }
1224
1225        (keys, hashmap)
1226    };
1227
1228    let result: Vec<T> = keys
1229        .iter()
1230        .map(|key: &ValueTuple| hashmap.get(key).cloned().unwrap_or_default())
1231        .collect();
1232
1233    Ok(result)
1234}
1235
1236fn cmp_table_ref(left: &TableRef, right: &TableRef) -> bool {
1237    left == right
1238}
1239
1240fn extract_col_type<Model>(
1241    left: &Identity,
1242    right: &Identity,
1243) -> Result<Vec<dynamic::FieldType>, DbErr>
1244where
1245    Model: ModelTrait,
1246{
1247    use itertools::Itertools;
1248
1249    if left.arity() != right.arity() {
1250        return Err(DbErr::Type(format!(
1251            "Identity mismatch: left: {} != right: {}",
1252            left.arity(),
1253            right.arity()
1254        )));
1255    }
1256
1257    let vec = left
1258        .iter()
1259        .zip_eq(right.iter())
1260        .map(|(l, r)| {
1261            let col_a =
1262                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(
1263                    &l.inner(),
1264                )
1265                .map_err(|_| DbErr::Type(format!("Failed at mapping '{l}'")))?;
1266            Ok(dynamic::FieldType::new(
1267                r.clone(),
1268                Model::get_value_type(col_a),
1269            ))
1270        })
1271        .collect::<Result<Vec<_>, DbErr>>()?;
1272
1273    Ok(vec)
1274}
1275
1276#[allow(clippy::unwrap_used)]
1277fn dyn_model_to_key(dyn_model: dynamic::Model) -> Result<ValueTuple, DbErr> {
1278    Ok(match dyn_model.fields.len() {
1279        0 => return Err(DbErr::Type("Identity zero?".into())),
1280        1 => ValueTuple::One(dyn_model.fields.into_iter().next().unwrap().value),
1281        2 => {
1282            let mut iter = dyn_model.fields.into_iter();
1283            ValueTuple::Two(iter.next().unwrap().value, iter.next().unwrap().value)
1284        }
1285        3 => {
1286            let mut iter = dyn_model.fields.into_iter();
1287            ValueTuple::Three(
1288                iter.next().unwrap().value,
1289                iter.next().unwrap().value,
1290                iter.next().unwrap().value,
1291            )
1292        }
1293        _ => ValueTuple::Many(dyn_model.fields.into_iter().map(|v| v.value).collect()),
1294    })
1295}
1296
1297fn arity_mismatch(expected: usize, actual: &ValueTuple) -> DbErr {
1298    DbErr::Type(format!(
1299        "Loader: arity mismatch: expected {expected}, got {} in {actual:?}",
1300        actual.arity()
1301    ))
1302}
1303
1304fn prepare_condition<Model>(
1305    table: &TableRef,
1306    from: &Identity,
1307    to: &Identity,
1308    keys: &[ValueTuple],
1309    db_backend: DbBackend,
1310) -> Result<Condition, DbErr>
1311where
1312    Model: ModelTrait,
1313{
1314    if matches!(db_backend, DbBackend::Postgres) {
1315        prepare_condition_with_save_as::<Model>(table, from, to, keys)
1316    } else {
1317        column_tuple_in_condition(table, to, keys, db_backend)
1318    }
1319}
1320
1321fn prepare_condition_with_save_as<Model>(
1322    table: &TableRef,
1323    from: &Identity,
1324    to: &Identity,
1325    keys: &[ValueTuple],
1326) -> Result<Condition, DbErr>
1327where
1328    Model: ModelTrait,
1329{
1330    use itertools::Itertools;
1331
1332    let keys = keys.iter().unique();
1333    let (from_cols, to_cols) = resolve_column_pairs::<Model>(table, from, to)?;
1334
1335    if from_cols.is_empty() || to_cols.is_empty() {
1336        return Err(DbErr::Type(format!(
1337            "Loader: resolved zero columns for identities {from:?} -> {to:?}"
1338        )));
1339    }
1340
1341    let arity = from_cols.len();
1342
1343    let value_tuples = keys
1344        .map(|key| {
1345            let key_arity = key.arity();
1346            if arity != key_arity {
1347                return Err(arity_mismatch(arity, key));
1348            }
1349
1350            // For Postgres, we need to use `AS` to cast the value to the correct type
1351            Ok(apply_save_as::<Model>(&from_cols, key.clone()))
1352        })
1353        .collect::<Result<Vec<_>, DbErr>>()?;
1354
1355    // Build `(c1, c2, ...) IN ((v11, v12, ...), (v21, v22, ...), ...)`
1356    let expr = Expr::tuple(create_table_columns(table, to)).is_in(value_tuples);
1357
1358    Ok(expr.into())
1359}
1360
1361type ModelColumn<M> = <<M as ModelTrait>::Entity as EntityTrait>::Column;
1362
1363type ColumnPairs<M> = (Vec<ModelColumn<M>>, Vec<ColumnRef>);
1364
1365fn resolve_column_pairs<Model>(
1366    table: &TableRef,
1367    from: &Identity,
1368    to: &Identity,
1369) -> Result<ColumnPairs<Model>, DbErr>
1370where
1371    Model: ModelTrait,
1372    ModelColumn<Model>: ColumnTrait,
1373{
1374    let from_columns = parse_identity_columns::<Model>(from)?;
1375    let to_columns = column_refs_from_identity(table, to);
1376
1377    if from_columns.len() != to_columns.len() {
1378        return Err(DbErr::Type(format!(
1379            "Loader: identity column count mismatch between {from:?} and {to:?}"
1380        )));
1381    }
1382
1383    Ok((from_columns, to_columns))
1384}
1385
1386fn check_self_ref(rel_def: &RelationDef) -> Result<(), DbErr> {
1387    if rel_def.from_tbl != rel_def.to_tbl {
1388        return Err(query_err("Relation must be self referencing"));
1389    }
1390    if rel_def.is_owner {
1391        return Err(query_err("Relation must be belongs_to"));
1392    }
1393    Ok(())
1394}
1395
1396fn check_self_ref_many(rel_def: &RelationDef) -> Result<(), DbErr> {
1397    if rel_def.from_tbl != rel_def.to_tbl {
1398        return Err(query_err("Relation must be self referencing"));
1399    }
1400    if !rel_def.is_owner {
1401        return Err(query_err("Relation must not be belongs_to"));
1402    }
1403    Ok(())
1404}
1405
1406fn column_refs_from_identity(table: &TableRef, identity: &Identity) -> Vec<ColumnRef> {
1407    identity
1408        .iter()
1409        .map(|col| table_column(table, col))
1410        .collect()
1411}
1412
1413fn parse_identity_columns<Model>(identity: &Identity) -> Result<Vec<ModelColumn<Model>>, DbErr>
1414where
1415    Model: ModelTrait,
1416{
1417    identity
1418        .iter()
1419        .map(|from_col| try_conv_ident_to_column::<Model>(from_col))
1420        .collect()
1421}
1422
1423fn try_conv_ident_to_column<Model>(ident: &DynIden) -> Result<ModelColumn<Model>, DbErr>
1424where
1425    Model: ModelTrait,
1426{
1427    let column_name = ident.inner();
1428    ModelColumn::<Model>::from_str(&column_name)
1429        .map_err(|_| DbErr::Type(format!("Failed at mapping '{column_name}' to column")))
1430}
1431
1432fn table_column(tbl: &TableRef, col: &DynIden) -> ColumnRef {
1433    (tbl.sea_orm_table().to_owned(), col.clone()).into_column_ref()
1434}
1435
1436/// Create a vector of `Expr::col` from the table and identity, e.g. [Expr::col((table, col1)), Expr::col((table, col2)), ...]
1437fn create_table_columns(table: &TableRef, cols: &Identity) -> Vec<Expr> {
1438    cols.iter()
1439        .map(|col| table_column(table, col))
1440        .map(Expr::col)
1441        .collect()
1442}
1443
1444/// Apply `save_as` to each value in the tuple, e.g. `(Cast(val1 as type1), Cast(val2 as type2), ...)`
1445fn apply_save_as<M: ModelTrait>(cols: &[ModelColumn<M>], values: ValueTuple) -> Expr {
1446    let values_expr_iter = values.into_iter().map(Expr::val);
1447
1448    let tuple_exprs: Vec<_> = cols
1449        .iter()
1450        .zip(values_expr_iter)
1451        .map(|(model_column, value)| model_column.save_as(value))
1452        .collect();
1453
1454    Expr::tuple(tuple_exprs)
1455}
1456
1457#[cfg(test)]
1458mod tests {
1459    fn cake_model(id: i32) -> sea_orm::tests_cfg::cake::Model {
1460        let name = match id {
1461            1 => "apple cake",
1462            2 => "orange cake",
1463            3 => "fruit cake",
1464            4 => "chocolate cake",
1465            _ => "",
1466        }
1467        .to_string();
1468        sea_orm::tests_cfg::cake::Model { id, name }
1469    }
1470
1471    fn fruit_model(id: i32, cake_id: Option<i32>) -> sea_orm::tests_cfg::fruit::Model {
1472        let name = match id {
1473            1 => "apple",
1474            2 => "orange",
1475            3 => "grape",
1476            4 => "strawberry",
1477            _ => "",
1478        }
1479        .to_string();
1480        sea_orm::tests_cfg::fruit::Model { id, name, cake_id }
1481    }
1482
1483    fn filling_model(id: i32) -> sea_orm::tests_cfg::filling::Model {
1484        let name = match id {
1485            1 => "apple juice",
1486            2 => "orange jam",
1487            3 => "chocolate crust",
1488            4 => "strawberry jam",
1489            _ => "",
1490        }
1491        .to_string();
1492        sea_orm::tests_cfg::filling::Model {
1493            id,
1494            name,
1495            vendor_id: Some(1),
1496            ignored_attr: 0,
1497        }
1498    }
1499
1500    fn cake_filling_model(
1501        cake_id: i32,
1502        filling_id: i32,
1503    ) -> sea_orm::tests_cfg::cake_filling::Model {
1504        sea_orm::tests_cfg::cake_filling::Model {
1505            cake_id,
1506            filling_id,
1507        }
1508    }
1509
1510    #[test]
1511    fn test_load_one() {
1512        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1513
1514        let db = MockDatabase::new(DbBackend::Postgres)
1515            .append_query_results([[cake_model(1), cake_model(2)]])
1516            .into_connection();
1517
1518        let fruits = vec![fruit_model(1, Some(1))];
1519
1520        let cakes = fruits
1521            .load_one(cake::Entity::find(), &db)
1522            .expect("Should return something");
1523
1524        assert_eq!(cakes, [Some(cake_model(1))]);
1525    }
1526
1527    #[test]
1528    fn test_load_one_same_cake() {
1529        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1530
1531        let db = MockDatabase::new(DbBackend::Postgres)
1532            .append_query_results([[cake_model(1), cake_model(2)]])
1533            .into_connection();
1534
1535        let fruits = vec![fruit_model(1, Some(1)), fruit_model(2, Some(1))];
1536
1537        let cakes = fruits
1538            .load_one(cake::Entity::find(), &db)
1539            .expect("Should return something");
1540
1541        assert_eq!(cakes, [Some(cake_model(1)), Some(cake_model(1))]);
1542    }
1543
1544    #[test]
1545    fn test_load_one_empty() {
1546        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1547
1548        let db = MockDatabase::new(DbBackend::Postgres)
1549            .append_query_results([[cake_model(1), cake_model(2)]])
1550            .into_connection();
1551
1552        let fruits: Vec<fruit::Model> = vec![];
1553
1554        let cakes = fruits
1555            .load_one(cake::Entity::find(), &db)
1556            .expect("Should return something");
1557
1558        assert_eq!(cakes, []);
1559    }
1560
1561    #[test]
1562    fn test_load_many() {
1563        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1564
1565        let db = MockDatabase::new(DbBackend::Postgres)
1566            .append_query_results([[fruit_model(1, Some(1))]])
1567            .into_connection();
1568
1569        let cakes = vec![cake_model(1), cake_model(2)];
1570
1571        let fruits = cakes
1572            .load_many(fruit::Entity::find(), &db)
1573            .expect("Should return something");
1574
1575        assert_eq!(fruits, [vec![fruit_model(1, Some(1))], vec![]]);
1576    }
1577
1578    #[test]
1579    fn test_load_many_same_fruit() {
1580        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1581
1582        let db = MockDatabase::new(DbBackend::Postgres)
1583            .append_query_results([[fruit_model(1, Some(1)), fruit_model(2, Some(1))]])
1584            .into_connection();
1585
1586        let cakes = vec![cake_model(1), cake_model(2)];
1587
1588        let fruits = cakes
1589            .load_many(fruit::Entity::find(), &db)
1590            .expect("Should return something");
1591
1592        assert_eq!(
1593            fruits,
1594            [
1595                vec![fruit_model(1, Some(1)), fruit_model(2, Some(1))],
1596                vec![]
1597            ]
1598        );
1599    }
1600
1601    #[test]
1602    fn test_load_many_empty() {
1603        use sea_orm::{DbBackend, MockDatabase, entity::prelude::*, tests_cfg::*};
1604
1605        let db = MockDatabase::new(DbBackend::Postgres).into_connection();
1606
1607        let cakes: Vec<cake::Model> = vec![];
1608
1609        let fruits = cakes
1610            .load_many(fruit::Entity::find(), &db)
1611            .expect("Should return something");
1612
1613        let empty_vec: Vec<Vec<fruit::Model>> = vec![];
1614
1615        assert_eq!(fruits, empty_vec);
1616    }
1617
1618    #[test]
1619    fn test_load_many_to_many_base() {
1620        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
1621
1622        let db = MockDatabase::new(DbBackend::Postgres)
1623            .append_query_results([
1624                [cake_filling_model(1, 1).into_mock_row()],
1625                [filling_model(1).into_mock_row()],
1626            ])
1627            .into_connection();
1628
1629        let cakes = vec![cake_model(1)];
1630
1631        let fillings = cakes
1632            .load_many_to_many(Filling, CakeFilling, &db)
1633            .expect("Should return something");
1634
1635        assert_eq!(fillings, vec![vec![filling_model(1)]]);
1636    }
1637
1638    #[test]
1639    fn test_load_many_to_many_complex() {
1640        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
1641
1642        let db = MockDatabase::new(DbBackend::Postgres)
1643            .append_query_results([
1644                [
1645                    cake_filling_model(1, 1).into_mock_row(),
1646                    cake_filling_model(1, 2).into_mock_row(),
1647                    cake_filling_model(1, 3).into_mock_row(),
1648                    cake_filling_model(2, 1).into_mock_row(),
1649                    cake_filling_model(2, 2).into_mock_row(),
1650                ],
1651                [
1652                    filling_model(1).into_mock_row(),
1653                    filling_model(2).into_mock_row(),
1654                    filling_model(3).into_mock_row(),
1655                    filling_model(4).into_mock_row(),
1656                    filling_model(5).into_mock_row(),
1657                ],
1658            ])
1659            .into_connection();
1660
1661        let cakes = vec![cake_model(1), cake_model(2), cake_model(3)];
1662
1663        let fillings = cakes
1664            .load_many_to_many(Filling, CakeFilling, &db)
1665            .expect("Should return something");
1666
1667        assert_eq!(
1668            fillings,
1669            vec![
1670                vec![filling_model(1), filling_model(2), filling_model(3)],
1671                vec![filling_model(1), filling_model(2)],
1672                vec![],
1673            ]
1674        );
1675    }
1676
1677    #[test]
1678    fn test_load_many_to_many_empty() {
1679        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
1680
1681        let db = MockDatabase::new(DbBackend::Postgres)
1682            .append_query_results([
1683                [cake_filling_model(1, 1).into_mock_row()],
1684                [filling_model(1).into_mock_row()],
1685            ])
1686            .into_connection();
1687
1688        let cakes: Vec<cake::Model> = vec![];
1689
1690        let fillings = cakes
1691            .load_many_to_many(Filling, CakeFilling, &db)
1692            .expect("Should return something");
1693
1694        let empty_vec: Vec<Vec<filling::Model>> = vec![];
1695
1696        assert_eq!(fillings, empty_vec);
1697    }
1698
1699    #[test]
1700    fn test_load_one_duplicate_keys() {
1701        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1702
1703        let db = MockDatabase::new(DbBackend::Postgres)
1704            .append_query_results([[cake_model(1), cake_model(2)]])
1705            .into_connection();
1706
1707        let fruits = vec![
1708            fruit_model(1, Some(1)),
1709            fruit_model(2, Some(1)),
1710            fruit_model(3, Some(1)),
1711            fruit_model(4, Some(1)),
1712        ];
1713
1714        let cakes = fruits
1715            .load_one(cake::Entity::find(), &db)
1716            .expect("Should return something");
1717
1718        assert_eq!(cakes.len(), 4);
1719        for cake in &cakes {
1720            assert_eq!(cake, &Some(cake_model(1)));
1721        }
1722        let logs = db.into_transaction_log();
1723        let sql = format!("{:?}", logs[0]);
1724
1725        let values_count = sql.matches("$1").count();
1726        assert_eq!(values_count, 1, "Duplicate values were not removed");
1727    }
1728
1729    #[test]
1730    fn test_load_many_duplicate_keys() {
1731        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
1732
1733        let db = MockDatabase::new(DbBackend::Postgres)
1734            .append_query_results([[
1735                fruit_model(1, Some(1)),
1736                fruit_model(2, Some(1)),
1737                fruit_model(3, Some(2)),
1738            ]])
1739            .into_connection();
1740
1741        let cakes = vec![cake_model(1), cake_model(1), cake_model(2), cake_model(2)];
1742
1743        let fruits = cakes
1744            .load_many(fruit::Entity::find(), &db)
1745            .expect("Should return something");
1746
1747        assert_eq!(fruits.len(), 4);
1748
1749        let logs = db.into_transaction_log();
1750        let sql = format!("{:?}", logs[0]);
1751
1752        let values_count = sql.matches("$1").count() + sql.matches("$2").count();
1753        assert_eq!(values_count, 2, "Duplicate values were not removed");
1754    }
1755
1756    #[test]
1757    fn test_assemble_vectors() {
1758        use super::assemble_vectors;
1759
1760        assert_eq!(
1761            assemble_vectors(&[vec![1], vec![], vec![2, 3], vec![]], vec![11, 22, 33]),
1762            [vec![11], vec![], vec![22, 33], vec![]]
1763        );
1764    }
1765}
sea_orm/query/loader.rs

sea_orm/query/
loader.rs
