use std::fmt; use std::sync::Arc; #[cfg(feature = "postgres")] use super::PostgresBackend; use crate::repository::announcements::AnnouncementWriteRepository; use crate::repository::auth::AuthApiKeyWriteRepository; use crate::repository::auth_modules::AuthModuleWriteRepository; use crate::repository::background_tasks::BackgroundTaskWriteRepository; use crate::repository::candidates::RequestCandidateWriteRepository; use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository; use crate::repository::global_models::GlobalModelWriteRepository; use crate::repository::management_tokens::ManagementTokenWriteRepository; use crate::repository::oauth_providers::OAuthProviderWriteRepository; use crate::repository::pool_scores::PoolMemberScoreWriteRepository; use crate::repository::provider_catalog::ProviderCatalogWriteRepository; use crate::repository::proxy_nodes::ProxyNodeWriteRepository; use crate::repository::quota::ProviderQuotaWriteRepository; use crate::repository::routing_profiles::RoutingGroupWriteRepository; use crate::repository::settlement::SettlementWriteRepository; use crate::repository::usage::UsageWriteRepository; use crate::repository::video_tasks::VideoTaskWriteRepository; use crate::repository::wallet::WalletWriteRepository; #[derive(Clone, Default)] pub struct DataWriteRepositories { announcements: Option>, auth_api_keys: Option>, auth_modules: Option>, background_tasks: Option>, request_candidates: Option>, gemini_file_mappings: Option>, global_models: Option>, management_tokens: Option>, oauth_providers: Option>, pool_scores: Option>, proxy_nodes: Option>, provider_catalog: Option>, provider_quotas: Option>, routing_groups: Option>, settlement: Option>, usage: Option>, video_tasks: Option>, wallets: Option>, } impl fmt::Debug for DataWriteRepositories { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("DataWriteRepositories") .field("has_announcements", &self.announcements.is_some()) .field("has_auth_api_keys", &self.auth_api_keys.is_some()) .field("has_auth_modules", &self.auth_modules.is_some()) .field("has_background_tasks", &self.background_tasks.is_some()) .field("has_request_candidates", &self.request_candidates.is_some()) .field( "has_gemini_file_mappings", &self.gemini_file_mappings.is_some(), ) .field("has_global_models", &self.global_models.is_some()) .field("has_management_tokens", &self.management_tokens.is_some()) .field("has_oauth_providers", &self.oauth_providers.is_some()) .field("has_pool_scores", &self.pool_scores.is_some()) .field("has_proxy_nodes", &self.proxy_nodes.is_some()) .field("has_provider_catalog", &self.provider_catalog.is_some()) .field("has_provider_quotas", &self.provider_quotas.is_some()) .field("has_routing_groups", &self.routing_groups.is_some()) .field("has_settlement", &self.settlement.is_some()) .field("has_usage", &self.usage.is_some()) .field("has_video_tasks", &self.video_tasks.is_some()) .field("has_wallets", &self.wallets.is_some()) .finish() } } impl DataWriteRepositories { pub(crate) fn from_backends( #[cfg(feature = "postgres")] postgres: Option<&PostgresBackend>, ) -> Self { let mut repositories = Self::default(); #[cfg(feature = "postgres")] if let Some(postgres) = postgres { repositories.install_postgres(postgres); } repositories } #[cfg(feature = "postgres")] fn install_postgres(&mut self, backend: &PostgresBackend) { if self.announcements.is_none() { self.announcements = Some(PostgresBackend::announcement_write_repository(backend)); } if self.auth_api_keys.is_none() { self.auth_api_keys = Some(PostgresBackend::auth_api_key_write_repository(backend)); } if self.auth_modules.is_none() { self.auth_modules = Some(PostgresBackend::auth_module_write_repository(backend)); } if self.background_tasks.is_none() { self.background_tasks = Some(PostgresBackend::background_task_write_repository(backend)); } if self.request_candidates.is_none() { self.request_candidates = Some(PostgresBackend::request_candidate_write_repository(backend)); } if self.gemini_file_mappings.is_none() { self.gemini_file_mappings = Some( PostgresBackend::gemini_file_mapping_write_repository(backend), ); } if self.global_models.is_none() { self.global_models = Some(PostgresBackend::global_model_write_repository(backend)); } if self.management_tokens.is_none() { self.management_tokens = Some(PostgresBackend::management_token_write_repository(backend)); } if self.oauth_providers.is_none() { self.oauth_providers = Some(PostgresBackend::oauth_provider_write_repository(backend)); } if self.pool_scores.is_none() { self.pool_scores = Some(PostgresBackend::pool_score_write_repository(backend)); } if self.proxy_nodes.is_none() { self.proxy_nodes = Some(PostgresBackend::proxy_node_write_repository(backend)); } if self.provider_catalog.is_none() { self.provider_catalog = Some(PostgresBackend::provider_catalog_write_repository(backend)); } if self.provider_quotas.is_none() { self.provider_quotas = Some(PostgresBackend::provider_quota_write_repository(backend)); } if self.routing_groups.is_none() { self.routing_groups = Some(PostgresBackend::routing_group_write_repository(backend)); } if self.settlement.is_none() { self.settlement = Some(PostgresBackend::settlement_write_repository(backend)); } if self.usage.is_none() { self.usage = Some(PostgresBackend::usage_write_repository(backend)); } if self.video_tasks.is_none() { self.video_tasks = Some(PostgresBackend::video_task_write_repository(backend)); } if self.wallets.is_none() { self.wallets = Some(PostgresBackend::wallet_write_repository(backend)); } } #[cfg(test)] #[cfg(feature = "postgres")] pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self { Self::from_backends(postgres) } pub fn announcements(&self) -> Option> { self.announcements.clone() } pub fn auth_api_keys(&self) -> Option> { self.auth_api_keys.clone() } pub fn auth_modules(&self) -> Option> { self.auth_modules.clone() } pub fn background_tasks(&self) -> Option> { self.background_tasks.clone() } pub fn usage(&self) -> Option> { self.usage.clone() } pub fn request_candidates(&self) -> Option> { self.request_candidates.clone() } pub fn gemini_file_mappings(&self) -> Option> { self.gemini_file_mappings.clone() } pub fn global_models(&self) -> Option> { self.global_models.clone() } pub fn management_tokens(&self) -> Option> { self.management_tokens.clone() } pub fn oauth_providers(&self) -> Option> { self.oauth_providers.clone() } pub fn pool_scores(&self) -> Option> { self.pool_scores.clone() } pub fn proxy_nodes(&self) -> Option> { self.proxy_nodes.clone() } pub fn provider_quotas(&self) -> Option> { self.provider_quotas.clone() } pub fn routing_groups(&self) -> Option> { self.routing_groups.clone() } pub fn provider_catalog(&self) -> Option> { self.provider_catalog.clone() } pub fn settlement(&self) -> Option> { self.settlement.clone() } pub fn video_tasks(&self) -> Option> { self.video_tasks.clone() } pub fn wallets(&self) -> Option> { self.wallets.clone() } pub fn has_any(&self) -> bool { self.announcements.is_some() || self.auth_api_keys.is_some() || self.auth_modules.is_some() || self.background_tasks.is_some() || self.request_candidates.is_some() || self.gemini_file_mappings.is_some() || self.global_models.is_some() || self.management_tokens.is_some() || self.oauth_providers.is_some() || self.pool_scores.is_some() || self.proxy_nodes.is_some() || self.provider_catalog.is_some() || self.provider_quotas.is_some() || self.routing_groups.is_some() || self.settlement.is_some() || self.usage.is_some() || self.video_tasks.is_some() || self.wallets.is_some() } } #[cfg(all(test, feature = "postgres"))] mod tests { use super::DataWriteRepositories; use crate::backend::PostgresBackend; use crate::driver::postgres::PostgresPoolConfig; #[tokio::test] async fn builds_write_repositories_from_postgres_backend() { let backend = PostgresBackend::from_config(PostgresPoolConfig { database_url: "postgres://localhost/aether".to_string(), min_connections: 1, max_connections: 4, acquire_timeout_ms: 1_000, idle_timeout_ms: 5_000, max_lifetime_ms: 30_000, statement_cache_capacity: 64, require_ssl: false, }) .expect("postgres backend should build"); let write = DataWriteRepositories::from_postgres(Some(&backend)); assert!(write.has_any()); assert!(write.announcements().is_some()); assert!(write.auth_api_keys().is_some()); assert!(write.auth_modules().is_some()); assert!(write.request_candidates().is_some()); assert!(write.gemini_file_mappings().is_some()); assert!(write.global_models().is_some()); assert!(write.management_tokens().is_some()); assert!(write.oauth_providers().is_some()); assert!(write.proxy_nodes().is_some()); assert!(write.provider_catalog().is_some()); assert!(write.provider_quotas().is_some()); assert!(write.settlement().is_some()); assert!(write.usage().is_some()); assert!(write.video_tasks().is_some()); assert!(write.wallets().is_some()); } }