magnetar/ext_model/src/model_ext.rs

use chrono::{DateTime, Utc};
use ext_model_migration::{Alias, Expr, IntoIden, Quote, SelectExpr, SelectStatement, TableRef};
use magnetar_sdk::types::SpanFilter;
use sea_orm::{
    ColumnTrait, Condition, ConnectionTrait, Cursor, DbErr, DynIden, EntityTrait, FromQueryResult,
    Iden, IntoIdentity, Iterable, JoinType, QueryTrait, RelationDef, RelationTrait, Select,
    SelectModel, SelectorTrait,
};
use serde::{Deserialize, Serialize};
use std::fmt::Write;

#[derive(Clone)]
pub enum MagIden {
    DynIden(DynIden),
    Alias(Alias),
}

impl MagIden {
    pub fn alias(alias: &str) -> MagIden {
        Self::Alias(Alias::new(alias))
    }
}

impl From<DynIden> for MagIden {
    fn from(value: DynIden) -> Self {
        Self::DynIden(value)
    }
}

impl Iden for MagIden {
    fn prepare(&self, s: &mut dyn Write, q: Quote) {
        match self {
            MagIden::DynIden(di) => di.prepare(s, q),
            MagIden::Alias(a) => a.prepare(s, q),
        }
    }

    fn quoted(&self, q: Quote) -> String {
        match self {
            MagIden::DynIden(di) => di.quoted(q),
            MagIden::Alias(a) => a.quoted(q),
        }
    }

    fn to_string(&self) -> String {
        match self {
            MagIden::DynIden(di) => di.to_string(),
            MagIden::Alias(a) => a.to_string(),
        }
    }

    fn unquoted(&self, s: &mut dyn Write) {
        match self {
            MagIden::DynIden(di) => di.unquoted(s),
            MagIden::Alias(a) => a.unquoted(s),
        }
    }
}

pub(crate) trait SelectColumnsExt {
    fn add_aliased_columns<T: EntityTrait>(&mut self, alias: &MagIden) -> &mut Self;

    fn join_columns(&mut self, join: JoinType, rel: RelationDef, alias: &MagIden) -> &mut Self;

    fn join_columns_on(
        &mut self,
        join: JoinType,
        rel: RelationDef,
        alias: &MagIden,
        condition: Condition,
    ) -> &mut Self;
}

pub(crate) fn join_columns_default(rel: RelationDef) -> Condition {
    let tbl_id = |tbl: &TableRef| {
        match tbl {
            TableRef::Table(id) => id,
            TableRef::TableAlias(_, id) => id,
            _ => unreachable!(),
        }
        .clone()
    };

    let mut cond = Condition::all();
    for (owner_key, foreign_key) in rel.from_col.into_iter().zip(rel.to_col.into_iter()) {
        cond = cond.add(
            Expr::col((tbl_id(&rel.from_tbl), owner_key))
                .equals((tbl_id(&rel.to_tbl), foreign_key)),
        );
    }
    cond
}

impl SelectColumnsExt for SelectStatement {
    fn add_aliased_columns<T: EntityTrait>(&mut self, iden: &MagIden) -> &mut Self {
        for col in T::Column::iter() {
            let column: &T::Column = &col;
            let alias = iden.join(&col);
            let column_ref = iden.col(column.as_column_ref().1);

            self.expr(SelectExpr {
                expr: col.select_as(column_ref),
                alias: Some(alias.into_iden()),
                window: None,
            });
        }

        self
    }

    fn join_columns(&mut self, join: JoinType, mut rel: RelationDef, alias: &MagIden) -> &mut Self {
        rel.to_tbl = rel.to_tbl.alias(alias.clone().into_iden());
        self.join(join, rel.to_tbl.clone(), join_columns_default(rel));
        self
    }

    fn join_columns_on(
        &mut self,
        join: JoinType,
        mut rel: RelationDef,
        alias: &MagIden,
        condition: Condition,
    ) -> &mut Self {
        rel.to_tbl = rel.to_tbl.alias(alias.clone().into_iden());
        self.join(join, rel.to_tbl.clone(), condition);
        self
    }
}

pub trait AliasColumnExt {
    fn col_tuple(&self, col: impl IntoIden) -> (MagIden, MagIden);

    fn col(&self, col: impl IntoIden) -> Expr;
}

impl<T: Iden + ?Sized> AliasColumnExt for T {
    fn col_tuple(&self, col: impl IntoIden) -> (MagIden, MagIden) {
        (
            MagIden::alias(&self.to_string()),
            MagIden::DynIden(col.into_iden()),
        )
    }

    fn col(&self, col: impl IntoIden) -> Expr {
        Expr::col(self.col_tuple(col))
    }
}

pub trait AliasSuffixExt {
    fn join(&self, suffix: &dyn Iden) -> MagIden;

    fn join_str(&self, suffix: &str) -> MagIden;

    fn join_as_str(&self, suffix: &dyn Iden) -> String;

    fn join_str_as_str(&self, suffix: &str) -> String;
}

