magnetar/ext_calckey_model/src/note_model.rs

use lazy_static::lazy_static;
use sea_orm::sea_query::{
    Alias, Asterisk, Expr, IntoCondition, IntoIden, Query, SelectExpr, SimpleExpr,
};
use sea_orm::{
    ColumnTrait, DbErr, EntityName, EntityTrait, FromQueryResult, Iden, Iterable, JoinType,
    QueryFilter, QueryOrder, QueryResult, QuerySelect, QueryTrait, RelationTrait, Select,
};
use serde::{Deserialize, Serialize};

use ck::{drive_file, note, note_reaction, user, user_note_pining};
use magnetar_sdk::types::RangeFilter;

use crate::{AliasSourceExt, CalckeyDbError, CalckeyModel};

pub mod sub_interaction_renote {
    use sea_orm::{DeriveColumn, EnumIter, FromQueryResult};
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
    pub struct Model {
        pub renotes: Option<i64>,
    }

    #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
    pub enum Column {
        Renotes,
    }
}

pub mod sub_interaction_reaction {
    use sea_orm::{DeriveColumn, EnumIter, FromQueryResult};
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
    pub struct Model {
        pub reaction_name: Option<String>,
    }

    #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
    pub enum Column {
        ReactionName,
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NoteData {
    pub note: note::Model,
    pub interaction_user_renote: Option<sub_interaction_renote::Model>,
    pub interaction_user_reaction: Option<sub_interaction_reaction::Model>,
    pub user: user::Model,
    pub avatar: Option<drive_file::Model>,
    pub banner: Option<drive_file::Model>,
    pub reply: Option<Box<NoteData>>,
    pub renote: Option<Box<NoteData>>,
}

const PIN: &str = "pin.";
const INTERACTION_REACTION: &str = "interaction.reaction.";
const INTERACTION_RENOTE: &str = "interaction.renote.";
const USER: &str = "user.";
const USER_AVATAR: &str = "user.avatar.";
const USER_BANNER: &str = "user.banner.";
const REPLY: &str = "reply.";
const REPLY_INTERACTION_REACTION: &str = "reply.interaction.reaction.";
const REPLY_INTERACTION_RENOTE: &str = "reply.interaction.renote.";
const REPLY_USER: &str = "reply.user.";
const REPLY_USER_AVATAR: &str = "reply.user.avatar.";
const REPLY_USER_BANNER: &str = "reply.user.banner.";
const RENOTE: &str = "renote.";
const RENOTE_INTERACTION_REACTION: &str = "renote.interaction.reaction.";
const RENOTE_INTERACTION_RENOTE: &str = "renote.interaction.renote.";
const RENOTE_USER: &str = "renote.user.";
const RENOTE_USER_AVATAR: &str = "renote.user.avatar.";
const RENOTE_USER_BANNER: &str = "renote.user.banner.";

impl FromQueryResult for NoteData {
    fn from_query_result(res: &QueryResult, _pre: &str) -> Result<Self, DbErr> {
        let reply = note::Model::from_query_result_optional(res, REPLY)?
            .map::<Result<_, DbErr>, _>(|r| {
                Ok(Box::new(NoteData {
                    note: r,
                    interaction_user_renote:
                        sub_interaction_renote::Model::from_query_result_optional(
                            res,
                            REPLY_INTERACTION_RENOTE,
                        )?,
                    interaction_user_reaction:
                        sub_interaction_reaction::Model::from_query_result_optional(
                            res,
                            REPLY_INTERACTION_REACTION,
                        )?,
                    user: user::Model::from_query_result(res, REPLY_USER)?,
                    avatar: drive_file::Model::from_query_result_optional(res, REPLY_USER_AVATAR)?,
                    banner: drive_file::Model::from_query_result_optional(res, REPLY_USER_BANNER)?,
                    reply: None,
                    renote: None,
                }))
            })
            .transpose()?;

        let renote = note::Model::from_query_result_optional(res, RENOTE)?
            .map::<Result<_, DbErr>, _>(|r| {
                Ok(Box::new(NoteData {
                    note: r,
                    interaction_user_renote:
                        sub_interaction_renote::Model::from_query_result_optional(
                            res,
                            RENOTE_INTERACTION_RENOTE,
                        )?,
                    interaction_user_reaction:
                        sub_interaction_reaction::Model::from_query_result_optional(
                            res,
                            RENOTE_INTERACTION_REACTION,
                        )?,
                    user: user::Model::from_query_result(res, RENOTE_USER)?,
                    avatar: drive_file::Model::from_query_result_optional(res, RENOTE_USER_AVATAR)?,
                    banner: drive_file::Model::from_query_result_optional(res, RENOTE_USER_BANNER)?,
                    reply: None,
                    renote: None,
                }))
            })
            .transpose()?;

        Ok(NoteData {
            note: note::Model::from_query_result(res, "")?,
            interaction_user_renote: sub_interaction_renote::Model::from_query_result_optional(
                res,
                INTERACTION_RENOTE,
            )?,
            interaction_user_reaction: sub_interaction_reaction::Model::from_query_result_optional(
                res,
                INTERACTION_REACTION,
            )?,
            user: user::Model::from_query_result(res, USER)?,
            avatar: drive_file::Model::from_query_result_optional(res, USER_AVATAR)?,
            banner: drive_file::Model::from_query_result_optional(res, USER_BANNER)?,
            reply,
            renote,
        })
    }
}

pub struct NoteResolver {
    db: CalckeyModel,
}

pub trait NoteVisibilityFilterFactory: Send + Sync {
    fn with_note_and_user_tables(&self, note: Option<Alias>) -> SimpleExpr;
}

pub struct NoteResolveOptions {
    pub ids: Option<Vec<String>>,
    pub visibility_filter: Box<dyn NoteVisibilityFilterFactory>,
    pub time_range: Option<RangeFilter>,
    pub with_user: bool,
    pub with_reply_target: bool,
    pub with_renote_target: bool,
    pub with_interactions_from: Option<String>, // User ID
    pub only_pins_from: Option<String>,         // User ID
}

trait SelectColumnsExt {
    fn add_aliased_columns<T: EntityTrait>(self, alias: Option<&str>, entity: T) -> Self;

