use sea_orm::{
ColumnTrait, EntityTrait, Order, PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, Select,
};
use std::future::Future;
use crate::database::DB;
use crate::error::FrameworkError;
pub struct QueryBuilder<E>
where
E: EntityTrait,
{
select: Select<E>,
}
impl<E> QueryBuilder<E>
where
E: EntityTrait,
E::Model: Send + Sync,
{
pub fn new() -> Self {
Self { select: E::find() }
}
pub fn filter<F>(mut self, filter: F) -> Self
where
F: sea_orm::sea_query::IntoCondition,
{
self.select = self.select.filter(filter);
self
}
pub fn order_by_asc<C>(mut self, col: C) -> Self
where
C: ColumnTrait,
{
self.select = self.select.order_by(col, Order::Asc);
self
}
pub fn order_by_desc<C>(mut self, col: C) -> Self
where
C: ColumnTrait,
{
self.select = self.select.order_by(col, Order::Desc);
self
}
pub fn order_by<C>(mut self, col: C, order: Order) -> Self
where
C: ColumnTrait,
{
self.select = self.select.order_by(col, order);
self
}
pub fn limit(mut self, limit: u64) -> Self {
self.select = self.select.limit(limit);
self
}
pub fn offset(mut self, offset: u64) -> Self {
self.select = self.select.offset(offset);
self
}
pub async fn all(self) -> Result<Vec<E::Model>, FrameworkError> {
let db = DB::connection()?;
self.select
.all(db.inner())
.await
.map_err(|e| FrameworkError::database(e.to_string()))
}
pub async fn first(self) -> Result<Option<E::Model>, FrameworkError> {
let db = DB::connection()?;
self.select
.one(db.inner())
.await
.map_err(|e| FrameworkError::database(e.to_string()))
}
pub async fn first_or_fail(self) -> Result<E::Model, FrameworkError> {
self.first().await?.ok_or_else(|| {
FrameworkError::database(format!("{} not found", std::any::type_name::<E::Model>()))
})
}
pub async fn count(self) -> Result<u64, FrameworkError> {
let db = DB::connection()?;
self.select
.count(db.inner())
.await
.map_err(|e| FrameworkError::database(e.to_string()))
}
pub async fn exists(self) -> Result<bool, FrameworkError> {
Ok(self.count().await? > 0)
}
pub fn into_select(self) -> Select<E> {
self.select
}
pub async fn all_with<R, F, Fut>(self, loader: F) -> Result<(Vec<E::Model>, R), FrameworkError>
where
F: FnOnce(&[E::Model]) -> Fut,
Fut: Future<Output = Result<R, FrameworkError>>,
{
let models = self.all().await?;
let related = loader(&models).await?;
Ok((models, related))
}
pub async fn all_with2<R1, R2, F1, F2, Fut1, Fut2>(
self,
loader1: F1,
loader2: F2,
) -> Result<(Vec<E::Model>, (R1, R2)), FrameworkError>
where
F1: FnOnce(&[E::Model]) -> Fut1,
F2: FnOnce(&[E::Model]) -> Fut2,
Fut1: Future<Output = Result<R1, FrameworkError>>,
Fut2: Future<Output = Result<R2, FrameworkError>>,
{
let models = self.all().await?;
let (r1, r2) = tokio::try_join!(loader1(&models), loader2(&models))?;
Ok((models, (r1, r2)))
}
pub async fn all_with3<R1, R2, R3, F1, F2, F3, Fut1, Fut2, Fut3>(
self,
loader1: F1,
loader2: F2,
loader3: F3,
) -> Result<(Vec<E::Model>, (R1, R2, R3)), FrameworkError>
where
F1: FnOnce(&[E::Model]) -> Fut1,
F2: FnOnce(&[E::Model]) -> Fut2,
F3: FnOnce(&[E::Model]) -> Fut3,
Fut1: Future<Output = Result<R1, FrameworkError>>,
Fut2: Future<Output = Result<R2, FrameworkError>>,
Fut3: Future<Output = Result<R3, FrameworkError>>,
{
let models = self.all().await?;
let (r1, r2, r3) = tokio::try_join!(loader1(&models), loader2(&models), loader3(&models))?;
Ok((models, (r1, r2, r3)))
}
}
impl<E> Default for QueryBuilder<E>
where
E: EntityTrait,
E::Model: Send + Sync,
{
fn default() -> Self {
Self::new()
}
}