use std::sync::Arc; use crate::database::SqlDatabaseConfig; use crate::driver::sqlite::{SqlitePool, SqlitePoolFactory}; use crate::repository::announcements::{ AnnouncementReadRepository, AnnouncementWriteRepository, SqliteAnnouncementRepository, }; use crate::repository::audit::{AuditLogReadRepository, SqliteAuditLogReadRepository}; use crate::repository::auth::{ AuthApiKeyReadRepository, AuthApiKeyWriteRepository, SqliteAuthApiKeyReadRepository, }; use crate::repository::auth_modules::{ AuthModuleReadRepository, AuthModuleWriteRepository, SqliteAuthModuleReadRepository, SqliteAuthModuleRepository, }; use crate::repository::billing::{BillingReadRepository, SqliteBillingReadRepository}; use crate::repository::candidate_selection::{ MinimalCandidateSelectionReadRepository, SqliteMinimalCandidateSelectionReadRepository, }; use crate::repository::candidates::{ RequestCandidateReadRepository, RequestCandidateWriteRepository, SqliteRequestCandidateRepository, }; use crate::repository::gemini_file_mappings::{ GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository, SqliteGeminiFileMappingRepository, }; use crate::repository::global_models::{ GlobalModelReadRepository, GlobalModelWriteRepository, SqliteGlobalModelReadRepository, }; use crate::repository::management_tokens::{ ManagementTokenReadRepository, ManagementTokenWriteRepository, SqliteManagementTokenRepository, }; use crate::repository::oauth_providers::{ OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository, }; use crate::repository::provider_catalog::{ ProviderCatalogReadRepository, ProviderCatalogWriteRepository, SqliteProviderCatalogReadRepository, }; use crate::repository::proxy_nodes::{ ProxyNodeReadRepository, ProxyNodeWriteRepository, SqliteProxyNodeReadRepository, }; use crate::repository::quota::{ ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqliteProviderQuotaRepository, }; use crate::repository::settlement::{SettlementWriteRepository, SqliteSettlementRepository}; use crate::repository::usage::{ SqliteUsageReadRepository, SqliteUsageWriteRepository, UsageReadRepository, UsageWriteRepository, }; use crate::repository::users::{SqliteUserReadRepository, UserReadRepository}; use crate::repository::video_tasks::{ SqliteVideoTaskRepository, VideoTaskReadRepository, VideoTaskWriteRepository, }; use crate::repository::wallet::{ SqliteWalletReadRepository, WalletReadRepository, WalletWriteRepository, }; use crate::DataLayerError; #[derive(Debug, Clone)] pub struct SqliteBackend { config: SqlDatabaseConfig, pool: SqlitePool, } impl SqliteBackend { pub fn from_config(config: SqlDatabaseConfig) -> Result { let factory = SqlitePoolFactory::new(config.clone())?; let pool = factory.connect_lazy()?; Ok(Self { config, pool }) } pub fn config(&self) -> &SqlDatabaseConfig { &self.config } pub fn pool(&self) -> &SqlitePool { &self.pool } pub fn pool_clone(&self) -> SqlitePool { self.pool.clone() } pub fn auth_api_key_read_repository(&self) -> Arc { Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone())) } pub fn announcement_read_repository(&self) -> Arc { Arc::new(SqliteAnnouncementRepository::new(self.pool_clone())) } pub fn audit_log_read_repository(&self) -> Arc { Arc::new(SqliteAuditLogReadRepository::new(self.pool_clone())) } pub fn announcement_write_repository(&self) -> Arc { Arc::new(SqliteAnnouncementRepository::new(self.pool_clone())) } pub fn auth_api_key_write_repository(&self) -> Arc { Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone())) } pub fn management_token_read_repository(&self) -> Arc { Arc::new(SqliteManagementTokenRepository::new(self.pool_clone())) } pub fn management_token_write_repository(&self) -> Arc { Arc::new(SqliteManagementTokenRepository::new(self.pool_clone())) } pub fn auth_module_read_repository(&self) -> Arc { Arc::new(SqliteAuthModuleReadRepository::new(self.pool_clone())) } pub fn auth_module_write_repository(&self) -> Arc { Arc::new(SqliteAuthModuleRepository::new(self.pool_clone())) } pub fn billing_read_repository(&self) -> Arc { Arc::new(SqliteBillingReadRepository::new(self.pool_clone())) } pub fn request_candidate_read_repository(&self) -> Arc { Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone())) } pub fn request_candidate_write_repository(&self) -> Arc { Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone())) } pub fn minimal_candidate_selection_read_repository( &self, ) -> Arc { Arc::new(SqliteMinimalCandidateSelectionReadRepository::new( self.pool_clone(), )) } pub fn gemini_file_mapping_read_repository(&self) -> Arc { Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone())) } pub fn gemini_file_mapping_write_repository( &self, ) -> Arc { Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone())) } pub fn global_model_read_repository(&self) -> Arc { Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone())) } pub fn global_model_write_repository(&self) -> Arc { Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone())) } pub fn user_read_repository(&self) -> Arc { Arc::new(SqliteUserReadRepository::new(self.pool_clone())) } pub fn video_task_read_repository(&self) -> Arc { Arc::new(SqliteVideoTaskRepository::new(self.pool_clone())) } pub fn video_task_write_repository(&self) -> Arc { Arc::new(SqliteVideoTaskRepository::new(self.pool_clone())) } pub fn oauth_provider_read_repository(&self) -> Arc { Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone())) } pub fn oauth_provider_write_repository(&self) -> Arc { Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone())) } pub fn provider_catalog_read_repository(&self) -> Arc { Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone())) } pub fn provider_catalog_write_repository(&self) -> Arc { Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone())) } pub fn proxy_node_read_repository(&self) -> Arc { Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone())) } pub fn proxy_node_write_repository(&self) -> Arc { Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone())) } pub fn provider_quota_read_repository(&self) -> Arc { Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone())) } pub fn provider_quota_write_repository(&self) -> Arc { Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone())) } pub fn settlement_write_repository(&self) -> Arc { Arc::new(SqliteSettlementRepository::new(self.pool_clone())) } pub fn usage_write_repository(&self) -> Arc { Arc::new(SqliteUsageWriteRepository::new(self.pool_clone())) } pub fn usage_read_repository(&self) -> Arc { Arc::new(SqliteUsageReadRepository::new(self.pool_clone())) } pub fn wallet_read_repository(&self) -> Arc { Arc::new(SqliteWalletReadRepository::new(self.pool_clone())) } pub fn wallet_write_repository(&self) -> Arc { Arc::new(SqliteWalletReadRepository::new(self.pool_clone())) } } #[cfg(test)] mod tests { use super::SqliteBackend; use crate::lifecycle::migrate::run_sqlite_migrations; use crate::{ DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, StatsDailyAggregationInput, StatsHourlyAggregationInput, WalletDailyUsageAggregationInput, }; #[tokio::test] async fn backend_retains_config_and_pool() { let config = SqlDatabaseConfig { driver: DatabaseDriver::Sqlite, url: "sqlite://./data/aether.db".to_string(), pool: SqlPoolConfig::default(), }; let backend = SqliteBackend::from_config(config.clone()).expect("backend should build"); assert_eq!(backend.config(), &config); let _pool = backend.pool(); let _pool_clone = backend.pool_clone(); } #[tokio::test] async fn system_config_round_trips_after_sqlite_migrations() { let config = SqlDatabaseConfig { driver: DatabaseDriver::Sqlite, url: "sqlite::memory:".to_string(), pool: SqlPoolConfig { max_connections: 1, ..SqlPoolConfig::default() }, }; let backend = SqliteBackend::from_config(config).expect("backend should build"); run_sqlite_migrations(backend.pool()) .await .expect("sqlite migrations should run"); let value = serde_json::json!({"enabled": true}); let stored = backend .upsert_system_config_entry("feature.local", &value, Some("local flag")) .await .expect("system config should upsert"); assert_eq!(stored.value, value); assert_eq!( backend .find_system_config_value("feature.local") .await .expect("system config should read"), Some(value) ); assert_eq!( backend .list_system_config_entries() .await .expect("system config should list") .len(), 1 ); assert!(backend .delete_system_config_value("feature.local") .await .expect("system config should delete")); } #[tokio::test] async fn table_maintenance_runs_after_sqlite_migrations() { let config = SqlDatabaseConfig { driver: DatabaseDriver::Sqlite, url: "sqlite::memory:".to_string(), pool: SqlPoolConfig { max_connections: 1, ..SqlPoolConfig::default() }, }; let backend = SqliteBackend::from_config(config).expect("backend should build"); run_sqlite_migrations(backend.pool()) .await .expect("sqlite migrations should run"); let summary = backend .run_table_maintenance(&["usage", "request_candidates", "audit_logs"]) .await .expect("sqlite table maintenance should run"); assert_eq!(summary.attempted, 3); assert_eq!(summary.succeeded, 3); } #[tokio::test] async fn wallet_daily_usage_aggregation_uses_settlement_wallets_after_sqlite_migrations() { let config = SqlDatabaseConfig { driver: DatabaseDriver::Sqlite, url: "sqlite::memory:".to_string(), pool: SqlPoolConfig { max_connections: 1, ..SqlPoolConfig::default() }, }; let backend = SqliteBackend::from_config(config).expect("backend should build"); run_sqlite_migrations(backend.pool()) .await .expect("sqlite migrations should run"); sqlx::query( r#" INSERT INTO wallets (id, user_id, balance, gift_balance, limit_mode, created_at, updated_at) VALUES ('wallet-1', 'user-1', 10.0, 2.0, 'finite', 1, 1), ('wallet-stale', 'user-stale', 0.0, 0.0, 'finite', 1, 1) "#, ) .execute(backend.pool()) .await .expect("wallets should seed"); sqlx::query( r#" INSERT INTO "usage" ( request_id, wallet_id, provider_name, model, status, billing_status, total_cost_usd, input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens, finalized_at, created_at_unix_ms, updated_at_unix_secs ) VALUES ('request-1', 'wrong-wallet', 'provider', 'model', 'completed', 'pending', 1.25, 10, 20, 3, 4, 900, 900000, 900), ('request-2', NULL, 'provider', 'model', 'completed', 'pending', 2.00, 5, 7, 1, 2, 901, 901000, 901), ('request-zero', NULL, 'provider', 'model', 'completed', 'pending', 0.00, 100, 100, 0, 0, 902, 902000, 902), ('request-outside', NULL, 'provider', 'model', 'completed', 'pending', 9.00, 50, 50, 0, 0, 903, 903000, 903) "#, ) .execute(backend.pool()) .await .expect("usage should seed"); sqlx::query( r#" INSERT INTO usage_settlement_snapshots ( request_id, billing_status, wallet_id, finalized_at, created_at, updated_at ) VALUES ('request-1', 'settled', 'wallet-1', 1000, 1000, 1000), ('request-2', 'settled', 'wallet-1', 1100, 1100, 1100), ('request-zero', 'settled', 'wallet-1', 1150, 1150, 1150), ('request-outside', 'settled', 'wallet-1', 1200, 1200, 1200) "#, ) .execute(backend.pool()) .await .expect("settlement snapshots should seed"); sqlx::query( r#" INSERT INTO wallet_daily_usage_ledgers ( id, wallet_id, billing_date, billing_timezone, total_cost_usd, total_requests, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, aggregated_at, created_at, updated_at ) VALUES ( 'stale-ledger', 'wallet-stale', '2026-05-03', 'Asia/Shanghai', 7.0, 3, 1, 1, 0, 0, 999, 999, 999 ) "#, ) .execute(backend.pool()) .await .expect("stale ledger should seed"); let summary = backend .aggregate_wallet_daily_usage(&WalletDailyUsageAggregationInput { billing_date: "2026-05-03".to_string(), billing_timezone: "Asia/Shanghai".to_string(), window_start_unix_secs: 1000, window_end_unix_secs: 1200, aggregated_at_unix_secs: 1300, }) .await .expect("wallet daily usage aggregation should run"); assert_eq!(summary.aggregated_wallets, 1); assert_eq!(summary.deleted_stale_ledgers, 1); let ledger = sqlx::query_as::< _, ( String, String, f64, i64, i64, i64, i64, i64, Option, Option, i64, ), >( r#" SELECT id, wallet_id, total_cost_usd, total_requests, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, first_finalized_at, last_finalized_at, aggregated_at FROM wallet_daily_usage_ledgers WHERE billing_date = '2026-05-03' AND billing_timezone = 'Asia/Shanghai' "#, ) .fetch_one(backend.pool()) .await .expect("aggregated ledger should load"); assert_eq!(ledger.0.len(), 64); assert_eq!(ledger.1, "wallet-1"); assert!((ledger.2 - 3.25).abs() < f64::EPSILON); assert_eq!(ledger.3, 2); assert_eq!(ledger.4, 15); assert_eq!(ledger.5, 27); assert_eq!(ledger.6, 4); assert_eq!(ledger.7, 6); assert_eq!(ledger.8, Some(1000)); assert_eq!(ledger.9, Some(1100)); assert_eq!(ledger.10, 1300); let stale_count: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM wallet_daily_usage_ledgers WHERE id = 'stale-ledger'", ) .fetch_one(backend.pool()) .await .expect("stale ledger count should load"); assert_eq!(stale_count, 0); } #[tokio::test] async fn stats_aggregation_runs_after_sqlite_migrations() { let config = SqlDatabaseConfig { driver: DatabaseDriver::Sqlite, url: "sqlite::memory:".to_string(), pool: SqlPoolConfig { max_connections: 1, ..SqlPoolConfig::default() }, }; let backend = SqliteBackend::from_config(config).expect("backend should build"); run_sqlite_migrations(backend.pool()) .await .expect("sqlite migrations should run"); sqlx::query( r#" INSERT INTO "usage" ( request_id, user_id, api_key_id, provider_name, model, status, billing_status, status_code, error_category, input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens, total_cost_usd, actual_total_cost_usd, response_time_ms, created_at_unix_ms, updated_at_unix_secs ) VALUES ('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'completed', 'settled', 200, NULL, 10, 20, 1, 2, 0.30, 0.25, 100, 3600000, 3600), ('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'failed', 'void', 500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 300, 3610000, 3610), ('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'pending', 'pending', NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3620000, 3620), ('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'completed', 'settled', 200, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3630000, 3630) "#, ) .execute(backend.pool()) .await .expect("usage stats rows should seed"); let target_hour = chrono::DateTime::::from_timestamp(3600, 0) .expect("target hour should be valid"); let aggregated_at = chrono::DateTime::::from_timestamp(7200, 0) .expect("aggregation time should be valid"); let hourly = backend .aggregate_stats_hourly(&StatsHourlyAggregationInput { target_hour_utc: target_hour, aggregated_at, }) .await .expect("hourly stats aggregation should run") .expect("hourly bucket should aggregate"); assert_eq!(hourly.hour_utc, target_hour); assert_eq!(hourly.total_requests, 2); assert_eq!(hourly.user_rows, 2); assert_eq!(hourly.user_model_rows, 2); assert_eq!(hourly.model_rows, 2); assert_eq!(hourly.provider_rows, 2); let hourly_row = sqlx::query_as::<_, (i64, i64, i64, i64, f64)>( r#" SELECT total_requests, success_requests, error_requests, input_tokens, total_cost FROM stats_hourly WHERE hour_utc = 3600 "#, ) .fetch_one(backend.pool()) .await .expect("hourly stats row should load"); assert_eq!(hourly_row.0, 2); assert_eq!(hourly_row.1, 1); assert_eq!(hourly_row.2, 1); assert_eq!(hourly_row.3, 15); assert!((hourly_row.4 - 0.50).abs() < f64::EPSILON); let second_hourly = backend .aggregate_stats_hourly(&StatsHourlyAggregationInput { target_hour_utc: target_hour, aggregated_at, }) .await .expect("second hourly aggregation should run"); assert!(second_hourly.is_none()); let target_day = chrono::DateTime::::from_timestamp(0, 0) .expect("target day should be valid"); let daily = backend .aggregate_stats_daily(&StatsDailyAggregationInput { target_day_utc: target_day, aggregated_at, }) .await .expect("daily stats aggregation should run") .expect("daily bucket should aggregate"); assert_eq!(daily.day_start_utc, target_day); assert_eq!(daily.total_requests, 2); assert_eq!(daily.model_rows, 2); assert_eq!(daily.provider_rows, 2); assert_eq!(daily.api_key_rows, 2); assert_eq!(daily.error_rows, 1); assert_eq!(daily.user_rows, 2); let daily_row = sqlx::query_as::<_, (i64, i64, i64, i64)>( r#" SELECT total_requests, success_requests, error_requests, unique_models FROM stats_daily WHERE "date" = 0 "#, ) .fetch_one(backend.pool()) .await .expect("daily stats row should load"); assert_eq!(daily_row, (2, 1, 1, 2)); } }