use std::collections::HashSet; use async_trait::async_trait; use sea_orm::{DeriveIden, Statement}; use sea_orm_migration::{ prelude::{extension::postgres::IntoTypeRef, *}, schema::timestamp_with_time_zone, }; use crate::migrations::extension::postgres::Type; #[derive(DeriveIden)] pub enum GeneralIds { CreatedAt, UpdatedAt, } #[derive(DeriveIden)] pub enum Subscribers { Table, Id, DisplayName, DownloaderId, BangumiConf, } #[derive(DeriveIden)] pub enum Subscriptions { Table, Id, DisplayName, SubscriberId, Category, SourceUrl, Enabled, CredentialId, } #[derive(DeriveIden)] pub enum Bangumi { Table, Id, MikanBangumiId, DisplayName, SubscriberId, OriginName, Season, SeasonRaw, Fansub, MikanFansubId, Filter, RssLink, PosterLink, OriginPosterLink, /** * @deprecated */ SavePath, Homepage, BangumiType, } #[derive(DeriveIden)] pub enum SubscriptionBangumi { Table, Id, SubscriberId, SubscriptionId, BangumiId, } #[derive(DeriveIden)] pub enum Episodes { Table, Id, MikanEpisodeId, OriginName, DisplayName, BangumiId, SubscriberId, DownloadId, SavePath, Resolution, Season, SeasonRaw, Fansub, PosterLink, OriginPosterLink, EpisodeIndex, Homepage, Subtitle, Source, EpisodeType, EnclosureTorrentLink, EnclosureMagnetLink, EnclosurePubDate, EnclosureContentLength, } #[derive(DeriveIden)] pub enum SubscriptionEpisode { Table, Id, SubscriberId, SubscriptionId, EpisodeId, } #[derive(DeriveIden)] pub enum Downloads { Table, Id, OriginName, DisplayName, SubscriberId, DownloaderId, EpisodeId, Status, CurrSize, AllSize, Mime, Url, Homepage, SavePath, } #[derive(DeriveIden)] pub enum Downloaders { Table, Id, Category, Endpoint, Password, Username, SubscriberId, SavePath, } #[derive(DeriveIden)] pub enum Auth { Table, Id, Pid, SubscriberId, AuthType, } #[derive(DeriveIden)] pub enum Credential3rd { Table, Id, SubscriberId, CredentialType, Cookies, Username, Password, UserAgent, } #[derive(DeriveIden)] pub enum Feeds { Table, Id, Token, FeedType, FeedSource, SubscriberId, SubscriptionId, } macro_rules! create_postgres_enum_for_active_enum { ($manager: expr, $active_enum: expr, $($enum_value:expr),+) => { { use sea_orm::ActiveEnum; let values = [$($enum_value,)+].map(|v| ActiveEnum::to_value(&v)); ($manager).create_postgres_enum_for_active_enum($active_enum, values) } }; } pub fn timestamps_z(t: TableCreateStatement) -> TableCreateStatement { let mut t = t; t.col(timestamp_with_time_zone(GeneralIds::CreatedAt).default(Expr::current_timestamp())) .col(timestamp_with_time_zone(GeneralIds::UpdatedAt).default(Expr::current_timestamp())) .take() } pub fn table_auto_z(name: T) -> TableCreateStatement { timestamps_z(Table::create().table(name).if_not_exists().take()) } #[async_trait] pub trait CustomSchemaManagerExt { async fn create_postgres_auto_update_ts_fn(&self, col_name: &str) -> Result<(), DbErr>; async fn create_postgres_auto_update_ts_fn_for_col( &self, col: C, ) -> Result<(), DbErr> { let column_ident = col.into_iden(); self.create_postgres_auto_update_ts_fn(&column_ident.to_string()) .await?; Ok(()) } async fn create_postgres_auto_update_ts_trigger( &self, tab_name: &str, col_name: &str, ) -> Result<(), DbErr>; async fn create_postgres_auto_update_ts_trigger_for_col< T: IntoIden + 'static + Send, C: IntoIden + 'static + Send, >( &self, tab: T, col: C, ) -> Result<(), DbErr> { let column_ident = col.into_iden(); let table_ident = tab.into_iden(); self.create_postgres_auto_update_ts_trigger( &table_ident.to_string(), &column_ident.to_string(), ) .await?; Ok(()) } async fn drop_postgres_auto_update_ts_fn(&self, col_name: &str) -> Result<(), DbErr>; async fn drop_postgres_auto_update_ts_fn_for_col( &self, col: C, ) -> Result<(), DbErr> { let column_ident = col.into_iden(); self.drop_postgres_auto_update_ts_fn(&column_ident.to_string()) .await?; Ok(()) } async fn drop_postgres_auto_update_ts_trigger( &self, tab_name: &str, col_name: &str, ) -> Result<(), DbErr>; async fn drop_postgres_auto_update_ts_trigger_for_col< T: IntoIden + 'static + Send, C: IntoIden + 'static + Send, >( &self, tab: T, col: C, ) -> Result<(), DbErr> { let column_ident = col.into_iden(); let table_ident = tab.into_iden(); self.drop_postgres_auto_update_ts_trigger( &table_ident.to_string(), &column_ident.to_string(), ) .await?; Ok(()) } async fn create_postgres_enum_for_active_enum< E: IntoTypeRef + IntoIden + Send + Clone, I: IntoIterator + Send, >( &self, enum_name: E, values: I, ) -> Result<(), DbErr>; async fn add_postgres_enum_values_for_active_enum< E: IntoTypeRef + IntoIden + Send + Clone, I: IntoIterator + Send, >( &self, enum_name: E, values: I, ) -> Result<(), DbErr>; async fn drop_postgres_enum_for_active_enum( &self, enum_name: E, ) -> Result<(), DbErr>; async fn if_postgres_enum_exists( &self, enum_name: E, ) -> Result; async fn get_postgres_enum_values( &self, enum_name: E, ) -> Result, DbErr>; } #[async_trait] impl CustomSchemaManagerExt for SchemaManager<'_> { async fn create_postgres_auto_update_ts_fn(&self, col_name: &str) -> Result<(), DbErr> { let sql = format!( "CREATE OR REPLACE FUNCTION update_{col_name}_column() RETURNS TRIGGER AS $$ BEGIN \ NEW.{col_name} = current_timestamp; RETURN NEW; END; $$ language 'plpgsql';" ); self.get_connection() .execute(Statement::from_string(self.get_database_backend(), sql)) .await?; Ok(()) } async fn create_postgres_auto_update_ts_trigger( &self, tab_name: &str, col_name: &str, ) -> Result<(), DbErr> { let sql = format!( "CREATE OR REPLACE TRIGGER update_{tab_name}_{col_name}_column_trigger BEFORE UPDATE \ ON {tab_name} FOR EACH ROW EXECUTE PROCEDURE update_{col_name}_column();" ); self.get_connection() .execute(Statement::from_string(self.get_database_backend(), sql)) .await?; Ok(()) } async fn drop_postgres_auto_update_ts_fn(&self, col_name: &str) -> Result<(), DbErr> { let sql = format!("DROP FUNCTION IF EXISTS update_{col_name}_column();"); self.get_connection() .execute(Statement::from_string(self.get_database_backend(), sql)) .await?; Ok(()) } async fn drop_postgres_auto_update_ts_trigger( &self, tab_name: &str, col_name: &str, ) -> Result<(), DbErr> { let sql = format!( "DROP TRIGGER IF EXISTS update_{tab_name}_{col_name}_column_trigger ON {tab_name};" ); self.get_connection() .execute(Statement::from_string(self.get_database_backend(), sql)) .await?; Ok(()) } async fn create_postgres_enum_for_active_enum< E: IntoTypeRef + IntoIden + Send + Clone, I: IntoIterator + Send, >( &self, enum_name: E, values: I, ) -> Result<(), DbErr> { let existed = self.if_postgres_enum_exists(enum_name.clone()).await?; if !existed { let idents = values.into_iter().map(Alias::new).collect::>(); self.create_type(Type::create().as_enum(enum_name).values(idents).to_owned()) .await?; } else { self.add_postgres_enum_values_for_active_enum(enum_name, values) .await?; } Ok(()) } async fn add_postgres_enum_values_for_active_enum< E: IntoTypeRef + IntoIden + Send + Clone, I: IntoIterator + Send, >( &self, enum_name: E, values: I, ) -> Result<(), DbErr> { let exists_values = self.get_postgres_enum_values(enum_name.clone()).await?; let to_add_values = values .into_iter() .filter(|v| !exists_values.contains(v as &str)) .collect::>(); if to_add_values.is_empty() { return Ok(()); } let mut type_alter = Type::alter().name(enum_name); for v in to_add_values { type_alter = type_alter.add_value(Alias::new(v)); } self.alter_type(type_alter.to_owned()).await?; Ok(()) } async fn drop_postgres_enum_for_active_enum( &self, enum_name: E, ) -> Result<(), DbErr> { if self.if_postgres_enum_exists(enum_name.clone()).await? { self.drop_type(Type::drop().name(enum_name).to_owned()) .await?; } Ok(()) } async fn if_postgres_enum_exists( &self, enum_name: E, ) -> Result { let enum_name: String = enum_name.into_iden().to_string(); let sql = format!("SELECT 1 FROM pg_type WHERE typname = '{enum_name}'"); let result = self .get_connection() .query_one(Statement::from_string(self.get_database_backend(), sql)) .await?; Ok(result.is_some()) } async fn get_postgres_enum_values( &self, enum_name: E, ) -> Result, DbErr> { let enum_name: String = enum_name.into_iden().to_string(); let sql = format!( "SELECT pg_enum.enumlabel AS enumlabel FROM pg_type JOIN pg_enum ON pg_enum.enumtypid \ = pg_type.oid WHERE pg_type.typname = '{enum_name}';" ); let results = self .get_connection() .query_all(Statement::from_string(self.get_database_backend(), sql)) .await?; let mut items = HashSet::new(); for r in results { items.insert(r.try_get::("", "enumlabel")?); } Ok(items) } }