Add multi-database data layer

Introduce aether-data-schema and driver-specific schema generation for Postgres, MySQL, and SQLite.

Split data backends, lifecycle, repositories, and gateway runtime integration across database drivers.

Verified with cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace.

apps/aether-gateway/src/data/state/core.rs

@@ -1,5 +1,5 @@
-use aether_data::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner};
-use aether_data::{DataBackends, DataLayerError};
+use aether_data::driver::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner};
+use aether_data::{DataBackends, DataLayerError, DatabaseDriver};
 
 use super::{GatewayDataConfig, GatewayDataState, StoredSystemConfigEntry};
 
@@ -138,6 +138,36 @@ impl GatewayDataState {
         self.backends.is_some()
     }
 
+    pub(crate) fn has_database_maintenance_backend(&self) -> bool {
+        self.backends
+            .as_ref()
+            .is_some_and(|backends| backends.has_database_maintenance_backend())
+    }
+
+    pub(crate) fn has_database_pool_summary(&self) -> bool {
+        self.backends
+            .as_ref()
+            .is_some_and(|backends| backends.has_database_pool_summary())
+    }
+
+    pub(crate) fn has_wallet_daily_usage_aggregation_backend(&self) -> bool {
+        self.backends
+            .as_ref()
+            .is_some_and(|backends| backends.has_wallet_daily_usage_aggregation_backend())
+    }
+
+    pub(crate) fn has_stats_hourly_aggregation_backend(&self) -> bool {
+        self.backends
+            .as_ref()
+            .is_some_and(|backends| backends.has_stats_hourly_aggregation_backend())
+    }
+
+    pub(crate) fn has_stats_daily_aggregation_backend(&self) -> bool {
+        self.backends
+            .as_ref()
+            .is_some_and(|backends| backends.has_stats_daily_aggregation_backend())
+    }
+
     pub(crate) fn has_auth_api_key_reader(&self) -> bool {
         self.auth_api_key_reader.is_some()
     }
@@ -158,6 +188,13 @@ impl GatewayDataState {
         self.announcement_writer.is_some()
     }
 
+    pub(crate) fn has_audit_log_reader(&self) -> bool {
+        self.backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+            .is_some()
+    }
+
     pub(crate) fn has_management_token_reader(&self) -> bool {
         self.management_token_reader.is_some()
     }
