i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add multi-database data layer

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

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

Verified with cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace.
This commit is contained in:
fawney19
2026-05-05 18:27:36 +08:00
parent 099653f732
commit fce7e959e5
372 changed files with 86217 additions and 21160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4610
apps/aether-gateway/src/maintenance/runtime.rs
View File

File diff suppressed because it is too large Load Diff

16
apps/aether-gateway/src/maintenance/runtime/audit_cleanup.rs
View File

@@ -3,26 +3,14 @@ use chrono::{DateTime, Utc};
use crate::data::GatewayDataState;
use super::{
postgres_error, system_config_bool, system_config_u64, system_config_usize,
DELETE_AUDIT_LOGS_BEFORE_SQL,
};
use super::{system_config_bool, system_config_u64, system_config_usize};
pub(crate) async fn cleanup_audit_logs_once(
data: &GatewayDataState,
) -> Result<usize, DataLayerError> {
cleanup_audit_logs_with(data, |cutoff_time, delete_limit| async move {
let Some(pool) = data.postgres_pool() else {
return Ok(0);
};
let deleted = sqlx::query(DELETE_AUDIT_LOGS_BEFORE_SQL)
.bind(cutoff_time)
.bind(i64::try_from(delete_limit).unwrap_or(i64::MAX))
.execute(&pool)
data.delete_audit_logs_before(cutoff_time.timestamp().max(0) as u64, delete_limit)
.await
.map_err(postgres_error)?
.rows_affected();
Ok(usize::try_from(deleted).unwrap_or(usize::MAX))
})
.await
}

28
apps/aether-gateway/src/maintenance/runtime/db_maintenance.rs
View File

