magnetar/ext_calckey_model/src/note_model.rs

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

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

use crate::model_ext::{
    joined_prefix, joined_prefix_alias, joined_prefix_str, EntityPrefixExt, SelectColumnsExt,
};
use crate::paginated::PaginatedModel;
use crate::user_model::{UserData, UserResolver};
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: UserData,
    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 REPLY: &str = "reply.";
const RENOTE: &str = "renote.";

impl FromQueryResult for NoteData {
    fn from_query_result(res: &QueryResult, prefix: &str) -> Result<Self, DbErr> {
        let fallback = format!("{}.", note::Entity.table_name());
        let prefix = if prefix.is_empty() { &fallback } else { prefix };

        Ok(NoteData {
            note: note::Model::from_query_result(res, prefix)?,
            interaction_user_renote: sub_interaction_renote::Model::from_query_result_optional(
                res,
                &joined_prefix_str(prefix, INTERACTION_RENOTE),
            )?,
            interaction_user_reaction: sub_interaction_reaction::Model::from_query_result_optional(
                res,
                &joined_prefix_str(prefix, INTERACTION_REACTION),
            )?,
            user: UserData::from_query_result(res, &joined_prefix_str(prefix, USER))?,
            reply: NoteData::from_query_result_optional(res, &joined_prefix_str(prefix, REPLY))?
                .map(Box::new),
            renote: NoteData::from_query_result_optional(res, &joined_prefix_str(prefix, RENOTE))?
                .map(Box::new),
        })
    }
}

pub struct NoteResolver {
    db: CalckeyModel,
    user_resolver: UserResolver,
}

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<SpanFilter>,
    pub limit: Option<u64>,
    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 SelectNoteInteractionsExt {
    fn add_sub_select_interaction_reaction(
        &mut self,
        source_note_alias: &str,
        alias: &str,
        user_id: &str,
    ) -> &mut Self;

    fn add_sub_select_interaction_renote(
        &mut self,
        source_note_alias: &str,
        alias: &str,
        user_id: &str,
    ) -> &mut Self;
}

impl SelectNoteInteractionsExt for SelectStatement {
    fn add_sub_select_interaction_reaction(
        &mut self,
        iden: &str,
        prefix_alias: &str,
        user_id: &str,
    ) -> &mut Self {
        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();

        self.expr(SimpleExpr::SubQuery(None, Box::new(sub_query)));
        self
    }

    fn add_sub_select_interaction_renote(
        &mut self,
        iden: &str,
        prefix_alias: &str,
        user_id: &str,
    ) -> &mut Self {
        let note_id_col = Expr::col((Alias::new(iden), note::Column::Id));

        let renote_note_tbl = joined_prefix_alias(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();

        self.expr(SimpleExpr::SubQuery(None, Box::new(sub_query)));
        self
    }
}

fn ids_into_expr(ids: &Vec<String>) -> SimpleExpr {
    let id_col = Expr::col((note::Entity.base_prefix(), note::Column::Id));

    if ids.len() == 1 {
        id_col.eq(&ids[0])
    } else {
        id_col.is_in(ids)
    }
}

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

    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()
            .and_then(note::Model::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()
            .and_then(note::Model::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((
                    joined_prefix(&note::Entity.base_prefix_str(), PIN),
                    user_note_pining::Column::CreatedAt,
                )))
            })
            .apply_if(options.limit, Select::<note::Entity>::limit)
            .into_model::<NoteData>()
            .all(self.db.inner())
            .await?;

        Ok(notes)
    }

    fn attach_note(
        &self,
        q: &mut SelectStatement,
        prefix: &str,
        reply_depth: usize,
        renote_depth: usize,
        options: &NoteResolveOptions,
        user_resolver: &UserResolver,
    ) {
        q.add_aliased_columns::<note::Entity>(prefix);

        // Add the note's author
        q.add_aliased_columns::<user::Entity>(&joined_prefix_str(prefix, USER));
        q.join_columns(
            JoinType::LeftJoin,
            note::Relation::User.with_alias(Alias::new(prefix).into_iden()),
            joined_prefix(prefix, USER),
        );

        // Interactions like renotes or reactions from the specified user
        if let Some(user_id) = &options.with_interactions_from {
            q.add_sub_select_interaction_reaction(
                prefix,
                &joined_prefix_str(prefix, INTERACTION_REACTION),
                user_id,
            );
            q.add_sub_select_interaction_renote(
                prefix,
                &joined_prefix_str(prefix, INTERACTION_RENOTE),
                user_id,
            );
        }

        // Recursively attach reply parents
        if reply_depth > 0 {
            q.join_columns(
                JoinType::LeftJoin,
                note::Relation::SelfRef2.with_alias(Alias::new(prefix).into_iden()),
                joined_prefix(prefix, REPLY),
            );

            self.attach_note(
                q,
                &joined_prefix_str(prefix, REPLY),
                reply_depth - 1,
                renote_depth,
                options,
                user_resolver,
            );
        }

        // Recursively attach renote/quote targets
        if renote_depth > 0 {
            q.join_columns(
                JoinType::LeftJoin,
                note::Relation::SelfRef1.with_alias(Alias::new(prefix).into_iden()),
                joined_prefix(prefix, RENOTE),
            );

            self.attach_note(
                q,
                &joined_prefix_str(prefix, RENOTE),
                reply_depth,
                renote_depth - 1,
                options,
                user_resolver,
            );
        }
    }

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

        let prefix = note::Entity.base_prefix_str();

        let query = QuerySelect::query(&mut select);
        query.clear_selects();
        query.from_as(note::Entity, Alias::new(&prefix));

        if let Some(pins_user) = &options.only_pins_from {
            select = select
                .join_as(
                    JoinType::InnerJoin,
                    note::Relation::UserNotePining.with_alias(Alias::new(&prefix)),
                    joined_prefix(&prefix, PIN),
                )
                .filter(
                    Expr::col((
                        joined_prefix(&prefix, PIN),
                        user_note_pining::Column::UserId,
                    ))
                    .eq(pins_user)
                    .into_condition(),
                )
        }

        self.attach_note(
            QuerySelect::query(&mut select),
            &prefix,
            options.with_reply_target.then_some(1).unwrap_or_default(),
            options.with_renote_target.then_some(1).unwrap_or_default(),
            options,
            &self.user_resolver,
        );

        select
    }
}