use std::marker::PhantomData;
use std::sync::Arc;
use rainier_events::Dispatcher;
use rainier_orm::sea_query::Value;
use rainier_support::{Error, Result};
use crate::connection::Database;
use crate::criteria::Criteria;
use crate::model::{Created, Creating, Deleted, Deleting, Model, Updated, Updating};
use crate::pagination::Paginated;
use crate::relation::{PivotQuery, RelationKey};
use crate::row::Cell;
use crate::statement;
#[async_trait::async_trait]
pub trait Repository<M: Model>: Send + Sync + 'static {
async fn all(&self) -> Result<Vec<M>>;
async fn find(&self, key: Value) -> Result<Option<M>>;
async fn find_by(&self, column: &str, value: Value) -> Result<Vec<M>>;
async fn first_by(&self, column: &str, value: Value) -> Result<Option<M>>;
async fn matching(&self, criteria: Criteria) -> Result<Vec<M>>;
async fn first_matching(&self, criteria: Criteria) -> Result<Option<M>>;
async fn count_matching(&self, criteria: Criteria) -> Result<u64>;
async fn paginate_matching(
&self,
criteria: Criteria,
page: u64,
per_page: u64,
) -> Result<Paginated<M>>;
async fn create(&self, model: M) -> Result<M>;
async fn update(&self, model: &M) -> Result<u64>;
async fn upsert(
&self,
model: &M,
conflict_columns: &[&str],
update_columns: &[&str],
) -> Result<u64>;
async fn delete(&self, key: Value) -> Result<u64>;
async fn delete_matching(&self, criteria: Criteria) -> Result<u64>;
async fn count_grouped(
&self,
column: &str,
criteria: Criteria,
) -> Result<Vec<(RelationKey, u64)>> {
let mut counts: std::collections::HashMap<RelationKey, u64> =
std::collections::HashMap::new();
for model in self.matching(criteria).await? {
let Some(value) = model.value_of(column) else { continue };
*counts.entry(RelationKey::new(&value)).or_insert(0) += 1;
}
Ok(counts.into_iter().collect())
}
async fn pivot_links(&self, query: PivotQuery) -> Result<Vec<(Value, Value)>> {
let _ = query;
Ok(Vec::new())
}
async fn find_key(&self, key: impl Into<Value> + Send) -> Result<Option<M>>
where
Self: Sized,
{
self.find(key.into()).await
}
async fn find_or_fail(&self, key: impl Into<Value> + Send) -> Result<M>
where
Self: Sized,
{
self.find(key.into()).await?.ok_or_else(|| {
Error::not_found(format!("No {} matches the given key.", M::model_name()))
})
}
async fn find_by_route_key(&self, value: impl Into<Value> + Send) -> Result<Option<M>>
where
Self: Sized,
{
self.first_by(M::route_key_name(), value.into()).await
}
async fn where_eq(&self, column: &str, value: impl Into<Value> + Send) -> Result<Vec<M>>
where
Self: Sized,
{
self.find_by(column, value.into()).await
}
async fn count(&self) -> Result<u64>
where
Self: Sized,
{
self.count_matching(Criteria::new()).await
}
async fn paginate(&self, page: u64, per_page: u64) -> Result<Paginated<M>>
where
Self: Sized,
{
self.paginate_matching(Criteria::new(), page, per_page).await
}
async fn exists(&self, criteria: Criteria) -> Result<bool>
where
Self: Sized,
{
Ok(self.count_matching(criteria).await? > 0)
}
}
pub struct EntityRepository<M> {
db: Database,
events: Option<Arc<Dispatcher>>,
_model: PhantomData<fn() -> M>,
}
impl<M> Clone for EntityRepository<M> {
fn clone(&self) -> Self {
Self { db: self.db.clone(), events: self.events.clone(), _model: PhantomData }
}
}
impl<M: Model> EntityRepository<M> {
pub fn new(db: Database) -> Self {
Self { db, events: None, _model: PhantomData }
}
pub fn with_events(mut self, events: Arc<Dispatcher>) -> Self {
self.events = Some(events);
self
}
pub fn database(&self) -> &Database {
&self.db
}
async fn fire<E>(&self, build: impl FnOnce() -> E) -> Result<()>
where
E: Send + Sync + 'static,
{
let Some(events) = &self.events else { return Ok(()) };
if !events.has_listeners::<E>() {
return Ok(());
}
events.dispatch(build()).await
}
}
#[async_trait::async_trait]
impl<M: Model> Repository<M> for EntityRepository<M> {
async fn all(&self) -> Result<Vec<M>> {
self.db.fetch_all(statement::select_all::<M>(self.db.dialect())).await
}
async fn find(&self, key: Value) -> Result<Option<M>> {
self.db.fetch_one(statement::select_by_pk::<M>(self.db.dialect(), key)).await
}
async fn find_by(&self, column: &str, value: Value) -> Result<Vec<M>> {
let prepared = statement::select_by_column::<M>(self.db.dialect(), column, value, None);
self.db.fetch_all(prepared).await
}
async fn first_by(&self, column: &str, value: Value) -> Result<Option<M>> {
let prepared = statement::select_by_column::<M>(self.db.dialect(), column, value, Some(1));
self.db.fetch_one(prepared).await
}
async fn matching(&self, criteria: Criteria) -> Result<Vec<M>> {
self.db.fetch_all(statement::select_matching::<M>(self.db.dialect(), &criteria)).await
}
async fn first_matching(&self, criteria: Criteria) -> Result<Option<M>> {
let criteria = criteria.limit(1);
self.db.fetch_one(statement::select_matching::<M>(self.db.dialect(), &criteria)).await
}
async fn count_matching(&self, criteria: Criteria) -> Result<u64> {
self.db.fetch_count(statement::count_matching::<M>(self.db.dialect(), &criteria)).await
}
async fn count_grouped(
&self,
column: &str,
criteria: Criteria,
) -> Result<Vec<(RelationKey, u64)>> {
let Some(spec) = M::columns().iter().find(|spec| spec.name == column) else {
return Err(Error::internal(format!(
"`{}` has no column `{column}` to group by",
M::table()
)));
};
let prepared = statement::count_grouped::<M>(self.db.dialect(), column, &criteria);
let columns = vec![
crate::row::ColumnRequest::new(column, spec.ty),
crate::row::ColumnRequest::new("cnt", rainier_orm::ColumnType::BigInt),
];
let rows = self.db.fetch(prepared, columns).await?;
Ok(rows
.iter()
.filter_map(|row| {
let key = row.cell(column).map(RelationKey::from_cell)?;
let count = row.cell("cnt").and_then(Cell::as_u64).unwrap_or(0);
Some((key, count))
})
.collect())
}
async fn pivot_links(&self, query: PivotQuery) -> Result<Vec<(Value, Value)>> {
if query.parent_keys.is_empty() {
return Ok(Vec::new());
}
let prepared = statement::select_pivot(self.db.dialect(), &query);
let key_type = M::columns()
.iter()
.find(|spec| spec.pk)
.map(|spec| spec.ty)
.unwrap_or(rainier_orm::ColumnType::BigInt);
let columns = vec![
crate::row::ColumnRequest::new(&query.parent_column, key_type),
crate::row::ColumnRequest::new(&query.related_column, key_type),
];
let rows = self.db.fetch(prepared, columns).await?;
Ok(rows
.iter()
.filter_map(|row| {
let parent = row.cell(&query.parent_column).and_then(Cell::to_value)?;
let related = row.cell(&query.related_column).and_then(Cell::to_value)?;
Some((parent, related))
})
.collect())
}
async fn paginate_matching(
&self,
criteria: Criteria,
page: u64,
per_page: u64,
) -> Result<Paginated<M>> {
let page = page.max(1);
let per_page = per_page.max(1);
let total = self.count_matching(criteria.clone()).await?;
if total == 0 {
return Ok(Paginated::empty(page, per_page));
}
let window = criteria.limit(per_page).offset((page - 1) * per_page);
let data = self.matching(window).await?;
Ok(Paginated::new(data, total, page, per_page))
}
async fn create(&self, model: M) -> Result<M> {
self.fire(|| Creating { model: model.clone() }).await?;
let assigned = statement::shard_key_for_insert::<M>(&model)
.and_then(|shard_key| self.db.connection().allocate_id(shard_key));
let prepared = statement::insert::<M>(self.db.dialect(), &model, assigned);
let outcome = self.db.execute(prepared).await?;
let id = assigned.map(|id| id as i64).unwrap_or(outcome.last_insert_id);
let stored = if assigned.is_none() && id > 0 && has_auto_increment_key::<M>() {
self.find(id.into()).await?.unwrap_or(model)
} else {
model
};
self.fire(|| Created { model: stored.clone(), id }).await?;
Ok(stored)
}
async fn update(&self, model: &M) -> Result<u64> {
self.fire(|| Updating { model: model.clone() }).await?;
let prepared = statement::update::<M>(self.db.dialect(), model);
let rows_affected = self.db.execute(prepared).await?.rows_affected;
self.fire(|| Updated { model: model.clone(), rows_affected }).await?;
Ok(rows_affected)
}
async fn upsert(
&self,
model: &M,
conflict_columns: &[&str],
update_columns: &[&str],
) -> Result<u64> {
let prepared =
statement::upsert::<M>(self.db.dialect(), model, conflict_columns, update_columns);
Ok(self.db.execute(prepared).await?.rows_affected)
}
async fn delete(&self, key: Value) -> Result<u64> {
let described = describe(&key);
self.fire(|| Deleting::<M> { key: described.clone(), model: PhantomData }).await?;
let prepared = statement::delete_by_pk::<M>(self.db.dialect(), key);
let rows_affected = self.db.execute(prepared).await?.rows_affected;
self.fire(|| Deleted::<M> { key: described, rows_affected, model: PhantomData }).await?;
Ok(rows_affected)
}
async fn delete_matching(&self, criteria: Criteria) -> Result<u64> {
let prepared = statement::delete_matching::<M>(self.db.dialect(), &criteria);
Ok(self.db.execute(prepared).await?.rows_affected)
}
}
fn has_auto_increment_key<M: Model>() -> bool {
M::columns().iter().any(|column| column.name == M::primary_key() && column.auto_increment)
}
fn describe(key: &Value) -> String {
match key {
Value::String(Some(text)) => text.to_string(),
Value::BigInt(Some(n)) => n.to_string(),
Value::Int(Some(n)) => n.to_string(),
Value::BigUnsigned(Some(n)) => n.to_string(),
Value::Unsigned(Some(n)) => n.to_string(),
other => format!("{other:?}"),
}
}
impl<M> std::fmt::Debug for EntityRepository<M> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EntityRepository")
.field("model", &std::any::type_name::<M>())
.field("hooks", &self.events.is_some())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::row::OwnedRow;
use crate::testing::{fake_database, MemoryConnection};
use rainier_orm::{Dialect, ShardRoute};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
#[derive(rainier_orm::Entity, Clone, Debug, PartialEq)]
#[orm(table = "posts")]
struct Post {
#[orm(pk, auto_increment)]
id: u64,
title: String,
published: bool,
}
impl Model for Post {}
#[derive(rainier_orm::Entity, Clone, Debug, PartialEq)]
#[orm(table = "tokens")]
struct Token {
#[orm(pk, shard_key)]
id: u64,
#[orm(shard_key)]
user_id: u64,
hash: String,
}
impl Model for Token {}
fn row(id: u64, title: &str) -> OwnedRow {
OwnedRow::new().with("id", id).with("title", title).with("published", true)
}
fn count(n: i64) -> OwnedRow {
OwnedRow::new().with("cnt", n)
}
fn repository(connection: MemoryConnection) -> (EntityRepository<Post>, Arc<MemoryConnection>) {
let (db, handle) = fake_database(connection);
(EntityRepository::<Post>::new(db), handle)
}
#[test]
fn repository_calls_produce_send_futures() {
fn assert_send<T: Send>(_: T) {}
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite));
assert_send(async move {
let _ = posts.all().await;
});
}
#[tokio::test]
async fn all_decodes_every_row() {
let (posts, _) = repository(
MemoryConnection::new(Dialect::Sqlite).returning([row(1, "a"), row(2, "b")]),
);
let found = posts.all().await.unwrap();
assert_eq!(found.len(), 2);
assert_eq!(found[0], Post { id: 1, title: "a".into(), published: true });
}
#[tokio::test]
async fn find_returns_none_when_nothing_matches() {
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite));
assert!(posts.find(1_i64.into()).await.unwrap().is_none());
}
#[tokio::test]
async fn find_or_fail_reports_a_404_naming_the_model() {
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite));
let err = posts.find_or_fail(1_i64).await.unwrap_err();
assert_eq!(err.status(), 404);
assert!(err.message().contains("Post"), "{}", err.message());
}
#[tokio::test]
async fn find_or_fail_returns_the_row_when_it_exists() {
let (posts, _) =
repository(MemoryConnection::new(Dialect::Sqlite).returning([row(1, "found")]));
assert_eq!(posts.find_or_fail(1_i64).await.unwrap().title, "found");
}
#[tokio::test]
async fn first_by_limits_to_one_row() {
let (posts, connection) =
repository(MemoryConnection::new(Dialect::Sqlite).returning([row(1, "a")]));
posts.first_by("title", "a".into()).await.unwrap();
assert!(connection.last_statement().unwrap().contains("LIMIT"));
}
#[tokio::test]
async fn criteria_reach_the_generated_sql() {
let (posts, connection) = repository(MemoryConnection::new(Dialect::Sqlite));
posts
.matching(Criteria::new().where_eq("published", true).order_by_desc("id").limit(5))
.await
.unwrap();
let sql = connection.last_statement().unwrap();
assert!(sql.contains("published"), "{sql}");
assert!(sql.contains("ORDER BY"), "{sql}");
assert!(sql.contains("LIMIT"), "{sql}");
}
#[tokio::test]
async fn counting_ignores_the_criterias_paging() {
let (posts, connection) =
repository(MemoryConnection::new(Dialect::Sqlite).returning([count(42)]));
let total = posts
.count_matching(Criteria::new().where_eq("published", true).limit(5).offset(10))
.await
.unwrap();
assert_eq!(total, 42);
let sql = connection.last_statement().unwrap();
assert!(!sql.contains("LIMIT"), "the count must not be limited: {sql}");
assert!(sql.contains("published"), "but it must keep the filters: {sql}");
}
#[tokio::test]
async fn paginate_reports_the_window() {
let connection = MemoryConnection::new(Dialect::Sqlite)
.returning([count(25)])
.returning([row(11, "k"), row(12, "l")]);
let (posts, _) = repository(connection);
let page = posts.paginate(2, 10).await.unwrap();
assert_eq!(page.total, 25);
assert_eq!(page.current_page, 2);
assert_eq!(page.len(), 2);
assert_eq!(page.last_page(), 3);
assert_eq!(page.from(), Some(11));
}
#[tokio::test]
async fn paginating_no_results_skips_the_row_query() {
let (posts, connection) =
repository(MemoryConnection::new(Dialect::Sqlite).returning([count(0)]));
let page = posts.paginate(1, 10).await.unwrap();
assert!(page.is_empty());
assert_eq!(connection.statement_count(), 1);
}
#[tokio::test]
async fn exists_is_a_count() {
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite).returning([count(3)]));
assert!(posts.exists(Criteria::new()).await.unwrap());
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite).returning([count(0)]));
assert!(!posts.exists(Criteria::new()).await.unwrap());
}
#[tokio::test]
async fn create_reads_the_row_back_when_the_database_assigned_a_key() {
let connection =
MemoryConnection::new(Dialect::Sqlite).with_outcome(1, 7).returning([row(7, "stored")]);
let (posts, _) = repository(connection);
let created =
posts.create(Post { id: 0, title: "sent".into(), published: false }).await.unwrap();
assert_eq!(created.id, 7);
assert_eq!(created.title, "stored", "the stored row wins over the struct we sent");
}
#[tokio::test]
async fn create_does_not_re_read_a_model_with_an_application_assigned_key() {
#[derive(rainier_orm::Entity, Clone, Debug, PartialEq)]
#[orm(table = "sessions")]
struct Session {
#[orm(pk)]
id: String,
user_id: u64,
}
impl Model for Session {}
let (db, connection) =
fake_database(MemoryConnection::new(Dialect::Sqlite).with_outcome(1, 99));
let sessions = EntityRepository::<Session>::new(db);
let created = sessions.create(Session { id: "abc".into(), user_id: 1 }).await.unwrap();
assert_eq!(created.id, "abc");
assert_eq!(connection.statement_count(), 1, "the insert only — no re-read");
assert!(connection.last_statement().unwrap().starts_with("INSERT INTO"));
}
#[tokio::test]
async fn a_sharded_insert_asks_the_connector_for_an_id_and_routes_to_it() {
let (db, connection) =
fake_database(MemoryConnection::new(Dialect::Sqlite).sharded("users").allocating(4242));
let tokens = EntityRepository::<Token>::new(db);
let created = tokens.create(Token { id: 0, user_id: 42, hash: "h".into() }).await.unwrap();
assert_eq!(created.id, 0);
assert_eq!(connection.last_route(), Some(ShardRoute::Key(4242)));
assert_eq!(connection.statement_count(), 1);
}
#[tokio::test]
async fn a_database_failure_surfaces() {
let (posts, _) = repository(MemoryConnection::new(Dialect::Sqlite).failing("disk full"));
let err = posts.all().await.unwrap_err();
assert!(err.message().contains("disk full"), "{}", err.message());
}
#[allow(clippy::type_complexity)]
fn with_hooks() -> (EntityRepository<Post>, Arc<Dispatcher>, Arc<MemoryConnection>) {
let (db, connection) = fake_database(
MemoryConnection::new(Dialect::Sqlite).with_outcome(1, 7).returning([row(7, "stored")]),
);
let events = Arc::new(Dispatcher::new());
(EntityRepository::<Post>::new(db).with_events(Arc::clone(&events)), events, connection)
}
#[tokio::test]
async fn create_fires_creating_then_created() {
let (posts, events, _) = with_hooks();
let log = Arc::new(Mutex::new(Vec::new()));
let sink = Arc::clone(&log);
events.listen(move |event: Arc<Creating<Post>>| {
let sink = Arc::clone(&sink);
async move {
sink.lock().unwrap().push(format!("creating:{}", event.model.title));
Ok(())
}
});
let sink = Arc::clone(&log);
events.listen(move |event: Arc<Created<Post>>| {
let sink = Arc::clone(&sink);
async move {
sink.lock().unwrap().push(format!("created:{}", event.id));
Ok(())
}
});
posts.create(Post { id: 0, title: "sent".into(), published: false }).await.unwrap();
assert_eq!(*log.lock().unwrap(), vec!["creating:sent", "created:7"]);
}
#[tokio::test]
async fn a_creating_listener_can_veto_the_insert() {
let (posts, events, connection) = with_hooks();
events.listen(|_: Arc<Creating<Post>>| async { Err(Error::unauthorized("read only")) });
let err = posts
.create(Post { id: 0, title: "blocked".into(), published: false })
.await
.unwrap_err();
assert_eq!(err.status(), 403);
assert_eq!(
connection.statement_count(),
0,
"a vetoed insert must never reach the database"
);
}
#[tokio::test]
async fn update_and_delete_fire_their_hooks() {
let (posts, events, _) = with_hooks();
let seen = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&seen);
events.listen(move |_: Arc<Updating<Post>>| {
let counter = Arc::clone(&counter);
async move {
counter.fetch_add(1, Ordering::SeqCst);
Ok(())
}
});
let counter = Arc::clone(&seen);
events.listen(move |_: Arc<Updated<Post>>| {
let counter = Arc::clone(&counter);
async move {
counter.fetch_add(10, Ordering::SeqCst);
Ok(())
}
});
let counter = Arc::clone(&seen);
events.listen(move |_: Arc<Deleting<Post>>| {
let counter = Arc::clone(&counter);
async move {
counter.fetch_add(100, Ordering::SeqCst);
Ok(())
}
});
let counter = Arc::clone(&seen);
events.listen(move |event: Arc<Deleted<Post>>| {
let counter = Arc::clone(&counter);
async move {
assert_eq!(event.key, "1");
counter.fetch_add(1000, Ordering::SeqCst);
Ok(())
}
});
posts.update(&Post { id: 1, title: "x".into(), published: true }).await.unwrap();
posts.delete(1_i64.into()).await.unwrap();
assert_eq!(seen.load(Ordering::SeqCst), 1111);
}
#[tokio::test]
async fn hooks_are_skipped_entirely_when_nothing_listens() {
let (posts, _, connection) = with_hooks();
posts.create(Post { id: 0, title: "x".into(), published: false }).await.unwrap();
assert!(connection.statement_count() >= 1);
}
#[tokio::test]
async fn the_contract_is_object_safe() {
let (db, _) = fake_database(MemoryConnection::new(Dialect::Sqlite));
let boxed: Arc<dyn Repository<Post>> = Arc::new(EntityRepository::<Post>::new(db));
assert!(boxed.all().await.unwrap().is_empty());
}
}