sea_orm/query/
loader.rs

1use crate::{
2    Condition, ConnectionTrait, DbErr, EntityTrait, Identity, ModelTrait, QueryFilter, Related,
3    RelationType, Select, error::*,
4};
5use async_trait::async_trait;
6use sea_query::{
7    ColumnRef, DynIden, Expr, ExprTrait, IntoColumnRef, SimpleExpr, TableRef, ValueTuple,
8};
9use std::{
10    collections::{HashMap, HashSet},
11    str::FromStr,
12};
13
14/// Entity, or a Select<Entity>; to be used as parameters in [`LoaderTrait`]
15pub trait EntityOrSelect<E: EntityTrait>: Send {
16    /// If self is Entity, use Entity::find()
17    fn select(self) -> Select<E>;
18}
19
20/// This trait implements the Data Loader API
21#[async_trait]
22pub trait LoaderTrait {
23    /// Source model
24    type Model: ModelTrait;
25
26    /// Used to eager load has_one relations
27    async fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
28    where
29        C: ConnectionTrait,
30        R: EntityTrait,
31        R::Model: Send + Sync,
32        S: EntityOrSelect<R>,
33        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>;
34
35    /// Used to eager load has_many relations
36    async fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
37    where
38        C: ConnectionTrait,
39        R: EntityTrait,
40        R::Model: Send + Sync,
41        S: EntityOrSelect<R>,
42        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>;
43
44    /// Used to eager load many_to_many relations
45    async fn load_many_to_many<R, S, V, C>(
46        &self,
47        stmt: S,
48        via: V,
49        db: &C,
50    ) -> Result<Vec<Vec<R::Model>>, DbErr>
51    where
52        C: ConnectionTrait,
53        R: EntityTrait,
54        R::Model: Send + Sync,
55        S: EntityOrSelect<R>,
56        V: EntityTrait,
57        V::Model: Send + Sync,
58        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>;
59}
60
61impl<E> EntityOrSelect<E> for E
62where
63    E: EntityTrait,
64{
65    fn select(self) -> Select<E> {
66        E::find()
67    }
68}
69
70impl<E> EntityOrSelect<E> for Select<E>
71where
72    E: EntityTrait,
73{
74    fn select(self) -> Select<E> {
75        self
76    }
77}
78
79#[async_trait]
80impl<M> LoaderTrait for Vec<M>
81where
82    M: ModelTrait + Sync,
83{
84    type Model = M;
85
86    async fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
87    where
88        C: ConnectionTrait,
89        R: EntityTrait,
90        R::Model: Send + Sync,
91        S: EntityOrSelect<R>,
92        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
93    {
94        self.as_slice().load_one(stmt, db).await
95    }
96
97    async fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
98    where
99        C: ConnectionTrait,
100        R: EntityTrait,
101        R::Model: Send + Sync,
102        S: EntityOrSelect<R>,
103        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
104    {
105        self.as_slice().load_many(stmt, db).await
106    }
107
108    async fn load_many_to_many<R, S, V, C>(
109        &self,
110        stmt: S,
111        via: V,
112        db: &C,
113    ) -> Result<Vec<Vec<R::Model>>, DbErr>
114    where
115        C: ConnectionTrait,
116        R: EntityTrait,
117        R::Model: Send + Sync,
118        S: EntityOrSelect<R>,
119        V: EntityTrait,
120        V::Model: Send + Sync,
121        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
122    {
123        self.as_slice().load_many_to_many(stmt, via, db).await
124    }
125}
126
127#[async_trait]
128impl<M> LoaderTrait for &[M]
129where
130    M: ModelTrait + Sync,
131{
132    type Model = M;
133
134    async fn load_one<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Option<R::Model>>, DbErr>
135    where
136        C: ConnectionTrait,
137        R: EntityTrait,
138        R::Model: Send + Sync,
139        S: EntityOrSelect<R>,
140        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
141    {
142        // we verify that is HasOne relation
143        if <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::via().is_some() {
144            return Err(query_err("Relation is ManytoMany instead of HasOne"));
145        }
146        let rel_def = <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::to();
147        if rel_def.rel_type == RelationType::HasMany {
148            return Err(query_err("Relation is HasMany instead of HasOne"));
149        }
150
151        if self.is_empty() {
152            return Ok(Vec::new());
153        }
154
155        let keys = self
156            .iter()
157            .map(|model| extract_key(&rel_def.from_col, model))
158            .collect::<Result<Vec<_>, _>>()?;
159
160        let condition = prepare_condition(&rel_def.to_tbl, &rel_def.to_col, &keys);
161
162        let stmt = <Select<R> as QueryFilter>::filter(stmt.select(), condition);
163
164        let data = stmt.all(db).await?;
165
166        let hashmap = data.into_iter().try_fold(
167            HashMap::<ValueTuple, <R as EntityTrait>::Model>::new(),
168            |mut acc, value| {
169                extract_key(&rel_def.to_col, &value).map(|key| {
170                    acc.insert(key, value);
171
172                    acc
173                })
174            },
175        )?;
176
177        let result: Vec<Option<<R as EntityTrait>::Model>> =
178            keys.iter().map(|key| hashmap.get(key).cloned()).collect();
179
180        Ok(result)
181    }
182
183    async fn load_many<R, S, C>(&self, stmt: S, db: &C) -> Result<Vec<Vec<R::Model>>, DbErr>
184    where
185        C: ConnectionTrait,
186        R: EntityTrait,
187        R::Model: Send + Sync,
188        S: EntityOrSelect<R>,
189        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
190    {
191        // we verify that is HasMany relation
192
193        if <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::via().is_some() {
194            return Err(query_err("Relation is ManyToMany instead of HasMany"));
195        }
196        let rel_def = <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::to();
197        if rel_def.rel_type == RelationType::HasOne {
198            return Err(query_err("Relation is HasOne instead of HasMany"));
199        }
200
201        if self.is_empty() {
202            return Ok(Vec::new());
203        }
204
205        let keys = self
206            .iter()
207            .map(|model| extract_key(&rel_def.from_col, model))
208            .collect::<Result<Vec<_>, _>>()?;
209
210        let condition = prepare_condition(&rel_def.to_tbl, &rel_def.to_col, &keys);
211
212        let stmt = <Select<R> as QueryFilter>::filter(stmt.select(), condition);
213
214        let data = stmt.all(db).await?;
215
216        let mut hashmap: HashMap<ValueTuple, Vec<<R as EntityTrait>::Model>> =
217            keys.iter()
218                .fold(HashMap::new(), |mut acc, key: &ValueTuple| {
219                    acc.insert(key.clone(), Vec::new());
220                    acc
221                });
222
223        for value in data {
224            let key = extract_key(&rel_def.to_col, &value)?;
225
226            let vec = hashmap.get_mut(&key).ok_or_else(|| {
227                DbErr::RecordNotFound(format!("Loader: failed to find model for {key:?}"))
228            })?;
229
230            vec.push(value);
231        }
232
233        let result: Vec<Vec<R::Model>> = keys
234            .iter()
235            .map(|key: &ValueTuple| hashmap.get(key).cloned().unwrap_or_default())
236            .collect();
237
238        Ok(result)
239    }
240
241    async fn load_many_to_many<R, S, V, C>(
242        &self,
243        stmt: S,
244        via: V,
245        db: &C,
246    ) -> Result<Vec<Vec<R::Model>>, DbErr>
247    where
248        C: ConnectionTrait,
249        R: EntityTrait,
250        R::Model: Send + Sync,
251        S: EntityOrSelect<R>,
252        V: EntityTrait,
253        V::Model: Send + Sync,
254        <<Self as LoaderTrait>::Model as ModelTrait>::Entity: Related<R>,
255    {
256        if let Some(via_rel) =
257            <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::via()
258        {
259            let rel_def =
260                <<<Self as LoaderTrait>::Model as ModelTrait>::Entity as Related<R>>::to();
261            if rel_def.rel_type != RelationType::HasOne {
262                return Err(query_err("Relation to is not HasOne"));
263            }
264
265            if !cmp_table_ref(&via_rel.to_tbl, &via.table_ref()) {
266                return Err(query_err(format!(
267                    "The given via Entity is incorrect: expected: {:?}, given: {:?}",
268                    via_rel.to_tbl,
269                    via.table_ref()
270                )));
271            }
272
273            if self.is_empty() {
274                return Ok(Vec::new());
275            }
276
277            let pkeys = self
278                .iter()
279                .map(|model| extract_key(&via_rel.from_col, model))
280                .collect::<Result<Vec<_>, _>>()?;
281
282            // Map of M::PK -> Vec<R::PK>
283            let mut keymap: HashMap<ValueTuple, Vec<ValueTuple>> = Default::default();
284
285            let keys: Vec<ValueTuple> = {
286                let condition = prepare_condition(&via_rel.to_tbl, &via_rel.to_col, &pkeys);
287                let stmt = V::find().filter(condition);
288                let data = stmt.all(db).await?;
289                for model in data {
290                    let pk = extract_key(&via_rel.to_col, &model)?;
291                    let entry = keymap.entry(pk).or_default();
292
293                    let fk = extract_key(&rel_def.from_col, &model)?;
294                    entry.push(fk);
295                }
296
297                keymap.values().flatten().cloned().collect()
298            };
299
300            let condition = prepare_condition(&rel_def.to_tbl, &rel_def.to_col, &keys);
301
302            let stmt = <Select<R> as QueryFilter>::filter(stmt.select(), condition);
303
304            let models = stmt.all(db).await?;
305
306            // Map of R::PK -> R::Model
307            let data = models.into_iter().try_fold(
308                HashMap::<ValueTuple, <R as EntityTrait>::Model>::new(),
309                |mut acc, model| {
310                    extract_key(&rel_def.to_col, &model).map(|key| {
311                        acc.insert(key, model);
312
313                        acc
314                    })
315                },
316            )?;
317
318            let result: Vec<Vec<R::Model>> = pkeys
319                .into_iter()
320                .map(|pkey| {
321                    let fkeys = keymap.get(&pkey).cloned().unwrap_or_default();
322
323                    let models: Vec<_> = fkeys
324                        .into_iter()
325                        .filter_map(|fkey| data.get(&fkey).cloned())
326                        .collect();
327
328                    models
329                })
330                .collect();
331
332            Ok(result)
333        } else {
334            return Err(query_err("Relation is not ManyToMany"));
335        }
336    }
337}
338
339fn cmp_table_ref(left: &TableRef, right: &TableRef) -> bool {
340    // not ideal; but
341    format!("{left:?}") == format!("{right:?}")
342}
343
344fn extract_key<Model>(target_col: &Identity, model: &Model) -> Result<ValueTuple, DbErr>
345where
346    Model: ModelTrait,
347{
348    Ok(match target_col {
349        Identity::Unary(a) => {
350            let a = a.to_string();
351            let column_a =
352                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(&a)
353                    .map_err(|_| DbErr::Type(format!("Failed at mapping '{a}' to column A:1")))?;
354            ValueTuple::One(model.get(column_a))
355        }
356        Identity::Binary(a, b) => {
357            let a = a.to_string();
358            let b = b.to_string();
359            let column_a =
360                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(&a)
361                    .map_err(|_| DbErr::Type(format!("Failed at mapping '{a}' to column A:2")))?;
362            let column_b =
363                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(&b)
364                    .map_err(|_| DbErr::Type(format!("Failed at mapping '{b}' to column B:2")))?;
365            ValueTuple::Two(model.get(column_a), model.get(column_b))
366        }
367        Identity::Ternary(a, b, c) => {
368            let a = a.to_string();
369            let b = b.to_string();
370            let c = c.to_string();
371            let column_a =
372                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(
373                    &a.to_string(),
374                )
375                .map_err(|_| DbErr::Type(format!("Failed at mapping '{a}' to column A:3")))?;
376            let column_b =
377                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(
378                    &b.to_string(),
379                )
380                .map_err(|_| DbErr::Type(format!("Failed at mapping '{b}' to column B:3")))?;
381            let column_c =
382                <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(
383                    &c.to_string(),
384                )
385                .map_err(|_| DbErr::Type(format!("Failed at mapping '{c}' to column C:3")))?;
386            ValueTuple::Three(
387                model.get(column_a),
388                model.get(column_b),
389                model.get(column_c),
390            )
391        }
392        Identity::Many(cols) => {
393            let mut values = Vec::new();
394            for col in cols {
395                let col_name = col.to_string();
396                let column =
397                    <<<Model as ModelTrait>::Entity as EntityTrait>::Column as FromStr>::from_str(
398                        &col_name,
399                    )
400                    .map_err(|_| DbErr::Type(format!("Failed at mapping '{col_name}' to colum")))?;
401                values.push(model.get(column))
402            }
403            ValueTuple::Many(values)
404        }
405    })
406}
407
408fn prepare_condition(table: &TableRef, col: &Identity, keys: &[ValueTuple]) -> Condition {
409    let keys = if !keys.is_empty() {
410        let set: HashSet<_> = keys.iter().cloned().collect();
411        set.into_iter().collect()
412    } else {
413        Vec::new()
414    };
415
416    match col {
417        Identity::Unary(column_a) => {
418            let column_a = table_column(table, column_a);
419            Condition::all().add(Expr::col(column_a).is_in(keys.into_iter().flatten()))
420        }
421        Identity::Binary(column_a, column_b) => Condition::all().add(
422            Expr::tuple([
423                SimpleExpr::Column(table_column(table, column_a)),
424                SimpleExpr::Column(table_column(table, column_b)),
425            ])
426            .in_tuples(keys),
427        ),
428        Identity::Ternary(column_a, column_b, column_c) => Condition::all().add(
429            Expr::tuple([
430                SimpleExpr::Column(table_column(table, column_a)),
431                SimpleExpr::Column(table_column(table, column_b)),
432                SimpleExpr::Column(table_column(table, column_c)),
433            ])
434            .in_tuples(keys),
435        ),
436        Identity::Many(cols) => {
437            let columns = cols
438                .iter()
439                .map(|col| SimpleExpr::Column(table_column(table, col)));
440            Condition::all().add(Expr::tuple(columns).in_tuples(keys))
441        }
442    }
443}
444
445fn table_column(tbl: &TableRef, col: &DynIden) -> ColumnRef {
446    (tbl.sea_orm_table().to_owned(), col.clone()).into_column_ref()
447}
448
449#[cfg(test)]
450mod tests {
451    fn cake_model(id: i32) -> sea_orm::tests_cfg::cake::Model {
452        let name = match id {
453            1 => "apple cake",
454            2 => "orange cake",
455            3 => "fruit cake",
456            4 => "chocolate cake",
457            _ => "",
458        }
459        .to_string();
460        sea_orm::tests_cfg::cake::Model { id, name }
461    }
462
463    fn fruit_model(id: i32, cake_id: Option<i32>) -> sea_orm::tests_cfg::fruit::Model {
464        let name = match id {
465            1 => "apple",
466            2 => "orange",
467            3 => "grape",
468            4 => "strawberry",
469            _ => "",
470        }
471        .to_string();
472        sea_orm::tests_cfg::fruit::Model { id, name, cake_id }
473    }
474
475    fn filling_model(id: i32) -> sea_orm::tests_cfg::filling::Model {
476        let name = match id {
477            1 => "apple juice",
478            2 => "orange jam",
479            3 => "chocolate crust",
480            4 => "strawberry jam",
481            _ => "",
482        }
483        .to_string();
484        sea_orm::tests_cfg::filling::Model {
485            id,
486            name,
487            vendor_id: Some(1),
488            ignored_attr: 0,
489        }
490    }
491
492    fn cake_filling_model(
493        cake_id: i32,
494        filling_id: i32,
495    ) -> sea_orm::tests_cfg::cake_filling::Model {
496        sea_orm::tests_cfg::cake_filling::Model {
497            cake_id,
498            filling_id,
499        }
500    }
501
502    #[tokio::test]
503    async fn test_load_one() {
504        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
505
506        let db = MockDatabase::new(DbBackend::Postgres)
507            .append_query_results([[cake_model(1), cake_model(2)]])
508            .into_connection();
509
510        let fruits = vec![fruit_model(1, Some(1))];
511
512        let cakes = fruits
513            .load_one(cake::Entity::find(), &db)
514            .await
515            .expect("Should return something");
516
517        assert_eq!(cakes, [Some(cake_model(1))]);
518    }
519
520    #[tokio::test]
521    async fn test_load_one_same_cake() {
522        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
523
524        let db = MockDatabase::new(DbBackend::Postgres)
525            .append_query_results([[cake_model(1), cake_model(2)]])
526            .into_connection();
527
528        let fruits = vec![fruit_model(1, Some(1)), fruit_model(2, Some(1))];
529
530        let cakes = fruits
531            .load_one(cake::Entity::find(), &db)
532            .await
533            .expect("Should return something");
534
535        assert_eq!(cakes, [Some(cake_model(1)), Some(cake_model(1))]);
536    }
537
538    #[tokio::test]
539    async fn test_load_one_empty() {
540        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
541
542        let db = MockDatabase::new(DbBackend::Postgres)
543            .append_query_results([[cake_model(1), cake_model(2)]])
544            .into_connection();
545
546        let fruits: Vec<fruit::Model> = vec![];
547
548        let cakes = fruits
549            .load_one(cake::Entity::find(), &db)
550            .await
551            .expect("Should return something");
552
553        assert_eq!(cakes, []);
554    }
555
556    #[tokio::test]
557    async fn test_load_many() {
558        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
559
560        let db = MockDatabase::new(DbBackend::Postgres)
561            .append_query_results([[fruit_model(1, Some(1))]])
562            .into_connection();
563
564        let cakes = vec![cake_model(1), cake_model(2)];
565
566        let fruits = cakes
567            .load_many(fruit::Entity::find(), &db)
568            .await
569            .expect("Should return something");
570
571        assert_eq!(fruits, [vec![fruit_model(1, Some(1))], vec![]]);
572    }
573
574    #[tokio::test]
575    async fn test_load_many_same_fruit() {
576        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
577
578        let db = MockDatabase::new(DbBackend::Postgres)
579            .append_query_results([[fruit_model(1, Some(1)), fruit_model(2, Some(1))]])
580            .into_connection();
581
582        let cakes = vec![cake_model(1), cake_model(2)];
583
584        let fruits = cakes
585            .load_many(fruit::Entity::find(), &db)
586            .await
587            .expect("Should return something");
588
589        assert_eq!(
590            fruits,
591            [
592                vec![fruit_model(1, Some(1)), fruit_model(2, Some(1))],
593                vec![]
594            ]
595        );
596    }
597
598    #[tokio::test]
599    async fn test_load_many_empty() {
600        use sea_orm::{DbBackend, MockDatabase, entity::prelude::*, tests_cfg::*};
601
602        let db = MockDatabase::new(DbBackend::Postgres).into_connection();
603
604        let cakes: Vec<cake::Model> = vec![];
605
606        let fruits = cakes
607            .load_many(fruit::Entity::find(), &db)
608            .await
609            .expect("Should return something");
610
611        let empty_vec: Vec<Vec<fruit::Model>> = vec![];
612
613        assert_eq!(fruits, empty_vec);
614    }
615
616    #[tokio::test]
617    async fn test_load_many_to_many_base() {
618        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
619
620        let db = MockDatabase::new(DbBackend::Postgres)
621            .append_query_results([
622                [cake_filling_model(1, 1).into_mock_row()],
623                [filling_model(1).into_mock_row()],
624            ])
625            .into_connection();
626
627        let cakes = vec![cake_model(1)];
628
629        let fillings = cakes
630            .load_many_to_many(Filling, CakeFilling, &db)
631            .await
632            .expect("Should return something");
633
634        assert_eq!(fillings, vec![vec![filling_model(1)]]);
635    }
636
637    #[tokio::test]
638    async fn test_load_many_to_many_complex() {
639        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
640
641        let db = MockDatabase::new(DbBackend::Postgres)
642            .append_query_results([
643                [
644                    cake_filling_model(1, 1).into_mock_row(),
645                    cake_filling_model(1, 2).into_mock_row(),
646                    cake_filling_model(1, 3).into_mock_row(),
647                    cake_filling_model(2, 1).into_mock_row(),
648                    cake_filling_model(2, 2).into_mock_row(),
649                ],
650                [
651                    filling_model(1).into_mock_row(),
652                    filling_model(2).into_mock_row(),
653                    filling_model(3).into_mock_row(),
654                    filling_model(4).into_mock_row(),
655                    filling_model(5).into_mock_row(),
656                ],
657            ])
658            .into_connection();
659
660        let cakes = vec![cake_model(1), cake_model(2), cake_model(3)];
661
662        let fillings = cakes
663            .load_many_to_many(Filling, CakeFilling, &db)
664            .await
665            .expect("Should return something");
666
667        assert_eq!(
668            fillings,
669            vec![
670                vec![filling_model(1), filling_model(2), filling_model(3)],
671                vec![filling_model(1), filling_model(2)],
672                vec![],
673            ]
674        );
675    }
676
677    #[tokio::test]
678    async fn test_load_many_to_many_empty() {
679        use sea_orm::{DbBackend, IntoMockRow, LoaderTrait, MockDatabase, tests_cfg::*};
680
681        let db = MockDatabase::new(DbBackend::Postgres)
682            .append_query_results([
683                [cake_filling_model(1, 1).into_mock_row()],
684                [filling_model(1).into_mock_row()],
685            ])
686            .into_connection();
687
688        let cakes: Vec<cake::Model> = vec![];
689
690        let fillings = cakes
691            .load_many_to_many(Filling, CakeFilling, &db)
692            .await
693            .expect("Should return something");
694
695        let empty_vec: Vec<Vec<filling::Model>> = vec![];
696
697        assert_eq!(fillings, empty_vec);
698    }
699
700    #[tokio::test]
701    async fn test_load_one_duplicate_keys() {
702        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
703
704        let db = MockDatabase::new(DbBackend::Postgres)
705            .append_query_results([[cake_model(1), cake_model(2)]])
706            .into_connection();
707
708        let fruits = vec![
709            fruit_model(1, Some(1)),
710            fruit_model(2, Some(1)),
711            fruit_model(3, Some(1)),
712            fruit_model(4, Some(1)),
713        ];
714
715        let cakes = fruits
716            .load_one(cake::Entity::find(), &db)
717            .await
718            .expect("Should return something");
719
720        assert_eq!(cakes.len(), 4);
721        for cake in &cakes {
722            assert_eq!(cake, &Some(cake_model(1)));
723        }
724        let logs = db.into_transaction_log();
725        let sql = format!("{:?}", logs[0]);
726
727        let values_count = sql.matches("$1").count();
728        assert_eq!(values_count, 1, "Duplicate values were not removed");
729    }
730
731    #[tokio::test]
732    async fn test_load_many_duplicate_keys() {
733        use sea_orm::{DbBackend, LoaderTrait, MockDatabase, entity::prelude::*, tests_cfg::*};
734
735        let db = MockDatabase::new(DbBackend::Postgres)
736            .append_query_results([[
737                fruit_model(1, Some(1)),
738                fruit_model(2, Some(1)),
739                fruit_model(3, Some(2)),
740            ]])
741            .into_connection();
742
743        let cakes = vec![cake_model(1), cake_model(1), cake_model(2), cake_model(2)];
744
745        let fruits = cakes
746            .load_many(fruit::Entity::find(), &db)
747            .await
748            .expect("Should return something");
749
750        assert_eq!(fruits.len(), 4);
751
752        let logs = db.into_transaction_log();
753        let sql = format!("{:?}", logs[0]);
754
755        let values_count = sql.matches("$1").count() + sql.matches("$2").count();
756        assert_eq!(values_count, 2, "Duplicate values were not removed");
757    }
758}