sea_orm/query/
loader.rs

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