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, } #[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, } #[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, pub interaction_user_reaction: Option, pub user: user::Model, pub avatar: Option, pub banner: Option, pub reply: Option>, pub renote: Option>, } 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 { let reply = note::Model::from_query_result_optional(res, REPLY)? .map::, _>(|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::, _>(|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) -> SimpleExpr; } pub struct NoteResolveOptions { pub ids: Option>, pub visibility_filter: Box, pub time_range: Option, pub with_user: bool, pub with_reply_target: bool, pub with_renote_target: bool, pub with_interactions_from: Option, // User ID pub only_pins_from: Option, // User ID } trait SelectColumnsExt { fn add_aliased_columns(self, alias: Option<&str>, entity: T) -> Self; fn add_sub_select_reaction( self, source_note_alias: Option<&str>, alias: &str, user_id: &str, ) -> Select; fn add_sub_select_renote( self, source_note_alias: Option<&str>, alias: &str, user_id: &str, ) -> Select; } 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 } fn add_sub_select_reaction( mut self: Select, source_note_alias: Option<&str>, prefix_alias: &str, user_id: &str, ) -> Select { 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, source_note_alias: Option<&str>, prefix_alias: &str, user_id: &str, ) -> Select { 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) -> 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, 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::::filter) .apply_if(time_filter, Select::::filter) .into_model::() .one(self.db.inner()) .await?; Ok(notes) } pub async fn get_many( &self, options: &NoteResolveOptions, ) -> Result, 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::::filter) .apply_if(time_filter, Select::::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::() .all(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 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 } }