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