magnetar/ext_calckey_model/src/note_model/mod.rs

pub mod data;

use ext_calckey_model_migration::SelectStatement;
use sea_orm::sea_query::{Asterisk, Expr, IntoIden, Query, SelectExpr, SimpleExpr};
use sea_orm::{
    Condition, EntityTrait, Iden, JoinType, QueryFilter, QueryOrder, QuerySelect, QueryTrait,
    Select,
};

use ck::{note, note_reaction, user, user_note_pining};
use data::{sub_interaction_reaction, sub_interaction_renote, NoteData};
use magnetar_sdk::types::SpanFilter;

use crate::model_ext::{
    join_columns_default, AliasColumnExt, AliasSourceExt, AliasSuffixExt, CursorPaginationExt,
    EntityPrefixExt, MagIden, SelectColumnsExt,
};
use crate::user_model::{UserResolveOptions, UserResolver};
use crate::{CalckeyDbError, CalckeyModel};

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.";

#[derive(Clone)]
pub struct NoteResolver {
    db: CalckeyModel,
    user_resolver: UserResolver,
}

pub trait NoteVisibilityFilterFactory: Send + Sync {
    fn with_note_and_user_tables(&self, note: &dyn Iden) -> 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_reply_target: bool,
    pub with_renote_target: bool,
    pub with_interactions_from: Option<String>, // User ID
    pub only_pins_from: Option<String>,         // User ID
    pub user_options: UserResolveOptions,
}

trait SelectNoteInteractionsExt {
    fn add_sub_select_interaction_reaction(
        &mut self,
        note_tbl: &MagIden,
        user_id: &str,
    ) -> &mut Self;

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

impl SelectNoteInteractionsExt for SelectStatement {
    fn add_sub_select_interaction_reaction(
        &mut self,
        note_tbl: &MagIden,
        user_id: &str,
    ) -> &mut Self {
        let note_id_col = note_tbl.col(note::Column::Id);
        let alias = note_tbl
            .join_str(INTERACTION_REACTION)
            .join(&sub_interaction_reaction::Column::ReactionName);

        let sub_query = Query::select()
            .expr(SelectExpr {
                expr: Expr::col(note_reaction::Column::Reaction).into(),
                alias: Some(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,
        note_tbl: &MagIden,
        user_id: &str,
    ) -> &mut Self {
        let interaction_tbl_prefix = note_tbl.join_str(INTERACTION_RENOTE);
        let renote_note_tbl = interaction_tbl_prefix.join_str("note");
        let alias = interaction_tbl_prefix.join(&sub_interaction_renote::Column::Renotes);

        let sub_query = Query::select()
            .expr(SelectExpr {
                expr: Expr::count(Expr::col(Asterisk)),
                alias: Some(alias.into_iden()),
                window: None,
            })
            .from_as(note::Entity, renote_note_tbl.clone().into_iden())
            .cond_where(
                renote_note_tbl
                    .col(note::Column::RenoteId)
                    .eq(note_tbl.col(note::Column::Id)),
            )
            .and_where(renote_note_tbl.col(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 = note::Entity.base_prefix().col(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(&note::Entity.base_prefix());

        let id_filter = options.ids.as_ref().map(ids_into_expr);

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

        Ok(note)
    }

    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(&note::Entity.base_prefix());
        let id_filter = options.ids.as_ref().map(ids_into_expr);

        let notes_select = select
            .filter(visibility_filter)
            .apply_if(id_filter, Select::<note::Entity>::filter)
            .apply_if(options.only_pins_from.as_deref(), |s, _| {
                s.order_by_desc(Expr::col((
                    note::Entity.base_prefix().into_iden().join_str(PIN),
                    user_note_pining::Column::CreatedAt,
                )))
            });

        let notes = if let Some(pagination) = &options.time_range {
            notes_select
                .get_paginated_model::<NoteData, _, _>(
                    self.db.inner(),
                    (note::Column::CreatedAt, note::Column::Id),
                    pagination,
                    &mut None,
                    &mut None,
                    options.limit.unwrap_or(20),
                )
                .await?
        } else {
            notes_select
                .apply_if(options.limit, Select::<note::Entity>::limit)
                .into_model::<NoteData>()
                .all(self.db.inner())
                .await?
        };

        Ok(notes)
    }

    /*
    pub async fn get_children(&self, options: &NoteResolveOptions) -> Select<note::Entity> {
        let max_breadth: usize = 10;

        let nth_child = Alias::new("nth_child");

        let mut select = Query::select();
        select.distinct();
        select.columns([Alias::new("id")]);
        select.from(cte);
        select.and_where(Expr::col(nth_child).lt(max_breadth))
    }

     */

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

        // Add the note's author
        let user_tbl = note_tbl.join_str(USER);
        q.add_aliased_columns::<user::Entity>(&user_tbl);
        q.join_columns(
            JoinType::LeftJoin,
            note::Relation::User.with_from_alias(note_tbl),
            &user_tbl,
        );

        // user_resolver.resolve(q, &user_tbl, &options.user_options);

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

        // Recursively attach reply parents
        if reply_depth > 0 {
            let reply_tbl = note_tbl.join_str(REPLY);
            let visibility_filter = options
                .visibility_filter
                .with_note_and_user_tables(&reply_tbl);
            let condition = Condition::all()
                .add(join_columns_default(
                    note::Relation::SelfRef2.with_alias(note_tbl, &reply_tbl),
                ))
                .add(visibility_filter);
            q.join_columns_on(
                JoinType::LeftJoin,
                note::Relation::SelfRef2.with_from_alias(note_tbl),
                &reply_tbl,
                condition,
            );

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

        // Recursively attach renote/quote targets
        if renote_depth > 0 {
            let renote_tbl = note_tbl.join_str(RENOTE);
            let visibility_filter = options
                .visibility_filter
                .with_note_and_user_tables(&renote_tbl);
            let condition = Condition::all()
                .add(join_columns_default(
                    note::Relation::SelfRef1.with_alias(note_tbl, &renote_tbl),
                ))
                .add(visibility_filter);
            q.join_columns_on(
                JoinType::LeftJoin,
                note::Relation::SelfRef1.with_from_alias(note_tbl),
                &renote_tbl,
                condition,
            );

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

    pub fn resolve(&self, options: &NoteResolveOptions) -> Select<note::Entity> {
        let note_tbl = note::Entity.base_prefix();

        let mut select = Query::select();
        select.from_as(note::Entity, note_tbl.clone().into_iden());

        if let Some(pins_user) = &options.only_pins_from {
            select
                .join_columns(
                    JoinType::InnerJoin,
                    note::Relation::UserNotePining.with_from_alias(&note_tbl),
                    &note_tbl.join_str(PIN),
                )
                .and_where(note_tbl.col(user_note_pining::Column::UserId).eq(pins_user));
        }

        self.attach_note(
            &mut select,
            &note_tbl,
            options.with_reply_target.then_some(1).unwrap_or_default(),
            options.with_renote_target.then_some(2).unwrap_or_default(),
            options,
            &self.user_resolver,
        );

        let mut ret_select = note::Entity::find();
        *QuerySelect::query(&mut ret_select) = select;
        ret_select
    }
}