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

refactor: 大规模模块拆分与代码精简，新增 ai-pipeline/data-contracts 独立 crate

Browse Source 
- 新增 aether-ai-pipeline 和 aether-data-contracts crate，将 pipeline 逻辑与数据契约从 gateway 中解耦
- 重构 admin handlers：拆分单体模块为 auth/billing/endpoint/features/model/observability/provider/system 等独立子模块
- 合并 chat/cli 重复代码路径：精简 conversion、finalize、planner 中的 sync/chat/cli 分支
- 重构 scheduler/executor/data 层，引入 facade 模式降低模块间耦合
- 移除冗余的 intent 模块，将 plan_fallback/policy/stream_path/sync_path 迁移至 executor
- 前端适配：调整 admin API 调用和 provider 模型测试对话框
This commit is contained in:
fawney19
2026-04-07 02:50:19 +08:00
parent 763ff03a7b
commit 5d96d6673b
732 changed files with 28593 additions and 20666 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,7 +3,7 @@ use std::io::Write;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data::repository::provider_catalog::{
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use chrono::{DateTime, Datelike, TimeZone, Timelike, Utc, Weekday};
@@ -16,7 +16,7 @@ use tracing::{debug, info, warn};
use uuid::Uuid;
use crate::data::GatewayDataState;
use crate::handlers::admin::provider_ops::admin_provider_ops_local_action_response;
use crate::handlers::admin::provider::ops::admin_provider_ops_local_action_response;
use crate::{AppState, GatewayError};
#[path = "runtime/audit_cleanup.rs"]
@@ -62,6 +62,12 @@ use usage_cleanup::*;
use wallet_daily_usage::*;
pub(crate) use workers::*;
pub(super) fn postgres_error(
error: impl std::fmt::Display,
) -> aether_data_contracts::DataLayerError {
aether_data_contracts::DataLayerError::postgres(error)
}
const AUDIT_LOG_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const GEMINI_FILE_MAPPING_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60);
const PENDING_CLEANUP_INTERVAL: Duration = Duration::from_secs(5 * 60);

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

