use crate::data::BaseEntity;
use crate::logging::correlation::CorrelationContext;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use chrono::Utc;
use sea_orm::prelude::DateTimeWithTimeZone;
use sea_orm::{entity::EntityLoaderTrait, DatabaseTransaction, EntityTrait};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct NoUser;
impl BaseEntity for NoUser {
type LoaderType = ();
fn load() -> Self::LoaderType {}
fn id(&self) -> i64 {
0
}
fn uid(&self) -> uuid::Uuid {
uuid::Uuid::nil()
}
fn created_at(&self) -> DateTimeWithTimeZone {
Utc::now().into()
}
fn updated_at(&self) -> DateTimeWithTimeZone {
Utc::now().into()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CursorData {
pub limit: usize,
pub cursor: i64,
}
impl CursorData {
pub fn encode(&self) -> String {
let json = serde_json::to_string(self).unwrap_or_default();
BASE64.encode(json.as_bytes())
}
pub fn decode(s: &str) -> Option<Self> {
let bytes = BASE64.decode(s).ok()?;
let json = String::from_utf8(bytes).ok()?;
serde_json::from_str(&json).ok()
}
}
#[derive(Debug, Clone)]
pub struct RepositoryOptions<U = NoUser>
where
U: BaseEntity + Clone + Send + Sync + 'static,
{
pub limit: usize,
pub cursor: Option<String>,
pub distinct: bool,
pub detach: bool,
pub user: Option<U>,
pub user_id: Option<i64>,
pub correlation: Option<Arc<CorrelationContext>>,
pub txn: Option<Arc<DatabaseTransaction>>,
}
impl<U> Default for RepositoryOptions<U>
where
U: BaseEntity + Clone + Send + Sync + 'static,
{
fn default() -> Self {
Self {
limit: 15,
cursor: None,
distinct: false,
detach: true,
user: None,
user_id: None,
correlation: None,
txn: None,
}
}
}
impl<U> RepositoryOptions<U>
where
U: BaseEntity + Clone + Send + Sync + 'static,
{
pub fn new() -> Self {
Self::default()
}
pub fn with_limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
pub fn with_cursor(mut self, cursor: impl Into<String>) -> Self {
self.cursor = Some(cursor.into());
self
}
pub fn with_user(mut self, user: U) -> Self {
self.user_id = Some(user.id());
self.user = Some(user);
self
}
pub fn with_user_id(mut self, user_id: i64) -> Self {
self.user_id = Some(user_id);
self
}
pub fn with_correlation(mut self, ctx: Arc<CorrelationContext>) -> Self {
self.correlation = Some(ctx);
self
}
pub fn distinct(mut self) -> Self {
self.distinct = true;
self
}
pub fn with_detach(mut self, detach: bool) -> Self {
self.detach = detach;
self
}
pub fn with_transaction(mut self, txn: Arc<DatabaseTransaction>) -> Self {
self.txn = Some(txn);
self
}
pub fn with_txn(mut self, txn: Arc<DatabaseTransaction>) -> Self {
self.txn = Some(txn);
self
}
pub fn without_transaction(mut self) -> Self {
self.txn = None;
self
}
pub fn transaction(&self) -> Option<&DatabaseTransaction> {
self.txn.as_deref()
}
pub fn has_transaction(&self) -> bool {
self.txn.is_some()
}
pub fn user_id(&self) -> Option<i64> {
self.user_id.or_else(|| self.user.as_ref().map(|u| u.id()))
}
pub fn join(mut self, other: Self) -> Self {
self.limit = other.limit;
self.distinct = other.distinct;
self.detach = other.detach;
self.cursor = other.cursor.or(self.cursor);
if other.user.is_some() {
self.user = other.user;
if other.user_id.is_none() {
self.user_id = None;
} else {
self.user_id = other.user_id;
}
} else {
self.user_id = other.user_id.or(self.user_id);
}
self.correlation = other.correlation.or(self.correlation);
self.txn = other.txn.or(self.txn);
self
}
}
impl RepositoryOptions<NoUser> {
pub fn from_ctx(ctx: Arc<CorrelationContext>) -> Self {
Self {
limit: ctx.pagination_limit(),
cursor: ctx.pagination_cursor(),
user_id: ctx.user_id().and_then(|s| s.parse().ok()),
correlation: Some(ctx),
..Self::default()
}
}
pub fn with_user_entity<U>(self, user: U) -> RepositoryOptions<U>
where
U: BaseEntity + Clone + Send + Sync + 'static,
{
RepositoryOptions {
limit: self.limit,
cursor: self.cursor,
distinct: self.distinct,
detach: self.detach,
user: Some(user),
user_id: None,
correlation: self.correlation,
txn: self.txn,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PageResult<T> {
#[serde(rename = "items")]
pub data: Vec<T>,
#[serde(rename = "hasNext")]
pub has_next: bool,
#[serde(rename = "cursor", skip_serializing_if = "Option::is_none")]
pub cursor: Option<String>,
#[serde(rename = "next", skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
#[serde(rename = "count")]
pub total: i64,
#[serde(rename = "limit")]
pub limit: usize,
}
impl<T> PageResult<T> {
pub fn items(&self) -> &[T] {
&self.data
}
pub fn limit(&self) -> Option<usize> {
self.next_cursor
.as_ref()
.and_then(|c| CursorData::decode(c))
.map(|d| d.limit)
}
}
pub struct QueryData<E>
where
E: EntityTrait,
E::ModelEx: BaseEntity,
<E::ModelEx as BaseEntity>::LoaderType: EntityLoaderTrait<E, ModelEx = E::ModelEx>,
{
pub loader: <E::ModelEx as BaseEntity>::LoaderType,
pub opts: RepositoryOptions<NoUser>,
pub order_asc: Option<bool>,
}
impl<E> QueryData<E>
where
E: EntityTrait,
E::ModelEx: BaseEntity,
<E::ModelEx as BaseEntity>::LoaderType: EntityLoaderTrait<E, ModelEx = E::ModelEx>,
{
pub fn new(loader: <E::ModelEx as BaseEntity>::LoaderType) -> Self {
Self {
loader,
opts: RepositoryOptions::default(),
order_asc: None,
}
}
pub fn with_options<U>(mut self, opts: RepositoryOptions<U>) -> Self
where
U: BaseEntity + Clone + Send + Sync + 'static,
{
self.opts.limit = opts.limit;
self.opts.cursor = opts.cursor;
self.opts.distinct = opts.distinct;
self.opts.detach = opts.detach;
self.opts.correlation = opts.correlation;
self.opts.txn = opts.txn;
self
}
pub fn order_by(mut self, asc: bool) -> Self {
if asc {
self.loader = self.loader.order_by_id_asc();
} else {
self.loader = self.loader.order_by_id_desc();
}
self.order_asc = Some(asc);
self
}
pub fn order_by_asc(mut self) -> Self {
self.loader = self.loader.order_by_id_asc();
self.order_asc = Some(true);
self
}
pub fn order_by_desc(mut self) -> Self {
self.loader = self.loader.order_by_id_desc();
self.order_asc = Some(false);
self
}
pub fn distinct(mut self) -> Self {
self.opts.distinct = true;
self
}
pub fn filter<F>(mut self, f: F) -> Self
where
F: FnOnce(<E::ModelEx as BaseEntity>::LoaderType) -> <E::ModelEx as BaseEntity>::LoaderType,
{
self.loader = f(self.loader);
self
}
pub fn with_limit(mut self, limit: u64) -> Self {
self.opts.limit = limit as usize;
self
}
}
pub struct DeleteQueryData<E: EntityTrait> {
pub delete: sea_orm::DeleteMany<E>,
}
impl<E: EntityTrait> DeleteQueryData<E> {
pub fn new(delete: sea_orm::DeleteMany<E>) -> Self {
Self { delete }
}
pub fn filter<F>(mut self, f: F) -> Self
where
F: FnOnce(sea_orm::DeleteMany<E>) -> sea_orm::DeleteMany<E>,
{
self.delete = f(self.delete);
self
}
}