    fn add_sub_select_reaction(
        self,
        source_note_alias: Option<&str>,
        alias: &str,
        user_id: &str,
    ) -> Select<note::Entity>;

    fn add_sub_select_renote(
        self,
        source_note_alias: Option<&str>,
        alias: &str,
        user_id: &str,
    ) -> Select<note::Entity>;
}

impl SelectColumnsExt for Select<note::Entity> {
    fn add_aliased_columns<T: EntityTrait>(
        mut self: Select<note::Entity>,
        alias: Option<&str>,
        entity: T,
    ) -> Select<note::Entity> {
        for col in T::Column::iter() {
            let column: &T::Column = &col;

            let iden = alias.unwrap_or_else(|| entity.table_name());
            let alias = format!("{}{}", iden, col.to_string());

            let column_ref = Expr::col((Alias::new(iden), column.as_column_ref().1));

            QuerySelect::query(&mut self).expr(SelectExpr {
                expr: col.select_as(column_ref),
                alias: Some(Alias::new(&alias).into_iden()),
                window: None,
            });
        }

        self
    }

    fn add_sub_select_reaction(
        mut self: Select<note::Entity>,
        source_note_alias: Option<&str>,
        prefix_alias: &str,
        user_id: &str,
    ) -> Select<note::Entity> {
        let iden = source_note_alias.unwrap_or_else(|| note::Entity.table_name());
        let note_id_col = Expr::col((Alias::new(iden), note::Column::Id));

        let column = sub_interaction_reaction::Column::ReactionName;
        let alias = format!("{}{}", prefix_alias, column.to_string());

        let sub_query = Query::select()
            .expr(SelectExpr {
                expr: Expr::col(note_reaction::Column::Reaction).into(),
                alias: Some(Alias::new(alias).into_iden()),
                window: None,
            })
            .from(note_reaction::Entity)
            .cond_where(Expr::col(note_reaction::Column::NoteId).eq(note_id_col))
            .and_where(Expr::col(note_reaction::Column::UserId).eq(user_id))
            .take()
            .into_sub_query_statement();

        QuerySelect::query(&mut self).expr(SimpleExpr::SubQuery(None, Box::new(sub_query)));

        self
    }