impl<T: Iden + ?Sized> AliasSuffixExt for T {
    fn join(&self, suffix: &dyn Iden) -> MagIden {
        MagIden::alias(&self.join_as_str(suffix))
    }

    fn join_str(&self, suffix: &str) -> MagIden {
        MagIden::alias(&self.join_str_as_str(suffix))
    }

    fn join_as_str(&self, suffix: &dyn Iden) -> String {
        format!("{}{}", self.to_string(), suffix.to_string())
    }

    fn join_str_as_str(&self, suffix: &str) -> String {
        format!("{}{}", self.to_string(), suffix)
    }
}

pub(crate) trait EntityPrefixExt {
    fn base_prefix_str(&self) -> String;
    fn base_prefix(&self) -> MagIden;
}

impl<T: EntityTrait> EntityPrefixExt for T {
    fn base_prefix_str(&self) -> String {
        format!("{}.", self.table_name())
    }

    fn base_prefix(&self) -> MagIden {
        MagIden::alias(&self.base_prefix_str())
    }
}

pub trait AliasSourceExt {
    fn with_from_alias(&self, alias: &MagIden) -> RelationDef;

    fn with_to_alias(&self, alias: &MagIden) -> RelationDef;

    fn with_alias(&self, from: &MagIden, to: &MagIden) -> RelationDef;
}

impl<T: RelationTrait> AliasSourceExt for T {
    fn with_from_alias(&self, alias: &MagIden) -> RelationDef {
        let mut def = self.def();
        def.from_tbl = def.from_tbl.alias(Alias::new(alias.clone().to_string()));
        def
    }

    fn with_to_alias(&self, alias: &MagIden) -> RelationDef {
        let mut def = self.def();
        def.to_tbl = def.to_tbl.alias(Alias::new(alias.clone().to_string()));
        def
    }

    fn with_alias(&self, from: &MagIden, to: &MagIden) -> RelationDef {
        let mut def = self.def();
        def.from_tbl = def.from_tbl.alias(from.clone().into_iden());
        def.to_tbl = def.to_tbl.alias(to.clone().into_iden());
        def
    }
}

pub trait CursorPaginationExt<E> {
    type Selector: SelectorTrait + Send + Sync;

    fn cursor_by_columns_and_span<C>(
        self,
        cursor_prefix_alias: Option<MagIden>,
        order_columns: C,
        pagination: &SpanFilter,
        limit: Option<u64>,
    ) -> Cursor<Self::Selector>
    where
        C: IntoIdentity;

    async fn get_paginated_model<M, C, T>(
        self,
        db: &T,
        cursor_prefix_alias: Option<MagIden>,
        columns: C,
        curr: &SpanFilter,
        prev: &mut Option<SpanFilter>,
        next: &mut Option<SpanFilter>,
        limit: u64,
    ) -> Result<Vec<M>, DbErr>
    where
        M: FromQueryResult + ModelPagination,
        C: IntoIdentity,
        T: ConnectionTrait;
}

impl<E, M> CursorPaginationExt<E> for Select<E>
where
    E: EntityTrait<Model = M>,
    M: FromQueryResult + Sized + Send + Sync,
{
    type Selector = SelectModel<M>;

    fn cursor_by_columns_and_span<C>(
        self,
        cursor_prefix_alias: Option<MagIden>,
        order_columns: C,
        pagination: &SpanFilter,
        limit: Option<u64>,
    ) -> Cursor<Self::Selector>
    where
        C: IntoIdentity,
    {
        let mut cursor = Cursor::new(
            self.into_query(),
            cursor_prefix_alias.map_or_else(|| E::default().into_iden(), MagIden::into_iden),
            order_columns,
        );

        if let Some(start) = pagination.start() {
            cursor.after(start);
        }

        if let Some(end) = pagination.end() {
            cursor.before(end);
        }

        if let Some(lim) = limit {
            if pagination.is_desc() {
                cursor.last(lim);
            } else {
                cursor.first(lim);
            }
        }

        cursor
    }

    async fn get_paginated_model<Q, C, T>(
        self,
        db: &T,
        cursor_prefix_alias: Option<MagIden>,
        columns: C,
        curr: &SpanFilter,
        prev: &mut Option<SpanFilter>,
        next: &mut Option<SpanFilter>,
        limit: u64,
    ) -> Result<Vec<Q>, DbErr>
    where
        Q: FromQueryResult + ModelPagination,
        C: IntoIdentity,
        T: ConnectionTrait,
    {
        let mut result = self
            .cursor_by_columns_and_span(cursor_prefix_alias, columns, curr, Some(limit + 1))
            .into_model::<Q>()
            .all(db)
            .await?;

        if curr.is_desc() {
            result.reverse();
        }

        if result.len() > limit as usize {
            result.pop();
            let last = result.last();
            *next = last.and_then(|c| curr.next(c.time(), c.id()));
        }

        let first = result.first();
        *prev = first.and_then(|c| curr.prev(c.time(), c.id()));

        Ok(result)
    }
}

pub trait ModelPagination {
    fn id(&self) -> &str;
    fn time(&self) -> DateTime<Utc>;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IdShape {
    pub id: String,
}