magnetar/ext_calckey_model/src/note_model.rs

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

use ck::{drive_file, note, user};
use once_cell::unsync::Lazy;

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NoteData {
    pub note: note::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 USER: &str = "user";
const USER_AVATAR: &str = "user.avatar";
const USER_BANNER: &str = "user.banner";
const REPLY: &str = "reply";
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_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,
                    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,
                    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, "")?,
            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 {
    fn with_note_and_user_tables(&self, note: Option<Alias>, user: Option<Alias>) -> SimpleExpr;
}

pub struct NoteResolveOptions {
    visibility_filter: Box<dyn NoteVisibilityFilterFactory>,
    with_user: bool,
    with_reply_target: bool,
    with_renote_target: bool,
}

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

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
    }
}

const ALIAS_USER: Lazy<Alias> = Lazy::new(|| Alias::new(USER));
const ALIAS_USER_AVATAR: Lazy<Alias> = Lazy::new(|| Alias::new(USER_AVATAR));
const ALIAS_USER_BANNER: Lazy<Alias> = Lazy::new(|| Alias::new(USER_BANNER));
const ALIAS_REPLY: Lazy<Alias> = Lazy::new(|| Alias::new(REPLY));
const ALIAS_REPLY_USER: Lazy<Alias> = Lazy::new(|| Alias::new(REPLY_USER));
const ALIAS_REPLY_USER_AVATAR: Lazy<Alias> = Lazy::new(|| Alias::new(REPLY_USER_AVATAR));
const ALIAS_REPLY_USER_BANNER: Lazy<Alias> = Lazy::new(|| Alias::new(REPLY_USER_BANNER));
const ALIAS_RENOTE: Lazy<Alias> = Lazy::new(|| Alias::new(RENOTE));
const ALIAS_RENOTE_USER: Lazy<Alias> = Lazy::new(|| Alias::new(RENOTE_USER));
const ALIAS_RENOTE_USER_AVATAR: Lazy<Alias> = Lazy::new(|| Alias::new(RENOTE_USER_AVATAR));
const ALIAS_RENOTE_USER_BANNER: Lazy<Alias> = Lazy::new(|| Alias::new(RENOTE_USER_BANNER));

impl NoteResolver {
    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, Some(ALIAS_USER.clone()));
        let notes = select
            .filter(visibility_filter)
            .into_model::<NoteData>()
            .one(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 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 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 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::InnerJoin,
                    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
    }
}