use std::{collections::HashSet, fmt::Display}; use sea_orm::{DeriveIden, Statement}; use sea_orm_migration::prelude::{extension::postgres::IntoTypeRef, *}; use crate::migrations::extension::postgres::Type; #[derive(DeriveIden)] pub enum GeneralIds { CreatedAt, UpdatedAt, } #[derive(DeriveIden)] pub enum Subscribers { Table, Id, Pid, DisplayName, DownloaderId, BangumiConf, } #[derive(DeriveIden)] pub enum Subscriptions { Table, Id, SubscriberId, DisplayName, Category, SourceUrl, Aggregate, Enabled, } #[derive(DeriveIden)] pub enum Bangumi { Table, Id, SubscriptionId, DisplayName, OfficialTitle, Fansub, Season, Filter, PosterLink, SavePath, LastEp, BangumiConfOverride, } #[derive(DeriveIden)] pub enum Episodes { Table, Id, OriginTitle, OfficialTitle, DisplayName, NameZh, NameJp, NameEn, SNameZh, SNameJp, SNameEn, BangumiId, DownloadId, SavePath, Resolution, Season, SeasonRaw, Fansub, PosterLink, HomePage, Subtitle, Source, } #[derive(DeriveIden)] pub enum Downloads { Table, Id, SubscriptionId, OriginTitle, DisplayName, Status, CurrSize, AllSize, Mime, Url, HomePage, } #[derive(DeriveIden)] pub enum Downloaders { Table, Id, Category, Endpoint, Password, Username, SubscriberId, SavePath, } #[async_trait::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, T: Display + Send, 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, T: Display + Send, 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::async_trait] impl<'c> CustomSchemaManagerExt for SchemaManager<'c> { 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, T: Display + Send, 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(|v| Alias::new(v.to_string())) .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, T: Display + Send, 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.to_string())) .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.to_string())); } self.alter_type(type_alter.to_owned()).await?; Ok(()) } async fn drop_postgres_enum_for_active_enum( &self, enum_name: E, ) -> Result<(), DbErr> { 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) } }