use crate::route::route_for;
use crate::{Entity, Executor, Result, ShardRoute};
use sea_query::{Alias, Expr, OnConflict, Order, Query, SimpleExpr, Value};
fn col(name: &str) -> Alias {
Alias::new(name)
}
fn route_from_pairs<E: Entity>(pairs: &[(&'static str, Value)]) -> ShardRoute {
if E::shard_columns().is_empty() {
return ShardRoute::Global;
}
for (c, v) in pairs {
let r = route_for::<E>(c, v);
if !matches!(r, ShardRoute::Global) {
return r;
}
}
ShardRoute::Global
}
fn assign_shard_pk<E: Entity, X: Executor>(exec: &X, pairs: &[(&str, Value)]) -> Option<u64> {
let pk = E::primary_key();
if !E::shard_columns().contains(&pk) {
return None;
}
let pk_spec = E::columns().iter().find(|c| c.name == pk)?;
if pk_spec.auto_increment {
return None;
}
let pk_val = pairs.iter().find(|(c, _)| *c == pk).map(|(_, v)| v)?;
if crate::route::routing_key(pk_val) != Some(0) {
return None;
}
let (_, sk_val) = pairs.iter().find(|(c, v)| {
*c != pk && E::shard_columns().contains(c) && crate::route::routing_key(v).is_some()
})?;
let shard_key = crate::route::routing_key(sk_val)?;
exec.allocate_id(shard_key)
}
pub fn query<E: Entity>() -> crate::Query<E> {
crate::Query::new()
}
pub async fn insert<E, X>(exec: &X, entity: &E) -> Result<i64>
where
E: Entity,
X: Executor,
{
let (route, sql, params, assigned) = {
let mut pairs = entity.insert_values();
let assigned = assign_shard_pk::<E, X>(exec, &pairs);
if let Some(id) = assigned {
let pk = E::primary_key();
for (c, v) in pairs.iter_mut() {
if *c == pk {
*v = id.into();
}
}
}
let route = route_from_pairs::<E>(&pairs);
let mut stmt = Query::insert();
stmt.into_table(col(E::table()));
stmt.columns(pairs.iter().map(|(c, _)| col(c)));
stmt.values_panic(pairs.into_iter().map(|(_, v)| SimpleExpr::from(Expr::val(v))));
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params), assigned)
};
let outcome = exec.execute_routed(route, &sql, params).await?;
Ok(assigned.map(|i| i as i64).unwrap_or(outcome.last_insert_id))
}
pub async fn upsert<E, X>(
exec: &X,
entity: &E,
conflict_columns: &[&str],
update_columns: &[&str],
) -> Result<u64>
where
E: Entity,
X: Executor,
{
let (route, sql, params) = {
let pairs = entity.insert_values();
let route = route_from_pairs::<E>(&pairs);
let mut stmt = Query::insert();
stmt.into_table(col(E::table()));
stmt.columns(pairs.iter().map(|(c, _)| col(c)));
stmt.values_panic(pairs.into_iter().map(|(_, v)| SimpleExpr::from(Expr::val(v))));
let mut on_conflict = OnConflict::columns(conflict_columns.iter().map(|c| col(c)));
if !update_columns.is_empty() {
on_conflict.update_columns(update_columns.iter().map(|c| col(c)));
} else if let Some(c) = conflict_columns.first() {
on_conflict.value(col(c), Expr::col(col(c)));
} else {
on_conflict.do_nothing();
}
stmt.on_conflict(on_conflict);
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
Ok(exec.execute_routed(route, &sql, params).await?.rows_affected)
}
pub async fn find_by_pk<E, X, V>(exec: &X, pk: V) -> Result<Option<E>>
where
E: Entity,
X: Executor,
V: Into<Value>,
{
let (route, sql, params) = {
let pk_val: Value = pk.into();
let route = route_for::<E>(E::primary_key(), &pk_val);
let mut stmt = Query::select();
stmt.from(col(E::table()))
.columns(E::columns().iter().map(|c| col(c.name)))
.and_where(Expr::col(col(E::primary_key())).eq(pk_val))
.limit(1);
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let rows = exec.fetch_all_routed(route, &sql, params).await?;
match rows.first() {
Some(row) => Ok(Some(E::from_row(row.as_ref())?)),
None => Ok(None),
}
}
pub async fn find_by<E, X, V>(exec: &X, column: &str, value: V) -> Result<Vec<E>>
where
E: Entity,
X: Executor,
V: Into<Value>,
{
let (route, sql, params) = {
let val: Value = value.into();
let route = route_for::<E>(column, &val);
let mut stmt = Query::select();
stmt.from(col(E::table()))
.columns(E::columns().iter().map(|c| col(c.name)))
.and_where(Expr::col(col(column)).eq(val));
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let rows = exec.fetch_all_routed(route, &sql, params).await?;
rows.iter().map(|r| E::from_row(r.as_ref())).collect()
}
pub async fn find_one_by<E, X, V>(exec: &X, column: &str, value: V) -> Result<Option<E>>
where
E: Entity,
X: Executor,
V: Into<Value>,
{
let (route, sql, params) = {
let val: Value = value.into();
let route = route_for::<E>(column, &val);
let mut stmt = Query::select();
stmt.from(col(E::table()))
.columns(E::columns().iter().map(|c| col(c.name)))
.and_where(Expr::col(col(column)).eq(val))
.limit(1);
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let rows = exec.fetch_all_routed(route, &sql, params).await?;
match rows.first() {
Some(row) => Ok(Some(E::from_row(row.as_ref())?)),
None => Ok(None),
}
}
pub async fn all<E, X>(exec: &X) -> Result<Vec<E>>
where
E: Entity,
X: Executor,
{
let (sql, params) = {
let mut stmt = Query::select();
stmt.from(col(E::table())).columns(E::columns().iter().map(|c| col(c.name)));
let (sql, params) = exec.dialect().build_query(&stmt);
(sql, into_vec(params))
};
let rows = exec.fetch_all_routed(ShardRoute::Global, &sql, params).await?;
rows.iter().map(|r| E::from_row(r.as_ref())).collect()
}
pub async fn update<E, X>(exec: &X, entity: &E) -> Result<u64>
where
E: Entity,
X: Executor,
{
let (route, sql, params) = {
let pk_val = entity.pk_value();
let route = route_for::<E>(E::primary_key(), &pk_val);
let mut stmt = Query::update();
stmt.table(col(E::table()));
for (c, v) in entity.update_values() {
stmt.value(col(c), v);
}
stmt.and_where(Expr::col(col(E::primary_key())).eq(pk_val));
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let outcome = exec.execute_routed(route, &sql, params).await?;
Ok(outcome.rows_affected)
}
pub async fn delete_by_pk<E, X, V>(exec: &X, pk: V) -> Result<u64>
where
E: Entity,
X: Executor,
V: Into<Value>,
{
let (route, sql, params) = {
let pk_val: Value = pk.into();
let route = route_for::<E>(E::primary_key(), &pk_val);
let mut stmt = Query::delete();
stmt.from_table(col(E::table())).and_where(Expr::col(col(E::primary_key())).eq(pk_val));
let (sql, params) = exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let outcome = exec.execute_routed(route, &sql, params).await?;
Ok(outcome.rows_affected)
}
pub fn cursor<E, X>(exec: &X, page_size: u64) -> Cursor<'_, E, X>
where
E: Entity,
X: Executor,
{
Cursor::new(exec, page_size)
}
pub struct Cursor<'a, E, X> {
exec: &'a X,
page_size: u64,
after: Option<Value>,
done: bool,
_entity: core::marker::PhantomData<E>,
}
impl<'a, E, X> Cursor<'a, E, X>
where
E: Entity,
X: Executor,
{
fn new(exec: &'a X, page_size: u64) -> Self {
Self {
exec,
page_size: page_size.max(1),
after: None,
done: false,
_entity: core::marker::PhantomData,
}
}
pub fn after(mut self, pk: impl Into<Value>) -> Self {
self.after = Some(pk.into());
self
}
pub async fn next_page(&mut self) -> Result<Option<Vec<E>>> {
if self.done {
return Ok(None);
}
let pk = E::primary_key();
let (route, sql, params) = {
let mut stmt = Query::select();
stmt.from(col(E::table()))
.columns(E::columns().iter().map(|c| col(c.name)))
.order_by(col(pk), Order::Asc)
.limit(self.page_size);
if let Some(after) = &self.after {
stmt.and_where(Expr::col(col(pk)).gt(after.clone()));
}
let route = match &self.after {
Some(after) => route_for::<E>(pk, after),
None => ShardRoute::Global,
};
let (sql, params) = self.exec.dialect().build_query(&stmt);
(route, sql, into_vec(params))
};
let rows = self.exec.fetch_all_routed(route, &sql, params).await?;
if rows.is_empty() {
self.done = true;
return Ok(None);
}
let page: Vec<E> = rows.iter().map(|r| E::from_row(r.as_ref())).collect::<Result<_>>()?;
if let Some(last) = rows.last() {
self.after = Some(pk_from_row::<E>(last.as_ref())?);
}
if (rows.len() as u64) < self.page_size {
self.done = true;
}
Ok(Some(page))
}
pub fn checkpoint(&self) -> Option<&Value> {
self.after.as_ref()
}
}
fn pk_from_row<E: Entity>(row: &dyn crate::Row) -> Result<Value> {
use crate::ColumnType;
let pk = E::primary_key();
let spec = E::columns()
.iter()
.find(|c| c.name == pk)
.ok_or_else(|| crate::Error::msg("primary key column not in metadata"))?;
let null = || crate::Error::msg(format!("primary key `{pk}` was NULL"));
Ok(match spec.ty {
ColumnType::Bool => row.get_bool(pk)?.ok_or_else(null)?.into(),
ColumnType::Int => row.get_i32(pk)?.ok_or_else(null)?.into(),
ColumnType::BigInt => row.get_i64(pk)?.ok_or_else(null)?.into(),
ColumnType::Uint => row.get_u32(pk)?.ok_or_else(null)?.into(),
ColumnType::BigUint => row.get_u64(pk)?.ok_or_else(null)?.into(),
ColumnType::Double => row.get_f64(pk)?.ok_or_else(null)?.into(),
ColumnType::Text => row.get_string(pk)?.ok_or_else(null)?.into(),
ColumnType::Binary => row.get_bytes(pk)?.ok_or_else(null)?.into(),
ColumnType::Timestamp => row.get_datetime(pk)?.ok_or_else(null)?.into(),
ColumnType::Date => row.get_naive_date(pk)?.ok_or_else(null)?.into(),
})
}
fn into_vec(values: sea_query::Values) -> Vec<Value> {
values.0
}