@@ -223,8 +260,7 @@ impl GatewayDataState {
             || self
                 .backends
                 .as_ref()
-                .and_then(|backends| backends.postgres())
-                .is_some()
+                .is_some_and(|backends| backends.has_system_config_backend())
     }
 
     pub(crate) fn oauth_refresh_lock_runner(&self) -> Option<RedisLockRunner> {
@@ -240,15 +276,10 @@ impl GatewayDataState {
             .and_then(|backend| backend.kv_runner(RedisKvRunnerConfig::default()).ok())
     }
 
-    pub(crate) fn postgres_pool(&self) -> Option<aether_data::postgres::PostgresPool> {
+    pub(crate) fn database_driver(&self) -> Option<DatabaseDriver> {
         self.backends
             .as_ref()
-            .and_then(|backends| backends.postgres())
-            .map(|backend| backend.pool_clone())
-    }
-
-    pub(crate) fn postgres_max_connections(&self) -> Option<u32> {
-        self.config.postgres().map(|config| config.max_connections)
+            .and_then(|backends| backends.database_driver())
     }
 
     pub(crate) fn has_provider_quota_writer(&self) -> bool {
@@ -307,14 +338,10 @@ impl GatewayDataState {
                 .get(key)
                 .map(|entry| entry.value.clone()));
         }
-        match self
-            .backends
-            .as_ref()
-            .and_then(|backends| backends.postgres())
-        {
-            Some(backend) => backend.find_system_config_value(key).await,
-            None => Ok(None),
-        }
+        let Some(backends) = self.backends.as_ref() else {
+            return Ok(None);
+        };
+        backends.find_system_config_value(key).await
     }
 
     pub(crate) async fn upsert_system_config_value(
@@ -340,14 +367,10 @@ impl GatewayDataState {
                 .cloned()
                 .collect());
         }
-        match self
-            .backends
-            .as_ref()
-            .and_then(|backends| backends.postgres())
-        {
-            Some(backend) => backend.list_system_config_entries().await,
-            None => Ok(Vec::new()),
-        }
+        let Some(backends) = self.backends.as_ref() else {
+            return Ok(Vec::new());
+        };
+        backends.list_system_config_entries().await
     }
 
     pub(crate) async fn upsert_system_config_entry(
@@ -370,23 +393,20 @@ impl GatewayDataState {
             values.insert(key.to_string(), entry.clone());
             return Ok(entry);
         }
-        match self
-            .backends
-            .as_ref()
-            .and_then(|backends| backends.postgres())
-        {
-            Some(backend) => {
-                backend
-                    .upsert_system_config_entry(key, value, description)
-                    .await
+        if let Some(backends) = self.backends.as_ref() {
+            if let Some(entry) = backends
+                .upsert_system_config_entry(key, value, description)
+                .await?
+            {
+                return Ok(entry);
             }
-            None => Ok(StoredSystemConfigEntry {
-                key: key.to_string(),
-                value: value.clone(),
-                description: description.map(ToOwned::to_owned),
-                updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
-            }),
         }
+        Ok(StoredSystemConfigEntry {
+            key: key.to_string(),
+            value: value.clone(),
+            description: description.map(ToOwned::to_owned),
+            updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
+        })
     }
 
     pub(crate) async fn delete_system_config_value(
@@ -400,25 +420,17 @@ impl GatewayDataState {
                 .remove(key)
                 .is_some());
         }
-        match self
-            .backends
-            .as_ref()
-            .and_then(|backends| backends.postgres())
-        {
-            Some(backend) => backend.delete_system_config_value(key).await,
-            None => Ok(false),
-        }
+        let Some(backends) = self.backends.as_ref() else {
+            return Ok(false);
+        };
+        backends.delete_system_config_value(key).await
     }
 
     pub(crate) async fn read_admin_system_stats(
         &self,
     ) -> Result<super::AdminSystemStats, DataLayerError> {
-        match self
-            .backends
-            .as_ref()
-            .and_then(|backends| backends.postgres())
-        {
-            Some(backend) => backend.read_admin_system_stats().await,
+        match self.backends.as_ref() {
+            Some(backends) => backends.read_admin_system_stats().await,
             None => Ok(super::AdminSystemStats::default()),
         }
     }

apps/aether-gateway/src/data/state/integrations.rs

@@ -1,6 +1,6 @@
 use aether_billing::enrich_usage_event_with_billing;
 use aether_billing::BillingModelContextLookup;
-use aether_data::redis::RedisStreamRunner;
+use aether_data::driver::redis::RedisStreamRunner;
 use aether_data::repository::audit::RequestAuditReader;
 use aether_data::repository::auth::{
     AuthApiKeyLookupKey, ResolvedAuthApiKeySnapshotReader, StoredAuthApiKeySnapshot,

apps/aether-gateway/src/data/state/mod.rs

@@ -11,12 +11,17 @@ use crate::provider_transport::{
     read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
 };
 use crate::video_tasks::LocalVideoTaskReadResponse;
-use aether_data::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner, RedisStreamRunner};
+use aether_data::driver::redis::{
+    RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner, RedisStreamRunner,
+};
 use aether_data::repository::announcements::{
     AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
     CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
 };
-use aether_data::repository::audit::RequestAuditBundle;
+use aether_data::repository::audit::{
+    AuditLogListQuery, RequestAuditBundle, StoredAdminAuditLogPage, StoredSuspiciousActivity,
+    StoredUserAuditLogPage,
+};
 use aether_data::repository::auth::{
     AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
     StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
@@ -47,7 +52,8 @@ use aether_data::repository::proxy_nodes::{
 };
 pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
 use aether_data::repository::users::{
-    StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserReadRepository,
+    StoredUserAuthRecord, StoredUserExportRow, StoredUserOAuthLinkSummary, StoredUserSummary,
+    UserReadRepository,
 };
 pub(crate) use aether_data::repository::users::{
     StoredUserPreferenceRecord, StoredUserSessionRecord,
@@ -72,8 +78,13 @@ use aether_data::repository::wallet::{
     StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
     WalletReadRepository, WalletWriteRepository,
 };
-use aether_data::{DataBackends, DataLayerError};
+use aether_data::{
+    DataBackends, DataLayerError, DatabaseMaintenanceSummary, WalletDailyUsageAggregationInput,
+    WalletDailyUsageAggregationResult,
+};
 use aether_data_contracts::repository::billing::{
+    AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome,
+    AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput,
     BillingReadRepository, StoredBillingModelContext,
 };
 use aether_data_contracts::repository::candidate_selection::{
@@ -104,8 +115,8 @@ use aether_data_contracts::repository::settlement::{
     SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
 };
 use aether_data_contracts::repository::usage::{
-    StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
-    UsageWriteRepository,
+    PendingUsageCleanupSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
+    UpsertUsageRecord, UsageReadRepository, UsageWriteRepository,
 };
 use aether_data_contracts::repository::video_tasks::{
     StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,

apps/aether-gateway/src/data/state/runtime.rs

@@ -1,36 +1,139 @@
 use super::{
     read_decision_trace, read_provider_transport_snapshot, read_request_candidate_trace,
-    AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
+    AdjustWalletBalanceInput, AdminBillingCollectorRecord, AdminBillingCollectorWriteInput,
+    AdminBillingMutationOutcome, AdminBillingPresetApplyResult, AdminBillingRuleRecord,
+    AdminBillingRuleWriteInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
     AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
-    AdminWalletRefundRequestListQuery, AnnouncementListQuery, CompleteAdminWalletRefundInput,
-    CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord,
-    CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
-    CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
-    CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError, DecisionTrace,
-    DeleteAdminRedeemCodeBatchInput, DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput,
-    FailAdminWalletRefundInput, GatewayDataState, GatewayProviderTransportSnapshot,
-    LocalVideoTaskReadResponse, ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput,
-    ProcessPaymentCallbackOutcome, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
-    RedisStreamRunner, RequestAuditBundle, RequestCandidateTrace, StoredAdminPaymentCallbackPage,
+    AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
+    CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
+    CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
+    CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
+    CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
+    DataLayerError, DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
+    DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
+    GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse,
+    ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
+    RedeemWalletCodeInput, RedeemWalletCodeOutcome, RedisStreamRunner, RequestAuditBundle,
+    RequestCandidateTrace, StoredAdminAuditLogPage, StoredAdminPaymentCallbackPage,
     StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminRedeemCodeBatch,
     StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage,
     StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
     StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
     StoredAdminWalletTransactionPage, StoredAnnouncement, StoredAnnouncementPage,
     StoredBillingModelContext, StoredProviderQuotaSnapshot, StoredProviderUsageSummary,
-    StoredRequestUsageAudit, StoredUsageSettlement, StoredUserAuthRecord, StoredUserExportRow,
-    StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger,
-    StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord,
-    UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, VideoTaskLookupKey,
-    VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount, WalletLookupKey,
-    WalletMutationOutcome,
+    StoredRequestUsageAudit, StoredSuspiciousActivity, StoredUsageSettlement,
+    StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary,
+    StoredVideoTask, StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage,
+    StoredWalletSnapshot, UpdateAnnouncementRecord, UpsertUsageRecord, UpsertVideoTask,
+    UsageSettlementInput, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
+    VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
+    WalletLookupKey, WalletMutationOutcome,
 };
 use aether_data_contracts::repository::usage::{
-    StoredUsageDailySummary, UsageAuditListQuery, UsageDailyHeatmapQuery,
+    PendingUsageCleanupSummary, StoredUsageDailySummary, UsageAuditListQuery, UsageCleanupSummary,
+    UsageCleanupWindow, UsageDailyHeatmapQuery,
 };
 use aether_video_tasks_core::read_data_backed_video_task_response;
 
 impl GatewayDataState {
+    pub(crate) async fn run_database_maintenance(
+        &self,
+        table_names: &[&str],
+    ) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
+        match &self.backends {
+            Some(backends) => backends.run_database_maintenance(table_names).await,
+            None => Ok(DatabaseMaintenanceSummary::default()),
+        }
+    }
+
+    pub(crate) async fn run_database_migrations(
+        &self,
+    ) -> Result<bool, sqlx::migrate::MigrateError> {
+        match &self.backends {
+            Some(backends) => backends.run_database_migrations().await,
+            None => Ok(false),
+        }
+    }
+
+    pub(crate) async fn run_database_backfills(&self) -> Result<bool, sqlx::migrate::MigrateError> {
+        match &self.backends {
+            Some(backends) => backends.run_database_backfills().await,
+            None => Ok(false),
+        }
+    }
+
+    pub(crate) async fn pending_database_migrations(
+        &self,
+    ) -> Result<
+        Option<Vec<aether_data::lifecycle::migrate::PendingMigrationInfo>>,
+        sqlx::migrate::MigrateError,
+    > {
+        match &self.backends {
+            Some(backends) => backends.pending_database_migrations().await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn prepare_database_for_startup(
+        &self,
+    ) -> Result<
+        Option<Vec<aether_data::lifecycle::migrate::PendingMigrationInfo>>,
+        sqlx::migrate::MigrateError,
+    > {
+        match &self.backends {
+            Some(backends) => backends.prepare_database_for_startup().await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn pending_database_backfills(
+        &self,
+    ) -> Result<
+        Option<Vec<aether_data::lifecycle::backfill::PendingBackfillInfo>>,
+        sqlx::migrate::MigrateError,
+    > {
+        match &self.backends {
+            Some(backends) => backends.pending_database_backfills().await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) fn database_pool_summary(&self) -> Option<aether_data::DatabasePoolSummary> {
+        self.backends
+            .as_ref()
+            .and_then(|backends| backends.database_pool_summary())
+    }
+
+    pub(crate) async fn aggregate_wallet_daily_usage(
+        &self,
+        input: &WalletDailyUsageAggregationInput,
+    ) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
+        match &self.backends {
+            Some(backends) => backends.aggregate_wallet_daily_usage(input).await,
+            None => Ok(WalletDailyUsageAggregationResult::default()),
+        }
+    }
+
+    pub(crate) async fn aggregate_stats_hourly(
+        &self,
+        input: &aether_data::StatsHourlyAggregationInput,
+    ) -> Result<Option<aether_data::StatsHourlyAggregationSummary>, DataLayerError> {
+        match &self.backends {
+            Some(backends) => backends.aggregate_stats_hourly(input).await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn aggregate_stats_daily(
+        &self,
+        input: &aether_data::StatsDailyAggregationInput,
+    ) -> Result<Option<aether_data::StatsDailyAggregationSummary>, DataLayerError> {
+        match &self.backends {
+            Some(backends) => backends.aggregate_stats_daily(input).await,
+            None => Ok(None),
+        }
+    }
+
     pub(crate) async fn list_announcements(
         &self,
         query: &AnnouncementListQuery,
@@ -51,6 +154,91 @@ impl GatewayDataState {
         }
     }
 
+    pub(crate) async fn list_admin_audit_logs(
+        &self,
+        query: &AuditLogListQuery,
+    ) -> Result<StoredAdminAuditLogPage, DataLayerError> {
+        let Some(repository) = self
+            .backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+        else {
+            return Ok(StoredAdminAuditLogPage {
+                items: Vec::new(),
+                total: 0,
+            });
+        };
+        repository.list_admin_audit_logs(query).await
+    }
+
+    pub(crate) async fn list_admin_suspicious_activities(
+        &self,
+        cutoff_unix_secs: u64,
+    ) -> Result<Vec<StoredSuspiciousActivity>, DataLayerError> {
+        let Some(repository) = self
+            .backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+        else {
+            return Ok(Vec::new());
+        };
+        repository
+            .list_admin_suspicious_activities(cutoff_unix_secs)
+            .await
+    }
+
+    pub(crate) async fn read_admin_user_behavior_event_counts(
+        &self,
+        user_id: &str,
+        cutoff_unix_secs: u64,
+    ) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
+        let Some(repository) = self
+            .backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+        else {
+            return Ok(std::collections::BTreeMap::new());
+        };
+        repository
+            .read_admin_user_behavior_event_counts(user_id, cutoff_unix_secs)
+            .await
+    }
+
+    pub(crate) async fn list_user_audit_logs(
+        &self,
+        user_id: &str,
+        query: &AuditLogListQuery,
+    ) -> Result<StoredUserAuditLogPage, DataLayerError> {
+        let Some(repository) = self
+            .backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+        else {
+            return Ok(StoredUserAuditLogPage {
+                items: Vec::new(),
+                total: 0,
+            });
+        };
+        repository.list_user_audit_logs(user_id, query).await
+    }
+
+    pub(crate) async fn delete_audit_logs_before(
+        &self,
+        cutoff_unix_secs: u64,
+        limit: usize,
+    ) -> Result<usize, DataLayerError> {
+        let Some(repository) = self
+            .backends
+            .as_ref()
+            .and_then(|backends| backends.read().audit_logs())
+        else {
+            return Ok(0);
+        };
+        repository
+            .delete_audit_logs_before(cutoff_unix_secs, limit)
+            .await
+    }
+
     pub(crate) async fn count_unread_active_announcements(
         &self,
         user_id: &str,
@@ -748,6 +936,44 @@ impl GatewayDataState {
         }
     }
 
+    pub(crate) async fn cleanup_stale_pending_requests(
+        &self,
+        cutoff_unix_secs: u64,
+        now_unix_secs: u64,
+        timeout_minutes: u64,
+        batch_size: usize,
+    ) -> Result<PendingUsageCleanupSummary, DataLayerError> {
+        match &self.usage_writer {
+            Some(repository) => {
+                repository
+                    .cleanup_stale_pending_requests(
+                        cutoff_unix_secs,
+                        now_unix_secs,
+                        timeout_minutes,
+                        batch_size,
+                    )
+                    .await
+            }
+            None => Ok(PendingUsageCleanupSummary::default()),
+        }
+    }
+
+    pub(crate) async fn cleanup_usage(
+        &self,
+        window: &UsageCleanupWindow,
+        batch_size: usize,
+        auto_delete_expired_keys: bool,
+    ) -> Result<UsageCleanupSummary, DataLayerError> {
+        match &self.usage_writer {
+            Some(repository) => {
+                repository
+                    .cleanup_usage(window, batch_size, auto_delete_expired_keys)
+                    .await
+            }
+            None => Ok(UsageCleanupSummary::default()),
+        }
+    }
+
     pub(crate) async fn find_request_usage_by_request_id(
         &self,
         request_id: &str,
@@ -1257,6 +1483,153 @@ impl GatewayDataState {
         }
     }
 
+    pub(crate) async fn admin_billing_enabled_default_value_exists(
+        &self,
+        api_format: &str,
+        task_type: &str,
+        dimension_name: &str,
+        existing_id: Option<&str>,
+    ) -> Result<Option<bool>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => {
+                repository
+                    .admin_billing_enabled_default_value_exists(
+                        api_format,
+                        task_type,
+                        dimension_name,
+                        existing_id,
+                    )
+                    .await
+            }
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn create_admin_billing_rule(
+        &self,
+        input: &AdminBillingRuleWriteInput,
+    ) -> Result<AdminBillingMutationOutcome<AdminBillingRuleRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => repository.create_admin_billing_rule(input).await,
+            None => Ok(AdminBillingMutationOutcome::Unavailable),
+        }
+    }
+
+    pub(crate) async fn list_admin_billing_rules(
+        &self,
+        task_type: Option<&str>,
+        is_enabled: Option<bool>,
+        page: u32,
+        page_size: u32,
+    ) -> Result<Option<(Vec<AdminBillingRuleRecord>, u64)>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => {
+                repository
+                    .list_admin_billing_rules(task_type, is_enabled, page, page_size)
+                    .await
+            }
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn find_admin_billing_rule(
+        &self,
+        rule_id: &str,
+    ) -> Result<Option<AdminBillingRuleRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => repository.find_admin_billing_rule(rule_id).await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn update_admin_billing_rule(
+        &self,
+        rule_id: &str,
+        input: &AdminBillingRuleWriteInput,
+    ) -> Result<AdminBillingMutationOutcome<AdminBillingRuleRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => repository.update_admin_billing_rule(rule_id, input).await,
+            None => Ok(AdminBillingMutationOutcome::Unavailable),
+        }
+    }
+
+    pub(crate) async fn create_admin_billing_collector(
+        &self,
+        input: &AdminBillingCollectorWriteInput,
+    ) -> Result<AdminBillingMutationOutcome<AdminBillingCollectorRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => repository.create_admin_billing_collector(input).await,
+            None => Ok(AdminBillingMutationOutcome::Unavailable),
+        }
+    }
+
+    pub(crate) async fn list_admin_billing_collectors(
+        &self,
+        api_format: Option<&str>,
+        task_type: Option<&str>,
+        dimension_name: Option<&str>,
+        is_enabled: Option<bool>,
+        page: u32,
+        page_size: u32,
+    ) -> Result<Option<(Vec<AdminBillingCollectorRecord>, u64)>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => {
+                repository
+                    .list_admin_billing_collectors(
+                        api_format,
+                        task_type,
+                        dimension_name,
+                        is_enabled,
+                        page,
+                        page_size,
+                    )
+                    .await
+            }
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn find_admin_billing_collector(
+        &self,
+        collector_id: &str,
+    ) -> Result<Option<AdminBillingCollectorRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => repository.find_admin_billing_collector(collector_id).await,
+            None => Ok(None),
+        }
+    }
+
+    pub(crate) async fn update_admin_billing_collector(
+        &self,
+        collector_id: &str,
+        input: &AdminBillingCollectorWriteInput,
+    ) -> Result<AdminBillingMutationOutcome<AdminBillingCollectorRecord>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => {
+                repository
+                    .update_admin_billing_collector(collector_id, input)
+                    .await
+            }
+            None => Ok(AdminBillingMutationOutcome::Unavailable),
+        }
+    }
+
+    pub(crate) async fn apply_admin_billing_preset(
+        &self,
+        preset: &str,
+        mode: &str,
+        collectors: &[AdminBillingCollectorWriteInput],
+    ) -> Result<AdminBillingMutationOutcome<AdminBillingPresetApplyResult>, DataLayerError> {
+        match &self.billing_reader {
+            Some(repository) => {
+                repository
+                    .apply_admin_billing_preset(preset, mode, collectors)
+                    .await
+            }
+            None => Ok(AdminBillingMutationOutcome::Unavailable),
+        }
+    }
+
     pub(crate) async fn read_request_candidate_trace(
         &self,
         request_id: &str,