    fn add_sub_select_renote(
        mut self: Select<note::Entity>,
        source_note_alias: Option<&str>,
        prefix_alias: &str,
        user_id: &str,
    ) -> Select<note::Entity> {
        let iden = source_note_alias.unwrap_or_else(|| note::Entity.table_name());
        let note_id_col = Expr::col((Alias::new(iden), note::Column::Id));

        let renote_note_tbl = Alias::new(format!("{}{}", prefix_alias, "note"));

        let column = sub_interaction_renote::Column::Renotes;
        let alias = format!("{}{}", prefix_alias, column.to_string());

        let sub_query = Query::select()
            .expr(SelectExpr {
                expr: Expr::count(Expr::col(Asterisk)),
                alias: Some(Alias::new(alias).into_iden()),
                window: None,
            })
            .from_as(note::Entity, renote_note_tbl.clone())
            .cond_where(
                Expr::col((renote_note_tbl.clone(), note::Column::RenoteId)).eq(note_id_col),
            )
            .and_where(Expr::col((renote_note_tbl.clone(), note::Column::UserId)).eq(user_id))
            .take()
            .into_sub_query_statement();

        QuerySelect::query(&mut self).expr(SimpleExpr::SubQuery(None, Box::new(sub_query)));

        self
    }
}

lazy_static! {
    static ref ALIAS_PIN: Alias = Alias::new(PIN);
    static ref ALIAS_INTERACTION_RENOTE: Alias = Alias::new(INTERACTION_RENOTE);
    static ref ALIAS_INTERACTION_REACTION: Alias = Alias::new(INTERACTION_REACTION);
    static ref ALIAS_USER: Alias = Alias::new(USER);
    static ref ALIAS_USER_AVATAR: Alias = Alias::new(USER_AVATAR);
    static ref ALIAS_USER_BANNER: Alias = Alias::new(USER_BANNER);
    static ref ALIAS_REPLY: Alias = Alias::new(REPLY);
    static ref ALIAS_REPLY_INTERACTION_RENOTE: Alias = Alias::new(REPLY_INTERACTION_RENOTE);
    static ref ALIAS_REPLY_INTERACTION_REACTION: Alias = Alias::new(REPLY_INTERACTION_REACTION);
    static ref ALIAS_REPLY_USER: Alias = Alias::new(REPLY_USER);
    static ref ALIAS_REPLY_USER_AVATAR: Alias = Alias::new(REPLY_USER_AVATAR);
    static ref ALIAS_REPLY_USER_BANNER: Alias = Alias::new(REPLY_USER_BANNER);
    static ref ALIAS_RENOTE: Alias = Alias::new(RENOTE);
    static ref ALIAS_RENOTE_INTERACTION_RENOTE: Alias = Alias::new(RENOTE_INTERACTION_RENOTE);
    static ref ALIAS_RENOTE_INTERACTION_REACTION: Alias = Alias::new(RENOTE_INTERACTION_REACTION);
    static ref ALIAS_RENOTE_USER: Alias = Alias::new(RENOTE_USER);
    static ref ALIAS_RENOTE_USER_AVATAR: Alias = Alias::new(RENOTE_USER_AVATAR);
    static ref ALIAS_RENOTE_USER_BANNER: Alias = Alias::new(RENOTE_USER_BANNER);
}

fn range_into_expr(filter: &RangeFilter) -> SimpleExpr {
    match filter {
        RangeFilter::TimeStart(start) => note::Column::CreatedAt.gte(*start),
        RangeFilter::TimeRange(range) => {
            note::Column::CreatedAt.between(*range.start(), *range.end())
        }
        RangeFilter::TimeEnd(end) => note::Column::CreatedAt.lt(*end),
    }
}

fn ids_into_expr(ids: &Vec<String>) -> SimpleExpr {
    if ids.len() == 1 {
        note::Column::Id.eq(&ids[0])
    } else {
        note::Column::Id.is_in(ids)
    }
}

impl NoteResolver {
    pub fn new(db: CalckeyModel) -> Self {
        NoteResolver { db }
    }

    pub async fn get_one(
        &self,
        options: &NoteResolveOptions,
    ) -> Result<Option<NoteData>, CalckeyDbError> {
        let select = self.resolve(options);
        let visibility_filter = options.visibility_filter.with_note_and_user_tables(None);
        let time_filter = options.time_range.as_ref().map(range_into_expr);
        let id_filter = options.ids.as_ref().map(ids_into_expr);

        let notes = select
            .filter(visibility_filter)
            .apply_if(id_filter, Select::<note::Entity>::filter)
            .apply_if(time_filter, Select::<note::Entity>::filter)
            .into_model::<NoteData>()
            .one(self.db.inner())
            .await?;

        Ok(notes)
    }

    pub async fn get_many(
        &self,
        options: &NoteResolveOptions,
    ) -> Result<Vec<NoteData>, CalckeyDbError> {
        let select = self.resolve(options);
        let visibility_filter = options.visibility_filter.with_note_and_user_tables(None);
        let time_filter = options.time_range.as_ref().map(range_into_expr);
        let id_filter = options.ids.as_ref().map(ids_into_expr);

        let notes = select
            .filter(visibility_filter)
            .apply_if(id_filter, Select::<note::Entity>::filter)
            .apply_if(time_filter, Select::<note::Entity>::filter)
            .apply_if(options.only_pins_from.as_deref(), |s, _| {
                s.order_by_desc(Expr::col((
                    ALIAS_PIN.clone(),
                    user_note_pining::Column::CreatedAt,
                )))
            })
            .into_model::<NoteData>()
            .all(self.db.inner())
            .await?;

        Ok(notes)
    }

