use sea_orm::sea_query::{Alias, Expr, IntoIden, SelectExpr, SimpleExpr}; use sea_orm::{ ColumnTrait, DbErr, EntityTrait, FromQueryResult, Iden, Iterable, JoinType, QueryFilter, QueryResult, QuerySelect, QueryTrait, RelationTrait, Select, }; use serde::{Deserialize, Serialize}; use ck::{drive_file, note, user}; use magnetar_sdk::types::RangeFilter; 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, pub banner: Option, pub reply: Option>, pub renote: Option>, } 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 { let reply = note::Model::from_query_result_optional(res, REPLY)? .map::, _>(|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::, _>(|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, user: Option) -> SimpleExpr; } pub struct NoteResolveOptions { visibility_filter: Box, time_range: Option, with_user: bool, with_reply_target: bool, with_renote_target: bool, } trait SelectColumnsExt { fn add_aliased_columns(self, alias: Option<&str>, entity: T) -> Self; } impl SelectColumnsExt for Select { fn add_aliased_columns( mut self: Select, alias: Option<&str>, entity: T, ) -> Select { 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 = Lazy::new(|| Alias::new(USER)); const ALIAS_USER_AVATAR: Lazy = Lazy::new(|| Alias::new(USER_AVATAR)); const ALIAS_USER_BANNER: Lazy = Lazy::new(|| Alias::new(USER_BANNER)); const ALIAS_REPLY: Lazy = Lazy::new(|| Alias::new(REPLY)); const ALIAS_REPLY_USER: Lazy = Lazy::new(|| Alias::new(REPLY_USER)); const ALIAS_REPLY_USER_AVATAR: Lazy = Lazy::new(|| Alias::new(REPLY_USER_AVATAR)); const ALIAS_REPLY_USER_BANNER: Lazy = Lazy::new(|| Alias::new(REPLY_USER_BANNER)); const ALIAS_RENOTE: Lazy = Lazy::new(|| Alias::new(RENOTE)); const ALIAS_RENOTE_USER: Lazy = Lazy::new(|| Alias::new(RENOTE_USER)); const ALIAS_RENOTE_USER_AVATAR: Lazy = Lazy::new(|| Alias::new(RENOTE_USER_AVATAR)); const ALIAS_RENOTE_USER_BANNER: Lazy = Lazy::new(|| Alias::new(RENOTE_USER_BANNER)); impl NoteResolver { pub async fn get_one( &self, options: &NoteResolveOptions, ) -> Result, CalckeyDbError> { let select = self.resolve(options); let visibility_filter = options .visibility_filter .with_note_and_user_tables(None, Some(ALIAS_USER.clone())); let time_filter = options.time_range.as_ref().map(|f| match f { RangeFilter::TimeStart(start) => note::Column::CreatedAt.gte(start.clone()), RangeFilter::TimeRange(range) => { note::Column::CreatedAt.between(range.start().clone(), range.end().clone()) } RangeFilter::TimeEnd(end) => note::Column::CreatedAt.lt(end.clone()), }); let notes = select .filter(visibility_filter) .apply_if(time_filter, Select::::filter) .into_model::() .one(self.db.inner()) .await?; Ok(notes) } pub fn resolve(&self, options: &NoteResolveOptions) -> Select { 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 } }