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

refactor(runtime): 优化管理端摘要查询与维护聚合链路

Browse Source 
- 为 provider catalog key 和 video task 列表增加 summary/page 查询与排序能力，减少列表场景读取重字段
- 将多处 SQL 结果读取改为流式收集，降低 `fetch_all` 的内存占用
- 把日/小时统计、钱包日用量等维护任务改为数据库侧 `CTE + upsert` 聚合
- 修复视频任务轮询更新时从本地 snapshot 回填稀疏字段，避免 `prompt` 和请求体信息丢失
This commit is contained in:
AAEE86
2026-04-12 20:45:47 +08:00
committed by fawney19
parent 6aa16ec792
commit b2d85d70ca
55 changed files with 1954 additions and 1382 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,6 +1,7 @@
use std::collections::HashSet;
use chrono::Utc;
use futures_util::TryStreamExt;
use sqlx::Row;
use crate::data::GatewayDataState;
@@ -59,13 +60,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
loop {
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
.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;
@@ -93,14 +99,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
let completed_request_ids = if request_ids.is_empty() {
HashSet::new()
} else {
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(request_ids)
.fetch_all(&mut *tx)
.await
.map_err(postgres_error)?
.into_iter()
.filter_map(|row| row.try_get::<String, _>("request_id").ok())
.collect::<HashSet<_>>()
{
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();

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

@@ -1,5 +1,3 @@
use std::collections::HashMap;
use chrono::{DateTime, Utc};
use sqlx::Row;
use uuid::Uuid;
@@ -10,15 +8,13 @@ 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_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,
SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL, SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL, SELECT_STATS_SUMMARY_TOTALS_SQL,
SELECT_STATS_USER_DAILY_AGGREGATES_SQL, UPDATE_STATS_SUMMARY_SQL,
UPSERT_STATS_DAILY_API_KEY_SQL, UPSERT_STATS_DAILY_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL,
UPSERT_STATS_DAILY_SQL, UPSERT_STATS_USER_DAILY_SQL,
INSERT_STATS_SUMMARY_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_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_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
UPSERT_STATS_USER_DAILY_SQL,
};
pub(super) async fn perform_stats_aggregation_once(
@@ -226,31 +222,15 @@ async fn upsert_stats_daily_model_rows(
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows = sqlx::query(SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_all(&mut **tx)
.await?;
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
for row in &rows {
sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
.bind(Uuid::new_v4().to_string())
.bind(day_start_utc)
.bind(row.try_get::<String, _>("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::<i64, _>("cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(rows.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_provider_rows(
@@ -259,30 +239,15 @@ async fn upsert_stats_daily_provider_rows(
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows = sqlx::query(SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_all(&mut **tx)
.await?;
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
for row in &rows {
sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
.bind(Uuid::new_v4().to_string())
.bind(day_start_utc)
.bind(row.try_get::<String, _>("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::<i64, _>("cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(rows.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_api_key_rows(
@@ -291,35 +256,15 @@ async fn upsert_stats_daily_api_key_rows(
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows = sqlx::query(SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_all(&mut **tx)
.await?;
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
for row in &rows {
let total_requests = row.try_get::<i64, _>("total_requests")?;
let error_requests = row.try_get::<i64, _>("error_requests")?;
sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(row.try_get::<String, _>("api_key_id")?)
.bind(row.try_get::<Option<String>, _>("api_key_name")?)
.bind(day_start_utc)
.bind(total_requests)
.bind(total_requests.saturating_sub(error_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(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(rows.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn refresh_stats_daily_error_rows(
@@ -332,27 +277,15 @@ async fn refresh_stats_daily_error_rows(
.bind(day_start_utc)
.execute(&mut **tx)
.await?;
let rows = sqlx::query(SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL)
let rows_affected = sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_all(&mut **tx)
.await?;
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
for row in &rows {
sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
.bind(Uuid::new_v4().to_string())
.bind(day_start_utc)
.bind(row.try_get::<String, _>("error_category")?)
.bind(row.try_get::<Option<String>, _>("provider_name")?)
.bind(row.try_get::<Option<String>, _>("model")?)
.bind(row.try_get::<i64, _>("total_count")?)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(rows.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_rows(
@@ -361,87 +294,15 @@ async fn upsert_stats_user_daily_rows(
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let active_user_ids = sqlx::query(SELECT_ACTIVE_USER_IDS_SQL)
.fetch_all(&mut **tx)
.await?
.into_iter()
.map(|row| row.try_get::<String, _>("id"))
.collect::<Result<Vec<_>, _>>()?;
if active_user_ids.is_empty() {
return Ok(0);
}
let aggregated_rows = sqlx::query(SELECT_STATS_USER_DAILY_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_all(&mut **tx)
.await?;
let mut aggregated_by_user = HashMap::with_capacity(aggregated_rows.len());
for row in aggregated_rows {
let user_id = row.try_get::<String, _>("user_id")?;
aggregated_by_user.insert(user_id, row);
}
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
for user_id in &active_user_ids {
let aggregated = aggregated_by_user.get(user_id);
let total_requests = aggregated
.map(|row| row.try_get::<i64, _>("total_requests"))
.transpose()?
.unwrap_or_default();
let error_requests = aggregated
.map(|row| row.try_get::<i64, _>("error_requests"))
.transpose()?
.unwrap_or_default();
sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(user_id)
.bind(
aggregated
.map(|row| row.try_get::<Option<String>, _>("username"))
.transpose()?
.flatten(),
)
.bind(day_start_utc)
.bind(total_requests)
.bind(total_requests.saturating_sub(error_requests))
.bind(error_requests)
.bind(
aggregated
.map(|row| row.try_get::<i64, _>("input_tokens"))
.transpose()?
.unwrap_or_default(),
)
.bind(
aggregated
.map(|row| row.try_get::<i64, _>("output_tokens"))
.transpose()?
.unwrap_or_default(),
)
.bind(
aggregated
.map(|row| row.try_get::<i64, _>("cache_creation_tokens"))
.transpose()?
.unwrap_or_default(),
)
.bind(
aggregated
.map(|row| row.try_get::<i64, _>("cache_read_tokens"))
.transpose()?
.unwrap_or_default(),
)
.bind(
aggregated
.map(|row| row.try_get::<f64, _>("total_cost"))
.transpose()?
.unwrap_or_default(),
)
.bind(now_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?;
}
Ok(active_user_ids.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn refresh_stats_summary_row(

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

@@ -7,9 +7,8 @@ use aether_data_contracts::DataLayerError;
use super::{
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_STATS_HOURLY_AGGREGATE_SQL,
SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL, SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL,
SELECT_STATS_HOURLY_USER_AGGREGATES_SQL, UPSERT_STATS_HOURLY_MODEL_SQL,
UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL, UPSERT_STATS_HOURLY_USER_SQL,
UPSERT_STATS_HOURLY_MODEL_SQL, UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL,
UPSERT_STATS_HOURLY_USER_SQL,
};
#[derive(Debug, Clone, PartialEq)]
@@ -115,51 +114,16 @@ async fn upsert_stats_hourly_user_rows(
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_USER_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?;
.map_err(postgres_error)?
.rows_affected();
for row in &rows {
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())
.bind(hour_utc)
.bind(user_id)
.bind(total_requests)
.bind(success_requests)
.bind(error_requests)
.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
.map_err(postgres_error)?;
}
Ok(rows.len())
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_model_rows(
@@ -168,54 +132,16 @@ async fn upsert_stats_hourly_model_rows(
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?;
let mut inserted = 0usize;
.map_err(postgres_error)?
.rows_affected();
for row in &rows {
let model = row
.try_get::<Option<String>, _>("model")
.map_err(postgres_error)?;
let Some(model) = model.filter(|value| !value.is_empty()) else {
continue;
};
sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(model)
.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
.map_err(postgres_error)?;
inserted += 1;
}
Ok(inserted)
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_provider_rows(
@@ -224,50 +150,16 @@ async fn upsert_stats_hourly_provider_rows(
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
let rows = sqlx::query(SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL)
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_all(&mut **tx)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?;
let mut inserted = 0usize;
.map_err(postgres_error)?
.rows_affected();
for row in &rows {
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;
};
sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(provider_name)
.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
.map_err(postgres_error)?;
inserted += 1;
}
Ok(inserted)
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
fn postgres_error(error: sqlx::Error) -> DataLayerError {

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

@@ -29,7 +29,7 @@ use super::{
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, SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL,
SELECT_STALE_PENDING_USAGE_BATCH_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
UPDATE_FAILED_VOID_STALE_USAGE_SQL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
@@ -82,9 +82,9 @@ async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
#[test]
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
assert!(
SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL.contains("JOIN usage_settlement_snapshots")
UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL.contains("JOIN usage_settlement_snapshots")
);
assert!(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL
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

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

@@ -6,6 +6,7 @@ use aether_data_contracts::repository::usage::{
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;
@@ -21,7 +22,8 @@ use super::{
DISABLE_EXPIRED_API_KEY_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_HEADER_BATCH_SQL, SELECT_USAGE_LEGACY_BODY_REF_METADATA_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,
@@ -211,23 +213,20 @@ async fn cleanup_usage_header_fields(
let mut total_cleaned = 0usize;
loop {
let rows = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
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_all(pool)
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| {
Ok(UsageBodyCleanupRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
})
})
.collect::<Result<Vec<_>, DataLayerError>>()?;
.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;
}
@@ -279,23 +278,20 @@ async fn cleanup_usage_stale_body_fields(
let mut total_cleaned = 0usize;
loop {
let rows = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
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_all(pool)
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| {
Ok(UsageBodyCleanupRow {
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
request_id: row
.try_get::<String, _>("request_id")
.map_err(postgres_error)?,
})
})
.collect::<Result<Vec<_>, DataLayerError>>()?;
.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;
}
@@ -354,53 +350,59 @@ async fn compress_usage_body_fields(
let mut no_progress_count = 0usize;
let batch_size = batch_size.clamp(1, 25);
loop {
let rows = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
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_all(pool)
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| {
Ok(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)?,
})
})
.collect::<Result<Vec<_>, DataLayerError>>()?;
if rows.is_empty() {
.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 row in rows {
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)?;
@@ -638,12 +640,9 @@ async fn cleanup_expired_api_keys(
pool: &aether_data::postgres::PostgresPool,
auto_delete_expired_keys: bool,
) -> Result<usize, DataLayerError> {
let expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL)
.fetch_all(pool)
.await
.map_err(postgres_error)?;
let mut expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL).fetch(pool);
let mut cleaned = 0usize;
for row in &expired_keys {
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(),

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

@@ -1,14 +1,11 @@
use chrono::{DateTime, Utc};
use sqlx::Row;
use uuid::Uuid;
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, SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL,
UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
};
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -43,60 +40,16 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
};
let mut tx = pool.begin().await.map_err(postgres_error)?;
let rows = sqlx::query(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL)
let aggregated_wallets = sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
.bind(target.window_start_utc)
.bind(target.window_end_utc)
.fetch_all(&mut *tx)
.bind(target.billing_date)
.bind(target.billing_timezone.as_str())
.bind(now_utc)
.execute(&mut *tx)
.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")
.map_err(postgres_error)?,
)
.bind(target.billing_date)
.bind(target.billing_timezone.as_str())
.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
.map_err(postgres_error)?;
}
.map_err(postgres_error)?
.rows_affected();
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
.bind(target.billing_date)
@@ -112,7 +65,7 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
Ok(WalletDailyUsageAggregationSummary {
billing_date: target.billing_date,
billing_timezone: target.billing_timezone,
aggregated_wallets: rows.len(),
aggregated_wallets: usize::try_from(aggregated_wallets).unwrap_or(usize::MAX),
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
})
}