mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +08:00
Integrate upstream updates while preserving the local analytics dashboards and schema-only migration changes. Combine user account analysis with upstream user/group usage statistics in separate tabs, retain all migration versions, and keep the deleted audit document removed. Validation: gateway all-target cargo check, frontend type check and 57 focused tests, 48 migration tests, schema composition checks, and diff whitespace checks.
3742 lines
129 KiB
Rust
3742 lines
129 KiB
Rust
use super::{
|
|
read_decision_trace, read_provider_transport_snapshot, read_request_candidate_trace,
|
|
AdjustWalletBalanceInBatchInput, AdjustWalletBalanceInput, AdminBillingCollectorRecord,
|
|
AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, AdminBillingPresetApplyResult,
|
|
AdminBillingRuleRecord, AdminBillingRuleWriteInput, AdminPaymentOrderListQuery,
|
|
AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery,
|
|
AdminUserWalletBalanceBatchUserOutcome, AdminWalletLedgerQuery, AdminWalletListQuery,
|
|
AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
|
|
BackgroundTaskListQuery, BackgroundTaskSummary, BillingModelContextCacheKey,
|
|
BillingModelContextCacheState, BillingModelContextInflightState, BillingPlanRecord,
|
|
BillingPlanWriteInput, CompareAndSwapPaymentOrderStripeClientSecretInput,
|
|
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
|
|
CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
|
|
CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput,
|
|
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
|
|
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError,
|
|
DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
|
|
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
|
|
FailWalletRechargeCheckoutInput, GatewayDataState, GatewayProviderTransportSnapshot,
|
|
LocalVideoTaskReadResponse, PaymentGatewayConfigCasWriteInput, PaymentGatewayConfigRecord,
|
|
PaymentGatewayConfigWriteInput, PaymentGatewaySecretCasUpdate,
|
|
PrepareAdminUserWalletBalanceBatchInput, PrepareAdminUserWalletBalanceBatchOutcome,
|
|
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
|
|
ReclaimWalletRechargeCheckoutInput, ReconcileUsagePolicyCostInput, RedeemWalletCodeInput,
|
|
RedeemWalletCodeOutcome, ReleaseUsagePolicyRequestAdmissionInput, RequestAuditBundle,
|
|
RequestCandidateTrace, ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome,
|
|
ReserveUsagePolicyRequestInput, ReserveUsagePolicyRequestOutcome, StoredAdminAuditLogPage,
|
|
StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage,
|
|
StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage,
|
|
StoredAdminUserWalletBalanceBatch, StoredAdminWalletLedgerPage, StoredAdminWalletListPage,
|
|
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
|
|
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredAnnouncement,
|
|
StoredAnnouncementPage, StoredBackgroundTaskEvent, StoredBackgroundTaskRun,
|
|
StoredBackgroundTaskRunPage, StoredBillingModelContext, StoredProviderQuotaSnapshot,
|
|
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredSuspiciousActivity,
|
|
StoredUsagePolicyCostReservation, StoredUsagePolicyRequestAdmission, StoredUsageSettlement,
|
|
StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary,
|
|
StoredVideoTask, StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage,
|
|
StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput, UpdateAnnouncementRecord,
|
|
UpdateWalletRechargeCheckoutInput, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
|
|
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, UserDailyQuotaAvailabilityRecord,
|
|
UserPlanEntitlementRecord, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
|
|
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
|
|
WalletLookupKey, WalletMutationOutcome,
|
|
};
|
|
use aether_data_contracts::repository::billing::{
|
|
ProviderExpenseInput, ProviderExpensePage, ProviderExpenseQuery, ProviderExpenseRecord,
|
|
};
|
|
use aether_data_contracts::repository::usage::{
|
|
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
|
|
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
|
|
UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
|
|
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot,
|
|
UsageDailyHeatmapQuery,
|
|
};
|
|
use aether_runtime_state::RuntimeQueueStore;
|
|
use aether_video_tasks_core::{
|
|
read_data_backed_video_task_response, read_data_backed_video_task_response_for_user,
|
|
};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::time::timeout;
|
|
|
|
fn normalize_billing_context_cache_part(value: &str) -> String {
|
|
value.trim().to_string()
|
|
}
|
|
|
|
fn normalize_optional_billing_context_cache_part(value: Option<&str>) -> Option<String> {
|
|
value
|
|
.map(str::trim)
|
|
.filter(|value| !value.is_empty())
|
|
.map(ToOwned::to_owned)
|
|
}
|
|
|
|
enum BillingModelContextInflightRegistration<'a> {
|
|
Leader(BillingModelContextInflightGuard<'a>),
|
|
Follower(std::sync::Arc<BillingModelContextInflightState>),
|
|
Saturated,
|
|
}
|
|
|
|
struct BillingModelContextInflightGuard<'a> {
|
|
state: &'a GatewayDataState,
|
|
key: Option<BillingModelContextCacheKey>,
|
|
inflight_state: std::sync::Arc<BillingModelContextInflightState>,
|
|
admission: Option<tokio::sync::OwnedSemaphorePermit>,
|
|
}
|
|
|
|
impl<'a> BillingModelContextInflightGuard<'a> {
|
|
fn new(
|
|
state: &'a GatewayDataState,
|
|
key: BillingModelContextCacheKey,
|
|
inflight_state: std::sync::Arc<BillingModelContextInflightState>,
|
|
admission: tokio::sync::OwnedSemaphorePermit,
|
|
) -> Self {
|
|
Self {
|
|
state,
|
|
key: Some(key),
|
|
inflight_state,
|
|
admission: Some(admission),
|
|
}
|
|
}
|
|
|
|
fn epoch(&self) -> u64 {
|
|
self.inflight_state.epoch
|
|
}
|
|
|
|
fn finish(&mut self, error: Option<DataLayerError>) {
|
|
let removed = self.key.take().and_then(|key| {
|
|
self.state.finish_billing_model_context_inflight(
|
|
&key,
|
|
&self.inflight_state,
|
|
self.admission.take(),
|
|
)
|
|
});
|
|
self.admission.take();
|
|
if let Some(removed) = removed {
|
|
removed.complete(error.map_or(Ok(()), Err));
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for BillingModelContextInflightGuard<'_> {
|
|
fn drop(&mut self) {
|
|
self.finish(None);
|
|
}
|
|
}
|
|
|
|
impl BillingModelContextInflightState {
|
|
fn complete(&self, result: Result<(), DataLayerError>) {
|
|
if self.completion.set(result).is_ok() {
|
|
self.notify.notify_waiters();
|
|
}
|
|
}
|
|
|
|
async fn wait(&self) -> Result<(), DataLayerError> {
|
|
loop {
|
|
if let Some(result) = self.completion.get() {
|
|
return result.clone();
|
|
}
|
|
|
|
let mut notified = Box::pin(self.notify.notified());
|
|
notified.as_mut().enable();
|
|
if let Some(result) = self.completion.get() {
|
|
return result.clone();
|
|
}
|
|
notified.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for BillingModelContextCacheState {
|
|
fn default() -> Self {
|
|
Self {
|
|
entries: aether_cache::ExpiringMap::default(),
|
|
inflight: std::sync::Mutex::new(std::collections::HashMap::new()),
|
|
epoch: std::sync::atomic::AtomicU64::new(0),
|
|
mutation: std::sync::Mutex::new(()),
|
|
admission: std::sync::Arc::new(tokio::sync::Semaphore::new(
|
|
GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT,
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl GatewayDataState {
|
|
const MAINTENANCE_POOL_IDLE_RESERVE_ENV: &'static str =
|
|
"AETHER_GATEWAY_MAINTENANCE_POOL_IDLE_RESERVE";
|
|
const MAINTENANCE_POOL_PRESSURE_MAX_DEFER: Duration = Duration::from_secs(30);
|
|
const BILLING_MODEL_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(30);
|
|
const BILLING_MODEL_CONTEXT_CACHE_MAX_ENTRIES: usize = 4096;
|
|
const BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT: usize = 4096;
|
|
#[cfg(not(test))]
|
|
const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
|
|
#[cfg(test)]
|
|
const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_millis(100);
|
|
|
|
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 warm_database_pool(&self) -> Result<(), DataLayerError> {
|
|
match &self.backends {
|
|
Some(backends) => backends.warm_database_pool().await,
|
|
None => Ok(()),
|
|
}
|
|
}
|
|
|
|
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 postgres_observability_snapshot(
|
|
&self,
|
|
) -> Result<Option<aether_data::DatabasePostgresObservabilitySnapshot>, DataLayerError> {
|
|
match &self.backends {
|
|
Some(backends) => backends.postgres_observability_snapshot().await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn postgres_activity_groups(
|
|
&self,
|
|
limit: i64,
|
|
) -> Result<Vec<aether_data::DatabasePostgresActivityGroup>, DataLayerError> {
|
|
match &self.backends {
|
|
Some(backends) => backends.postgres_activity_groups(limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn database_pool_under_maintenance_pressure(&self) -> bool {
|
|
self.database_pool_summary()
|
|
.as_ref()
|
|
.is_some_and(Self::database_pool_summary_under_maintenance_pressure)
|
|
}
|
|
|
|
pub(crate) fn database_pool_summary_under_maintenance_pressure(
|
|
summary: &aether_data::DatabasePoolSummary,
|
|
) -> bool {
|
|
summary.checked_out > 0
|
|
&& Self::database_pool_available_capacity(summary)
|
|
<= Self::maintenance_pool_idle_reserve(summary)
|
|
}
|
|
|
|
pub(crate) fn database_pool_summary_under_usage_worker_pressure(
|
|
summary: &aether_data::DatabasePoolSummary,
|
|
) -> bool {
|
|
summary.checked_out > 0
|
|
&& Self::database_pool_available_capacity(summary)
|
|
<= Self::usage_worker_pool_idle_reserve(summary)
|
|
}
|
|
|
|
fn database_pool_available_capacity(summary: &aether_data::DatabasePoolSummary) -> usize {
|
|
let unopened = (summary.max_connections as usize).saturating_sub(summary.pool_size);
|
|
summary.idle.saturating_add(unopened)
|
|
}
|
|
|
|
pub(crate) fn maintenance_pool_idle_reserve(
|
|
summary: &aether_data::DatabasePoolSummary,
|
|
) -> usize {
|
|
if let Some(override_value) = std::env::var(Self::MAINTENANCE_POOL_IDLE_RESERVE_ENV)
|
|
.ok()
|
|
.and_then(|value| value.trim().parse::<usize>().ok())
|
|
{
|
|
return override_value;
|
|
}
|
|
|
|
let max_connections = summary.max_connections as usize;
|
|
if max_connections == 0 {
|
|
return 0;
|
|
}
|
|
|
|
let ten_percent_ceil = (max_connections + 9) / 10;
|
|
ten_percent_ceil.clamp(2, 10).min(max_connections)
|
|
}
|
|
|
|
fn usage_worker_pool_idle_reserve(summary: &aether_data::DatabasePoolSummary) -> usize {
|
|
if summary.max_connections <= 1 {
|
|
return 0;
|
|
}
|
|
1
|
|
}
|
|
|
|
pub(crate) fn should_defer_maintenance_for_database_pool_pressure(
|
|
&self,
|
|
deferred_since: &mut Option<Instant>,
|
|
) -> bool {
|
|
Self::should_defer_maintenance_for_pool_pressure_state(
|
|
self.database_pool_under_maintenance_pressure(),
|
|
deferred_since,
|
|
)
|
|
}
|
|
|
|
pub(crate) fn should_defer_maintenance_for_pool_pressure_state(
|
|
pool_under_pressure: bool,
|
|
deferred_since: &mut Option<Instant>,
|
|
) -> bool {
|
|
if !pool_under_pressure {
|
|
*deferred_since = None;
|
|
return false;
|
|
}
|
|
|
|
let now = Instant::now();
|
|
let since = deferred_since.get_or_insert(now);
|
|
if now.duration_since(*since) >= Self::MAINTENANCE_POOL_PRESSURE_MAX_DEFER {
|
|
*deferred_since = None;
|
|
return false;
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
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 rebuild_overview_buckets(
|
|
&self,
|
|
input: &aether_data::StatsHourlyAggregationInput,
|
|
) -> Result<usize, DataLayerError> {
|
|
match &self.backends {
|
|
Some(backends) => backends.rebuild_overview_buckets(input).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn drain_overview_dirty_events(
|
|
&self,
|
|
now: chrono::DateTime<chrono::Utc>,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.backends {
|
|
Some(backends) => backends.drain_overview_dirty_events(now).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
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,
|
|
) -> Result<StoredAnnouncementPage, DataLayerError> {
|
|
match &self.announcement_reader {
|
|
Some(repository) => repository.list_announcements(query).await,
|
|
None => Ok(StoredAnnouncementPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_user_announcements(
|
|
&self,
|
|
user_id: &str,
|
|
query: &aether_data::repository::announcements::UserAnnouncementListQuery,
|
|
) -> Result<aether_data::repository::announcements::StoredUserAnnouncementPage, DataLayerError>
|
|
{
|
|
match &self.announcement_reader {
|
|
Some(repository) => repository.list_user_announcements(user_id, query).await,
|
|
None => Ok(Default::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_announcement_by_id(
|
|
&self,
|
|
announcement_id: &str,
|
|
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
|
match &self.announcement_reader {
|
|
Some(repository) => repository.find_by_id(announcement_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
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,
|
|
now_unix_secs: u64,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.announcement_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.count_unread_active_announcements(user_id, now_unix_secs)
|
|
.await
|
|
}
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_required_unread_active_announcements(
|
|
&self,
|
|
user_id: &str,
|
|
now_unix_secs: u64,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredAnnouncement>, DataLayerError> {
|
|
match &self.announcement_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_required_unread_active_announcements(user_id, now_unix_secs, limit)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_announcement(
|
|
&self,
|
|
record: CreateAnnouncementRecord,
|
|
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
|
match &self.announcement_writer {
|
|
Some(repository) => repository.create_announcement(record).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn update_announcement(
|
|
&self,
|
|
record: UpdateAnnouncementRecord,
|
|
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
|
match &self.announcement_writer {
|
|
Some(repository) => repository.update_announcement(record).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn delete_announcement(
|
|
&self,
|
|
announcement_id: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
match &self.announcement_writer {
|
|
Some(repository) => repository.delete_announcement(announcement_id).await,
|
|
None => Ok(false),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn mark_announcement_as_read(
|
|
&self,
|
|
user_id: &str,
|
|
announcement_id: &str,
|
|
read_at_unix_secs: u64,
|
|
) -> Result<bool, DataLayerError> {
|
|
match &self.announcement_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.mark_announcement_as_read(user_id, announcement_id, read_at_unix_secs)
|
|
.await
|
|
}
|
|
None => Ok(false),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_video_task(
|
|
&self,
|
|
key: VideoTaskLookupKey<'_>,
|
|
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.find(key).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_video_task_for_user(
|
|
&self,
|
|
key: VideoTaskLookupKey<'_>,
|
|
user_id: &str,
|
|
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.find_for_user(key, user_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_video_task_page(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
offset: usize,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.list_page(filter, offset, limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_video_task_page_summary(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
offset: usize,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.list_page_summary(filter, offset, limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_video_tasks(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.count(filter).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_video_tasks_by_status(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
) -> Result<Vec<VideoTaskStatusCount>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.count_by_status(filter).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_distinct_video_task_users(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.count_distinct_users(filter).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn top_video_task_models(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
limit: usize,
|
|
) -> Result<Vec<VideoTaskModelCount>, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => repository.top_models(filter, limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_video_tasks_created_since(
|
|
&self,
|
|
filter: &VideoTaskQueryFilter,
|
|
created_since_unix_secs: u64,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.video_task_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.count_created_since(filter, created_since_unix_secs)
|
|
.await
|
|
}
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn upsert_video_task(
|
|
&self,
|
|
task: UpsertVideoTask,
|
|
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_writer {
|
|
Some(repository) => repository.upsert(task).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn update_active_video_task(
|
|
&self,
|
|
task: UpsertVideoTask,
|
|
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_writer {
|
|
Some(repository) => repository.update_if_active(task).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn claim_due_video_tasks(
|
|
&self,
|
|
now_unix_secs: u64,
|
|
claim_until_unix_secs: u64,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
|
match &self.video_task_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.claim_due(now_unix_secs, claim_until_unix_secs, limit)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_wallet(
|
|
&self,
|
|
key: WalletLookupKey<'_>,
|
|
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.find(key).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_wallets_by_api_key_ids(
|
|
&self,
|
|
api_key_ids: &[String],
|
|
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_wallets_by_api_key_ids(api_key_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_wallets_by_user_ids(
|
|
&self,
|
|
user_ids: &[String],
|
|
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_wallets_by_user_ids(user_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_wallets(
|
|
&self,
|
|
query: &AdminWalletListQuery,
|
|
) -> Result<StoredAdminWalletListPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_wallets(query).await,
|
|
None => Ok(StoredAdminWalletListPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_wallet_ledger(
|
|
&self,
|
|
query: &AdminWalletLedgerQuery,
|
|
) -> Result<StoredAdminWalletLedgerPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_wallet_ledger(query).await,
|
|
None => Ok(StoredAdminWalletLedgerPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_wallet_refund_requests(
|
|
&self,
|
|
query: &AdminWalletRefundRequestListQuery,
|
|
) -> Result<StoredAdminWalletRefundRequestPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_wallet_refund_requests(query).await,
|
|
None => Ok(StoredAdminWalletRefundRequestPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_wallet_transactions(
|
|
&self,
|
|
wallet_id: &str,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<StoredAdminWalletTransactionPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_admin_wallet_transactions(wallet_id, limit, offset)
|
|
.await
|
|
}
|
|
None => Ok(StoredAdminWalletTransactionPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_wallet_today_usage(
|
|
&self,
|
|
wallet_id: &str,
|
|
billing_timezone: &str,
|
|
) -> Result<Option<StoredWalletDailyUsageLedger>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.find_wallet_today_usage(wallet_id, billing_timezone)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_wallet_daily_usage_history(
|
|
&self,
|
|
wallet_id: &str,
|
|
billing_timezone: &str,
|
|
limit: usize,
|
|
) -> Result<StoredWalletDailyUsageLedgerPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_wallet_daily_usage_history(wallet_id, billing_timezone, limit)
|
|
.await
|
|
}
|
|
None => Ok(StoredWalletDailyUsageLedgerPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_wallet_refunds(
|
|
&self,
|
|
wallet_id: &str,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<StoredAdminWalletRefundPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_admin_wallet_refunds(wallet_id, limit, offset)
|
|
.await
|
|
}
|
|
None => Ok(StoredAdminWalletRefundPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_payment_orders(
|
|
&self,
|
|
query: &AdminPaymentOrderListQuery,
|
|
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_payment_orders(query).await,
|
|
None => Ok(StoredAdminPaymentOrderPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_payment_callbacks(
|
|
&self,
|
|
payment_method: Option<&str>,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<StoredAdminPaymentCallbackPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_admin_payment_callbacks(payment_method, limit, offset)
|
|
.await
|
|
}
|
|
None => Ok(StoredAdminPaymentCallbackPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_redeem_code_batches(
|
|
&self,
|
|
query: &AdminRedeemCodeBatchListQuery,
|
|
) -> Result<StoredAdminRedeemCodeBatchPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_redeem_code_batches(query).await,
|
|
None => Ok(StoredAdminRedeemCodeBatchPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_admin_redeem_code_batch(
|
|
&self,
|
|
batch_id: &str,
|
|
) -> Result<Option<StoredAdminRedeemCodeBatch>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.find_admin_redeem_code_batch(batch_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_admin_redeem_codes(
|
|
&self,
|
|
query: &AdminRedeemCodeListQuery,
|
|
) -> Result<StoredAdminRedeemCodePage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.list_admin_redeem_codes(query).await,
|
|
None => Ok(StoredAdminRedeemCodePage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_admin_payment_order(
|
|
&self,
|
|
order_id: &str,
|
|
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.find_admin_payment_order(order_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_wallet_payment_orders_by_user_id(
|
|
&self,
|
|
user_id: &str,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_wallet_payment_orders_by_user_id(user_id, limit, offset)
|
|
.await
|
|
}
|
|
None => Ok(StoredAdminPaymentOrderPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_wallet_payment_order_by_user_id(
|
|
&self,
|
|
user_id: &str,
|
|
order_id: &str,
|
|
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.find_wallet_payment_order_by_user_id(user_id, order_id)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_wallet_recharge_order_by_order_no(
|
|
&self,
|
|
user_id: &str,
|
|
order_no: &str,
|
|
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.find_wallet_recharge_order_by_order_no(user_id, order_no)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_pending_plan_purchase_order_by_user_id(
|
|
&self,
|
|
user_id: &str,
|
|
product_id: &str,
|
|
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.find_pending_plan_purchase_order_by_user_id(user_id, product_id)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_payment_order_by_order_no(
|
|
&self,
|
|
order_no: &str,
|
|
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.find_payment_order_by_order_no(order_no).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_wallet_refund(
|
|
&self,
|
|
wallet_id: &str,
|
|
refund_id: &str,
|
|
) -> Result<Option<aether_data::repository::wallet::StoredAdminWalletRefund>, DataLayerError>
|
|
{
|
|
match &self.wallet_reader {
|
|
Some(repository) => repository.find_wallet_refund(wallet_id, refund_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_wallet_recharge_order(
|
|
&self,
|
|
input: CreateWalletRechargeOrderInput,
|
|
) -> Result<Option<CreateWalletRechargeOrderOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.create_wallet_recharge_order(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn update_wallet_recharge_checkout(
|
|
&self,
|
|
input: UpdateWalletRechargeCheckoutInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.update_wallet_recharge_checkout(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn compare_and_swap_payment_order_stripe_client_secret(
|
|
&self,
|
|
input: CompareAndSwapPaymentOrderStripeClientSecretInput,
|
|
) -> Result<Option<bool>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.compare_and_swap_payment_order_stripe_client_secret(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn fail_wallet_recharge_checkout(
|
|
&self,
|
|
input: FailWalletRechargeCheckoutInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.fail_wallet_recharge_checkout(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn reclaim_wallet_recharge_checkout(
|
|
&self,
|
|
input: ReclaimWalletRechargeCheckoutInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.reclaim_wallet_recharge_checkout(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_plan_purchase_order(
|
|
&self,
|
|
input: CreatePlanPurchaseOrderInput,
|
|
) -> Result<Option<CreatePlanPurchaseOrderOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.create_plan_purchase_order(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_wallet_refund_request(
|
|
&self,
|
|
input: CreateWalletRefundRequestInput,
|
|
) -> Result<Option<CreateWalletRefundRequestOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.create_wallet_refund_request(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn process_payment_callback(
|
|
&self,
|
|
input: ProcessPaymentCallbackInput,
|
|
) -> Result<Option<ProcessPaymentCallbackOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.process_payment_callback(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn adjust_wallet_balance(
|
|
&self,
|
|
input: AdjustWalletBalanceInput,
|
|
) -> Result<Option<(StoredWalletSnapshot, Option<StoredAdminWalletTransaction>)>, DataLayerError>
|
|
{
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.adjust_wallet_balance(input).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn prepare_admin_user_wallet_balance_batch(
|
|
&self,
|
|
input: PrepareAdminUserWalletBalanceBatchInput,
|
|
) -> Result<Option<PrepareAdminUserWalletBalanceBatchOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.prepare_admin_user_wallet_balance_batch(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn get_admin_user_wallet_balance_batch(
|
|
&self,
|
|
admin_user_id: &str,
|
|
idempotency_key: &str,
|
|
request_fingerprint: &str,
|
|
) -> Result<Option<PrepareAdminUserWalletBalanceBatchOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.get_admin_user_wallet_balance_batch(
|
|
admin_user_id,
|
|
idempotency_key,
|
|
request_fingerprint,
|
|
)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn adjust_admin_user_wallet_balance_batch_user(
|
|
&self,
|
|
input: AdjustWalletBalanceInBatchInput,
|
|
) -> Result<Option<AdminUserWalletBalanceBatchUserOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.adjust_admin_user_wallet_balance_batch_user(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn record_admin_user_wallet_balance_batch_failure(
|
|
&self,
|
|
admin_user_id: &str,
|
|
idempotency_key: &str,
|
|
user_id: &str,
|
|
reason: &str,
|
|
) -> Result<Option<AdminUserWalletBalanceBatchUserOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.record_admin_user_wallet_balance_batch_failure(
|
|
admin_user_id,
|
|
idempotency_key,
|
|
user_id,
|
|
reason,
|
|
)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_manual_wallet_recharge(
|
|
&self,
|
|
input: CreateManualWalletRechargeInput,
|
|
) -> Result<Option<(StoredWalletSnapshot, StoredAdminPaymentOrder)>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.create_manual_wallet_recharge(input).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn process_admin_wallet_refund(
|
|
&self,
|
|
input: ProcessAdminWalletRefundInput,
|
|
) -> Result<
|
|
Option<
|
|
WalletMutationOutcome<(
|
|
StoredWalletSnapshot,
|
|
StoredAdminWalletRefund,
|
|
StoredAdminWalletTransaction,
|
|
)>,
|
|
>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.process_admin_wallet_refund(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn update_admin_wallet_refund_gateway(
|
|
&self,
|
|
input: UpdateAdminWalletRefundGatewayInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminWalletRefund>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.update_admin_wallet_refund_gateway(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn complete_admin_wallet_refund(
|
|
&self,
|
|
input: CompleteAdminWalletRefundInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminWalletRefund>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.complete_admin_wallet_refund(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn fail_admin_wallet_refund(
|
|
&self,
|
|
input: FailAdminWalletRefundInput,
|
|
) -> Result<
|
|
Option<
|
|
WalletMutationOutcome<(
|
|
StoredWalletSnapshot,
|
|
StoredAdminWalletRefund,
|
|
Option<StoredAdminWalletTransaction>,
|
|
)>,
|
|
>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.fail_admin_wallet_refund(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn expire_admin_payment_order(
|
|
&self,
|
|
order_id: &str,
|
|
) -> Result<Option<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>>, DataLayerError>
|
|
{
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.expire_admin_payment_order(order_id)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn fail_admin_payment_order(
|
|
&self,
|
|
order_id: &str,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.fail_admin_payment_order(order_id)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn credit_admin_payment_order(
|
|
&self,
|
|
input: CreditAdminPaymentOrderInput,
|
|
) -> Result<Option<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>>, DataLayerError>
|
|
{
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.credit_admin_payment_order(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_admin_redeem_code_batch(
|
|
&self,
|
|
input: CreateAdminRedeemCodeBatchInput,
|
|
) -> Result<Option<CreateAdminRedeemCodeBatchResult>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.create_admin_redeem_code_batch(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn disable_admin_redeem_code_batch(
|
|
&self,
|
|
input: DisableAdminRedeemCodeBatchInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminRedeemCodeBatch>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.disable_admin_redeem_code_batch(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn delete_admin_redeem_code_batch(
|
|
&self,
|
|
input: DeleteAdminRedeemCodeBatchInput,
|
|
) -> Result<Option<WalletMutationOutcome<StoredAdminRedeemCodeBatch>>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository
|
|
.delete_admin_redeem_code_batch(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn disable_admin_redeem_code(
|
|
&self,
|
|
input: DisableAdminRedeemCodeInput,
|
|
) -> Result<
|
|
Option<WalletMutationOutcome<aether_data::repository::wallet::StoredAdminRedeemCode>>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.disable_admin_redeem_code(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn redeem_wallet_code(
|
|
&self,
|
|
input: RedeemWalletCodeInput,
|
|
) -> Result<Option<RedeemWalletCodeOutcome>, DataLayerError> {
|
|
match &self.wallet_writer {
|
|
Some(repository) => repository.redeem_wallet_code(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn settle_usage(
|
|
&self,
|
|
input: UsageSettlementInput,
|
|
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => repository.settle_usage(input).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn reserve_usage_policy_cost(
|
|
&self,
|
|
input: ReserveUsagePolicyCostInput,
|
|
) -> Result<Option<ReserveUsagePolicyCostOutcome>, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => repository.reserve_usage_policy_cost(input).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn reserve_usage_policy_request(
|
|
&self,
|
|
input: ReserveUsagePolicyRequestInput,
|
|
) -> Result<Option<ReserveUsagePolicyRequestOutcome>, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => repository
|
|
.reserve_usage_policy_request(input)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn release_usage_policy_request_admission(
|
|
&self,
|
|
input: ReleaseUsagePolicyRequestAdmissionInput,
|
|
) -> Result<Option<StoredUsagePolicyRequestAdmission>, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.release_usage_policy_request_admission(input)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn cleanup_usage_policy_request_admissions(
|
|
&self,
|
|
now_unix_secs: u64,
|
|
batch_size: usize,
|
|
) -> Result<usize, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.cleanup_usage_policy_request_admissions(now_unix_secs, batch_size)
|
|
.await
|
|
}
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn reconcile_usage_policy_cost(
|
|
&self,
|
|
input: ReconcileUsagePolicyCostInput,
|
|
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => repository.reconcile_usage_policy_cost(input).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn cleanup_usage_policy_cost_reservations(
|
|
&self,
|
|
now_unix_secs: u64,
|
|
batch_size: usize,
|
|
) -> Result<usize, DataLayerError> {
|
|
match &self.settlement_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.cleanup_usage_policy_cost_reservations(now_unix_secs, batch_size)
|
|
.await
|
|
}
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn reset_due_provider_quotas(
|
|
&self,
|
|
now_unix_secs: u64,
|
|
) -> Result<usize, DataLayerError> {
|
|
match &self.provider_quota_writer {
|
|
Some(repository) => repository.reset_due(now_unix_secs).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_provider_quota_by_provider_id(
|
|
&self,
|
|
provider_id: &str,
|
|
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
|
|
match &self.provider_quota_reader {
|
|
Some(repository) => repository.find_by_provider_id(provider_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_provider_quotas_by_provider_ids(
|
|
&self,
|
|
provider_ids: &[String],
|
|
) -> Result<Vec<StoredProviderQuotaSnapshot>, DataLayerError> {
|
|
match &self.provider_quota_reader {
|
|
Some(repository) => repository.find_by_provider_ids(provider_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
|
|
pub(crate) async fn upsert_usage(
|
|
&self,
|
|
usage: UpsertUsageRecord,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"usage_upsert",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.upsert(usage).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn upsert_first_byte_usage(
|
|
&self,
|
|
usage: UpsertUsageRecord,
|
|
) -> Result<(), DataLayerError> {
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"usage_first_byte_upsert",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.upsert_first_byte(usage).await,
|
|
None => Ok(()),
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn upsert_first_byte_usage_many(
|
|
&self,
|
|
usages: Vec<UpsertUsageRecord>,
|
|
) -> Result<(), DataLayerError> {
|
|
if usages.is_empty() {
|
|
return Ok(());
|
|
}
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"usage_first_byte_upsert_batch",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.upsert_first_byte_many(usages).await,
|
|
None => Ok(()),
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn upsert_pending_usage_many(
|
|
&self,
|
|
usages: Vec<UpsertUsageRecord>,
|
|
) -> Result<(), DataLayerError> {
|
|
if usages.is_empty() {
|
|
return Ok(());
|
|
}
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"usage_pending_upsert_batch",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.upsert_pending_many(usages).await,
|
|
None => Ok(()),
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub(crate) async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.rebuild_api_key_usage_stats().await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub(crate) async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.rebuild_provider_api_key_usage_stats().await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn flush_usage_counter_deltas(
|
|
&self,
|
|
batch_size: usize,
|
|
) -> Result<UsageCounterFlushSummary, DataLayerError> {
|
|
match &self.usage_writer {
|
|
Some(repository) => repository.flush_usage_counter_deltas(batch_size).await,
|
|
None => Ok(UsageCounterFlushSummary::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn cleanup_processed_usage_counter_deltas(
|
|
&self,
|
|
cutoff_unix_secs: u64,
|
|
batch_size: usize,
|
|
) -> Result<usize, DataLayerError> {
|
|
match &self.usage_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.cleanup_processed_usage_counter_deltas(cutoff_unix_secs, batch_size)
|
|
.await
|
|
}
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
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,
|
|
targets: UsageCleanupTargets,
|
|
mode: UsageCleanupExecutionMode,
|
|
) -> Result<UsageCleanupSummary, DataLayerError> {
|
|
match &self.usage_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.cleanup_usage(window, batch_size, auto_delete_expired_keys, targets, mode)
|
|
.await
|
|
}
|
|
None => Ok(UsageCleanupSummary::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn preview_usage_cleanup(
|
|
&self,
|
|
window: &UsageCleanupWindow,
|
|
targets: UsageCleanupTargets,
|
|
mode: UsageCleanupExecutionMode,
|
|
) -> Result<aether_data_contracts::repository::usage::UsageCleanupPreviewCounts, DataLayerError>
|
|
{
|
|
match &self.usage_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.preview_usage_cleanup(window, targets, mode)
|
|
.await
|
|
}
|
|
None => {
|
|
Ok(aether_data_contracts::repository::usage::UsageCleanupPreviewCounts::default())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_request_usage_by_request_id(
|
|
&self,
|
|
request_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.find_by_request_id(request_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_request_usage_by_request_id_shallow(
|
|
&self,
|
|
request_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.find_by_request_id_shallow(request_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_request_usage_by_id(
|
|
&self,
|
|
usage_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.find_by_id(usage_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_request_usage_by_ids(
|
|
&self,
|
|
usage_ids: &[String],
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_by_ids(usage_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn resolve_request_usage_body_ref(
|
|
&self,
|
|
body_ref: &str,
|
|
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.resolve_body_ref(body_ref).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn read_request_usage_body_payload(
|
|
&self,
|
|
body_ref: &str,
|
|
) -> Result<
|
|
Option<aether_data_contracts::repository::usage::StoredUsageBodyPayload>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.read_body_payload(body_ref).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_usage_audits(
|
|
&self,
|
|
query: &UsageAuditListQuery,
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_usage_audits(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_usage_audits(
|
|
&self,
|
|
query: &UsageAuditListQuery,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.count_usage_audits(query).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_usage_audits_by_keyword_search(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageAuditKeywordSearchQuery,
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_usage_audits_by_keyword_search(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_usage_audits_by_keyword_search(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageAuditKeywordSearchQuery,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.count_usage_audits_by_keyword_search(query).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn aggregate_usage_audits(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageAuditAggregationQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageAuditAggregation>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.aggregate_usage_audits(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn query_dashboard_summary(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::StoredDashboardSummary, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.query_dashboard_summary(query).await,
|
|
None => Err(DataLayerError::InvalidInput(
|
|
"dashboard summary repository is unavailable".into(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn query_dashboard_analytics(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageDashboardAnalytics,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.query_dashboard_analytics(query).await,
|
|
None => Err(DataLayerError::InvalidInput(
|
|
"usage analytics repository is unavailable".into(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn query_usage_analytics(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageAnalyticsQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::StoredUsageAnalytics, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.query_usage_analytics(query).await,
|
|
None => Err(DataLayerError::InvalidInput(
|
|
"usage analytics repository is unavailable".into(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_health_observations(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::HealthObservationQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::HealthObservationSummary, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_health_observations(query).await,
|
|
None => Err(DataLayerError::InvalidInput(
|
|
"health observations repository is unavailable".into(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_audits(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageAuditSummaryQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::StoredUsageAuditSummary, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_audits(query).await,
|
|
None => {
|
|
Ok(aether_data_contracts::repository::usage::StoredUsageAuditSummary::default())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn read_usage_counter_health(
|
|
&self,
|
|
) -> Result<UsageCounterHealthSnapshot, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.read_usage_counter_health().await,
|
|
None => Ok(UsageCounterHealthSnapshot::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn read_usage_counter_pending_health(
|
|
&self,
|
|
) -> Result<UsageCounterPendingHealthSnapshot, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.read_usage_counter_pending_health().await,
|
|
None => Ok(UsageCounterPendingHealthSnapshot::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_totals_by_user_ids(
|
|
&self,
|
|
user_ids: &[String],
|
|
) -> Result<Vec<aether_data_contracts::repository::usage::StoredUsageUserTotals>, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_usage_totals_by_user_ids(user_ids)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_cache_hit_summary(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageCacheHitSummaryQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::StoredUsageCacheHitSummary, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_cache_hit_summary(query).await,
|
|
None => {
|
|
Ok(aether_data_contracts::repository::usage::StoredUsageCacheHitSummary::default())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_settled_cost(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageSettledCostSummary,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_settled_cost(query).await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageSettledCostSummary::default(),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_cache_affinity_hit_summary(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageCacheAffinityHitSummaryQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageCacheAffinityHitSummary,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository
|
|
.summarize_usage_cache_affinity_hit_summary(query)
|
|
.await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageCacheAffinityHitSummary::default(),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_usage_cache_affinity_intervals(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageCacheAffinityIntervalQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageCacheAffinityIntervalRow>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_usage_cache_affinity_intervals(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_dashboard_usage(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardSummaryQuery,
|
|
) -> Result<aether_data_contracts::repository::usage::StoredUsageDashboardSummary, DataLayerError>
|
|
{
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_dashboard_usage(query).await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageDashboardSummary::default(),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_dashboard_stats(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardSummaryQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_dashboard_stats(query).await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary::default(
|
|
),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_dashboard_daily_breakdown(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardDailyBreakdownQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageDashboardDailyBreakdownRow>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_dashboard_daily_breakdown(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_dashboard_provider_counts(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageDashboardProviderCountsQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageDashboardProviderCount>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_dashboard_provider_counts(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_breakdown(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageBreakdownSummaryQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageBreakdownSummaryRow>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_breakdown(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_monitoring_usage_errors(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageMonitoringErrorCountQuery,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.count_monitoring_usage_errors(query).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_monitoring_usage_errors(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageMonitoringErrorListQuery,
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_monitoring_usage_errors(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_error_distribution(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageErrorDistributionQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageErrorDistributionRow>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_error_distribution(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_performance_percentiles(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsagePerformancePercentilesQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsagePerformancePercentilesRow>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_usage_performance_percentiles(query)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_provider_performance(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageProviderPerformanceQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageProviderPerformance,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_provider_performance(query).await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageProviderPerformance::default(),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_cost_savings(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageCostSavingsSummaryQuery,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::StoredUsageCostSavingsSummary,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_cost_savings(query).await,
|
|
None => Ok(
|
|
aether_data_contracts::repository::usage::StoredUsageCostSavingsSummary::default(),
|
|
),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_time_series(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageTimeSeriesQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageTimeSeriesBucket>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_time_series(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_leaderboard(
|
|
&self,
|
|
query: &aether_data_contracts::repository::usage::UsageLeaderboardQuery,
|
|
) -> Result<
|
|
Vec<aether_data_contracts::repository::usage::StoredUsageLeaderboardSummary>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_leaderboard(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_daily_heatmap(
|
|
&self,
|
|
query: &UsageDailyHeatmapQuery,
|
|
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.summarize_usage_daily_heatmap(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_recent_usage_audits(
|
|
&self,
|
|
user_id: Option<&str>,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => repository.list_recent_usage_audits(user_id, limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_total_tokens_by_api_key_ids(
|
|
&self,
|
|
api_key_ids: &[String],
|
|
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_total_tokens_by_api_key_ids(api_key_ids)
|
|
.await
|
|
}
|
|
None => Ok(std::collections::BTreeMap::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_by_provider_api_key_ids(
|
|
&self,
|
|
provider_api_key_ids: &[String],
|
|
) -> Result<
|
|
std::collections::BTreeMap<
|
|
String,
|
|
aether_data_contracts::repository::usage::StoredProviderApiKeyUsageSummary,
|
|
>,
|
|
DataLayerError,
|
|
> {
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_usage_by_provider_api_key_ids(provider_api_key_ids)
|
|
.await
|
|
}
|
|
None => Ok(std::collections::BTreeMap::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_usage_by_provider_api_key_windows(
|
|
&self,
|
|
requests: &[ProviderApiKeyWindowUsageRequest],
|
|
) -> Result<Vec<StoredProviderApiKeyWindowUsageSummary>, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_usage_by_provider_api_key_windows(requests)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_users_by_ids(
|
|
&self,
|
|
user_ids: &[String],
|
|
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.list_users_by_ids(user_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_users_by_username_search(
|
|
&self,
|
|
username_search: &str,
|
|
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_users_by_username_search(username_search)
|
|
.await
|
|
}
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_export_users(
|
|
&self,
|
|
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.list_export_users().await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_export_users_page(
|
|
&self,
|
|
query: &aether_data::repository::users::UserExportListQuery,
|
|
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.list_export_users_page(query).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn count_export_users(
|
|
&self,
|
|
query: &aether_data::repository::users::UserExportListQuery,
|
|
) -> Result<u64, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.count_export_users(query).await,
|
|
None => Ok(0),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_export_users(
|
|
&self,
|
|
) -> Result<aether_data::repository::users::UserExportSummary, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.summarize_export_users().await,
|
|
None => Ok(aether_data::repository::users::UserExportSummary::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_export_user_by_id(
|
|
&self,
|
|
user_id: &str,
|
|
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.find_export_user_by_id(user_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn read_user_feature_settings(
|
|
&self,
|
|
user_id: &str,
|
|
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
|
let user_id = user_id.trim();
|
|
if user_id.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
if let Some(user) = self.find_export_user_by_id(user_id).await? {
|
|
return Ok(user.feature_settings);
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
pub(crate) async fn list_non_admin_export_users(
|
|
&self,
|
|
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.list_non_admin_export_users().await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_user_auth_by_ids(
|
|
&self,
|
|
user_ids: &[String],
|
|
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
|
|
match &self.user_reader {
|
|
Some(repository) => repository.list_user_auth_by_ids(user_ids).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_provider_usage_since(
|
|
&self,
|
|
provider_id: &str,
|
|
since_unix_secs: u64,
|
|
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
|
match &self.usage_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.summarize_provider_usage_since(provider_id, since_unix_secs)
|
|
.await
|
|
}
|
|
None => Ok(StoredProviderUsageSummary::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn usage_worker_queue(&self) -> Option<std::sync::Arc<dyn RuntimeQueueStore>> {
|
|
self.usage_worker_queue.clone()
|
|
}
|
|
|
|
pub(crate) async fn find_billing_model_context(
|
|
&self,
|
|
provider_id: &str,
|
|
provider_api_key_id: Option<&str>,
|
|
global_model_name: &str,
|
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
|
let key = BillingModelContextCacheKey::ByGlobalModelName {
|
|
provider_id: normalize_billing_context_cache_part(provider_id),
|
|
provider_api_key_id: normalize_optional_billing_context_cache_part(provider_api_key_id),
|
|
global_model_name: normalize_billing_context_cache_part(global_model_name),
|
|
};
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
loop {
|
|
match self.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Saturated => {
|
|
return Err(DataLayerError::TimedOut(format!(
|
|
"billing model context cache admission saturated for {key:?}"
|
|
)));
|
|
}
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => {
|
|
match timeout(
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT,
|
|
inflight_state.wait(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(error)) => return Err(error),
|
|
Err(_) => self.expire_billing_model_context_inflight(&key, &inflight_state),
|
|
}
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
continue;
|
|
}
|
|
BillingModelContextInflightRegistration::Leader(mut guard) => {
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
let load_epoch = guard.epoch();
|
|
let result = match timeout(
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT,
|
|
self.load_billing_model_context_by_name(
|
|
key,
|
|
provider_id,
|
|
provider_api_key_id,
|
|
global_model_name,
|
|
load_epoch,
|
|
&guard.inflight_state,
|
|
),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(_) => Err(DataLayerError::TimedOut(
|
|
"billing model context load timed out".to_string(),
|
|
)),
|
|
};
|
|
guard.finish(result.as_ref().err().cloned());
|
|
return result;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_billing_model_context_by_model_id(
|
|
&self,
|
|
provider_id: &str,
|
|
provider_api_key_id: Option<&str>,
|
|
model_id: &str,
|
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
|
let key = BillingModelContextCacheKey::ByModelId {
|
|
provider_id: normalize_billing_context_cache_part(provider_id),
|
|
provider_api_key_id: normalize_optional_billing_context_cache_part(provider_api_key_id),
|
|
model_id: normalize_billing_context_cache_part(model_id),
|
|
};
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
loop {
|
|
match self.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Saturated => {
|
|
return Err(DataLayerError::TimedOut(format!(
|
|
"billing model context cache admission saturated for {key:?}"
|
|
)));
|
|
}
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => {
|
|
match timeout(
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT,
|
|
inflight_state.wait(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(error)) => return Err(error),
|
|
Err(_) => self.expire_billing_model_context_inflight(&key, &inflight_state),
|
|
}
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
continue;
|
|
}
|
|
BillingModelContextInflightRegistration::Leader(mut guard) => {
|
|
if let Some(value) = self.cached_billing_model_context(&key) {
|
|
return Ok(value);
|
|
}
|
|
let load_epoch = guard.epoch();
|
|
let result = match timeout(
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT,
|
|
self.load_billing_model_context_by_model_id(
|
|
key,
|
|
provider_id,
|
|
provider_api_key_id,
|
|
model_id,
|
|
load_epoch,
|
|
&guard.inflight_state,
|
|
),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(_) => Err(DataLayerError::TimedOut(
|
|
"billing model context load timed out".to_string(),
|
|
)),
|
|
};
|
|
guard.finish(result.as_ref().err().cloned());
|
|
return result;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn load_billing_model_context_by_name(
|
|
&self,
|
|
key: BillingModelContextCacheKey,
|
|
provider_id: &str,
|
|
provider_api_key_id: Option<&str>,
|
|
global_model_name: &str,
|
|
load_epoch: u64,
|
|
load_flight: &std::sync::Arc<BillingModelContextInflightState>,
|
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"billing_model_context",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
let value = repository
|
|
.find_model_context(provider_id, provider_api_key_id, global_model_name)
|
|
.await?;
|
|
self.remember_billing_model_context(
|
|
key,
|
|
value.clone(),
|
|
load_epoch,
|
|
load_flight,
|
|
);
|
|
Ok(value)
|
|
}
|
|
None => {
|
|
self.remember_billing_model_context(key, None, load_epoch, load_flight);
|
|
Ok(None)
|
|
}
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn load_billing_model_context_by_model_id(
|
|
&self,
|
|
key: BillingModelContextCacheKey,
|
|
provider_id: &str,
|
|
provider_api_key_id: Option<&str>,
|
|
model_id: &str,
|
|
load_epoch: u64,
|
|
load_flight: &std::sync::Arc<BillingModelContextInflightState>,
|
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
|
crate::request_diagnostics::observe_db_operation(
|
|
"billing_model_context",
|
|
self.database_pool_summary(),
|
|
async {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
let value = repository
|
|
.find_model_context_by_model_id(
|
|
provider_id,
|
|
provider_api_key_id,
|
|
model_id,
|
|
)
|
|
.await?;
|
|
self.remember_billing_model_context(
|
|
key,
|
|
value.clone(),
|
|
load_epoch,
|
|
load_flight,
|
|
);
|
|
Ok(value)
|
|
}
|
|
None => {
|
|
self.remember_billing_model_context(key, None, load_epoch, load_flight);
|
|
Ok(None)
|
|
}
|
|
}
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
fn register_billing_model_context_inflight(
|
|
&self,
|
|
key: &BillingModelContextCacheKey,
|
|
) -> BillingModelContextInflightRegistration<'_> {
|
|
let mut inflight = self
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
if let Some(inflight_state) = inflight.get(key) {
|
|
return BillingModelContextInflightRegistration::Follower(std::sync::Arc::clone(
|
|
inflight_state,
|
|
));
|
|
}
|
|
if inflight.len() >= Self::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT {
|
|
return BillingModelContextInflightRegistration::Saturated;
|
|
}
|
|
let Ok(admission) =
|
|
std::sync::Arc::clone(&self.billing_model_context_cache.admission).try_acquire_owned()
|
|
else {
|
|
return BillingModelContextInflightRegistration::Saturated;
|
|
};
|
|
let inflight_state = std::sync::Arc::new(BillingModelContextInflightState {
|
|
epoch: self
|
|
.billing_model_context_cache
|
|
.epoch
|
|
.load(std::sync::atomic::Ordering::Acquire),
|
|
completion: std::sync::OnceLock::new(),
|
|
notify: tokio::sync::Notify::new(),
|
|
});
|
|
inflight.insert(key.clone(), std::sync::Arc::clone(&inflight_state));
|
|
BillingModelContextInflightRegistration::Leader(BillingModelContextInflightGuard::new(
|
|
self,
|
|
key.clone(),
|
|
inflight_state,
|
|
admission,
|
|
))
|
|
}
|
|
|
|
fn finish_billing_model_context_inflight(
|
|
&self,
|
|
key: &BillingModelContextCacheKey,
|
|
inflight_state: &std::sync::Arc<BillingModelContextInflightState>,
|
|
admission: Option<tokio::sync::OwnedSemaphorePermit>,
|
|
) -> Option<std::sync::Arc<BillingModelContextInflightState>> {
|
|
let mut inflight = self
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
drop(admission);
|
|
if inflight
|
|
.get(key)
|
|
.is_some_and(|current| std::sync::Arc::ptr_eq(current, inflight_state))
|
|
{
|
|
inflight.remove(key)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
fn expire_billing_model_context_inflight(
|
|
&self,
|
|
key: &BillingModelContextCacheKey,
|
|
inflight_state: &std::sync::Arc<BillingModelContextInflightState>,
|
|
) {
|
|
let _mutation = self
|
|
.billing_model_context_cache
|
|
.mutation
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
let removed = {
|
|
let mut inflight = self
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
if inflight
|
|
.get(key)
|
|
.is_some_and(|current| std::sync::Arc::ptr_eq(current, inflight_state))
|
|
{
|
|
inflight.remove(key)
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
drop(_mutation);
|
|
if let Some(removed) = removed {
|
|
tracing::warn!(
|
|
event_name = "billing_model_context_cache_inflight_expired",
|
|
log_type = "ops",
|
|
cache_key = ?key,
|
|
wait_timeout_ms = Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64,
|
|
"gateway billing model context cache expired stale inflight load"
|
|
);
|
|
removed.complete(Ok(()));
|
|
}
|
|
}
|
|
|
|
fn cached_billing_model_context(
|
|
&self,
|
|
key: &BillingModelContextCacheKey,
|
|
) -> Option<Option<StoredBillingModelContext>> {
|
|
self.billing_model_context_cache
|
|
.entries
|
|
.get_fresh(key, Self::BILLING_MODEL_CONTEXT_CACHE_TTL)
|
|
}
|
|
|
|
fn remember_billing_model_context(
|
|
&self,
|
|
key: BillingModelContextCacheKey,
|
|
value: Option<StoredBillingModelContext>,
|
|
load_epoch: u64,
|
|
load_flight: &std::sync::Arc<BillingModelContextInflightState>,
|
|
) {
|
|
let _mutation = self
|
|
.billing_model_context_cache
|
|
.mutation
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
if load_epoch
|
|
!= self
|
|
.billing_model_context_cache
|
|
.epoch
|
|
.load(std::sync::atomic::Ordering::Acquire)
|
|
{
|
|
return;
|
|
}
|
|
let inflight = self
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
if !inflight
|
|
.get(&key)
|
|
.is_some_and(|current| std::sync::Arc::ptr_eq(current, load_flight))
|
|
{
|
|
return;
|
|
}
|
|
self.billing_model_context_cache.entries.insert(
|
|
key,
|
|
value,
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_TTL,
|
|
Self::BILLING_MODEL_CONTEXT_CACHE_MAX_ENTRIES,
|
|
);
|
|
}
|
|
|
|
pub(super) fn clear_billing_model_context_cache(&self) {
|
|
let _mutation = self
|
|
.billing_model_context_cache
|
|
.mutation
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
self.billing_model_context_cache
|
|
.epoch
|
|
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
|
|
self.billing_model_context_cache.entries.clear();
|
|
let inflight_states = self
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.drain()
|
|
.map(|(_, state)| state)
|
|
.collect::<Vec<_>>();
|
|
drop(_mutation);
|
|
if !inflight_states.is_empty() {
|
|
tracing::warn!(
|
|
event_name = "billing_model_context_cache_inflight_cleared",
|
|
log_type = "ops",
|
|
"gateway billing model context cache cleared in-flight loads"
|
|
);
|
|
for inflight_state in inflight_states {
|
|
inflight_state.complete(Ok(()));
|
|
}
|
|
}
|
|
}
|
|
|
|
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> {
|
|
let result = match &self.billing_reader {
|
|
Some(repository) => repository.create_admin_billing_rule(input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
};
|
|
if result.is_ok() {
|
|
self.clear_billing_model_context_cache();
|
|
}
|
|
result
|
|
}
|
|
|
|
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> {
|
|
let result = match &self.billing_reader {
|
|
Some(repository) => repository.update_admin_billing_rule(rule_id, input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
};
|
|
if result.is_ok() {
|
|
self.clear_billing_model_context_cache();
|
|
}
|
|
result
|
|
}
|
|
|
|
pub(crate) async fn create_admin_billing_collector(
|
|
&self,
|
|
input: &AdminBillingCollectorWriteInput,
|
|
) -> Result<AdminBillingMutationOutcome<AdminBillingCollectorRecord>, DataLayerError> {
|
|
let result = match &self.billing_reader {
|
|
Some(repository) => repository.create_admin_billing_collector(input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
};
|
|
if result.is_ok() {
|
|
self.clear_billing_model_context_cache();
|
|
}
|
|
result
|
|
}
|
|
|
|
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> {
|
|
let result = match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.update_admin_billing_collector(collector_id, input)
|
|
.await
|
|
}
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
};
|
|
if result.is_ok() {
|
|
self.clear_billing_model_context_cache();
|
|
}
|
|
result
|
|
}
|
|
|
|
pub(crate) async fn apply_admin_billing_preset(
|
|
&self,
|
|
preset: &str,
|
|
mode: &str,
|
|
collectors: &[AdminBillingCollectorWriteInput],
|
|
) -> Result<AdminBillingMutationOutcome<AdminBillingPresetApplyResult>, DataLayerError> {
|
|
let result = match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.apply_admin_billing_preset(preset, mode, collectors)
|
|
.await
|
|
}
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
};
|
|
if result.is_ok() {
|
|
self.clear_billing_model_context_cache();
|
|
}
|
|
result
|
|
}
|
|
|
|
pub(crate) async fn find_payment_gateway_config(
|
|
&self,
|
|
provider: &str,
|
|
) -> Result<Option<PaymentGatewayConfigRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.find_payment_gateway_config(provider).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_payment_gateway_config_strong(
|
|
&self,
|
|
provider: &str,
|
|
) -> Result<Option<PaymentGatewayConfigRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.find_payment_gateway_config_strong(provider)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn compare_and_swap_payment_gateway_secret(
|
|
&self,
|
|
update: &PaymentGatewaySecretCasUpdate,
|
|
) -> Result<bool, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.compare_and_swap_payment_gateway_secret(update)
|
|
.await
|
|
}
|
|
None => Ok(false),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn compare_and_swap_payment_gateway_config(
|
|
&self,
|
|
input: &PaymentGatewayConfigCasWriteInput,
|
|
) -> Result<AdminBillingMutationOutcome<PaymentGatewayConfigRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.compare_and_swap_payment_gateway_config(input)
|
|
.await
|
|
}
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn upsert_payment_gateway_config(
|
|
&self,
|
|
input: &PaymentGatewayConfigWriteInput,
|
|
) -> Result<AdminBillingMutationOutcome<PaymentGatewayConfigRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.upsert_payment_gateway_config(input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_provider_expenses(
|
|
&self,
|
|
query: &ProviderExpenseQuery,
|
|
) -> Result<Option<ProviderExpensePage>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repo) => repo.list_provider_expenses(query).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
pub(crate) async fn create_provider_expense(
|
|
&self,
|
|
input: &ProviderExpenseInput,
|
|
) -> Result<AdminBillingMutationOutcome<ProviderExpenseRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repo) => repo.create_provider_expense(input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
pub(crate) async fn void_provider_expense(
|
|
&self,
|
|
id: &str,
|
|
operator: Option<&str>,
|
|
) -> Result<AdminBillingMutationOutcome<ProviderExpenseRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repo) => repo.void_provider_expense(id, operator).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_billing_plans(
|
|
&self,
|
|
include_disabled: bool,
|
|
) -> Result<Option<Vec<BillingPlanRecord>>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.list_billing_plans(include_disabled).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_billing_plan(
|
|
&self,
|
|
plan_id: &str,
|
|
) -> Result<Option<BillingPlanRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.find_billing_plan(plan_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn create_billing_plan(
|
|
&self,
|
|
input: &BillingPlanWriteInput,
|
|
) -> Result<AdminBillingMutationOutcome<BillingPlanRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.create_billing_plan(input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn update_billing_plan(
|
|
&self,
|
|
plan_id: &str,
|
|
input: &BillingPlanWriteInput,
|
|
) -> Result<AdminBillingMutationOutcome<BillingPlanRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.update_billing_plan(plan_id, input).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn set_billing_plan_enabled(
|
|
&self,
|
|
plan_id: &str,
|
|
enabled: bool,
|
|
) -> Result<AdminBillingMutationOutcome<BillingPlanRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.set_billing_plan_enabled(plan_id, enabled).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn delete_billing_plan(
|
|
&self,
|
|
plan_id: &str,
|
|
) -> Result<AdminBillingMutationOutcome<()>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.delete_billing_plan(plan_id).await,
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_user_plan_entitlements(
|
|
&self,
|
|
user_id: &str,
|
|
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.list_user_plan_entitlements(user_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_user_plan_entitlements_with_history(
|
|
&self,
|
|
user_id: &str,
|
|
include_inactive: bool,
|
|
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.list_user_plan_entitlements_with_history(user_id, include_inactive)
|
|
.await
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn revoke_user_plan_entitlement(
|
|
&self,
|
|
user_id: &str,
|
|
entitlement_id: &str,
|
|
) -> Result<AdminBillingMutationOutcome<()>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => {
|
|
repository
|
|
.revoke_user_plan_entitlement(user_id, entitlement_id)
|
|
.await
|
|
}
|
|
None => Ok(AdminBillingMutationOutcome::Unavailable),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn find_user_daily_quota_availability(
|
|
&self,
|
|
user_id: &str,
|
|
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, DataLayerError> {
|
|
match &self.billing_reader {
|
|
Some(repository) => repository.find_user_daily_quota_availability(user_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn read_request_candidate_trace(
|
|
&self,
|
|
request_id: &str,
|
|
attempted_only: bool,
|
|
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
|
|
read_request_candidate_trace(self, request_id, attempted_only).await
|
|
}
|
|
|
|
pub(crate) async fn read_decision_trace(
|
|
&self,
|
|
request_id: &str,
|
|
attempted_only: bool,
|
|
) -> Result<Option<DecisionTrace>, DataLayerError> {
|
|
read_decision_trace(self, request_id, attempted_only).await
|
|
}
|
|
|
|
pub(crate) async fn read_request_usage_audit(
|
|
&self,
|
|
request_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
self.find_request_usage_by_request_id(request_id).await
|
|
}
|
|
|
|
pub(crate) async fn read_request_usage_audit_shallow(
|
|
&self,
|
|
request_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
self.find_request_usage_by_request_id_shallow(request_id)
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn read_request_audit_bundle(
|
|
&self,
|
|
request_id: &str,
|
|
attempted_only: bool,
|
|
now_unix_secs: u64,
|
|
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
|
|
aether_data::repository::audit::read_request_audit_bundle(
|
|
self,
|
|
request_id,
|
|
attempted_only,
|
|
now_unix_secs,
|
|
)
|
|
.await
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub(crate) async fn read_provider_transport_snapshot(
|
|
&self,
|
|
provider_id: &str,
|
|
endpoint_id: &str,
|
|
key_id: &str,
|
|
) -> Result<Option<GatewayProviderTransportSnapshot>, DataLayerError> {
|
|
read_provider_transport_snapshot(self, provider_id, endpoint_id, key_id).await
|
|
}
|
|
|
|
pub(crate) async fn read_video_task_response(
|
|
&self,
|
|
route_family: Option<&str>,
|
|
request_path: &str,
|
|
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
|
read_data_backed_video_task_response(self, route_family, request_path).await
|
|
}
|
|
|
|
pub(crate) async fn read_video_task_response_for_user(
|
|
&self,
|
|
route_family: Option<&str>,
|
|
request_path: &str,
|
|
user_id: &str,
|
|
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
|
read_data_backed_video_task_response_for_user(self, route_family, request_path, user_id)
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn find_background_task_run(
|
|
&self,
|
|
run_id: &str,
|
|
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
|
match &self.background_task_reader {
|
|
Some(repository) => repository.find_run(run_id).await,
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_background_task_runs(
|
|
&self,
|
|
query: &BackgroundTaskListQuery,
|
|
) -> Result<StoredBackgroundTaskRunPage, DataLayerError> {
|
|
match &self.background_task_reader {
|
|
Some(repository) => repository.list_runs(query).await,
|
|
None => Ok(StoredBackgroundTaskRunPage::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn list_background_task_events(
|
|
&self,
|
|
run_id: &str,
|
|
offset: usize,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredBackgroundTaskEvent>, DataLayerError> {
|
|
match &self.background_task_reader {
|
|
Some(repository) => repository.list_events(run_id, offset, limit).await,
|
|
None => Ok(Vec::new()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn summarize_background_task_runs(
|
|
&self,
|
|
) -> Result<BackgroundTaskSummary, DataLayerError> {
|
|
match &self.background_task_reader {
|
|
Some(repository) => repository.summarize_runs().await,
|
|
None => Ok(BackgroundTaskSummary::default()),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn upsert_background_task_run(
|
|
&self,
|
|
run: UpsertBackgroundTaskRun,
|
|
) -> Result<Option<StoredBackgroundTaskRun>, DataLayerError> {
|
|
match &self.background_task_writer {
|
|
Some(repository) => repository.upsert_run(run).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn request_cancel_background_task_run(
|
|
&self,
|
|
run_id: &str,
|
|
updated_at_unix_secs: u64,
|
|
) -> Result<bool, DataLayerError> {
|
|
match &self.background_task_writer {
|
|
Some(repository) => {
|
|
repository
|
|
.request_cancel(run_id, updated_at_unix_secs)
|
|
.await
|
|
}
|
|
None => Ok(false),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn upsert_background_task_event(
|
|
&self,
|
|
event: UpsertBackgroundTaskEvent,
|
|
) -> Result<Option<StoredBackgroundTaskEvent>, DataLayerError> {
|
|
match &self.background_task_writer {
|
|
Some(repository) => repository.upsert_event(event).await.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::Duration;
|
|
|
|
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
|
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow};
|
|
use aether_data_contracts::repository::billing::{
|
|
BillingReadRepository, StoredBillingModelContext,
|
|
};
|
|
use aether_data_contracts::repository::global_models::{
|
|
StoredAdminGlobalModel, StoredPublicGlobalModel, UpdateAdminGlobalModelRecord,
|
|
};
|
|
use aether_data_contracts::DataLayerError;
|
|
use async_trait::async_trait;
|
|
use serde_json::json;
|
|
use tokio::sync::Barrier;
|
|
|
|
use super::{
|
|
BillingModelContextCacheKey, BillingModelContextInflightRegistration, GatewayDataState,
|
|
};
|
|
|
|
struct SlowBillingContextRepository {
|
|
calls: AtomicUsize,
|
|
context: StoredBillingModelContext,
|
|
}
|
|
|
|
struct BlockedBillingContextRepository {
|
|
calls: AtomicUsize,
|
|
context: Mutex<StoredBillingModelContext>,
|
|
first_read: Barrier,
|
|
release_first_read: Barrier,
|
|
}
|
|
|
|
struct ConcurrentBillingContextRepository {
|
|
calls: AtomicUsize,
|
|
context: StoredBillingModelContext,
|
|
entered: Barrier,
|
|
release: Barrier,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl BillingReadRepository for SlowBillingContextRepository {
|
|
async fn find_model_context(
|
|
&self,
|
|
_provider_id: &str,
|
|
_provider_api_key_id: Option<&str>,
|
|
_global_model_name: &str,
|
|
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
|
|
{
|
|
self.calls.fetch_add(1, Ordering::AcqRel);
|
|
tokio::time::sleep(Duration::from_millis(25)).await;
|
|
Ok(Some(self.context.clone()))
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl BillingReadRepository for BlockedBillingContextRepository {
|
|
async fn find_model_context(
|
|
&self,
|
|
_provider_id: &str,
|
|
_provider_api_key_id: Option<&str>,
|
|
_global_model_name: &str,
|
|
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
|
|
{
|
|
let call_index = self.calls.fetch_add(1, Ordering::AcqRel);
|
|
let context = self
|
|
.context
|
|
.lock()
|
|
.expect("mutable billing context lock")
|
|
.clone();
|
|
if call_index == 0 {
|
|
self.first_read.wait().await;
|
|
self.release_first_read.wait().await;
|
|
}
|
|
Ok(Some(context))
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl BillingReadRepository for ConcurrentBillingContextRepository {
|
|
async fn find_model_context(
|
|
&self,
|
|
_provider_id: &str,
|
|
_provider_api_key_id: Option<&str>,
|
|
_global_model_name: &str,
|
|
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
|
self.calls.fetch_add(1, Ordering::AcqRel);
|
|
self.entered.wait().await;
|
|
self.release.wait().await;
|
|
Ok(Some(self.context.clone()))
|
|
}
|
|
}
|
|
|
|
fn billing_context() -> StoredBillingModelContext {
|
|
StoredBillingModelContext::new(
|
|
"provider-1".to_string(),
|
|
Some("pay_as_you_go".to_string()),
|
|
Some("key-1".to_string()),
|
|
None,
|
|
None,
|
|
"global-model-1".to_string(),
|
|
"gpt-5".to_string(),
|
|
None,
|
|
Some(0.02),
|
|
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})),
|
|
Some("model-1".to_string()),
|
|
Some("gpt-5-upstream".to_string()),
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("billing context should build")
|
|
}
|
|
|
|
fn billing_cache_key(global_model_name: impl Into<String>) -> BillingModelContextCacheKey {
|
|
BillingModelContextCacheKey::ByGlobalModelName {
|
|
provider_id: "provider-1".to_string(),
|
|
provider_api_key_id: Some("key-1".to_string()),
|
|
global_model_name: global_model_name.into(),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_cancelled_leader_cannot_lose_follower_wakeup() {
|
|
let state = GatewayDataState::default();
|
|
let key = billing_cache_key("lost-wakeup");
|
|
let leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("first registration should lead"),
|
|
};
|
|
let follower = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state,
|
|
_ => panic!("second registration should follow"),
|
|
};
|
|
|
|
// Complete before wait() is constructed or polled. A bare global
|
|
// notify_waiters() broadcast loses this ordering.
|
|
drop(leader);
|
|
tokio::time::timeout(Duration::from_millis(100), follower.wait())
|
|
.await
|
|
.expect("cancelled flight must release an unpolled follower")
|
|
.expect("leader cancellation should allow a retry");
|
|
assert!(matches!(
|
|
state.register_billing_model_context_inflight(&key),
|
|
BillingModelContextInflightRegistration::Leader(_)
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_failed_flight_fans_out_error() {
|
|
let state = GatewayDataState::default();
|
|
let key = billing_cache_key("failed-flight");
|
|
let mut leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("first registration should lead"),
|
|
};
|
|
let followers = (0..2)
|
|
.map(
|
|
|_| match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => {
|
|
inflight_state
|
|
}
|
|
_ => panic!("same-key registration should follow"),
|
|
},
|
|
)
|
|
.collect::<Vec<_>>();
|
|
|
|
leader.finish(Some(DataLayerError::Sql(
|
|
"forced billing context load failure".to_string(),
|
|
)));
|
|
for follower in followers {
|
|
let error = tokio::time::timeout(Duration::from_millis(100), follower.wait())
|
|
.await
|
|
.expect("failed flight should release every follower")
|
|
.expect_err("follower should receive the leader failure");
|
|
assert_eq!(
|
|
error.to_string(),
|
|
"sql error: forced billing context load failure"
|
|
);
|
|
}
|
|
assert!(state
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap()
|
|
.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_clear_wakes_old_follower_without_removing_replacement() {
|
|
let state = GatewayDataState::default();
|
|
let key = billing_cache_key("clear-replacement");
|
|
let old_leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("first registration should lead"),
|
|
};
|
|
let old_follower = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state,
|
|
_ => panic!("second registration should follow"),
|
|
};
|
|
let old_epoch = old_leader.epoch();
|
|
|
|
state.clear_billing_model_context_cache();
|
|
let replacement_leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("clear should allow a replacement leader"),
|
|
};
|
|
let replacement_follower = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state,
|
|
_ => panic!("registration behind replacement should follow"),
|
|
};
|
|
assert_ne!(replacement_leader.epoch(), old_epoch);
|
|
|
|
drop(old_leader);
|
|
assert!(state
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap()
|
|
.get(&key)
|
|
.is_some_and(|current| std::sync::Arc::ptr_eq(
|
|
current,
|
|
&replacement_leader.inflight_state
|
|
)));
|
|
tokio::time::timeout(Duration::from_millis(100), old_follower.wait())
|
|
.await
|
|
.expect("clear should wake the invalidated flight")
|
|
.expect("clear should allow an immediate retry");
|
|
|
|
drop(replacement_leader);
|
|
tokio::time::timeout(Duration::from_millis(100), replacement_follower.wait())
|
|
.await
|
|
.expect("old guard must not strand the replacement follower")
|
|
.expect("replacement completion should succeed");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_timeout_expiration_allows_replacement() {
|
|
let state = GatewayDataState::default();
|
|
let key = billing_cache_key("timeout-replacement");
|
|
let old_leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("first registration should lead"),
|
|
};
|
|
let old_follower = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state,
|
|
_ => panic!("second registration should follow"),
|
|
};
|
|
|
|
state.expire_billing_model_context_inflight(&key, &old_follower);
|
|
old_follower
|
|
.wait()
|
|
.await
|
|
.expect("expired flight should permit a retry");
|
|
let replacement_leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("timeout should allow a replacement leader"),
|
|
};
|
|
drop(old_leader);
|
|
assert!(state
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap()
|
|
.get(&key)
|
|
.is_some_and(|current| std::sync::Arc::ptr_eq(
|
|
current,
|
|
&replacement_leader.inflight_state
|
|
)));
|
|
}
|
|
|
|
#[test]
|
|
fn billing_model_context_expired_leader_cannot_publish_over_replacement() {
|
|
let state = GatewayDataState::default();
|
|
let key = billing_cache_key("timeout-publication-replacement");
|
|
let old_leader = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("first registration should lead"),
|
|
};
|
|
let old_flight = std::sync::Arc::clone(&old_leader.inflight_state);
|
|
let load_epoch = old_leader.epoch();
|
|
state.expire_billing_model_context_inflight(&key, &old_flight);
|
|
let replacement = match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => guard,
|
|
_ => panic!("expiration should allow a replacement leader"),
|
|
};
|
|
assert_eq!(replacement.epoch(), load_epoch);
|
|
|
|
let mut fresh = billing_context();
|
|
fresh.default_price_per_request = Some(2.0);
|
|
state.remember_billing_model_context(
|
|
key.clone(),
|
|
Some(fresh),
|
|
replacement.epoch(),
|
|
&replacement.inflight_state,
|
|
);
|
|
|
|
let mut stale = billing_context();
|
|
stale.default_price_per_request = Some(1.0);
|
|
state.remember_billing_model_context(key.clone(), Some(stale), load_epoch, &old_flight);
|
|
|
|
let cached = state
|
|
.cached_billing_model_context(&key)
|
|
.expect("replacement should publish")
|
|
.expect("billing context should exist");
|
|
assert_eq!(cached.default_price_per_request, Some(2.0));
|
|
}
|
|
|
|
#[test]
|
|
fn billing_model_context_inflight_limit_rejects_only_new_keys() {
|
|
let state = GatewayDataState::default();
|
|
let mut leaders =
|
|
Vec::with_capacity(GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT);
|
|
for index in 0..GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT {
|
|
let key = billing_cache_key(format!("model-{index}"));
|
|
match state.register_billing_model_context_inflight(&key) {
|
|
BillingModelContextInflightRegistration::Leader(guard) => leaders.push(guard),
|
|
_ => panic!("unique key below the hard limit should lead"),
|
|
}
|
|
}
|
|
|
|
assert!(matches!(
|
|
state.register_billing_model_context_inflight(&billing_cache_key("overflow")),
|
|
BillingModelContextInflightRegistration::Saturated
|
|
));
|
|
assert!(matches!(
|
|
state.register_billing_model_context_inflight(&billing_cache_key("model-0")),
|
|
BillingModelContextInflightRegistration::Follower(_)
|
|
));
|
|
assert_eq!(
|
|
state
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap()
|
|
.len(),
|
|
GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT
|
|
);
|
|
|
|
drop(leaders);
|
|
assert!(state
|
|
.billing_model_context_cache
|
|
.inflight
|
|
.lock()
|
|
.unwrap()
|
|
.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_different_keys_load_concurrently() {
|
|
let repository = Arc::new(ConcurrentBillingContextRepository {
|
|
calls: AtomicUsize::new(0),
|
|
context: billing_context(),
|
|
entered: Barrier::new(3),
|
|
release: Barrier::new(3),
|
|
});
|
|
let state = Arc::new(GatewayDataState::with_billing_reader_for_tests(
|
|
repository.clone(),
|
|
));
|
|
let task_a = {
|
|
let state = Arc::clone(&state);
|
|
tokio::spawn(async move {
|
|
state
|
|
.find_billing_model_context("provider-1", Some("key-1"), "model-a")
|
|
.await
|
|
})
|
|
};
|
|
let task_b = {
|
|
let state = Arc::clone(&state);
|
|
tokio::spawn(async move {
|
|
state
|
|
.find_billing_model_context("provider-1", Some("key-1"), "model-b")
|
|
.await
|
|
})
|
|
};
|
|
|
|
tokio::time::timeout(Duration::from_secs(1), repository.entered.wait())
|
|
.await
|
|
.expect("different cache keys should enter the repository concurrently");
|
|
assert_eq!(repository.calls.load(Ordering::Acquire), 2);
|
|
repository.release.wait().await;
|
|
task_a
|
|
.await
|
|
.expect("first lookup should join")
|
|
.expect("first lookup should succeed");
|
|
task_b
|
|
.await
|
|
.expect("second lookup should join")
|
|
.expect("second lookup should succeed");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn billing_model_context_cache_coalesces_concurrent_loads() {
|
|
let repository = Arc::new(SlowBillingContextRepository {
|
|
calls: AtomicUsize::new(0),
|
|
context: billing_context(),
|
|
});
|
|
let state = Arc::new(GatewayDataState::with_billing_reader_for_tests(
|
|
repository.clone(),
|
|
));
|
|
let mut tasks = Vec::new();
|
|
|
|
for _ in 0..16 {
|
|
let state = Arc::clone(&state);
|
|
tasks.push(tokio::spawn(async move {
|
|
state
|
|
.find_billing_model_context("provider-1", Some("key-1"), "gpt-5")
|
|
.await
|
|
.expect("billing context lookup should succeed")
|
|
.expect("billing context should exist");
|
|
}));
|
|
}
|
|
|
|
for task in tasks {
|
|
task.await.expect("lookup task should complete");
|
|
}
|
|
|
|
assert_eq!(repository.calls.load(Ordering::Acquire), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn global_model_price_update_invalidates_inflight_billing_context_cache() {
|
|
let mut initial_context = billing_context();
|
|
initial_context.default_price_per_request = None;
|
|
initial_context.default_tiered_pricing = None;
|
|
let repository = Arc::new(BlockedBillingContextRepository {
|
|
calls: AtomicUsize::new(0),
|
|
context: Mutex::new(initial_context),
|
|
first_read: Barrier::new(2),
|
|
release_first_read: Barrier::new(2),
|
|
});
|
|
let mut state = GatewayDataState::with_billing_reader_for_tests(repository.clone());
|
|
let stored_global_model = StoredAdminGlobalModel::new(
|
|
"global-model-1".to_string(),
|
|
"gpt-5".to_string(),
|
|
"GPT-5".to_string(),
|
|
true,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
1,
|
|
1,
|
|
0,
|
|
Some(1_711_000_000),
|
|
Some(1_711_000_000),
|
|
)
|
|
.expect("stored global model should build");
|
|
state.global_model_writer = Some(Arc::new(
|
|
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
|
|
.with_admin_global_models([stored_global_model]),
|
|
));
|
|
let state = Arc::new(state);
|
|
let lookup_state = Arc::clone(&state);
|
|
let stale_lookup = tokio::spawn(async move {
|
|
lookup_state
|
|
.find_billing_model_context("provider-1", Some("key-1"), "gpt-5")
|
|
.await
|
|
});
|
|
repository.first_read.wait().await;
|
|
|
|
let updated_pricing =
|
|
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
|
|
let update = UpdateAdminGlobalModelRecord::new(
|
|
"global-model-1".to_string(),
|
|
"GPT-5".to_string(),
|
|
true,
|
|
None,
|
|
Some(updated_pricing.clone()),
|
|
None,
|
|
None,
|
|
)
|
|
.expect("global model update should build");
|
|
repository
|
|
.context
|
|
.lock()
|
|
.expect("mutable billing context lock")
|
|
.default_tiered_pricing = Some(updated_pricing.clone());
|
|
state
|
|
.update_admin_global_model(&update)
|
|
.await
|
|
.expect("global model price update should succeed");
|
|
repository.release_first_read.wait().await;
|
|
|
|
let before = stale_lookup
|
|
.await
|
|
.expect("initial billing lookup task should complete")
|
|
.expect("initial billing lookup should succeed")
|
|
.expect("initial billing context should exist");
|
|
assert_eq!(before.default_tiered_pricing, None);
|
|
assert_eq!(repository.calls.load(Ordering::Acquire), 1);
|
|
|
|
let after = state
|
|
.find_billing_model_context("provider-1", Some("key-1"), "gpt-5")
|
|
.await
|
|
.expect("updated billing lookup should succeed")
|
|
.expect("updated billing context should exist");
|
|
assert_eq!(after.default_tiered_pricing, Some(updated_pricing));
|
|
assert_eq!(repository.calls.load(Ordering::Acquire), 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn lists_non_admin_export_users_from_user_reader() {
|
|
let repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![
|
|
StoredUserExportRow::new(
|
|
"user-1".to_string(),
|
|
Some("[email protected]".to_string()),
|
|
true,
|
|
"alice".to_string(),
|
|
Some("hash".to_string()),
|
|
"user".to_string(),
|
|
"local".to_string(),
|
|
Some(serde_json::json!(["openai"])),
|
|
Some(serde_json::json!(["openai:chat"])),
|
|
Some(serde_json::json!(["gpt-4.1"])),
|
|
Some(60),
|
|
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})),
|
|
true,
|
|
)
|
|
.expect("user export row should build"),
|
|
]));
|
|
let state = GatewayDataState::with_user_reader_for_tests(repository);
|
|
|
|
let rows = state
|
|
.list_non_admin_export_users()
|
|
.await
|
|
.expect("export users should succeed");
|
|
|
|
assert_eq!(rows.len(), 1);
|
|
assert_eq!(rows[0].username, "alice");
|
|
assert!(rows[0].email_verified);
|
|
assert_eq!(rows[0].password_hash.as_deref(), Some("hash"));
|
|
assert_eq!(rows[0].allowed_models, Some(vec!["gpt-4.1".to_string()]));
|
|
assert_eq!(
|
|
rows[0].model_capability_settings,
|
|
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}}))
|
|
);
|
|
}
|
|
}
|