use std::collections::HashMap;
use std::marker::PhantomData;
use rainier_orm::sea_query::Value;
use rainier_support::{str, Result};
use crate::criteria::Criteria;
use crate::model::Model;
use crate::repository::Repository;
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct RelationKey(String);
impl RelationKey {
pub fn new(value: &Value) -> Self {
Self(normalise(value))
}
pub fn from_cell(cell: &crate::row::Cell) -> Self {
match cell.to_value() {
Some(value) => Self::new(&value),
None => Self(null()),
}
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for RelationKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
fn normalise(value: &Value) -> String {
macro_rules! int {
($v:expr) => {
$v.map(|n| (n as i128).to_string()).unwrap_or_else(null)
};
}
match value {
Value::TinyInt(v) => int!(*v),
Value::SmallInt(v) => int!(*v),
Value::Int(v) => int!(*v),
Value::BigInt(v) => int!(*v),
Value::TinyUnsigned(v) => int!(*v),
Value::SmallUnsigned(v) => int!(*v),
Value::Unsigned(v) => int!(*v),
Value::BigUnsigned(v) => int!(*v),
Value::Bool(v) => v.map(|b| b.to_string()).unwrap_or_else(null),
Value::String(v) => v.as_ref().map(|s| s.to_string()).unwrap_or_else(null),
Value::Char(v) => v.map(|c| c.to_string()).unwrap_or_else(null),
Value::Bytes(v) => v.as_ref().map(|b| hex(b)).unwrap_or_else(null),
other => format!("{other:?}"),
}
}
fn null() -> String {
"\0null".to_string()
}
fn hex(bytes: &[u8]) -> String {
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}
#[derive(Debug, Clone)]
pub struct Related<C> {
near_key: String,
groups: HashMap<RelationKey, Vec<C>>,
}
impl<C> Related<C> {
pub fn empty(near_key: impl Into<String>) -> Self {
Self { near_key: near_key.into(), groups: HashMap::new() }
}
pub fn of<N: Model>(&self, near: &N) -> &[C] {
self.key_of(near).and_then(|key| self.groups.get(&key)).map_or(&[], Vec::as_slice)
}
pub fn one<N: Model>(&self, near: &N) -> Option<&C> {
self.of(near).first()
}
pub fn count_of<N: Model>(&self, near: &N) -> usize {
self.of(near).len()
}
pub fn is_empty(&self) -> bool {
self.groups.values().all(Vec::is_empty)
}
pub fn len(&self) -> usize {
self.groups.values().map(Vec::len).sum()
}
pub fn near_key(&self) -> &str {
&self.near_key
}
pub fn zip<N: Model>(mut self, near: Vec<N>) -> Vec<(N, Vec<C>)> {
near.into_iter()
.map(|model| {
let related = model
.value_of(&self.near_key)
.and_then(|value| self.groups.remove(&RelationKey::new(&value)))
.unwrap_or_default();
(model, related)
})
.collect()
}
fn key_of<N: Model>(&self, near: &N) -> Option<RelationKey> {
near.value_of(&self.near_key).map(|value| RelationKey::new(&value))
}
fn group_by(near_key: &str, far_key: &str, rows: Vec<C>) -> Self
where
C: Model,
{
let mut groups: HashMap<RelationKey, Vec<C>> = HashMap::new();
for row in rows {
let Some(value) = row.value_of(far_key) else { continue };
groups.entry(RelationKey::new(&value)).or_default().push(row);
}
Self { near_key: near_key.to_string(), groups }
}
}
#[derive(Debug, Clone, Default)]
pub struct RelatedCounts {
near_key: String,
counts: HashMap<RelationKey, u64>,
}
impl RelatedCounts {
pub fn of<N: Model>(&self, near: &N) -> u64 {
near.value_of(&self.near_key)
.and_then(|value| self.counts.get(&RelationKey::new(&value)).copied())
.unwrap_or(0)
}
pub fn total(&self) -> u64 {
self.counts.values().sum()
}
}
#[async_trait::async_trait]
pub trait Relation<P: Model>: Send + Sync {
type Related: Model;
fn near_key(&self) -> &str;
fn far_key(&self) -> &str;
async fn load(
&self,
parents: &[P],
related: &dyn Repository<Self::Related>,
) -> Result<Related<Self::Related>>;
async fn count(
&self,
parents: &[P],
related: &dyn Repository<Self::Related>,
) -> Result<RelatedCounts>;
async fn for_one(
&self,
parent: &P,
related: &dyn Repository<Self::Related>,
) -> Result<Vec<Self::Related>>
where
Self: Sized,
{
let loaded = self.load(std::slice::from_ref(parent), related).await?;
Ok(loaded.of(parent).to_vec())
}
}
fn keys_of<P: Model>(parents: &[P], column: &str) -> Vec<Value> {
let mut seen = std::collections::HashSet::new();
let mut keys = Vec::with_capacity(parents.len());
for parent in parents {
let Some(value) = parent.value_of(column) else { continue };
if seen.insert(RelationKey::new(&value)) {
keys.push(value);
}
}
keys
}
fn conventional_foreign_key<M: Model>() -> String {
format!("{}_{}", str::snake(M::model_name()), M::primary_key())
}
pub struct HasMany<P, C> {
foreign_key: String,
local_key: String,
criteria: Criteria,
_sides: PhantomData<fn() -> (P, C)>,
}
impl<P: Model, C: Model> HasMany<P, C> {
pub fn new() -> Self {
Self {
foreign_key: conventional_foreign_key::<P>(),
local_key: P::primary_key().to_string(),
criteria: Criteria::new(),
_sides: PhantomData,
}
}
pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
self.foreign_key = column.into();
self
}
pub fn local_key(mut self, column: impl Into<String>) -> Self {
self.local_key = column.into();
self
}
pub fn matching(mut self, criteria: Criteria) -> Self {
self.criteria = criteria;
self
}
}
impl<P: Model, C: Model> Default for HasMany<P, C> {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl<P: Model, C: Model> Relation<P> for HasMany<P, C> {
type Related = C;
fn near_key(&self) -> &str {
&self.local_key
}
fn far_key(&self) -> &str {
&self.foreign_key
}
async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
let keys = keys_of(parents, &self.local_key);
if keys.is_empty() {
return Ok(Related::empty(&self.local_key));
}
let rows = related.matching(self.scoped(keys)).await?;
Ok(Related::group_by(&self.local_key, &self.foreign_key, rows))
}
async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
let keys = keys_of(parents, &self.local_key);
if keys.is_empty() {
return Ok(RelatedCounts::default());
}
let counts = related.count_grouped(&self.foreign_key, self.scoped(keys)).await?;
Ok(RelatedCounts { near_key: self.local_key.clone(), counts: counts.into_iter().collect() })
}
}
impl<P: Model, C: Model> HasMany<P, C> {
fn scoped(&self, keys: Vec<Value>) -> Criteria {
self.criteria.clone().where_in(&self.foreign_key, keys)
}
}
pub struct HasOne<P, C>(HasMany<P, C>);
impl<P: Model, C: Model> HasOne<P, C> {
pub fn new() -> Self {
Self(HasMany::new())
}
pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
self.0 = self.0.foreign_key(column);
self
}
pub fn local_key(mut self, column: impl Into<String>) -> Self {
self.0 = self.0.local_key(column);
self
}
pub fn matching(mut self, criteria: Criteria) -> Self {
self.0 = self.0.matching(criteria);
self
}
}
impl<P: Model, C: Model> Default for HasOne<P, C> {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl<P: Model, C: Model> Relation<P> for HasOne<P, C> {
type Related = C;
fn near_key(&self) -> &str {
self.0.near_key()
}
fn far_key(&self) -> &str {
self.0.far_key()
}
async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
self.0.load(parents, related).await
}
async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
self.0.count(parents, related).await
}
}
pub struct BelongsTo<C, P> {
foreign_key: String,
owner_key: String,
criteria: Criteria,
_sides: PhantomData<fn() -> (C, P)>,
}
impl<C: Model, P: Model> BelongsTo<C, P> {
pub fn new() -> Self {
Self {
foreign_key: conventional_foreign_key::<P>(),
owner_key: P::primary_key().to_string(),
criteria: Criteria::new(),
_sides: PhantomData,
}
}
pub fn foreign_key(mut self, column: impl Into<String>) -> Self {
self.foreign_key = column.into();
self
}
pub fn owner_key(mut self, column: impl Into<String>) -> Self {
self.owner_key = column.into();
self
}
pub fn matching(mut self, criteria: Criteria) -> Self {
self.criteria = criteria;
self
}
}
impl<C: Model, P: Model> Default for BelongsTo<C, P> {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl<C: Model, P: Model> Relation<C> for BelongsTo<C, P> {
type Related = P;
fn near_key(&self) -> &str {
&self.foreign_key
}
fn far_key(&self) -> &str {
&self.owner_key
}
async fn load(&self, children: &[C], related: &dyn Repository<P>) -> Result<Related<P>> {
let keys = keys_of(children, &self.foreign_key);
if keys.is_empty() {
return Ok(Related::empty(&self.foreign_key));
}
let criteria = self.criteria.clone().where_in(&self.owner_key, keys);
let rows = related.matching(criteria).await?;
Ok(Related::group_by(&self.foreign_key, &self.owner_key, rows))
}
async fn count(&self, children: &[C], related: &dyn Repository<P>) -> Result<RelatedCounts> {
let keys = keys_of(children, &self.foreign_key);
if keys.is_empty() {
return Ok(RelatedCounts::default());
}
let criteria = self.criteria.clone().where_in(&self.owner_key, keys);
let counts = related.count_grouped(&self.owner_key, criteria).await?;
Ok(RelatedCounts {
near_key: self.foreign_key.clone(),
counts: counts.into_iter().collect(),
})
}
}
pub struct BelongsToMany<P, C> {
pivot: String,
parent_pivot_key: String,
related_pivot_key: String,
parent_key: String,
related_key: String,
criteria: Criteria,
_sides: PhantomData<fn() -> (P, C)>,
}
impl<P: Model, C: Model> BelongsToMany<P, C> {
pub fn new(pivot: impl Into<String>) -> Self {
Self {
pivot: pivot.into(),
parent_pivot_key: conventional_foreign_key::<P>(),
related_pivot_key: conventional_foreign_key::<C>(),
parent_key: P::primary_key().to_string(),
related_key: C::primary_key().to_string(),
criteria: Criteria::new(),
_sides: PhantomData,
}
}
pub fn conventional_pivot() -> String {
let mut sides =
[str::snake(P::model_name()), str::snake(C::model_name())].map(|s| str::singular(&s));
sides.sort();
sides.join("_")
}
pub fn parent_pivot_key(mut self, column: impl Into<String>) -> Self {
self.parent_pivot_key = column.into();
self
}
pub fn related_pivot_key(mut self, column: impl Into<String>) -> Self {
self.related_pivot_key = column.into();
self
}
pub fn parent_key(mut self, column: impl Into<String>) -> Self {
self.parent_key = column.into();
self
}
pub fn related_key(mut self, column: impl Into<String>) -> Self {
self.related_key = column.into();
self
}
pub fn matching(mut self, criteria: Criteria) -> Self {
self.criteria = criteria;
self
}
pub fn pivot(&self) -> &str {
&self.pivot
}
}
#[async_trait::async_trait]
impl<P: Model, C: Model> Relation<P> for BelongsToMany<P, C> {
type Related = C;
fn near_key(&self) -> &str {
&self.parent_key
}
fn far_key(&self) -> &str {
&self.related_key
}
async fn load(&self, parents: &[P], related: &dyn Repository<C>) -> Result<Related<C>> {
let keys = keys_of(parents, &self.parent_key);
if keys.is_empty() {
return Ok(Related::empty(&self.parent_key));
}
let links = related.pivot_links(self.links_for(keys)).await?;
if links.is_empty() {
return Ok(Related::empty(&self.parent_key));
}
let wanted: Vec<Value> = dedup(links.iter().map(|(_, related)| related.clone()));
let criteria = self.criteria.clone().where_in(&self.related_key, wanted);
let rows = related.matching(criteria).await?;
let by_key: HashMap<RelationKey, C> = rows
.into_iter()
.filter_map(|row| {
row.value_of(&self.related_key).map(|key| (RelationKey::new(&key), row))
})
.collect();
let mut groups: HashMap<RelationKey, Vec<C>> = HashMap::new();
for (parent, related) in &links {
if let Some(row) = by_key.get(&RelationKey::new(related)) {
groups.entry(RelationKey::new(parent)).or_default().push(row.clone());
}
}
Ok(Related { near_key: self.parent_key.clone(), groups })
}
async fn count(&self, parents: &[P], related: &dyn Repository<C>) -> Result<RelatedCounts> {
let keys = keys_of(parents, &self.parent_key);
if keys.is_empty() {
return Ok(RelatedCounts::default());
}
let links = related.pivot_links(self.links_for(keys)).await?;
let mut counts: HashMap<RelationKey, u64> = HashMap::new();
for (parent, _) in &links {
*counts.entry(RelationKey::new(parent)).or_insert(0) += 1;
}
Ok(RelatedCounts { near_key: self.parent_key.clone(), counts })
}
}
impl<P: Model, C: Model> BelongsToMany<P, C> {
fn links_for(&self, keys: Vec<Value>) -> PivotQuery {
PivotQuery {
table: self.pivot.clone(),
parent_column: self.parent_pivot_key.clone(),
related_column: self.related_pivot_key.clone(),
parent_keys: keys,
}
}
}
#[derive(Debug, Clone)]
pub struct PivotQuery {
pub table: String,
pub parent_column: String,
pub related_column: String,
pub parent_keys: Vec<Value>,
}
fn dedup(values: impl Iterator<Item = Value>) -> Vec<Value> {
let mut seen = std::collections::HashSet::new();
let mut unique = Vec::new();
for value in values {
if seen.insert(RelationKey::new(&value)) {
unique.push(value);
}
}
unique
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing::{fake_database, MemoryConnection};
use crate::{Dialect, EntityRepository};
use rainier_orm::Entity;
#[derive(Debug, Clone, PartialEq, Entity)]
#[orm(table = "users")]
struct User {
#[orm(pk, auto_increment)]
id: u64,
name: String,
}
impl Model for User {}
#[derive(Debug, Clone, PartialEq, Entity)]
#[orm(table = "posts")]
struct Post {
#[orm(pk, auto_increment)]
id: u64,
title: String,
user_id: u64,
}
impl Model for Post {}
#[derive(Debug, Clone, PartialEq, Entity)]
#[orm(table = "tags")]
struct Tag {
#[orm(pk, auto_increment)]
id: u64,
name: String,
}
impl Model for Tag {}
fn user(id: u64) -> User {
User { id, name: format!("user-{id}") }
}
fn post(id: u64, user_id: u64) -> Post {
Post { id, title: format!("post-{id}"), user_id }
}
#[test]
fn the_conventional_keys_come_from_the_model_names() {
let has_many = HasMany::<User, Post>::new();
assert_eq!(has_many.far_key(), "user_id");
assert_eq!(has_many.near_key(), "id");
let belongs_to = BelongsTo::<Post, User>::new();
assert_eq!(belongs_to.near_key(), "user_id", "the column on the post");
assert_eq!(belongs_to.far_key(), "id", "the column on the user");
assert_eq!(BelongsToMany::<Post, Tag>::conventional_pivot(), "post_tag");
}
#[test]
fn keys_are_normalised_across_integer_widths() {
assert_eq!(
RelationKey::new(&Value::Int(Some(7))),
RelationKey::new(&Value::BigUnsigned(Some(7)))
);
assert_ne!(RelationKey::new(&Value::Int(Some(7))), RelationKey::new(&Value::Int(Some(8))));
let null_key = RelationKey::new(&Value::BigUnsigned(None));
assert_ne!(null_key, RelationKey::new(&Value::BigUnsigned(Some(0))));
}
#[test]
fn related_rows_are_grouped_by_the_foreign_key() {
let related = Related::group_by("id", "user_id", vec![post(1, 7), post(2, 7), post(3, 8)]);
assert_eq!(related.of(&user(7)).len(), 2);
assert_eq!(related.of(&user(8)).len(), 1);
assert_eq!(related.of(&user(9)), &[], "a parent with no children");
assert_eq!(related.one(&user(8)).unwrap().id, 3);
assert_eq!(related.len(), 3);
}
#[test]
fn zipping_pairs_each_parent_with_its_own() {
let related = Related::group_by("id", "user_id", vec![post(1, 7), post(2, 8), post(3, 7)]);
let pairs = related.zip(vec![user(7), user(8), user(9)]);
assert_eq!(pairs[0].1.len(), 2);
assert_eq!(pairs[1].1.len(), 1);
assert!(pairs[2].1.is_empty(), "still present, with nothing");
}
#[test]
fn duplicate_parent_keys_are_asked_for_once() {
let posts = vec![post(1, 7), post(2, 7), post(3, 8)];
assert_eq!(keys_of(&posts, "user_id").len(), 2);
}
#[tokio::test]
async fn loading_many_parents_is_one_query() {
let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
let posts = EntityRepository::<Post>::new(db);
let users: Vec<User> = (1..=50).map(user).collect();
HasMany::<User, Post>::new().load(&users, &posts).await.unwrap();
assert_eq!(connection.statement_count(), 1, "fifty parents, one query");
let sql = connection.last_statement().unwrap();
assert!(sql.contains("IN"), "{sql}");
}
#[tokio::test]
async fn loading_no_parents_asks_nothing() {
let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
let posts = EntityRepository::<Post>::new(db);
let loaded = HasMany::<User, Post>::new().load(&[], &posts).await.unwrap();
assert_eq!(connection.statement_count(), 0, "an empty IN () is not a query worth running");
assert!(loaded.is_empty());
}
#[tokio::test]
async fn a_constrained_load_keeps_its_filters() {
let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
let posts = EntityRepository::<Post>::new(db);
HasMany::<User, Post>::new()
.matching(Criteria::new().order_by_desc("id"))
.load(&[user(1)], &posts)
.await
.unwrap();
let sql = connection.last_statement().unwrap();
assert!(sql.contains("ORDER BY"), "{sql}");
}
#[tokio::test]
async fn belongs_to_many_reads_the_pivot_then_the_rows() {
let (db, connection) = fake_database(MemoryConnection::new(Dialect::Sqlite));
let tags = EntityRepository::<Tag>::new(db);
BelongsToMany::<Post, Tag>::new("post_tag")
.load(&[post(1, 7), post(2, 7)], &tags)
.await
.unwrap();
assert_eq!(connection.statement_count(), 1);
assert!(connection.last_statement().unwrap().contains("post_tag"));
}
}