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

sea_orm/query/
loader.rs
