1use std::collections::HashMap;
43use std::marker::PhantomData;
44
45use rainier_orm::sea_query::Value;
46use rainier_support::{str, Result};
47
48use crate::criteria::Criteria;
49use crate::model::Model;
50use crate::repository::Repository;
51
52#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
59pub struct RelationKey(String);
60
61impl RelationKey {
62 pub fn new(value: &Value) -> Self {
64 Self(normalise(value))
65 }
66
67 pub fn from_cell(cell: &crate::row::Cell) -> Self {
69 match cell.to_value() {
70 Some(value) => Self::new(&value),
71 None => Self(null()),
72 }
73 }
74
75 pub fn as_str(&self) -> &str {
77 &self.0
78 }
79}
80
81impl std::fmt::Display for RelationKey {
82 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83 f.write_str(&self.0)
84 }
85}
86
87fn normalise(value: &Value) -> String {
92 macro_rules! int {
93 ($v:expr) => {
94 $v.map(|n| (n as i128).to_string()).unwrap_or_else(null)
95 };
96 }
97
98 match value {
99 Value::TinyInt(v) => int!(*v),
100 Value::SmallInt(v) => int!(*v),
101 Value::Int(v) => int!(*v),
102 Value::BigInt(v) => int!(*v),
103 Value::TinyUnsigned(v) => int!(*v),
104 Value::SmallUnsigned(v) => int!(*v),
105 Value::Unsigned(v) => int!(*v),
106 Value::BigUnsigned(v) => int!(*v),
107 Value::Bool(v) => v.map(|b| b.to_string()).unwrap_or_else(null),
108 Value::String(v) => v.as_ref().map(|s| s.to_string()).unwrap_or_else(null),
109 Value::Char(v) => v.map(|c| c.to_string()).unwrap_or_else(null),
110 Value::Bytes(v) => v.as_ref().map(|b| hex(b)).unwrap_or_else(null),
111 other => format!("{other:?}"),
114 }
115}
116
117fn null() -> String {
118 "\0null".to_string()
119}
120
121fn hex(bytes: &[u8]) -> String {
122 bytes.iter().map(|byte| format!("{byte:02x}")).collect()
123}
124
125#[derive(Debug, Clone)]
130pub struct Related<C> {
131 near_key: String,
134 groups: HashMap<RelationKey, Vec<C>>,
135}
136
137impl<C> Related<C> {
138 pub fn empty(near_key: impl Into<String>) -> Self {
140 Self { near_key: near_key.into(), groups: HashMap::new() }
141 }
142
143 pub fn of<N: Model>(&self, near: &N) -> &[C] {
145 self.key_of(near).and_then(|key| self.groups.get(&key)).map_or(&[], Vec::as_slice)
146 }
147
148 pub fn one<N: Model>(&self, near: &N) -> Option<&C> {
151 self.of(near).first()
152 }
153
154 pub fn count_of<N: Model>(&self, near: &N) -> usize {
159 self.of(near).len()
160 }
161
162 pub fn is_empty(&self) -> bool {
164 self.groups.values().all(Vec::is_empty)
165 }
166
167 pub fn len(&self) -> usize {
169 self.groups.values().map(Vec::len).sum()
170 }
171
172 pub fn near_key(&self) -> &str {
174 &self.near_key
175 }
176
177 pub fn zip<N: Model>(mut self, near: Vec<N>) -> Vec<(N, Vec<C>)> {
182 near.into_iter()
183 .map(|model| {
184 let related = model
185 .value_of(&self.near_key)
186 .and_then(|value| self.groups.remove(&RelationKey::new(&value)))
187 .unwrap_or_default();
188 (model, related)
189 })
190 .collect()
191 }
192
193 fn key_of<N: Model>(&self, near: &N) -> Option<RelationKey> {
194 near.value_of(&self.near_key).map(|value| RelationKey::new(&value))
195 }
196
197 fn group_by(near_key: &str, far_key: &str, rows: Vec<C>) -> Self
199 where
200 C: Model,
201 {
202 let mut groups: HashMap<RelationKey, Vec<C>> = HashMap::new();
203 for row in rows {
204 let Some(value) = row.value_of(far_key) else { continue };
205 groups.entry(RelationKey::new(&value)).or_default().push(row);
206 }
207 Self { near_key: near_key.to_string(), groups }
208 }
209}
210
211#[derive(Debug, Clone, Default)]
213pub struct RelatedCounts {
214 near_key: String,
215 counts: HashMap<RelationKey, u64>,
216}
217
218impl RelatedCounts {
219 pub fn of<N: Model>(&self, near: &N) -> u64 {
222 near.value_of(&self.near_key)
223 .and_then(|value| self.counts.get(&RelationKey::new(&value)).copied())
224 .unwrap_or(0)
225 }
226
227 pub fn total(&self) -> u64 {
229 self.counts.values().sum()
230 }
231}
232
233#[async_trait::async_trait]
238pub trait Relation<P: Model>: Send + Sync {
239 type Related: Model;
241
242 fn near_key(&self) -> &str;
244
245 fn far_key(&self) -> &str;
247
248 async fn load(
250 &self,
251 parents: &[P],
252 related: &dyn Repository<Self::Related>,
253 ) -> Result<Related<Self::Related>>;
254
255 async fn count(
258 &self,
259 parents: &[P],
260 related: &dyn Repository<Self::Related>,
261 ) -> Result<RelatedCounts>;
262
263 async fn for_one(
269 &self,
270 parent: &P,
271 related: &dyn Repository<Self::Related>,
272 ) -> Result<Vec<Self::Related>>
273 where
274 Self: Sized,
275 {
276 let loaded = self.load(std::slice::from_ref(parent), related).await?;
277 Ok(loaded.of(parent).to_vec())
278 }
279}
280
281fn keys_of<P: Model>(parents: &[P], column: &str) -> Vec<Value> {
283 let mut seen = std::collections::HashSet::new();
284 let mut keys = Vec::with_capacity(parents.len());
285
286 for parent in parents {
287 let Some(value) = parent.value_of(column) else { continue };
288 if seen.insert(RelationKey::new(&value)) {
289 keys.push(value);
290 }
291 }
292 keys
293}
294
295fn conventional_foreign_key<M: Model>() -> String {
299 format!("{}_{}", str::snake(M::model_name()), M::primary_key())
300}
301
302pub struct HasMany<P, C> {
313 foreign_key: String,
314 local_key: String,
315 criteria: Criteria,
316 _sides: PhantomData<fn() -> (P, C)>,
317}
318
319impl<P: Model, C: Model> HasMany<P, C> {
320 pub fn new() -> Self {
322 Self {
323 foreign_key: conventional_foreign_key::<P>(),
324 local_key: P::primary_key().to_string(),
325 criteria: Criteria::new(),
326 _sides: PhantomData,
327 }
328 }
329
330 pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
332 self.foreign_key = column.into();
333 self
334 }
335
336 pub fn local_key(mut self, column: impl Into<String>) -> Self {
338 self.local_key = column.into();
339 self
340 }
341
342 pub fn matching(mut self, criteria: Criteria) -> Self {
348 self.criteria = criteria;
349 self
350 }
351}
352
353impl<P: Model, C: Model> Default for HasMany<P, C> {
354 fn default() -> Self {
355 Self::new()
356 }
357}
358
359#[async_trait::async_trait]
360impl<P: Model, C: Model> Relation<P> for HasMany<P, C> {
361 type Related = C;
362
363 fn near_key(&self) -> &str {
364 &self.local_key
365 }
366
367 fn far_key(&self) -> &str {
368 &self.foreign_key
369 }
370
371 async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
372 let keys = keys_of(parents, &self.local_key);
373 if keys.is_empty() {
374 return Ok(Related::empty(&self.local_key));
375 }
376
377 let rows = related.matching(self.scoped(keys)).await?;
378 Ok(Related::group_by(&self.local_key, &self.foreign_key, rows))
379 }
380
381 async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
382 let keys = keys_of(parents, &self.local_key);
383 if keys.is_empty() {
384 return Ok(RelatedCounts::default());
385 }
386
387 let counts = related.count_grouped(&self.foreign_key, self.scoped(keys)).await?;
388 Ok(RelatedCounts { near_key: self.local_key.clone(), counts: counts.into_iter().collect() })
389 }
390}
391
392impl<P: Model, C: Model> HasMany<P, C> {
393 fn scoped(&self, keys: Vec<Value>) -> Criteria {
394 self.criteria.clone().where_in(&self.foreign_key, keys)
395 }
396}
397
398pub struct HasOne<P, C>(HasMany<P, C>);
409
410impl<P: Model, C: Model> HasOne<P, C> {
411 pub fn new() -> Self {
413 Self(HasMany::new())
414 }
415
416 pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
418 self.0 = self.0.foreign_key(column);
419 self
420 }
421
422 pub fn local_key(mut self, column: impl Into<String>) -> Self {
424 self.0 = self.0.local_key(column);
425 self
426 }
427
428 pub fn matching(mut self, criteria: Criteria) -> Self {
430 self.0 = self.0.matching(criteria);
431 self
432 }
433}
434
435impl<P: Model, C: Model> Default for HasOne<P, C> {
436 fn default() -> Self {
437 Self::new()
438 }
439}
440
441#[async_trait::async_trait]
442impl<P: Model, C: Model> Relation<P> for HasOne<P, C> {
443 type Related = C;
444
445 fn near_key(&self) -> &str {
446 self.0.near_key()
447 }
448
449 fn far_key(&self) -> &str {
450 self.0.far_key()
451 }
452
453 async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
454 self.0.load(parents, related).await
455 }
456
457 async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
458 self.0.count(parents, related).await
459 }
460}
461
462pub struct BelongsTo<C, P> {
476 foreign_key: String,
477 owner_key: String,
478 criteria: Criteria,
479 _sides: PhantomData<fn() -> (C, P)>,
480}
481
482impl<C: Model, P: Model> BelongsTo<C, P> {
483 pub fn new() -> Self {
485 Self {
486 foreign_key: conventional_foreign_key::<P>(),
487 owner_key: P::primary_key().to_string(),
488 criteria: Criteria::new(),
489 _sides: PhantomData,
490 }
491 }
492
493 pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
495 self.foreign_key = column.into();
496 self
497 }
498
499 pub fn owner_key(mut self, column: impl Into<String>) -> Self {
501 self.owner_key = column.into();
502 self
503 }
504
505 pub fn matching(mut self, criteria: Criteria) -> Self {
507 self.criteria = criteria;
508 self
509 }
510}
511
512impl<C: Model, P: Model> Default for BelongsTo<C, P> {
513 fn default() -> Self {
514 Self::new()
515 }
516}
517
518#[async_trait::async_trait]
519impl<C: Model, P: Model> Relation<C> for BelongsTo<C, P> {
520 type Related = P;
521
522 fn near_key(&self) -> &str {
523 &self.foreign_key
524 }
525
526 fn far_key(&self) -> &str {
527 &self.owner_key
528 }
529
530 async fn load(&self, children: &[C], related: &dyn Repository<P>) -> Result<Related<P>> {
531 let keys = keys_of(children, &self.foreign_key);
532 if keys.is_empty() {
533 return Ok(Related::empty(&self.foreign_key));
534 }
535
536 let criteria = self.criteria.clone().where_in(&self.owner_key, keys);
537 let rows = related.matching(criteria).await?;
538
539 Ok(Related::group_by(&self.foreign_key, &self.owner_key, rows))
540 }
541
542 async fn count(&self, children: &[C], related: &dyn Repository<P>) -> Result<RelatedCounts> {
543 let keys = keys_of(children, &self.foreign_key);
544 if keys.is_empty() {
545 return Ok(RelatedCounts::default());
546 }
547
548 let criteria = self.criteria.clone().where_in(&self.owner_key, keys);
549 let counts = related.count_grouped(&self.owner_key, criteria).await?;
550
551 Ok(RelatedCounts {
552 near_key: self.foreign_key.clone(),
553 counts: counts.into_iter().collect(),
554 })
555 }
556}
557
558pub struct BelongsToMany<P, C> {
578 pivot: String,
579 parent_pivot_key: String,
580 related_pivot_key: String,
581 parent_key: String,
582 related_key: String,
583 criteria: Criteria,
584 _sides: PhantomData<fn() -> (P, C)>,
585}
586
587impl<P: Model, C: Model> BelongsToMany<P, C> {
588 pub fn new(pivot: impl Into<String>) -> Self {
591 Self {
592 pivot: pivot.into(),
593 parent_pivot_key: conventional_foreign_key::<P>(),
594 related_pivot_key: conventional_foreign_key::<C>(),
595 parent_key: P::primary_key().to_string(),
596 related_key: C::primary_key().to_string(),
597 criteria: Criteria::new(),
598 _sides: PhantomData,
599 }
600 }
601
602 pub fn conventional_pivot() -> String {
605 let mut sides =
606 [str::snake(P::model_name()), str::snake(C::model_name())].map(|s| str::singular(&s));
607 sides.sort();
608 sides.join("_")
609 }
610
611 pub fn parent_pivot_key(mut self, column: impl Into<String>) -> Self {
613 self.parent_pivot_key = column.into();
614 self
615 }
616
617 pub fn related_pivot_key(mut self, column: impl Into<String>) -> Self {
619 self.related_pivot_key = column.into();
620 self
621 }
622
623 pub fn parent_key(mut self, column: impl Into<String>) -> Self {
625 self.parent_key = column.into();
626 self
627 }
628
629 pub fn related_key(mut self, column: impl Into<String>) -> Self {
631 self.related_key = column.into();
632 self
633 }
634
635 pub fn matching(mut self, criteria: Criteria) -> Self {
637 self.criteria = criteria;
638 self
639 }
640
641 pub fn pivot(&self) -> &str {
643 &self.pivot
644 }
645}
646
647#[async_trait::async_trait]
648impl<P: Model, C: Model> Relation<P> for BelongsToMany<P, C> {
649 type Related = C;
650
651 fn near_key(&self) -> &str {
652 &self.parent_key
653 }
654
655 fn far_key(&self) -> &str {
656 &self.related_key
657 }
658
659 async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
660 let keys = keys_of(parents, &self.parent_key);
661 if keys.is_empty() {
662 return Ok(Related::empty(&self.parent_key));
663 }
664
665 let links = related.pivot_links(self.links_for(keys)).await?;
666 if links.is_empty() {
667 return Ok(Related::empty(&self.parent_key));
668 }
669
670 let wanted: Vec<Value> = dedup(links.iter().map(|(_, related)| related.clone()));
672 let criteria = self.criteria.clone().where_in(&self.related_key, wanted);
673 let rows = related.matching(criteria).await?;
674
675 let by_key: HashMap<RelationKey, C> = rows
678 .into_iter()
679 .filter_map(|row| {
680 row.value_of(&self.related_key).map(|key| (RelationKey::new(&key), row))
681 })
682 .collect();
683
684 let mut groups: HashMap<RelationKey, Vec<C>> = HashMap::new();
685 for (parent, related) in &links {
686 if let Some(row) = by_key.get(&RelationKey::new(related)) {
687 groups.entry(RelationKey::new(parent)).or_default().push(row.clone());
688 }
689 }
690
691 Ok(Related { near_key: self.parent_key.clone(), groups })
692 }
693
694 async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
695 let keys = keys_of(parents, &self.parent_key);
696 if keys.is_empty() {
697 return Ok(RelatedCounts::default());
698 }
699
700 let links = related.pivot_links(self.links_for(keys)).await?;
704
705 let mut counts: HashMap<RelationKey, u64> = HashMap::new();
706 for (parent, _) in &links {
707 *counts.entry(RelationKey::new(parent)).or_insert(0) += 1;
708 }
709
710 Ok(RelatedCounts { near_key: self.parent_key.clone(), counts })
711 }
712}
713
714impl<P: Model, C: Model> BelongsToMany<P, C> {
715 fn links_for(&self, keys: Vec<Value>) -> PivotQuery {
716 PivotQuery {
717 table: self.pivot.clone(),
718 parent_column: self.parent_pivot_key.clone(),
719 related_column: self.related_pivot_key.clone(),
720 parent_keys: keys,
721 }
722 }
723}
724
725#[derive(Debug, Clone)]
730pub struct PivotQuery {
731 pub table: String,
733 pub parent_column: String,
735 pub related_column: String,
737 pub parent_keys: Vec<Value>,
739}
740
741fn dedup(values: impl Iterator<Item = Value>) -> Vec<Value> {
742 let mut seen = std::collections::HashSet::new();
743 let mut unique = Vec::new();
744 for value in values {
745 if seen.insert(RelationKey::new(&value)) {
746 unique.push(value);
747 }
748 }
749 unique
750}
751
752#[cfg(test)]
753mod tests {
754 use super::*;
755 use crate::testing::{fake_database, MemoryConnection};
756 use crate::{Dialect, EntityRepository};
757 use rainier_orm::Entity;
758
759 #[derive(Debug, Clone, PartialEq, Entity)]
760 #[orm(table = "users")]
761 struct User {
762 #[orm(pk, auto_increment)]
763 id: u64,
764 name: String,
765 }
766
767 impl Model for User {}
768
769 #[derive(Debug, Clone, PartialEq, Entity)]
770 #[orm(table = "posts")]
771 struct Post {
772 #[orm(pk, auto_increment)]
773 id: u64,
774 title: String,
775 user_id: u64,
776 }
777
778 impl Model for Post {}
779
780 #[derive(Debug, Clone, PartialEq, Entity)]
781 #[orm(table = "tags")]
782 struct Tag {
783 #[orm(pk, auto_increment)]
784 id: u64,
785 name: String,
786 }
787
788 impl Model for Tag {}
789
790 fn user(id: u64) -> User {
791 User { id, name: format!("user-{id}") }
792 }
793
794 fn post(id: u64, user_id: u64) -> Post {
795 Post { id, title: format!("post-{id}"), user_id }
796 }
797
798 #[test]
799 fn the_conventional_keys_come_from_the_model_names() {
800 let has_many = HasMany::<User, Post>::new();
801 assert_eq!(has_many.far_key(), "user_id");
802 assert_eq!(has_many.near_key(), "id");
803
804 let belongs_to = BelongsTo::<Post, User>::new();
805 assert_eq!(belongs_to.near_key(), "user_id", "the column on the post");
806 assert_eq!(belongs_to.far_key(), "id", "the column on the user");
807
808 assert_eq!(BelongsToMany::<Post, Tag>::conventional_pivot(), "post_tag");
809 }
810
811 #[test]
812 fn keys_are_normalised_across_integer_widths() {
813 assert_eq!(
816 RelationKey::new(&Value::Int(Some(7))),
817 RelationKey::new(&Value::BigUnsigned(Some(7)))
818 );
819 assert_ne!(RelationKey::new(&Value::Int(Some(7))), RelationKey::new(&Value::Int(Some(8))));
820
821 let null_key = RelationKey::new(&Value::BigUnsigned(None));
823 assert_ne!(null_key, RelationKey::new(&Value::BigUnsigned(Some(0))));
824 }
825
826 #[test]
827 fn related_rows_are_grouped_by_the_foreign_key() {
828 let related = Related::group_by("id", "user_id", vec![post(1, 7), post(2, 7), post(3, 8)]);
829
830 assert_eq!(related.of(&user(7)).len(), 2);
831 assert_eq!(related.of(&user(8)).len(), 1);
832 assert_eq!(related.of(&user(9)), &[], "a parent with no children");
833 assert_eq!(related.one(&user(8)).unwrap().id, 3);
834 assert_eq!(related.len(), 3);
835 }
836
837 #[test]
838 fn zipping_pairs_each_parent_with_its_own() {
839 let related = Related::group_by("id", "user_id", vec![post(1, 7), post(2, 8), post(3, 7)]);
840 let pairs = related.zip(vec![user(7), user(8), user(9)]);
841
842 assert_eq!(pairs[0].1.len(), 2);
843 assert_eq!(pairs[1].1.len(), 1);
844 assert!(pairs[2].1.is_empty(), "still present, with nothing");
845 }
846
847 #[test]
848 fn duplicate_parent_keys_are_asked_for_once() {
849 let posts = vec![post(1, 7), post(2, 7), post(3, 8)];
852 assert_eq!(keys_of(&posts, "user_id").len(), 2);
853 }
854
855 #[tokio::test]
856 async fn loading_many_parents_is_one_query() {
857 let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
858 let posts = EntityRepository::<Post>::new(db);
859
860 let users: Vec<User> = (1..=50).map(user).collect();
861 HasMany::<User, Post>::new().load(&users, &posts).await.unwrap();
862
863 assert_eq!(connection.statement_count(), 1, "fifty parents, one query");
864 let sql = connection.last_statement().unwrap();
865 assert!(sql.contains("IN"), "{sql}");
866 }
867
868 #[tokio::test]
869 async fn loading_no_parents_asks_nothing() {
870 let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
871 let posts = EntityRepository::<Post>::new(db);
872
873 let loaded = HasMany::<User, Post>::new().load(&[], &posts).await.unwrap();
874
875 assert_eq!(connection.statement_count(), 0, "an empty IN () is not a query worth running");
876 assert!(loaded.is_empty());
877 }
878
879 #[tokio::test]
880 async fn a_constrained_load_keeps_its_filters() {
881 let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
882 let posts = EntityRepository::<Post>::new(db);
883
884 HasMany::<User, Post>::new()
885 .matching(Criteria::new().order_by_desc("id"))
886 .load(&[user(1)], &posts)
887 .await
888 .unwrap();
889
890 let sql = connection.last_statement().unwrap();
891 assert!(sql.contains("ORDER BY"), "{sql}");
892 }
893
894 #[tokio::test]
895 async fn belongs_to_many_reads_the_pivot_then_the_rows() {
896 let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
897 let tags = EntityRepository::<Tag>::new(db);
898
899 BelongsToMany::<Post, Tag>::new("post_tag")
900 .load(&[post(1, 7), post(2, 7)], &tags)
901 .await
902 .unwrap();
903
904 assert_eq!(connection.statement_count(), 1);
907 assert!(connection.last_statement().unwrap().contains("post_tag"));
908 }
909}