    pub fn resolve(&self, options: &NoteResolveOptions) -> Select<note::Entity> {
        let mut select = note::Entity::find().add_aliased_columns(Some(USER), user::Entity);

        if let Some(pins_user) = options.only_pins_from.clone() {
            select = select.join_as(
                JoinType::InnerJoin,
                note::Relation::UserNotePining
                    .def()
                    .on_condition(move |left, _right| {
                        Expr::col((left, note::Column::UserId))
                            .eq(&pins_user)
                            .into_condition()
                    }),
                ALIAS_PIN.clone(),
            )
        }

        if let Some(user_id) = &options.with_interactions_from {
            select = select
                .add_sub_select_reaction(None, INTERACTION_REACTION, user_id)
                .add_sub_select_renote(None, INTERACTION_RENOTE, user_id);
        }

        if options.with_user {
            select = select
                .add_aliased_columns(Some(USER_AVATAR), drive_file::Entity)
                .add_aliased_columns(Some(USER_BANNER), drive_file::Entity);
        }

        if options.with_reply_target {
            select = select
                .add_aliased_columns(Some(REPLY), note::Entity)
                .add_aliased_columns(Some(REPLY_USER), user::Entity);

            if let Some(user_id) = &options.with_interactions_from {
                select = select
                    .add_sub_select_reaction(Some(REPLY), REPLY_INTERACTION_REACTION, user_id)
                    .add_sub_select_renote(Some(REPLY), REPLY_INTERACTION_RENOTE, user_id);
            }

            if options.with_user {
                select = select
                    .add_aliased_columns(Some(REPLY_USER_AVATAR), drive_file::Entity)
                    .add_aliased_columns(Some(REPLY_USER_BANNER), drive_file::Entity);
            }
        }

        if options.with_renote_target {
            select = select
                .add_aliased_columns(Some(RENOTE), note::Entity)
                .add_aliased_columns(Some(RENOTE_USER), user::Entity);

            if let Some(user_id) = &options.with_interactions_from {
                select = select
                    .add_sub_select_reaction(Some(RENOTE), RENOTE_INTERACTION_REACTION, user_id)
                    .add_sub_select_renote(Some(RENOTE), RENOTE_INTERACTION_RENOTE, user_id);
            }

            if options.with_user {
                select = select
                    .add_aliased_columns(Some(RENOTE_USER_AVATAR), drive_file::Entity)
                    .add_aliased_columns(Some(RENOTE_USER_BANNER), drive_file::Entity)
            }
        }

        if options.with_reply_target {
            select = select
                .join_as(
                    JoinType::LeftJoin,
                    note::Relation::SelfRef2.def(),
                    ALIAS_REPLY.clone(),
                )
                .join_as(
                    JoinType::LeftJoin,
                    note::Relation::User.with_alias(ALIAS_REPLY.clone()),
                    ALIAS_REPLY_USER.clone(),
                );
        }

        if options.with_renote_target {
            select = select
                .join_as(
                    JoinType::LeftJoin,
                    note::Relation::SelfRef1.def(),
                    ALIAS_RENOTE.clone(),
                )
                .join_as(
                    JoinType::LeftJoin,
                    note::Relation::User.with_alias(ALIAS_RENOTE.clone()),
                    ALIAS_RENOTE_USER.clone(),
                );
        }

        select = select.join_as(
            JoinType::InnerJoin,
            note::Relation::User.def(),
            ALIAS_USER.clone(),
        );

        if options.with_user {
            select = select
                .join_as(
                    JoinType::LeftJoin,
                    user::Relation::DriveFile2.with_alias(ALIAS_USER.clone()),
                    ALIAS_USER_AVATAR.clone(),
                )
                .join_as(
                    JoinType::LeftJoin,
                    user::Relation::DriveFile1.with_alias(ALIAS_USER.clone()),
                    ALIAS_USER_BANNER.clone(),
                );

            if options.with_reply_target {
                select = select
                    .join_as(
                        JoinType::LeftJoin,
                        user::Relation::DriveFile2.with_alias(ALIAS_REPLY_USER.clone()),
                        ALIAS_REPLY_USER_AVATAR.clone(),
                    )
                    .join_as(
                        JoinType::LeftJoin,
                        user::Relation::DriveFile1.with_alias(ALIAS_REPLY_USER.clone()),
                        ALIAS_REPLY_USER_BANNER.clone(),
                    );
            }

            if options.with_renote_target {
                select = select
                    .join_as(
                        JoinType::LeftJoin,
                        user::Relation::DriveFile2.with_alias(ALIAS_RENOTE_USER.clone()),
                        ALIAS_RENOTE_USER_AVATAR.clone(),
                    )
                    .join_as(
                        JoinType::LeftJoin,
                        user::Relation::DriveFile1.with_alias(ALIAS_RENOTE_USER.clone()),
                        ALIAS_RENOTE_USER_BANNER.clone(),
                    );
            }
        }

        select
    }
}