@@ -1,14 +1,16 @@
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Utc};
use crate::data::GatewayDataState;
use super::{
system_config_bool, system_config_u64, system_config_usize, DELETE_AUDIT_LOGS_BEFORE_SQL,
postgres_error, system_config_bool, system_config_u64, system_config_usize,
DELETE_AUDIT_LOGS_BEFORE_SQL,
};
pub(crate) async fn cleanup_audit_logs_once(
data: &GatewayDataState,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
cleanup_audit_logs_with(data, |cutoff_time, delete_limit| async move {
let Some(pool) = data.postgres_pool() else {
return Ok(0);
@@ -17,7 +19,8 @@ pub(crate) async fn cleanup_audit_logs_once(
.bind(cutoff_time)
.bind(i64::try_from(delete_limit).unwrap_or(i64::MAX))
.execute(&pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
Ok(usize::try_from(deleted).unwrap_or(usize::MAX))
})
@@ -27,10 +30,10 @@ pub(crate) async fn cleanup_audit_logs_once(
pub(super) async fn cleanup_audit_logs_with<F, Fut>(
data: &GatewayDataState,
mut delete_batch: F,
) -> Result<usize, aether_data::DataLayerError>
) -> Result<usize, DataLayerError>
where
F: FnMut(DateTime<Utc>, usize) -> Fut,
Fut: std::future::Future<Output = Result<usize, aether_data::DataLayerError>>,
Fut: std::future::Future<Output = Result<usize, DataLayerError>>,
{
if !system_config_bool(data, "enable_auto_cleanup", true).await? {
return Ok(0);

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

@@ -1,3 +1,4 @@
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Utc};
use serde_json::Value;
@@ -9,7 +10,7 @@ pub(super) async fn system_config_bool(
data: &GatewayDataState,
key: &str,
default: bool,
) -> Result<bool, aether_data::DataLayerError> {
) -> Result<bool, DataLayerError> {
Ok(match data.find_system_config_value(key).await? {
Some(Value::Bool(value)) => value,
Some(Value::Number(value)) => value.as_i64().map(|raw| raw != 0).unwrap_or(default),
@@ -26,7 +27,7 @@ pub(super) async fn system_config_u64(
data: &GatewayDataState,
key: &str,
default: u64,
) -> Result<u64, aether_data::DataLayerError> {
) -> Result<u64, DataLayerError> {
Ok(match data.find_system_config_value(key).await? {
Some(Value::Number(value)) => value
.as_u64()
@@ -41,7 +42,7 @@ pub(super) async fn system_config_usize(
data: &GatewayDataState,
key: &str,
default: usize,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
Ok(match data.find_system_config_value(key).await? {
Some(Value::Number(value)) => value
.as_u64()
@@ -62,7 +63,7 @@ pub(super) async fn system_config_string(
data: &GatewayDataState,
key: &str,
default: &str,
) -> Result<String, aether_data::DataLayerError> {
) -> Result<String, DataLayerError> {
Ok(match data.find_system_config_value(key).await? {
Some(Value::String(value)) => {
let trimmed = value.trim();
@@ -78,13 +79,13 @@ pub(super) async fn system_config_string(
pub(super) async fn pending_cleanup_timeout_minutes(
data: &GatewayDataState,
) -> Result<u64, aether_data::DataLayerError> {
) -> Result<u64, DataLayerError> {
system_config_u64(data, "pending_request_timeout_minutes", 10).await
}
pub(super) async fn pending_cleanup_batch_size(
data: &GatewayDataState,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
Ok(system_config_usize(data, "cleanup_batch_size", 1_000)
.await?
.max(1)
@@ -93,7 +94,7 @@ pub(super) async fn pending_cleanup_batch_size(
pub(super) async fn usage_cleanup_settings(
data: &GatewayDataState,
) -> Result<UsageCleanupSettings, aether_data::DataLayerError> {
) -> Result<UsageCleanupSettings, DataLayerError> {
Ok(UsageCleanupSettings {
detail_retention_days: system_config_u64(data, "detail_log_retention_days", 7).await?,
compressed_retention_days: system_config_u64(data, "compressed_log_retention_days", 30)

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

@@ -1,8 +1,9 @@
use aether_data_contracts::DataLayerError;
use tracing::warn;
use crate::data::GatewayDataState;
use super::{system_config_bool, DB_MAINTENANCE_TABLES};
use super::{postgres_error, system_config_bool, DB_MAINTENANCE_TABLES};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct DbMaintenanceRunSummary {
@@ -12,7 +13,7 @@ pub(super) struct DbMaintenanceRunSummary {
pub(super) async fn perform_db_maintenance_once(
data: &GatewayDataState,
) -> Result<DbMaintenanceRunSummary, aether_data::DataLayerError> {
) -> Result<DbMaintenanceRunSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
return Ok(DbMaintenanceRunSummary {
attempted: 0,
@@ -24,7 +25,10 @@ pub(super) async fn perform_db_maintenance_once(
let pool = pool.clone();
async move {
let statement = format!("VACUUM ANALYZE {table_name}");
sqlx::raw_sql(&statement).execute(&pool).await?;
sqlx::raw_sql(&statement)
.execute(&pool)
.await
.map_err(postgres_error)?;
Ok(())
}
})
@@ -34,10 +38,10 @@ pub(super) async fn perform_db_maintenance_once(
pub(super) async fn run_db_maintenance_with<F, Fut>(
data: &GatewayDataState,
mut vacuum_table: F,
) -> Result<DbMaintenanceRunSummary, aether_data::DataLayerError>
) -> Result<DbMaintenanceRunSummary, DataLayerError>
where
F: FnMut(&'static str) -> Fut,
Fut: std::future::Future<Output = Result<(), aether_data::DataLayerError>>,
Fut: std::future::Future<Output = Result<(), DataLayerError>>,
{
if !system_config_bool(data, "enable_db_maintenance", true).await? {
return Ok(DbMaintenanceRunSummary {

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

@@ -4,9 +4,10 @@ use chrono::Utc;
use sqlx::Row;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
pending_cleanup_batch_size, pending_cleanup_timeout_minutes,
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,
@@ -44,7 +45,7 @@ pub(super) struct PendingCleanupBatchPlan {
pub(crate) async fn cleanup_stale_pending_requests_once(
data: &GatewayDataState,
) -> Result<PendingCleanupSummary, aether_data::DataLayerError> {
) -> Result<PendingCleanupSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
return Ok(PendingCleanupSummary::default());
};
@@ -57,29 +58,34 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
let mut summary = PendingCleanupSummary::default();
loop {
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_err(postgres_error)?;
let stale_rows = 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_all(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
if stale_rows.is_empty() {
tx.rollback().await?;
tx.rollback().await.map_err(postgres_error)?;
break;
}
let stale_rows = stale_rows
.into_iter()
.map(|row| {
Ok(StalePendingUsageRow {
id: row.try_get::<String, _>("id")?,
request_id: row.try_get::<String, _>("request_id")?,
status: row.try_get::<String, _>("status")?,
billing_status: row.try_get::<String, _>("billing_status")?,
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<_>, sqlx::Error>>()?;
.collect::<Result<Vec<_>, DataLayerError>>()?;
let request_ids = stale_rows
.iter()
.map(|row| row.request_id.clone())
@@ -90,7 +96,8 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(request_ids)
.fetch_all(&mut *tx)
.await?
.await
.map_err(postgres_error)?
.into_iter()
.filter_map(|row| row.try_get::<String, _>("request_id").ok())
.collect::<HashSet<_>>()
@@ -102,7 +109,8 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
sqlx::query(UPDATE_RECOVERED_STALE_USAGE_SQL)
.bind(usage_id)
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
}
for failed_row in &plan.failed_usage_rows {
if failed_row.should_void_billing {
@@ -111,13 +119,15 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
.bind(&failed_row.error_message)
.bind(now)
.execute(&mut *tx)
.await?;
.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?;
.await
.map_err(postgres_error)?;
}
}
if !plan.recovered_request_ids.is_empty() {
@@ -125,7 +135,8 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
.bind(plan.recovered_request_ids.clone())
.bind(now)
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
}
if !plan.failed_request_ids.is_empty() {
sqlx::query(UPDATE_FAILED_PENDING_CANDIDATES_SQL)
@@ -133,10 +144,11 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
.bind(now)
.bind(active_statuses.clone())
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
}
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
summary.failed += plan.failed_usage_rows.len();
summary.recovered += plan.recovered_usage_ids.len();
}

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

@@ -1,12 +1,12 @@
use std::collections::HashMap;
use aether_data::repository::provider_catalog::{
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use futures_util::stream::{self, StreamExt};
use tracing::{debug, warn};
use crate::handlers::admin::provider_ops::admin_provider_ops_local_action_response;
use crate::handlers::admin::provider::ops::admin_provider_ops_local_action_response;
use crate::{AppState, GatewayError};
use super::{system_config_bool, PROVIDER_CHECKIN_CONCURRENCY};

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

@@ -1,3 +1,4 @@
use aether_data_contracts::DataLayerError;
use tracing::info;
use crate::data::GatewayDataState;
@@ -6,7 +7,7 @@ use super::{now_unix_secs, system_config_bool, system_config_u64, system_config_
pub(crate) async fn cleanup_request_candidates_once(
data: &GatewayDataState,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
if !system_config_bool(data, "enable_auto_cleanup", true).await? {
return Ok(0);
}
@@ -45,7 +46,7 @@ pub(crate) async fn cleanup_request_candidates_once(
pub(super) async fn run_request_candidate_cleanup_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let deleted = cleanup_request_candidates_once(data).await?;
if deleted > 0 {
info!(deleted, "gateway deleted expired request candidates");

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

@@ -1,3 +1,4 @@
use aether_data_contracts::DataLayerError;
use tracing::info;
use crate::data::GatewayDataState;
@@ -11,9 +12,7 @@ use super::{
perform_wallet_daily_usage_aggregation_once, summarize_postgres_pool,
};
pub(super) async fn run_audit_cleanup_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let deleted = cleanup_audit_logs_once(data).await?;
if deleted > 0 {
info!(
@@ -29,7 +28,7 @@ pub(super) async fn run_audit_cleanup_once(
pub(super) async fn run_gemini_file_mapping_cleanup_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let deleted = cleanup_expired_gemini_file_mappings_once(data).await?;
if deleted > 0 {
info!(
@@ -43,9 +42,7 @@ pub(super) async fn run_gemini_file_mapping_cleanup_once(
Ok(())
}
pub(super) async fn run_db_maintenance_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
pub(super) async fn run_db_maintenance_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let summary = perform_db_maintenance_once(data).await?;
if summary.attempted > 0 {
info!(
@@ -63,7 +60,7 @@ pub(super) async fn run_db_maintenance_once(
pub(super) async fn run_wallet_daily_usage_aggregation_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let summary = perform_wallet_daily_usage_aggregation_once(data).await?;
info!(
event_name = "wallet_daily_usage_aggregation_completed",
@@ -80,7 +77,7 @@ pub(super) async fn run_wallet_daily_usage_aggregation_once(
pub(super) async fn run_stats_aggregation_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let Some(summary) = perform_stats_aggregation_once(data).await? else {
return Ok(());
};
@@ -101,9 +98,7 @@ pub(super) async fn run_stats_aggregation_once(
Ok(())
}
pub(super) async fn run_usage_cleanup_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let summary = perform_usage_cleanup_once(data).await?;
if summary.body_compressed > 0
|| summary.body_cleaned > 0
@@ -146,7 +141,7 @@ pub(super) fn run_pool_monitor_once(data: &GatewayDataState) {
pub(super) async fn run_pending_cleanup_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let summary = cleanup_stale_pending_requests_once(data).await?;
if summary.failed > 0 || summary.recovered > 0 {
info!(
@@ -163,7 +158,7 @@ pub(super) async fn run_pending_cleanup_once(
pub(super) async fn run_stats_hourly_aggregation_once(
data: &GatewayDataState,
) -> Result<(), aether_data::DataLayerError> {
) -> Result<(), DataLayerError> {
let Some(summary) = perform_stats_hourly_aggregation_once(data).await? else {
return Ok(());
};

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

@@ -1,5 +1,6 @@
use std::time::Duration;
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Datelike, TimeZone, Timelike, Utc, Weekday};
use chrono_tz::Tz;
use tracing::warn;
@@ -15,7 +16,7 @@ use super::{
pub(super) async fn provider_checkin_schedule(
data: &GatewayDataState,
) -> Result<(u32, u32), aether_data::DataLayerError> {
) -> Result<(u32, u32), DataLayerError> {
let configured =
system_config_string(data, "provider_checkin_time", PROVIDER_CHECKIN_DEFAULT_TIME).await?;
Ok(parse_hhmm_time(&configured).unwrap_or_else(|| {

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

@@ -5,11 +5,12 @@ use sqlx::Row;
use uuid::Uuid;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
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_ACTIVE_USER_IDS_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
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_ACTIVE_USER_IDS_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL,
SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL, SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL,
@@ -22,7 +23,7 @@ use super::{
pub(super) async fn perform_stats_aggregation_once(
data: &GatewayDataState,
) -> Result<Option<StatsAggregationSummary>, aether_data::DataLayerError> {
) -> Result<Option<StatsAggregationSummary>, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
return Ok(None);
};
@@ -33,36 +34,45 @@ pub(super) async fn perform_stats_aggregation_once(
let now_utc = Utc::now();
let day_start_utc = stats_aggregation_target_day(now_utc);
let day_end_utc = day_start_utc + chrono::Duration::days(1);
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_err(postgres_error)?;
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")?;
.await
.map_err(postgres_error)?;
let total_requests = aggregate_row
.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?;
let error_requests = aggregate_row
.try_get::<i64, _>("error_requests")
.map_err(postgres_error)?;
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")?;
.await
.map_err(postgres_error)?
.try_get::<i64, _>("fallback_count")
.map_err(postgres_error)?;
let response_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await?;
.await
.map_err(postgres_error)?;
let first_byte_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await?;
.await
.map_err(postgres_error)?;
sqlx::query(UPSERT_STATS_DAILY_SQL)
.bind(Uuid::new_v4().to_string())
@@ -70,17 +80,61 @@ pub(super) async fn perform_stats_aggregation_once(
.bind(total_requests)
.bind(success_requests)
.bind(error_requests)
.bind(aggregate_row.try_get::<i64, _>("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_read_tokens")?)
.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, _>("avg_response_time_ms")?)
.bind(
aggregate_row
.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("cache_creation_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("cache_read_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("actual_total_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("input_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("output_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("cache_creation_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("cache_read_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("avg_response_time_ms")
.map_err(postgres_error)?,
)
.bind(response_percentiles.p50)
.bind(response_percentiles.p90)
.bind(response_percentiles.p99)
@@ -88,27 +142,45 @@ pub(super) async fn perform_stats_aggregation_once(
.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(
aggregate_row
.try_get::<i64, _>("unique_models")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("unique_providers")
.map_err(postgres_error)?,
)
.bind(true)
.bind(now_utc)
.bind(now_utc)
.bind(now_utc)
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
let model_rows =
upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let model_rows = upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let provider_rows =
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
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?;
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc).await?;
tx.commit().await?;
upsert_stats_daily_api_key_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let error_rows = refresh_stats_daily_error_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let user_rows = upsert_stats_user_daily_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
Ok(Some(StatsAggregationSummary {
day_start_utc,

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

@@ -3,6 +3,7 @@ use sqlx::Row;
use uuid::Uuid;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_STATS_HOURLY_AGGREGATE_SQL,
@@ -22,7 +23,7 @@ pub(super) struct StatsHourlyAggregationSummary {
pub(super) async fn perform_stats_hourly_aggregation_once(
data: &GatewayDataState,
) -> Result<Option<StatsHourlyAggregationSummary>, aether_data::DataLayerError> {
) -> Result<Option<StatsHourlyAggregationSummary>, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
return Ok(None);
};
@@ -34,15 +35,20 @@ pub(super) async fn perform_stats_hourly_aggregation_once(
let hour_utc = stats_hourly_aggregation_target_hour(now_utc);
let hour_end = hour_utc + chrono::Duration::hours(1);
let aggregated_at = now_utc;
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_err(postgres_error)?;
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")?;
.await
.map_err(postgres_error)?;
let total_requests = row
.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?;
let error_requests = row
.try_get::<i64, _>("error_requests")
.map_err(postgres_error)?;
let success_requests = total_requests.saturating_sub(error_requests);
sqlx::query(UPSERT_STATS_HOURLY_SQL)
.bind(Uuid::new_v4().to_string())
@@ -50,19 +56,41 @@ pub(super) async fn perform_stats_hourly_aggregation_once(
.bind(total_requests)
.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, _>("avg_response_time_ms")?)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_creation_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_read_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("actual_total_cost")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("avg_response_time_ms")
.map_err(postgres_error)?,
)
.bind(true)
.bind(aggregated_at)
.bind(aggregated_at)
.bind(aggregated_at)
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
let user_rows =
upsert_stats_hourly_user_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
@@ -70,7 +98,7 @@ pub(super) async fn perform_stats_hourly_aggregation_once(
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?;
tx.commit().await.map_err(postgres_error)?;
Ok(Some(StatsHourlyAggregationSummary {
hour_utc,
@@ -86,17 +114,24 @@ async fn upsert_stats_hourly_user_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_USER_AGGREGATES_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
for row in &rows {
let user_id = row.try_get::<String, _>("user_id")?;
let total_requests = row.try_get::<i64, _>("total_requests")?;
let error_requests = row.try_get::<i64, _>("error_requests")?;
let user_id = row
.try_get::<String, _>("user_id")
.map_err(postgres_error)?;
let total_requests = row
.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?;
let error_requests = row
.try_get::<i64, _>("error_requests")
.map_err(postgres_error)?;
let success_requests = total_requests.saturating_sub(error_requests);
sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
.bind(Uuid::new_v4().to_string())
@@ -105,13 +140,23 @@ async fn upsert_stats_hourly_user_rows(
.bind(total_requests)
.bind(success_requests)
.bind(error_requests)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
}
Ok(rows.len())
@@ -122,16 +167,19 @@ async fn upsert_stats_hourly_model_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
let mut inserted = 0usize;
for row in &rows {
let model = row.try_get::<Option<String>, _>("model")?;
let model = row
.try_get::<Option<String>, _>("model")
.map_err(postgres_error)?;
let Some(model) = model.filter(|value| !value.is_empty()) else {
continue;
};
@@ -139,15 +187,31 @@ async fn upsert_stats_hourly_model_rows(
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(model)
.bind(row.try_get::<i64, _>("total_requests")?)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(
row.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("avg_response_time_ms")
.map_err(postgres_error)?,
)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
inserted += 1;
}
@@ -159,16 +223,19 @@ async fn upsert_stats_hourly_provider_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
let mut inserted = 0usize;
for row in &rows {
let provider_name = row.try_get::<Option<String>, _>("provider_name")?;
let provider_name = row
.try_get::<Option<String>, _>("provider_name")
.map_err(postgres_error)?;
let Some(provider_name) = provider_name.filter(|value| !value.is_empty()) else {
continue;
};
@@ -176,16 +243,33 @@ async fn upsert_stats_hourly_provider_rows(
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(provider_name)
.bind(row.try_get::<i64, _>("total_requests")?)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(
row.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
.await
.map_err(postgres_error)?;
inserted += 1;
}
Ok(inserted)
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}

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

@@ -275,20 +275,18 @@ fn usage_cleanup_window_uses_non_overlapping_ranges() {
#[tokio::test]
async fn summarize_postgres_pool_uses_busy_connections_for_usage_rate() {
let data = GatewayDataState::from_config(
crate::data::GatewayDataConfig::from_postgres_config(
aether_data::postgres::PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 8,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
},
),
)
let data = GatewayDataState::from_config(crate::data::GatewayDataConfig::from_postgres_config(
aether_data::postgres::PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 8,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
},
))
.expect("gateway data state should build");
let summary = summarize_postgres_pool(&data).expect("pool summary should exist");

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

@@ -1,5 +1,6 @@
use std::io::Write;
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use serde_json::Value;
@@ -21,7 +22,7 @@ use super::{
pub(super) async fn perform_usage_cleanup_once(
data: &GatewayDataState,
) -> Result<UsageCleanupSummary, aether_data::DataLayerError> {
) -> Result<UsageCleanupSummary, DataLayerError> {
let Some(pool) = data.postgres_pool() else {
return Ok(UsageCleanupSummary::default());
};
@@ -76,14 +77,15 @@ async fn delete_old_usage_records(
pool: &aether_data::postgres::PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
) -> Result<usize, aether_data::DataLayerError> {
) -> 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?
.await
.map_err(postgres_error)?
.rows_affected();
let deleted = usize::try_from(deleted).unwrap_or(usize::MAX);
total_deleted += deleted;
@@ -99,7 +101,7 @@ async fn cleanup_usage_header_fields(
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
@@ -116,10 +118,11 @@ async fn cleanup_usage_header_fields(
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch_all(pool)
.await?
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| row.try_get::<String, _>("id"))
.collect::<Result<Vec<_>, sqlx::Error>>()?;
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
.collect::<Result<Vec<_>, DataLayerError>>()?;
if ids.is_empty() {
break;
}
@@ -127,7 +130,8 @@ async fn cleanup_usage_header_fields(
let cleaned = sqlx::query(CLEAR_USAGE_HEADER_FIELDS_SQL)
.bind(ids)
.execute(pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
total_cleaned += cleaned;
@@ -143,7 +147,7 @@ async fn cleanup_usage_stale_body_fields(
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
@@ -160,10 +164,11 @@ async fn cleanup_usage_stale_body_fields(
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch_all(pool)
.await?
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| row.try_get::<String, _>("id"))
.collect::<Result<Vec<_>, sqlx::Error>>()?;
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
.collect::<Result<Vec<_>, DataLayerError>>()?;
if ids.is_empty() {
break;
}
@@ -171,7 +176,8 @@ async fn cleanup_usage_stale_body_fields(
let cleaned = sqlx::query(CLEAR_USAGE_BODY_FIELDS_SQL)
.bind(ids)
.execute(pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
total_cleaned += cleaned;
@@ -187,7 +193,7 @@ async fn compress_usage_body_fields(
cutoff_time: DateTime<Utc>,
batch_size: usize,
newer_than: Option<DateTime<Utc>>,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
if matches!(newer_than, Some(value) if value >= cutoff_time) {
warn!(
cutoff_time = %cutoff_time,
@@ -206,20 +212,27 @@ async fn compress_usage_body_fields(
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch_all(pool)
.await?
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| {
Ok(UsageBodyCompressionRow {
id: row.try_get::<String, _>("id")?,
request_body: row.try_get::<Option<Value>, _>("request_body")?,
response_body: row.try_get::<Option<Value>, _>("response_body")?,
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_body: row
.try_get::<Option<Value>, _>("request_body")
.map_err(postgres_error)?,
response_body: row
.try_get::<Option<Value>, _>("response_body")
.map_err(postgres_error)?,
provider_request_body: row
.try_get::<Option<Value>, _>("provider_request_body")?,
.try_get::<Option<Value>, _>("provider_request_body")
.map_err(postgres_error)?,
client_response_body: row
.try_get::<Option<Value>, _>("client_response_body")?,
.try_get::<Option<Value>, _>("client_response_body")
.map_err(postgres_error)?,
})
})
.collect::<Result<Vec<_>, sqlx::Error>>()?;
.collect::<Result<Vec<_>, DataLayerError>>()?;
if rows.is_empty() {
break;
}
@@ -239,7 +252,8 @@ async fn compress_usage_body_fields(
.bind(compressed.2)
.bind(compressed.3)
.execute(pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
if updated > 0 {
batch_success += 1;
@@ -273,16 +287,19 @@ fn compress_usage_json_value(value: Option<&Value>) -> Option<Vec<u8>> {
async fn cleanup_expired_api_keys(
pool: &aether_data::postgres::PostgresPool,
auto_delete_expired_keys: bool,
) -> Result<usize, aether_data::DataLayerError> {
) -> Result<usize, DataLayerError> {
let expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL)
.fetch_all(pool)
.await?;
.await
.map_err(postgres_error)?;
let mut cleaned = 0usize;
for row in &expired_keys {
let api_key_id = row.try_get::<String, _>("id")?;
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")?,
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
@@ -293,7 +310,8 @@ async fn cleanup_expired_api_keys(
let deleted = sqlx::query(DELETE_EXPIRED_API_KEY_SQL)
.bind(key.id)
.execute(pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
if deleted > 0 {
cleaned += 1;
@@ -303,7 +321,8 @@ async fn cleanup_expired_api_keys(
.bind(key.id)
.bind(Utc::now())
.execute(pool)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
if updated > 0 {
cleaned += 1;
@@ -316,13 +335,14 @@ async fn cleanup_expired_api_keys(
async fn nullify_expired_api_key_usage_refs(
pool: &aether_data::postgres::PostgresPool,
api_key_id: &str,
) -> Result<(), aether_data::DataLayerError> {
) -> 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?
.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 {
@@ -335,13 +355,14 @@ async fn nullify_expired_api_key_usage_refs(
async fn nullify_expired_api_key_candidate_refs(
pool: &aether_data::postgres::PostgresPool,
api_key_id: &str,
) -> Result<(), aether_data::DataLayerError> {
) -> 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?
.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 {
@@ -350,3 +371,7 @@ async fn nullify_expired_api_key_candidate_refs(
}
Ok(())
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}

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

@@ -3,6 +3,7 @@ use sqlx::Row;
use uuid::Uuid;
use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
maintenance_timezone, wallet_daily_usage_aggregation_target,
@@ -28,7 +29,7 @@ pub(super) struct WalletDailyUsageAggregationTarget {
pub(super) async fn perform_wallet_daily_usage_aggregation_once(
data: &GatewayDataState,
) -> Result<WalletDailyUsageAggregationSummary, aether_data::DataLayerError> {
) -> Result<WalletDailyUsageAggregationSummary, DataLayerError> {
let timezone = maintenance_timezone();
let now_utc = Utc::now();
let target = wallet_daily_usage_aggregation_target(now_utc, timezone);
@@ -41,31 +42,60 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
});
};
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_err(postgres_error)?;
let rows = sqlx::query(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL)
.bind(target.window_start_utc)
.bind(target.window_end_utc)
.fetch_all(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
for row in &rows {
sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
.bind(Uuid::new_v4().to_string())
.bind(row.try_get::<String, _>("wallet_id")?)
.bind(
row.try_get::<String, _>("wallet_id")
.map_err(postgres_error)?,
)
.bind(target.billing_date)
.bind(target.billing_timezone.as_str())
.bind(row.try_get::<f64, _>("total_cost_usd")?)
.bind(row.try_get::<i64, _>("total_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::<Option<DateTime<Utc>>, _>("first_finalized_at")?)
.bind(row.try_get::<Option<DateTime<Utc>>, _>("last_finalized_at")?)
.bind(
row.try_get::<f64, _>("total_cost_usd")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_creation_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_read_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<Option<DateTime<Utc>>, _>("first_finalized_at")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<Option<DateTime<Utc>>, _>("last_finalized_at")
.map_err(postgres_error)?,
)
.bind(now_utc)
.bind(now_utc)
.bind(now_utc)
.execute(&mut *tx)
.await?;
.await
.map_err(postgres_error)?;
}
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
@@ -74,9 +104,10 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
.bind(target.window_start_utc)
.bind(target.window_end_utc)
.execute(&mut *tx)
.await?
.await
.map_err(postgres_error)?
.rows_affected();
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
Ok(WalletDailyUsageAggregationSummary {
billing_date: target.billing_date,
@@ -85,3 +116,7 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
})
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}

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

@@ -1,13 +1,13 @@
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
RequestCandidateWriteRepository, UpsertRequestCandidateRecord,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogProvider,
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
UpsertRequestCandidateRecord,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use axum::http::StatusCode;
use axum::routing::{get, post};
use axum::{Json, Router};
@@ -82,19 +82,18 @@ async fn gateway_background_request_candidate_cleanup_deletes_expired_entries_in
seed_candidate(&repository, "cand-expired-2", 2).await;
seed_candidate(&repository, "cand-active", now_unix_secs()).await;
let data_state =
crate::data::GatewayDataState::with_request_candidate_repository_for_tests(
Arc::clone(&repository),
)
.with_system_config_values_for_tests([
("enable_auto_cleanup".to_string(), json!(true)),
("cleanup_batch_size".to_string(), json!(1)),
(
"request_candidates_cleanup_batch_size".to_string(),
json!(1),
),
("request_candidates_retention_days".to_string(), json!(30)),
]);
let data_state = crate::data::GatewayDataState::with_request_candidate_repository_for_tests(
Arc::clone(&repository),
)
.with_system_config_values_for_tests([
("enable_auto_cleanup".to_string(), json!(true)),
("cleanup_batch_size".to_string(), json!(1)),
(
"request_candidates_cleanup_batch_size".to_string(),
json!(1),
),
("request_candidates_retention_days".to_string(), json!(30)),
]);
let gateway_state = AppState::new()
.expect("gateway state should build")
@@ -227,15 +226,13 @@ async fn gateway_provider_checkin_runs_local_query_balance_for_configured_provid
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
repository,
)
.with_system_config_values_for_tests([
("enable_provider_checkin".to_string(), json!(true)),
("provider_checkin_time".to_string(), json!("01:05")),
])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_system_config_values_for_tests([
("enable_provider_checkin".to_string(), json!(true)),
("provider_checkin_time".to_string(), json!("01:05")),
])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let summary = crate::maintenance::perform_provider_checkin_once(&gateway_state)
.await
@@ -266,14 +263,12 @@ async fn gateway_provider_checkin_skips_when_disabled_via_system_config() {
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
repository,
)
.with_system_config_values_for_tests([(
"enable_provider_checkin".to_string(),
json!(false),
)]),
);
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_system_config_values_for_tests([(
"enable_provider_checkin".to_string(),
json!(false),
)]),
);
let summary = crate::maintenance::perform_provider_checkin_once(&gateway_state)
.await