use crate::route::route_for;
use crate::{repo, Entity, Executor, Result, ShardRoute};
use core::future::Future;
use core::marker::PhantomData;
use sea_query::{
Alias, Asterisk, ColumnRef, Cond, Expr, Func, IntoColumnRef, JoinType, Order, Query as SqQuery,
SimpleExpr, Value,
};
pub struct Query<E> {
cond: Cond,
joins: Vec<(JoinType, String, SimpleExpr)>,
orders: Vec<(ColumnRef, Order)>,
limit: Option<u64>,
offset: Option<u64>,
route: ShardRoute,
_entity: PhantomData<E>,
}
impl<E: Entity> Default for Query<E> {
fn default() -> Self {
Self::new()
}
}
impl<E: Entity> Query<E> {
pub fn new() -> Self {
Self {
cond: Cond::all(),
joins: Vec::new(),
orders: Vec::new(),
limit: None,
offset: None,
route: ShardRoute::Global,
_entity: PhantomData,
}
}
fn add(mut self, expr: SimpleExpr) -> Self {
let cond = core::mem::replace(&mut self.cond, Cond::all());
self.cond = cond.add(expr);
self
}
pub fn where_eq(mut self, column: &str, value: impl Into<Value>) -> Self {
let val: Value = value.into();
if matches!(self.route, ShardRoute::Global) {
self.route = route_for::<E>(column, &val);
}
self.add(Expr::col(col_ref::<E>(column)).eq(val))
}
pub fn where_ne(self, column: &str, value: impl Into<Value>) -> Self {
self.add(Expr::col(col_ref::<E>(column)).ne(value.into()))
}
pub fn where_gt(self, column: &str, value: impl Into<Value>) -> Self {
self.add(Expr::col(col_ref::<E>(column)).gt(value.into()))
}
pub fn where_gte(self, column: &str, value: impl Into<Value>) -> Self {
self.add(Expr::col(col_ref::<E>(column)).gte(value.into()))
}
pub fn where_lt(self, column: &str, value: impl Into<Value>) -> Self {
self.add(Expr::col(col_ref::<E>(column)).lt(value.into()))
}
pub fn where_lte(self, column: &str, value: impl Into<Value>) -> Self {
self.add(Expr::col(col_ref::<E>(column)).lte(value.into()))
}
pub fn where_like(self, column: &str, pattern: &str) -> Self {
self.add(Expr::col(col_ref::<E>(column)).like(pattern))
}
pub fn where_not_like(self, column: &str, pattern: &str) -> Self {
self.add(Expr::col(col_ref::<E>(column)).not_like(pattern))
}
pub fn where_in<V: Into<Value>>(
self,
column: &str,
values: impl IntoIterator<Item = V>,
) -> Self {
let vals: Vec<Value> = values.into_iter().map(Into::into).collect();
self.add(Expr::col(col_ref::<E>(column)).is_in(vals))
}
pub fn where_not_in<V: Into<Value>>(
self,
column: &str,
values: impl IntoIterator<Item = V>,
) -> Self {
let vals: Vec<Value> = values.into_iter().map(Into::into).collect();
self.add(Expr::col(col_ref::<E>(column)).is_not_in(vals))
}
pub fn where_null(self, column: &str) -> Self {
self.add(Expr::col(col_ref::<E>(column)).is_null())
}
pub fn where_not_null(self, column: &str) -> Self {
self.add(Expr::col(col_ref::<E>(column)).is_not_null())
}
pub fn filter(mut self, cond: Cond) -> Self {
let existing = core::mem::replace(&mut self.cond, Cond::all());
self.cond = existing.add(cond);
self
}
pub fn join(self, foreign_table: &str, local_col: &str, foreign_col: &str) -> Self {
self.join_with(JoinType::InnerJoin, foreign_table, local_col, foreign_col)
}
pub fn left_join(self, foreign_table: &str, local_col: &str, foreign_col: &str) -> Self {
self.join_with(JoinType::LeftJoin, foreign_table, local_col, foreign_col)
}
fn join_with(
mut self,
kind: JoinType,
foreign_table: &str,
local_col: &str,
foreign_col: &str,
) -> Self {
let on = Expr::col((Alias::new(E::table()), Alias::new(local_col)))
.equals((Alias::new(foreign_table.to_owned()), Alias::new(foreign_col.to_owned())));
self.joins.push((kind, foreign_table.to_owned(), on));
self
}
pub fn order_by_asc(mut self, column: &str) -> Self {
self.orders.push((col_ref::<E>(column), Order::Asc));
self
}
pub fn order_by_desc(mut self, column: &str) -> Self {
self.orders.push((col_ref::<E>(column), Order::Desc));
self
}
pub fn limit(mut self, n: u64) -> Self {
self.limit = Some(n);
self
}
pub fn offset(mut self, n: u64) -> Self {
self.offset = Some(n);
self
}
fn render_select(self, dialect: crate::Dialect) -> (ShardRoute, String, Vec<Value>) {
let mut stmt = SqQuery::select();
stmt.from(Alias::new(E::table()));
for c in E::columns() {
stmt.column((Alias::new(E::table()), Alias::new(c.name)));
}
for (kind, table, on) in &self.joins {
stmt.join(*kind, Alias::new(table.clone()), on.clone());
}
stmt.cond_where(self.cond.clone());
for (col, ord) in &self.orders {
stmt.order_by(col.clone(), ord.clone());
}
if let Some(l) = self.limit {
stmt.limit(l);
}
if let Some(o) = self.offset {
stmt.offset(o);
}
let (sql, params) = dialect.build_query(&stmt);
(self.route, sql, params.0)
}
fn render_count(self, dialect: crate::Dialect) -> (ShardRoute, String, Vec<Value>) {
let mut stmt = SqQuery::select();
stmt.from(Alias::new(E::table()));
for (kind, table, on) in &self.joins {
stmt.join(*kind, Alias::new(table.clone()), on.clone());
}
stmt.cond_where(self.cond.clone());
stmt.expr_as(Func::count(Expr::col(Asterisk)), Alias::new("cnt"));
let (sql, params) = dialect.build_query(&stmt);
(self.route, sql, params.0)
}
fn render_update(
self,
dialect: crate::Dialect,
set: Vec<(&str, Value)>,
) -> (ShardRoute, String, Vec<Value>) {
let mut stmt = SqQuery::update();
stmt.table(Alias::new(E::table()));
for (c, v) in set {
stmt.value(Alias::new(c.to_owned()), v);
}
stmt.cond_where(self.cond.clone());
let (sql, params) = dialect.build_query(&stmt);
(self.route, sql, params.0)
}
fn render_increment(
self,
dialect: crate::Dialect,
column: &str,
by: i64,
) -> (ShardRoute, String, Vec<Value>) {
let mut stmt = SqQuery::update();
stmt.table(Alias::new(E::table()));
stmt.value(Alias::new(column.to_owned()), Expr::col(Alias::new(column.to_owned())).add(by));
stmt.cond_where(self.cond.clone());
let (sql, params) = dialect.build_query(&stmt);
(self.route, sql, params.0)
}
fn render_delete(self, dialect: crate::Dialect) -> (ShardRoute, String, Vec<Value>) {
let mut stmt = SqQuery::delete();
stmt.from_table(Alias::new(E::table()));
stmt.cond_where(self.cond.clone());
let (sql, params) = dialect.build_query(&stmt);
(self.route, sql, params.0)
}
pub fn all<'a, X: Executor>(self, exec: &'a X) -> impl Future<Output = Result<Vec<E>>> + 'a
where
E: 'a,
{
let (route, sql, params) = self.render_select(exec.dialect());
async move {
let rows = exec.fetch_all_routed(route, &sql, params).await?;
rows.iter().map(|r| E::from_row(r.as_ref())).collect()
}
}
pub fn first<'a, X: Executor>(self, exec: &'a X) -> impl Future<Output = Result<Option<E>>> + 'a
where
E: 'a,
{
let rows = self.limit(1).all(exec);
async move { Ok(rows.await?.into_iter().next()) }
}
pub fn count<'a, X: Executor>(self, exec: &'a X) -> impl Future<Output = Result<u64>> + 'a
where
E: 'a,
{
let (route, sql, params) = self.render_count(exec.dialect());
async move {
let rows = exec.fetch_all_routed(route, &sql, params).await?;
let n = match rows.first() {
Some(r) => r.get_i64("cnt")?.unwrap_or(0),
None => 0,
};
Ok(n.max(0) as u64)
}
}
pub fn update<'a, X: Executor>(
self,
exec: &'a X,
set: Vec<(&str, Value)>,
) -> impl Future<Output = Result<u64>> + 'a {
let (route, sql, params) = self.render_update(exec.dialect(), set);
async move { Ok(exec.execute_routed(route, &sql, params).await?.rows_affected) }
}
pub fn increment<'a, X: Executor>(
self,
exec: &'a X,
column: &str,
by: i64,
) -> impl Future<Output = Result<u64>> + 'a {
let (route, sql, params) = self.render_increment(exec.dialect(), column, by);
async move { Ok(exec.execute_routed(route, &sql, params).await?.rows_affected) }
}
pub fn delete<'a, X: Executor>(self, exec: &'a X) -> impl Future<Output = Result<u64>> + 'a {
let (route, sql, params) = self.render_delete(exec.dialect());
async move { Ok(exec.execute_routed(route, &sql, params).await?.rows_affected) }
}
pub fn first_or_create<'a, X: Executor>(
self,
exec: &'a X,
default: E,
) -> impl Future<Output = Result<E>> + 'a
where
E: 'a,
{
let existing = self.first(exec);
async move {
if let Some(found) = existing.await? {
return Ok(found);
}
let id = repo::insert(exec, &default).await?;
if id > 0 {
if let Some(created) = repo::find_by_pk(exec, id).await? {
return Ok(created);
}
}
Ok(default)
}
}
}
fn col_ref<E: Entity>(spec: &str) -> ColumnRef {
if let Some((t, c)) = spec.split_once('.') {
(Alias::new(t.to_owned()), Alias::new(c.to_owned())).into_column_ref()
} else {
(Alias::new(E::table()), Alias::new(spec.to_owned())).into_column_ref()
}
}