@@ -3,7 +3,7 @@ use tracing::warn;
use crate::data::GatewayDataState;
use super::{postgres_error, system_config_bool, DB_MAINTENANCE_TABLES};
use super::{system_config_bool, DB_MAINTENANCE_TABLES};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct DbMaintenanceRunSummary {
@@ -14,25 +14,25 @@ pub(super) struct DbMaintenanceRunSummary {
pub(super) async fn perform_db_maintenance_once(
data: &GatewayDataState,
) -> Result<DbMaintenanceRunSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
if !data.has_database_maintenance_backend() {
return Ok(DbMaintenanceRunSummary {
attempted: 0,
succeeded: 0,
});
};
}
run_db_maintenance_with(data, |table_name| {
let pool = pool.clone();
async move {
let statement = format!("VACUUM ANALYZE {table_name}");
sqlx::raw_sql(&statement)
.execute(&pool)
.await
.map_err(postgres_error)?;
Ok(())
}
if !system_config_bool(data, "enable_db_maintenance", true).await? {
return Ok(DbMaintenanceRunSummary {
attempted: 0,
succeeded: 0,
});
}
let summary = data.run_database_maintenance(DB_MAINTENANCE_TABLES).await?;
Ok(DbMaintenanceRunSummary {
attempted: summary.attempted,
succeeded: summary.succeeded,
})
.await
}
pub(super) async fn run_db_maintenance_with<F, Fut>(

138
apps/aether-gateway/src/maintenance/runtime/pending_cleanup.rs
View File

@@ -1,19 +1,11 @@
use std::collections::HashSet;
use chrono::Utc;
use futures_util::TryStreamExt;
use sqlx::Row;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
pending_cleanup_batch_size, pending_cleanup_timeout_minutes, postgres_error,
SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL, SELECT_STALE_PENDING_USAGE_BATCH_SQL,
UPDATE_FAILED_PENDING_CANDIDATES_SQL, UPDATE_FAILED_STALE_USAGE_SQL,
UPDATE_FAILED_VOID_STALE_USAGE_SQL, UPDATE_RECOVERED_STALE_USAGE_SQL,
UPDATE_RECOVERED_STREAMING_CANDIDATES_SQL,
};
use super::{pending_cleanup_batch_size, pending_cleanup_timeout_minutes};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub(crate) struct PendingCleanupSummary {
@@ -47,123 +39,27 @@ pub(super) struct PendingCleanupBatchPlan {
pub(crate) async fn cleanup_stale_pending_requests_once(
data: &GatewayDataState,
) -> Result<PendingCleanupSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
if !data.has_usage_writer() {
return Ok(PendingCleanupSummary::default());
};
}
let timeout_minutes = pending_cleanup_timeout_minutes(data).await?;
let batch_size = pending_cleanup_batch_size(data).await?;
let cutoff_time =
Utc::now() - chrono::Duration::minutes(i64::try_from(timeout_minutes).unwrap_or(i64::MAX));
let active_statuses = vec!["pending", "streaming"];
let mut summary = PendingCleanupSummary::default();
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let cutoff_unix_secs = now_unix_secs.saturating_sub(timeout_minutes.saturating_mul(60));
let summary = data
.cleanup_stale_pending_requests(
cutoff_unix_secs,
now_unix_secs,
timeout_minutes,
batch_size,
)
.await?;
loop {
let mut tx = pool.begin().await.map_err(postgres_error)?;
let stale_rows = {
let mut stale_rows_stream = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
.bind(active_statuses.clone())
.bind(cutoff_time)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch(&mut *tx);
let mut stale_rows = Vec::new();
while let Some(row) = stale_rows_stream.try_next().await.map_err(postgres_error)? {
stale_rows.push(row);
}
stale_rows
};
if stale_rows.is_empty() {
tx.rollback().await.map_err(postgres_error)?;
break;
}
let stale_rows = stale_rows
.into_iter()
.map(|row| {
Ok::<StalePendingUsageRow, DataLayerError>(StalePendingUsageRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
status: row.try_get::<String, _>("status").map_err(postgres_error)?,
billing_status: row
.try_get::<String, _>("billing_status")
.map_err(postgres_error)?,
})
})
.collect::<Result<Vec<_>, DataLayerError>>()?;
let request_ids = stale_rows
.iter()
.map(|row| row.request_id.clone())
.collect::<Vec<_>>();
let completed_request_ids = if request_ids.is_empty() {
HashSet::new()
} else {
{
let mut completed_rows = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(request_ids)
.fetch(&mut *tx);
let mut completed_request_ids = HashSet::new();
while let Some(row) = completed_rows.try_next().await.map_err(postgres_error)? {
if let Ok(request_id) = row.try_get::<String, _>("request_id") {
completed_request_ids.insert(request_id);
}
}
completed_request_ids
}
};
let plan = plan_pending_cleanup_batch(stale_rows, &completed_request_ids, timeout_minutes);
let now = Utc::now();
for usage_id in &plan.recovered_usage_ids {
sqlx::query(UPDATE_RECOVERED_STALE_USAGE_SQL)
.bind(usage_id)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
for failed_row in &plan.failed_usage_rows {
if failed_row.should_void_billing {
sqlx::query(UPDATE_FAILED_VOID_STALE_USAGE_SQL)
.bind(&failed_row.id)
.bind(&failed_row.error_message)
.bind(now)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
} else {
sqlx::query(UPDATE_FAILED_STALE_USAGE_SQL)
.bind(&failed_row.id)
.bind(&failed_row.error_message)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
}
if !plan.recovered_request_ids.is_empty() {
sqlx::query(UPDATE_RECOVERED_STREAMING_CANDIDATES_SQL)
.bind(plan.recovered_request_ids.clone())
.bind(now)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
if !plan.failed_request_ids.is_empty() {
sqlx::query(UPDATE_FAILED_PENDING_CANDIDATES_SQL)
.bind(plan.failed_request_ids.clone())
.bind(now)
.bind(active_statuses.clone())
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
tx.commit().await.map_err(postgres_error)?;
summary.failed += plan.failed_usage_rows.len();
summary.recovered += plan.recovered_usage_ids.len();
}
Ok(summary)
Ok(PendingCleanupSummary {
failed: summary.failed,
recovered: summary.recovered,
})
}
pub(super) fn plan_pending_cleanup_batch(

2
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
View File

@@ -1,7 +1,7 @@
use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data::redis::{RedisKvRunner, RedisLockLease, RedisLockRunner};
use aether_data::driver::redis::{RedisKvRunner, RedisLockLease, RedisLockRunner};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};

9
apps/aether-gateway/src/maintenance/runtime/runners.rs
View File

@@ -12,7 +12,7 @@ use super::{
perform_provider_checkin_once, perform_stats_aggregation_once,
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
perform_wallet_daily_usage_aggregation_once, record_proxy_upgrade_traffic_success,
summarize_postgres_pool,
summarize_database_pool,
};
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
@@ -212,20 +212,21 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
}
pub(super) fn run_pool_monitor_once(data: &GatewayDataState) {
let Some(summary) = summarize_postgres_pool(data) else {
let Some(summary) = summarize_database_pool(data) else {
return;
};
info!(
event_name = "postgres_pool_sampled",
event_name = "database_pool_sampled",
log_type = "ops",
worker = "pool_monitor",
driver = %summary.driver,
checked_out = summary.checked_out,
pool_size = summary.pool_size,
idle = summary.idle,
max_connections = summary.max_connections,
usage_rate = summary.usage_rate,
"gateway postgres pool status"
"gateway database pool status"
);
}

658
apps/aether-gateway/src/maintenance/runtime/stats_daily.rs
View File

@@ -1,662 +1,24 @@
use chrono::{DateTime, Utc};
use sqlx::Row;
use uuid::Uuid;
use aether_data::{DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary};
use chrono::Utc;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
postgres_error, stats_aggregation_target_day, system_config_bool, PercentileSummary,
StatsAggregationSummary, DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL, INSERT_STATS_DAILY_ERROR_SQL,
INSERT_STATS_SUMMARY_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
SELECT_LATEST_STATS_DAILY_DATE_SQL, SELECT_NEXT_STATS_DAILY_BUCKET_SQL,
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL, SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL,
SELECT_STATS_SUMMARY_TOTALS_SQL, UPDATE_STATS_SUMMARY_SQL, UPSERT_STATS_DAILY_API_KEY_SQL,
UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL, UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_SQL,
UPSERT_STATS_DAILY_COST_SAVINGS_PROVIDER_SQL, UPSERT_STATS_DAILY_COST_SAVINGS_SQL,
UPSERT_STATS_DAILY_MODEL_PROVIDER_SQL, UPSERT_STATS_DAILY_MODEL_SQL,
UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
UPSERT_STATS_USER_DAILY_API_FORMAT_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_COST_SAVINGS_SQL,
UPSERT_STATS_USER_DAILY_MODEL_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_MODEL_SQL,
UPSERT_STATS_USER_DAILY_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_SQL,
UPSERT_STATS_USER_SUMMARY_SQL,
};
use super::{stats_aggregation_target_day, system_config_bool};
pub(super) async fn perform_stats_aggregation_once(
data: &GatewayDataState,
) -> Result<Option<StatsAggregationSummary>, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
) -> Result<Option<StatsDailyAggregationSummary>, DataLayerError> {
if !data.has_stats_daily_aggregation_backend() {
return Ok(None);
};
}
if !system_config_bool(data, "enable_stats_aggregation", true).await? {
return Ok(None);
}
let now_utc = Utc::now();
let target_day_utc = stats_aggregation_target_day(now_utc);
let Some(day_start_utc) = next_stats_aggregation_day(&pool, target_day_utc)
.await
.map_err(postgres_error)?
else {
return Ok(None);
};
perform_stats_aggregation_for_day(&pool, day_start_utc, now_utc)
.await
.map(Some)
.map_err(postgres_error)
}
async fn next_stats_aggregation_day(
pool: &aether_data::postgres::PostgresPool,
target_day_utc: DateTime<Utc>,
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
let latest_row = sqlx::query(SELECT_LATEST_STATS_DAILY_DATE_SQL)
.fetch_one(pool)
.await?;
let latest_day = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_date")?;
let search_from = latest_day
.map(|value| value + chrono::Duration::days(1))
.unwrap_or_else(|| {
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
});
let search_until = target_day_utc + chrono::Duration::days(1);
if search_from >= search_until {
return Ok(None);
}
let next_row = sqlx::query(SELECT_NEXT_STATS_DAILY_BUCKET_SQL)
.bind(search_from)
.bind(search_until)
.fetch_one(pool)
.await?;
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
Ok(next_bucket.filter(|value| *value <= target_day_utc))
}
async fn perform_stats_aggregation_for_day(
pool: &aether_data::postgres::PostgresPool,
day_start_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<StatsAggregationSummary, sqlx::Error> {
let day_end_utc = day_start_utc + chrono::Duration::days(1);
let mut tx = pool.begin().await?;
let aggregate_row = sqlx::query(SELECT_STATS_DAILY_AGGREGATE_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_one(&mut *tx)
.await?;
let total_requests = aggregate_row.try_get::<i64, _>("total_requests")?;
let error_requests = aggregate_row.try_get::<i64, _>("error_requests")?;
let success_requests = total_requests.saturating_sub(error_requests);
let fallback_count = sqlx::query(SELECT_STATS_DAILY_FALLBACK_COUNT_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(vec!["success", "failed"])
.fetch_one(&mut *tx)
.await?
.try_get::<i64, _>("fallback_count")?;
let response_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await?;
let first_byte_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await?;
sqlx::query(UPSERT_STATS_DAILY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(day_start_utc)
.bind(total_requests)
.bind(aggregate_row.try_get::<i64, _>("cache_hit_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("cache_hit_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_hit_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_read_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_total_input_context")?)
.bind(aggregate_row.try_get::<f64, _>("completed_cache_creation_cost")?)
.bind(aggregate_row.try_get::<f64, _>("completed_cache_read_cost")?)
.bind(aggregate_row.try_get::<f64, _>("settled_total_cost")?)
.bind(aggregate_row.try_get::<i64, _>("settled_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("settled_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_output_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_cache_read_tokens")?)
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
.bind(success_requests)
.bind(error_requests)
.bind(aggregate_row.try_get::<i64, _>("input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("effective_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("output_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_5m_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_1h_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_read_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("total_input_context")?)
.bind(aggregate_row.try_get::<f64, _>("total_cost")?)
.bind(aggregate_row.try_get::<f64, _>("actual_total_cost")?)
.bind(aggregate_row.try_get::<f64, _>("input_cost")?)
.bind(aggregate_row.try_get::<f64, _>("output_cost")?)
.bind(aggregate_row.try_get::<f64, _>("cache_creation_cost")?)
.bind(aggregate_row.try_get::<f64, _>("cache_read_cost")?)
.bind(aggregate_row.try_get::<f64, _>("response_time_sum_ms")?)
.bind(aggregate_row.try_get::<i64, _>("response_time_samples")?)
.bind(aggregate_row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(response_percentiles.p50)
.bind(response_percentiles.p90)
.bind(response_percentiles.p99)
.bind(first_byte_percentiles.p50)
.bind(first_byte_percentiles.p90)
.bind(first_byte_percentiles.p99)
.bind(fallback_count)
.bind(aggregate_row.try_get::<i64, _>("unique_models")?)
.bind(aggregate_row.try_get::<i64, _>("unique_providers")?)
.bind(true)
.bind(now_utc)
.bind(now_utc)
.bind(now_utc)
.execute(&mut *tx)
.await?;
let model_rows =
upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let provider_rows =
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_cost_savings_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_daily_cost_savings_model_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
let api_key_rows =
upsert_stats_daily_api_key_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let error_rows =
refresh_stats_daily_error_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let user_rows =
upsert_stats_user_daily_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_user_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_cost_savings_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
upsert_stats_user_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_user_daily_cost_savings_model_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
upsert_stats_user_daily_api_format_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc).await?;
refresh_stats_user_summary_rows(&mut tx, day_end_utc, now_utc).await?;
tx.commit().await?;
Ok(StatsAggregationSummary {
day_start_utc,
total_requests,
model_rows,
provider_rows,
api_key_rows,
error_rows,
user_rows,
data.aggregate_stats_daily(&StatsDailyAggregationInput {
target_day_utc: stats_aggregation_target_day(now_utc),
aggregated_at: now_utc,
})
}
async fn fetch_stats_daily_percentiles(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
sql: &str,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
) -> Result<PercentileSummary, sqlx::Error> {
let row = sqlx::query(sql)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_one(&mut **tx)
.await?;
let sample_count = row.try_get::<i64, _>("sample_count")?;
if sample_count < 10 {
return Ok(PercentileSummary::default());
}
Ok(PercentileSummary {
p50: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p50")?),
p90: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p90")?),
p99: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p99")?),
})
}
fn percentile_ms_to_i64(value: Option<f64>) -> Option<i64> {
value.and_then(|raw| raw.is_finite().then_some(raw.floor() as i64))
}
async fn upsert_stats_daily_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_api_key_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn refresh_stats_daily_error_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
sqlx::query(DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL)
.bind(day_start_utc)
.execute(&mut **tx)
.await?;
let rows_affected = sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_api_format_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_API_FORMAT_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn refresh_stats_summary_row(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
cutoff_date: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<(), sqlx::Error> {
let totals_row = sqlx::query(SELECT_STATS_SUMMARY_TOTALS_SQL)
.bind(cutoff_date)
.fetch_one(&mut **tx)
.await?;
let entity_counts_row = sqlx::query(SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL)
.fetch_one(&mut **tx)
.await?;
let existing_summary_id = sqlx::query_scalar::<_, String>(SELECT_EXISTING_STATS_SUMMARY_ID_SQL)
.fetch_optional(&mut **tx)
.await?;
let all_time_requests = totals_row.try_get::<i64, _>("all_time_requests")?;
let all_time_success_requests = totals_row.try_get::<i64, _>("all_time_success_requests")?;
let all_time_error_requests = totals_row.try_get::<i64, _>("all_time_error_requests")?;
let all_time_input_tokens = totals_row.try_get::<i64, _>("all_time_input_tokens")?;
let all_time_output_tokens = totals_row.try_get::<i64, _>("all_time_output_tokens")?;
let all_time_cache_creation_tokens =
totals_row.try_get::<i64, _>("all_time_cache_creation_tokens")?;
let all_time_cache_read_tokens = totals_row.try_get::<i64, _>("all_time_cache_read_tokens")?;
let all_time_cost = totals_row.try_get::<f64, _>("all_time_cost")?;
let all_time_actual_cost = totals_row.try_get::<f64, _>("all_time_actual_cost")?;
let total_users = entity_counts_row.try_get::<i64, _>("total_users")?;
let active_users = entity_counts_row.try_get::<i64, _>("active_users")?;
let total_api_keys = entity_counts_row.try_get::<i64, _>("total_api_keys")?;
let active_api_keys = entity_counts_row.try_get::<i64, _>("active_api_keys")?;
if let Some(summary_id) = existing_summary_id {
sqlx::query(UPDATE_STATS_SUMMARY_SQL)
.bind(summary_id)
.bind(cutoff_date)
.bind(all_time_requests)
.bind(all_time_success_requests)
.bind(all_time_error_requests)
.bind(all_time_input_tokens)
.bind(all_time_output_tokens)
.bind(all_time_cache_creation_tokens)
.bind(all_time_cache_read_tokens)
.bind(all_time_cost)
.bind(all_time_actual_cost)
.bind(total_users)
.bind(active_users)
.bind(total_api_keys)
.bind(active_api_keys)
.bind(now_utc)
.execute(&mut **tx)
.await?;
} else {
sqlx::query(INSERT_STATS_SUMMARY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(cutoff_date)
.bind(all_time_requests)
.bind(all_time_success_requests)
.bind(all_time_error_requests)
.bind(all_time_input_tokens)
.bind(all_time_output_tokens)
.bind(all_time_cache_creation_tokens)
.bind(all_time_cache_read_tokens)
.bind(all_time_cost)
.bind(all_time_actual_cost)
.bind(total_users)
.bind(active_users)
.bind(total_api_keys)
.bind(active_api_keys)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(())
}
async fn refresh_stats_user_summary_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
cutoff_date: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<(), sqlx::Error> {
sqlx::query(UPSERT_STATS_USER_SUMMARY_SQL)
.bind(cutoff_date)
.bind(now_utc)
.execute(&mut **tx)
.await?;
Ok(())
.await
}

217
apps/aether-gateway/src/maintenance/runtime/stats_hourly.rs
View File

@@ -1,223 +1,24 @@
use chrono::{DateTime, Utc};
use sqlx::Row;
use uuid::Uuid;
use aether_data::{DataLayerError, StatsHourlyAggregationInput, StatsHourlyAggregationSummary};
use chrono::Utc;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_LATEST_STATS_HOURLY_HOUR_SQL,
SELECT_NEXT_STATS_HOURLY_BUCKET_SQL, SELECT_STATS_HOURLY_AGGREGATE_SQL,
UPSERT_STATS_HOURLY_MODEL_SQL, UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL,
UPSERT_STATS_HOURLY_USER_MODEL_SQL, UPSERT_STATS_HOURLY_USER_SQL,
};
#[derive(Debug, Clone, PartialEq)]
pub(super) struct StatsHourlyAggregationSummary {
pub(super) hour_utc: DateTime<Utc>,
pub(super) total_requests: i64,
pub(super) user_rows: usize,
pub(super) user_model_rows: usize,
pub(super) model_rows: usize,
pub(super) provider_rows: usize,
}
use super::{stats_hourly_aggregation_target_hour, system_config_bool};
pub(super) async fn perform_stats_hourly_aggregation_once(
data: &GatewayDataState,
) -> Result<Option<StatsHourlyAggregationSummary>, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
if !data.has_stats_hourly_aggregation_backend() {
return Ok(None);
};
}
if !system_config_bool(data, "enable_stats_aggregation", true).await? {
return Ok(None);
}
let now_utc = Utc::now();
let target_hour_utc = stats_hourly_aggregation_target_hour(now_utc);
let Some(hour_utc) = next_stats_hourly_bucket(&pool, target_hour_utc)
.await
.map_err(postgres_error)?
else {
return Ok(None);
};
perform_stats_hourly_aggregation_for_hour(&pool, hour_utc, now_utc)
.await
.map(Some)
.map_err(postgres_error)
}
async fn next_stats_hourly_bucket(
pool: &aether_data::postgres::PostgresPool,
target_hour_utc: DateTime<Utc>,
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
let latest_row = sqlx::query(SELECT_LATEST_STATS_HOURLY_HOUR_SQL)
.fetch_one(pool)
.await?;
let latest_hour = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_hour")?;
let search_from = latest_hour
.map(|value| value + chrono::Duration::hours(1))
.unwrap_or_else(|| {
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
});
let search_until = target_hour_utc + chrono::Duration::hours(1);
if search_from >= search_until {
return Ok(None);
}
let next_row = sqlx::query(SELECT_NEXT_STATS_HOURLY_BUCKET_SQL)
.bind(search_from)
.bind(search_until)
.fetch_one(pool)
.await?;
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
Ok(next_bucket.filter(|value| *value <= target_hour_utc))
}
async fn perform_stats_hourly_aggregation_for_hour(
pool: &aether_data::postgres::PostgresPool,
hour_utc: DateTime<Utc>,
aggregated_at: DateTime<Utc>,
) -> Result<StatsHourlyAggregationSummary, sqlx::Error> {
let hour_end = hour_utc + chrono::Duration::hours(1);
let mut tx = pool.begin().await?;
let row = sqlx::query(SELECT_STATS_HOURLY_AGGREGATE_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_one(&mut *tx)
.await?;
let total_requests = row.try_get::<i64, _>("total_requests")?;
let error_requests = row.try_get::<i64, _>("error_requests")?;
let success_requests = total_requests.saturating_sub(error_requests);
sqlx::query(UPSERT_STATS_HOURLY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(total_requests)
.bind(row.try_get::<i64, _>("cache_hit_total_requests")?)
.bind(row.try_get::<i64, _>("cache_hit_requests")?)
.bind(row.try_get::<i64, _>("completed_total_requests")?)
.bind(row.try_get::<i64, _>("completed_cache_hit_requests")?)
.bind(row.try_get::<i64, _>("completed_input_tokens")?)
.bind(row.try_get::<i64, _>("completed_cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("completed_cache_read_tokens")?)
.bind(row.try_get::<i64, _>("completed_total_input_context")?)
.bind(row.try_get::<f64, _>("completed_cache_creation_cost")?)
.bind(row.try_get::<f64, _>("completed_cache_read_cost")?)
.bind(row.try_get::<f64, _>("settled_total_cost")?)
.bind(row.try_get::<i64, _>("settled_total_requests")?)
.bind(row.try_get::<i64, _>("settled_input_tokens")?)
.bind(row.try_get::<i64, _>("settled_output_tokens")?)
.bind(row.try_get::<i64, _>("settled_cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("settled_cache_read_tokens")?)
.bind(row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
.bind(row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
.bind(success_requests)
.bind(error_requests)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(row.try_get::<f64, _>("actual_total_cost")?)
.bind(row.try_get::<f64, _>("response_time_sum_ms")?)
.bind(row.try_get::<i64, _>("response_time_samples")?)
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(true)
.bind(aggregated_at)
.bind(aggregated_at)
.bind(aggregated_at)
.execute(&mut *tx)
.await?;
let user_rows =
upsert_stats_hourly_user_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let user_model_rows =
upsert_stats_hourly_user_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let model_rows =
upsert_stats_hourly_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let provider_rows =
upsert_stats_hourly_provider_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
tx.commit().await?;
Ok(StatsHourlyAggregationSummary {
hour_utc,
total_requests,
user_rows,
user_model_rows,
model_rows,
provider_rows,
data.aggregate_stats_hourly(&StatsHourlyAggregationInput {
target_hour_utc: stats_hourly_aggregation_target_hour(now_utc),
aggregated_at: now_utc,
})
}
async fn upsert_stats_hourly_user_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_user_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_MODEL_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
.await
}

44
apps/aether-gateway/src/maintenance/runtime/tests.rs
View File

@@ -25,12 +25,10 @@ use super::{
spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_postgres_pool,
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
usage_cleanup_settings, usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState,
DbMaintenanceRunSummary, FailedPendingUsageRow, GatewayDataState,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings,
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, SELECT_STALE_PENDING_USAGE_BATCH_SQL,
UPDATE_FAILED_VOID_STALE_USAGE_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL, USAGE_CLEANUP_HOUR,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
@@ -41,12 +39,12 @@ async fn spawn_audit_cleanup_worker_skips_when_postgres_unavailable() {
}
#[tokio::test]
async fn spawn_db_maintenance_worker_skips_when_postgres_unavailable() {
async fn spawn_db_maintenance_worker_skips_when_database_maintenance_unavailable() {
assert!(spawn_db_maintenance_worker(Arc::new(GatewayDataState::disabled())).is_none());
}
#[tokio::test]
async fn spawn_pending_cleanup_worker_skips_when_postgres_unavailable() {
async fn spawn_pending_cleanup_worker_skips_when_usage_writer_unavailable() {
assert!(spawn_pending_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
}
@@ -89,24 +87,6 @@ async fn spawn_pool_quota_probe_worker_skips_when_provider_catalog_unavailable()
assert!(spawn_pool_quota_probe_worker(state).is_none());
}
#[test]
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
assert!(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL.contains("JOIN usage_settlement_snapshots"));
assert!(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL.contains("usage_settlement_snapshots.wallet_id"));
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL.contains("JOIN usage_settlement_snapshots"));
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL
.contains("usage_settlement_snapshots.wallet_id = ledgers.wallet_id"));
}
#[test]
fn pending_cleanup_queries_use_settlement_snapshots_for_billing_authority() {
assert!(SELECT_STALE_PENDING_USAGE_BATCH_SQL.contains("LEFT JOIN usage_settlement_snapshots"));
assert!(SELECT_STALE_PENDING_USAGE_BATCH_SQL
.contains("COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status)"));
assert!(UPDATE_FAILED_VOID_STALE_USAGE_SQL.contains("INSERT INTO usage_settlement_snapshots"));
assert!(UPDATE_FAILED_VOID_STALE_USAGE_SQL.contains("billing_status = EXCLUDED.billing_status"));
}
fn sample_connected_proxy_node(
node_id: &str,
heartbeat_interval: i32,
@@ -548,24 +528,25 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
}
#[tokio::test]
async fn spawn_stats_aggregation_worker_skips_when_postgres_unavailable() {
async fn spawn_stats_aggregation_worker_skips_when_stats_daily_backend_unavailable() {
assert!(spawn_stats_aggregation_worker(Arc::new(GatewayDataState::disabled())).is_none());
}
#[tokio::test]
async fn spawn_stats_hourly_aggregation_worker_skips_when_postgres_unavailable() {
async fn spawn_stats_hourly_aggregation_worker_skips_when_stats_hourly_backend_unavailable() {
assert!(
spawn_stats_hourly_aggregation_worker(Arc::new(GatewayDataState::disabled())).is_none()
);
}
#[tokio::test]
async fn spawn_usage_cleanup_worker_skips_when_postgres_unavailable() {
async fn spawn_usage_cleanup_worker_skips_when_usage_writer_unavailable() {
assert!(spawn_usage_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
}
#[tokio::test]
async fn spawn_wallet_daily_usage_aggregation_worker_skips_when_postgres_unavailable() {
async fn spawn_wallet_daily_usage_aggregation_worker_skips_when_wallet_daily_usage_backend_unavailable(
) {
assert!(
spawn_wallet_daily_usage_aggregation_worker(Arc::new(GatewayDataState::disabled()))
.is_none()
@@ -780,9 +761,9 @@ fn usage_cleanup_window_uses_non_overlapping_ranges() {
}
#[tokio::test]
async fn summarize_postgres_pool_uses_busy_connections_for_usage_rate() {
async fn summarize_database_pool_uses_busy_connections_for_usage_rate() {
let data = GatewayDataState::from_config(crate::data::GatewayDataConfig::from_postgres_config(
aether_data::postgres::PostgresPoolConfig {
aether_data::driver::postgres::PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 8,
@@ -795,8 +776,9 @@ async fn summarize_postgres_pool_uses_busy_connections_for_usage_rate() {
))
.expect("gateway data state should build");
let summary = summarize_postgres_pool(&data).expect("pool summary should exist");
let summary = summarize_database_pool(&data).expect("pool summary should exist");
assert_eq!(summary.driver, aether_data::DatabaseDriver::Postgres);
assert_eq!(summary.checked_out, 0);
assert_eq!(summary.pool_size, 0);
assert_eq!(summary.idle, 0);

856
apps/aether-gateway/src/maintenance/runtime/usage_cleanup.rs
View File

@@ -1,865 +1,27 @@
use std::io::Write;
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, UsageBodyField,
};
use aether_data_contracts::repository::usage::UsageCleanupSummary;
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use futures_util::TryStreamExt;
use serde_json::{Map, Value};
use sqlx::Row;
use tracing::warn;
use chrono::Utc;
use crate::data::GatewayDataState;
use super::{
system_config_bool, usage_cleanup_settings, usage_cleanup_window, ExpiredApiKeyRow,
UsageBodyCleanupRow, UsageBodyCompressionRow, UsageCleanupSummary, CLEAR_USAGE_BODY_FIELDS_SQL,
CLEAR_USAGE_HEADER_FIELDS_SQL, CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL, DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL,
DELETE_EXPIRED_API_KEY_SQL, DELETE_OLD_USAGE_RECORDS_SQL, DELETE_USAGE_BODY_BLOBS_SQL,
DISABLE_EXPIRED_API_KEY_SQL, DISABLE_EXPIRED_API_KEY_WALLET_SQL,
EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE, NULLIFY_REQUEST_CANDIDATE_API_KEY_BATCH_SQL,
NULLIFY_USAGE_API_KEY_BATCH_SQL, SELECT_EXPIRED_ACTIVE_API_KEYS_SQL,
SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL, SELECT_USAGE_BODY_COMPRESSION_ROW_SQL,
SELECT_USAGE_HEADER_BATCH_SQL, SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL,
SELECT_USAGE_STALE_BODY_BATCH_SQL, UPDATE_USAGE_BODY_COMPRESSION_SQL,
UPDATE_USAGE_REQUEST_METADATA_SQL, UPSERT_USAGE_BODY_BLOB_SQL,
UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
};
use super::{system_config_bool, usage_cleanup_settings, usage_cleanup_window};
pub(super) async fn perform_usage_cleanup_once(
data: &GatewayDataState,
) -> Result<UsageCleanupSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
if !data.has_usage_writer() {
return Ok(UsageCleanupSummary::default());
};
}
if !system_config_bool(data, "enable_auto_cleanup", true).await? {
return Ok(UsageCleanupSummary::default());
}
let settings = usage_cleanup_settings(data).await?;
let window = usage_cleanup_window(Utc::now(), settings);
let records_deleted =
delete_old_usage_records(&pool, window.log_cutoff, settings.batch_size).await?;
let header_cleaned = cleanup_usage_header_fields(
&pool,
window.header_cutoff,
data.cleanup_usage(
&window,
settings.batch_size,
Some(window.log_cutoff),
settings.auto_delete_expired_keys,
)
.await?;
let legacy_body_refs_migrated = migrate_legacy_usage_body_ref_metadata(
&pool,
window.detail_cutoff,
settings.batch_size,
Some(window.compressed_cutoff),
)
.await?;
let body_cleaned = cleanup_usage_stale_body_fields(
&pool,
window.compressed_cutoff,
settings.batch_size,
Some(window.log_cutoff),
)
.await?;
let body_externalized = compress_usage_body_fields(
&pool,
window.detail_cutoff,
settings.batch_size,
Some(window.compressed_cutoff),
)
.await?;
let keys_cleaned =
match cleanup_expired_api_keys(&pool, settings.auto_delete_expired_keys).await {
Ok(count) => count,
Err(err) => {
warn!(error = %err, "gateway expired api key cleanup failed");
0
}
};
Ok(UsageCleanupSummary {
body_externalized,
legacy_body_refs_migrated,
body_cleaned,
header_cleaned,
keys_cleaned,
records_deleted,
})
}
async fn migrate_legacy_usage_body_ref_metadata(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"gateway usage legacy body-ref migration skipped due to invalid window"
);
return Ok(0);
}
let mut total_migrated = 0usize;
loop {
let rows = sqlx::query(SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch_all(pool)
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| {
Ok(UsageLegacyBodyRefMetadataRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
request_metadata: row
.try_get::<Option<Value>, _>("request_metadata")
.map_err(postgres_error)?,
})
})
.collect::<Result<Vec<_>, DataLayerError>>()?;
if rows.is_empty() {
break;
}
let mut batch_migrated = 0usize;
for row in rows {
let Some(plan) =
migrate_legacy_body_ref_metadata_plan(&row.request_id, row.request_metadata)
else {
continue;
};
let mut tx = pool.begin().await.map_err(postgres_error)?;
if plan.refs.any_present() {
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&row.request_id)
.bind(plan.refs.request_body_ref.as_deref())
.bind(plan.refs.provider_request_body_ref.as_deref())
.bind(plan.refs.response_body_ref.as_deref())
.bind(plan.refs.client_response_body_ref.as_deref())
.bind("ref_backed")
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
let updated = sqlx::query(UPDATE_USAGE_REQUEST_METADATA_SQL)
.bind(&row.id)
.bind(plan.request_metadata)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
tx.commit().await.map_err(postgres_error)?;
if updated > 0 {
batch_migrated += 1;
}
}
total_migrated += batch_migrated;
if batch_migrated == 0 || batch_migrated < batch_size {
break;
}
}
Ok(total_migrated)
}
async fn delete_old_usage_records(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
) -> Result<usize, DataLayerError> {
let mut total_deleted = 0usize;
loop {
let deleted = sqlx::query(DELETE_OLD_USAGE_RECORDS_SQL)
.bind(cutoff_time)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
let deleted = usize::try_from(deleted).unwrap_or(usize::MAX);
total_deleted += deleted;
if deleted < batch_size {
break;
}
}
Ok(total_deleted)
}
async fn cleanup_usage_header_fields(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"gateway usage header cleanup skipped due to invalid window"
);
return Ok(0);
}
let mut total_cleaned = 0usize;
loop {
let mut stream = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch(pool);
let mut rows = Vec::new();
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
rows.push(UsageBodyCleanupRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
});
}
if rows.is_empty() {
break;
}
let ids = rows.iter().map(|row| row.id.clone()).collect::<Vec<_>>();
let request_ids = rows
.iter()
.map(|row| row.request_id.clone())
.collect::<Vec<_>>();
let cleaned = sqlx::query(CLEAR_USAGE_HEADER_FIELDS_SQL)
.bind(ids)
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL)
.bind(&request_ids)
.execute(pool)
.await
.map_err(postgres_error)?;
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
.bind(request_ids)
.execute(pool)
.await
.map_err(postgres_error)?;
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
total_cleaned += cleaned;
if cleaned == 0 || cleaned < batch_size {
break;
}
}
Ok(total_cleaned)
}
async fn cleanup_usage_stale_body_fields(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"gateway usage body cleanup skipped due to invalid window"
);
return Ok(0);
}
let mut total_cleaned = 0usize;
loop {
let mut stream = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch(pool);
let mut rows = Vec::new();
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
rows.push(UsageBodyCleanupRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
});
}
if rows.is_empty() {
break;
}
let ids = rows.iter().map(|row| row.id.clone()).collect::<Vec<_>>();
let request_ids = rows
.iter()
.map(|row| row.request_id.clone())
.collect::<Vec<_>>();
let cleaned = sqlx::query(CLEAR_USAGE_BODY_FIELDS_SQL)
.bind(ids)
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
sqlx::query(DELETE_USAGE_BODY_BLOBS_SQL)
.bind(&request_ids)
.execute(pool)
.await
.map_err(postgres_error)?;
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&request_ids)
.execute(pool)
.await
.map_err(postgres_error)?;
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
.bind(request_ids)
.execute(pool)
.await
.map_err(postgres_error)?;
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
total_cleaned += cleaned;
if cleaned == 0 || cleaned < batch_size {
break;
}
}
Ok(total_cleaned)
}
async fn compress_usage_body_fields(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"gateway usage body compression skipped due to invalid window"
);
return Ok(0);
}
let mut total_compressed = 0usize;
let mut no_progress_count = 0usize;
let batch_size = batch_size.clamp(1, 25);
loop {
let mut stream = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch(pool);
let mut ids = Vec::new();
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
ids.push(row.try_get::<String, _>("id").map_err(postgres_error)?);
}
if ids.is_empty() {
break;
}
let mut batch_success = 0usize;
for id in ids {
let row = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_ROW_SQL)
.bind(&id)
.fetch_optional(pool)
.await
.map_err(postgres_error)?;
let Some(row) = row else {
continue;
};
let row = UsageBodyCompressionRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
request_body: row
.try_get::<Option<Value>, _>("request_body")
.map_err(postgres_error)?,
request_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("request_body_compressed")
.map_err(postgres_error)?,
response_body: row
.try_get::<Option<Value>, _>("response_body")
.map_err(postgres_error)?,
response_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("response_body_compressed")
.map_err(postgres_error)?,
provider_request_body: row
.try_get::<Option<Value>, _>("provider_request_body")
.map_err(postgres_error)?,
provider_request_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
.map_err(postgres_error)?,
client_response_body: row
.try_get::<Option<Value>, _>("client_response_body")
.map_err(postgres_error)?,
client_response_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("client_response_body_compressed")
.map_err(postgres_error)?,
};
let detached = build_usage_body_externalization(&row)?;
if detached.refs.any_present() {
let mut tx = pool.begin().await.map_err(postgres_error)?;
for blob in &detached.blobs {
sqlx::query(UPSERT_USAGE_BODY_BLOB_SQL)
.bind(&blob.body_ref)
.bind(&row.request_id)
.bind(blob.body_field)
.bind(&blob.payload_gzip)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&row.request_id)
.bind(detached.refs.request_body_ref.as_deref())
.bind(detached.refs.provider_request_body_ref.as_deref())
.bind(detached.refs.response_body_ref.as_deref())
.bind(detached.refs.client_response_body_ref.as_deref())
.bind("ref_backed")
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
.bind(&row.id)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
tx.commit().await.map_err(postgres_error)?;
if updated > 0 {
batch_success += 1;
}
continue;
}
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
.bind(&row.id)
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
if updated > 0 {
batch_success += 1;
}
}
if batch_success == 0 {
no_progress_count += 1;
if no_progress_count >= 3 {
warn!(
"gateway usage body compression stopped after repeated zero-progress batches"
);
break;
}
} else {
no_progress_count = 0;
}
total_compressed += batch_success;
}
Ok(total_compressed)
}
fn compress_usage_json_value(value: &Value) -> Result<Vec<u8>, DataLayerError> {
let bytes = serde_json::to_vec(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to serialize usage json for gzip: {err}"))
})?;
let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
encoder.write_all(&bytes).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to gzip usage json: {err}"))
})?;
encoder.finish().map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to finish gzipped usage json: {err}"))
})
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct UsageDetachedBodyBlobWrite {
body_ref: String,
body_field: &'static str,
payload_gzip: Vec<u8>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageDetachedBodyRefs {
request_body_ref: Option<String>,
provider_request_body_ref: Option<String>,
response_body_ref: Option<String>,
client_response_body_ref: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
struct UsageLegacyBodyRefMetadataRow {
id: String,
request_id: String,
request_metadata: Option<Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageLegacyBodyRefMigrationPlan {
refs: UsageDetachedBodyRefs,
request_metadata: Option<Value>,
}
impl UsageDetachedBodyRefs {
fn any_present(&self) -> bool {
self.request_body_ref.is_some()
|| self.provider_request_body_ref.is_some()
|| self.response_body_ref.is_some()
|| self.client_response_body_ref.is_some()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageBodyExternalizationPlan {
blobs: Vec<UsageDetachedBodyBlobWrite>,
refs: UsageDetachedBodyRefs,
}
fn migrate_legacy_body_ref_metadata_plan(
request_id: &str,
request_metadata: Option<Value>,
) -> Option<UsageLegacyBodyRefMigrationPlan> {
let mut metadata = match request_metadata {
Some(Value::Object(object)) => object,
_ => return None,
};
let mut refs = UsageDetachedBodyRefs::default();
let mut removed_any = false;
for field in [
UsageBodyField::RequestBody,
UsageBodyField::ProviderRequestBody,
UsageBodyField::ResponseBody,
UsageBodyField::ClientResponseBody,
] {
let key = field.as_ref_key();
let Some(value) = metadata.remove(key) else {
continue;
};
removed_any = true;
let parsed = value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(parse_usage_body_ref)
.filter(|(parsed_request_id, parsed_field)| {
parsed_request_id == request_id && *parsed_field == field
})
.map(|(parsed_request_id, parsed_field)| {
usage_body_ref(&parsed_request_id, parsed_field)
});
match field {
UsageBodyField::RequestBody => refs.request_body_ref = parsed,
UsageBodyField::ProviderRequestBody => refs.provider_request_body_ref = parsed,
UsageBodyField::ResponseBody => refs.response_body_ref = parsed,
UsageBodyField::ClientResponseBody => refs.client_response_body_ref = parsed,
}
}
if !removed_any {
return None;
}
Some(UsageLegacyBodyRefMigrationPlan {
refs,
request_metadata: (!metadata.is_empty()).then_some(Value::Object(metadata)),
})
}
fn build_usage_body_externalization(
row: &UsageBodyCompressionRow,
) -> Result<UsageBodyExternalizationPlan, DataLayerError> {
let mut plan = UsageBodyExternalizationPlan::default();
maybe_externalize_usage_body_field(
&mut plan,
&row.request_id,
UsageBodyField::RequestBody,
row.request_body.as_ref(),
row.request_body_compressed.as_deref(),
)?;
maybe_externalize_usage_body_field(
&mut plan,
&row.request_id,
UsageBodyField::ProviderRequestBody,
row.provider_request_body.as_ref(),
row.provider_request_body_compressed.as_deref(),
)?;
maybe_externalize_usage_body_field(
&mut plan,
&row.request_id,
UsageBodyField::ResponseBody,
row.response_body.as_ref(),
row.response_body_compressed.as_deref(),
)?;
maybe_externalize_usage_body_field(
&mut plan,
&row.request_id,
UsageBodyField::ClientResponseBody,
row.client_response_body.as_ref(),
row.client_response_body_compressed.as_deref(),
)?;
Ok(plan)
}
fn maybe_externalize_usage_body_field(
plan: &mut UsageBodyExternalizationPlan,
request_id: &str,
field: UsageBodyField,
inline_body: Option<&Value>,
compressed_body: Option<&[u8]>,
) -> Result<(), DataLayerError> {
let Some(payload_gzip) = (match inline_body {
Some(value) => Some(compress_usage_json_value(value)?),
None => compressed_body.map(|value| value.to_vec()),
}) else {
return Ok(());
};
let body_ref = usage_body_ref(request_id, field);
plan.blobs.push(UsageDetachedBodyBlobWrite {
body_ref: body_ref.clone(),
body_field: field.as_storage_field(),
payload_gzip,
});
match field {
UsageBodyField::RequestBody => plan.refs.request_body_ref = Some(body_ref),
UsageBodyField::ProviderRequestBody => plan.refs.provider_request_body_ref = Some(body_ref),
UsageBodyField::ResponseBody => plan.refs.response_body_ref = Some(body_ref),
UsageBodyField::ClientResponseBody => plan.refs.client_response_body_ref = Some(body_ref),
}
Ok(())
}
async fn cleanup_expired_api_keys(
pool: &aether_data::postgres::PostgresPool,
auto_delete_expired_keys: bool,
) -> Result<usize, DataLayerError> {
let mut expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL).fetch(pool);
let mut cleaned = 0usize;
while let Some(row) = expired_keys.try_next().await.map_err(postgres_error)? {
let api_key_id = row.try_get::<String, _>("id").map_err(postgres_error)?;
let key = ExpiredApiKeyRow {
id: api_key_id.as_str(),
auto_delete_on_expiry: row
.try_get::<Option<bool>, _>("auto_delete_on_expiry")
.map_err(postgres_error)?,
};
let should_delete = key
.auto_delete_on_expiry
.unwrap_or(auto_delete_expired_keys);
if should_delete {
nullify_expired_api_key_usage_refs(pool, key.id).await?;
nullify_expired_api_key_candidate_refs(pool, key.id).await?;
sqlx::query(DISABLE_EXPIRED_API_KEY_WALLET_SQL)
.bind(key.id)
.execute(pool)
.await
.map_err(postgres_error)?;
let deleted = sqlx::query(DELETE_EXPIRED_API_KEY_SQL)
.bind(key.id)
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
if deleted > 0 {
cleaned += 1;
}
} else {
let updated = sqlx::query(DISABLE_EXPIRED_API_KEY_SQL)
.bind(key.id)
.bind(Utc::now())
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
if updated > 0 {
cleaned += 1;
}
}
}
Ok(cleaned)
}
async fn nullify_expired_api_key_usage_refs(
pool: &aether_data::postgres::PostgresPool,
api_key_id: &str,
) -> Result<(), DataLayerError> {
loop {
let updated = sqlx::query(NULLIFY_USAGE_API_KEY_BATCH_SQL)
.bind(api_key_id)
.bind(i64::try_from(EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE).unwrap_or(i64::MAX))
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
let updated = usize::try_from(updated).unwrap_or(usize::MAX);
if updated < EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE {
break;
}
}
Ok(())
}
async fn nullify_expired_api_key_candidate_refs(
pool: &aether_data::postgres::PostgresPool,
api_key_id: &str,
) -> Result<(), DataLayerError> {
loop {
let updated = sqlx::query(NULLIFY_REQUEST_CANDIDATE_API_KEY_BATCH_SQL)
.bind(api_key_id)
.bind(i64::try_from(EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE).unwrap_or(i64::MAX))
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
let updated = usize::try_from(updated).unwrap_or(usize::MAX);
if updated < EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE {
break;
}
}
Ok(())
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}
#[cfg(test)]
mod tests {
use std::io::Read;
use flate2::read::GzDecoder;
use serde_json::json;
use super::{
build_usage_body_externalization, compress_usage_json_value,
migrate_legacy_body_ref_metadata_plan, UsageBodyCompressionRow,
};
fn inflate_json(bytes: &[u8]) -> serde_json::Value {
let mut decoder = GzDecoder::new(bytes);
let mut decoded = Vec::new();
decoder
.read_to_end(&mut decoded)
.expect("gzip should decode");
serde_json::from_slice(&decoded).expect("json should decode")
}
#[test]
fn usage_body_externalization_moves_inline_json_into_ref_backed_blobs() {
let row = UsageBodyCompressionRow {
id: "usage-1".to_string(),
request_id: "req-1".to_string(),
request_body: Some(json!({"hello": "world"})),
request_body_compressed: None,
response_body: None,
response_body_compressed: None,
provider_request_body: Some(json!({"provider": true})),
provider_request_body_compressed: None,
client_response_body: None,
client_response_body_compressed: None,
};
let plan = build_usage_body_externalization(&row).expect("plan should build");
assert_eq!(plan.blobs.len(), 2);
assert_eq!(
plan.refs.request_body_ref.as_deref(),
Some("usage://request/req-1/request_body")
);
assert_eq!(
plan.refs.provider_request_body_ref.as_deref(),
Some("usage://request/req-1/provider_request_body")
);
assert_eq!(
inflate_json(&plan.blobs[0].payload_gzip),
json!({"hello": "world"})
);
assert_eq!(
inflate_json(&plan.blobs[1].payload_gzip),
json!({"provider": true})
);
}
#[test]
fn usage_body_externalization_reuses_existing_compressed_payloads() {
let compressed = compress_usage_json_value(&json!({"legacy": true}))
.expect("compressed payload should build");
let row = UsageBodyCompressionRow {
id: "usage-1".to_string(),
request_id: "req-legacy".to_string(),
request_body: None,
request_body_compressed: Some(compressed.clone()),
response_body: None,
response_body_compressed: None,
provider_request_body: None,
provider_request_body_compressed: None,
client_response_body: None,
client_response_body_compressed: None,
};
let plan = build_usage_body_externalization(&row).expect("plan should build");
assert_eq!(plan.blobs.len(), 1);
assert_eq!(plan.blobs[0].payload_gzip, compressed);
assert_eq!(
plan.refs.request_body_ref.as_deref(),
Some("usage://request/req-legacy/request_body")
);
}
#[test]
fn legacy_body_ref_metadata_migration_moves_matching_refs_and_strips_keys() {
let plan = migrate_legacy_body_ref_metadata_plan(
"req-1",
Some(json!({
"trace_id": "trace-1",
"request_body_ref": "usage://request/req-1/request_body",
"response_body_ref": "usage://request/req-1/response_body"
})),
)
.expect("migration plan should exist");
assert_eq!(
plan.refs.request_body_ref.as_deref(),
Some("usage://request/req-1/request_body")
);
assert_eq!(
plan.refs.response_body_ref.as_deref(),
Some("usage://request/req-1/response_body")
);
assert_eq!(
plan.request_metadata,
Some(json!({
"trace_id": "trace-1"
}))
);
}
#[test]
fn legacy_body_ref_metadata_migration_strips_invalid_and_cross_request_refs() {
let plan = migrate_legacy_body_ref_metadata_plan(
"req-1",
Some(json!({
"request_body_ref": "blob://legacy-request",
"provider_request_body_ref": "usage://request/req-other/provider_request_body",
"candidate_index": 2
})),
)
.expect("migration plan should exist");
assert!(!plan.refs.any_present());
assert_eq!(
plan.request_metadata,
Some(json!({
"candidate_index": 2
}))
);
}
.await
}

48
apps/aether-gateway/src/maintenance/runtime/wallet_daily_usage.rs
View File

@@ -3,10 +3,7 @@ use chrono::{DateTime, Utc};
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
maintenance_timezone, wallet_daily_usage_aggregation_target,
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
};
use super::{maintenance_timezone, wallet_daily_usage_aggregation_target};
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct WalletDailyUsageAggregationSummary {
@@ -30,46 +27,29 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
let timezone = maintenance_timezone();
let now_utc = Utc::now();
let target = wallet_daily_usage_aggregation_target(now_utc, timezone);
let Some(pool) = data.postgres_pool() else {
if !data.has_wallet_daily_usage_aggregation_backend() {
return Ok(WalletDailyUsageAggregationSummary {
billing_date: target.billing_date,
billing_timezone: target.billing_timezone,
aggregated_wallets: 0,
deleted_stale_ledgers: 0,
});
};
}
let mut tx = pool.begin().await.map_err(postgres_error)?;
let aggregated_wallets = sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
.bind(target.window_start_utc)
.bind(target.window_end_utc)
.bind(target.billing_date)
.bind(target.billing_timezone.as_str())
.bind(now_utc)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
.bind(target.billing_date)
.bind(target.billing_timezone.as_str())
.bind(target.window_start_utc)
.bind(target.window_end_utc)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
tx.commit().await.map_err(postgres_error)?;
let result = data
.aggregate_wallet_daily_usage(&aether_data::WalletDailyUsageAggregationInput {
billing_date: target.billing_date.to_string(),
billing_timezone: target.billing_timezone.clone(),
window_start_unix_secs: target.window_start_utc.timestamp().max(0) as u64,
window_end_unix_secs: target.window_end_utc.timestamp().max(0) as u64,
aggregated_at_unix_secs: now_utc.timestamp().max(0) as u64,
})
.await?;
Ok(WalletDailyUsageAggregationSummary {
billing_date: target.billing_date,
billing_timezone: target.billing_timezone,
aggregated_wallets: usize::try_from(aggregated_wallets).unwrap_or(usize::MAX),
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
aggregated_wallets: result.aggregated_wallets,
deleted_stale_ledgers: result.deleted_stale_ledgers,
})
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}

16
apps/aether-gateway/src/maintenance/runtime/workers.rs
View File

@@ -42,7 +42,7 @@ fn log_maintenance_worker_failure(
pub(crate) fn spawn_audit_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_audit_log_reader() {
return None;
}
@@ -65,7 +65,7 @@ pub(crate) fn spawn_audit_cleanup_worker(
pub(crate) fn spawn_db_maintenance_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_database_maintenance_backend() {
return None;
}
@@ -83,7 +83,7 @@ pub(crate) fn spawn_db_maintenance_worker(
pub(crate) fn spawn_wallet_daily_usage_aggregation_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_wallet_daily_usage_aggregation_backend() {
return None;
}
@@ -107,7 +107,7 @@ pub(crate) fn spawn_wallet_daily_usage_aggregation_worker(
pub(crate) fn spawn_stats_aggregation_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_stats_daily_aggregation_backend() {
return None;
}
@@ -137,7 +137,7 @@ pub(crate) fn spawn_stats_aggregation_worker(
pub(crate) fn spawn_usage_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_usage_writer() {
return None;
}
@@ -224,7 +224,7 @@ pub(crate) fn spawn_gemini_file_mapping_cleanup_worker(
pub(crate) fn spawn_pending_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_usage_writer() {
return None;
}
@@ -296,7 +296,7 @@ pub(crate) fn spawn_proxy_upgrade_rollout_worker(
pub(crate) fn spawn_pool_monitor_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_database_pool_summary() {
return None;
}
@@ -314,7 +314,7 @@ pub(crate) fn spawn_pool_monitor_worker(
pub(crate) fn spawn_stats_hourly_aggregation_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if data.postgres_pool().is_none() {
if !data.has_stats_hourly_aggregation_backend() {
return None;
}