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

feat(gateway): 重构 usage 数据层、迁移系统与系统导入

Browse Source 
数据库迁移:
- 引入 baseline v2 bootstrap，空库首次启动自动初始化
- 服务启动不再自动执行迁移，需显式 `--migrate` 运行
- 新增 pending migration 检测，schema 落后时拒绝启动

Usage 数据层:
- usage body 存储外部化为独立 blob 表
- 新增 HTTP audit 表拆分存储请求/响应头与 body ref
- 后台清理任务支持 legacy body ref 元数据迁移
- usage runtime 写入迁移到专用 tokio runtime（独立线程池, 8MB 栈）

系统导入/导出:
- 支持用户、API Keys、钱包数据的完整导入
- 兼容 legacy 与 v1.3+ 两种导出格式

其他改进:
- executor outcome 增加 runtime miss 诊断上下文
- 主 tokio runtime 栈大小调整为 8MB
- 前端 provider 管理支持 base URL 配置
- dev.sh 支持 --migrate 参数
This commit is contained in:
fawney19
2026-04-13 14:01:22 +08:00
parent 3698e5a833
commit 5bb08e6aa4
106 changed files with 21736 additions and 1529 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -111,22 +111,24 @@ const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
const DB_MAINTENANCE_TABLES: &[&str] = &["usage", "request_candidates", "audit_logs"];
const SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL: &str = r#"
SELECT
wallet_id,
COUNT(id) AS total_requests,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(cache_creation_input_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens,
MIN(finalized_at) AS first_finalized_at,
MAX(finalized_at) AS last_finalized_at
usage_settlement_snapshots.wallet_id,
COUNT(usage.id) AS total_requests,
CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
COALESCE(SUM(usage.input_tokens), 0) AS input_tokens,
COALESCE(SUM(usage.output_tokens), 0) AS output_tokens,
COALESCE(SUM(usage.cache_creation_input_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(usage.cache_read_input_tokens), 0) AS cache_read_tokens,
MIN(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS first_finalized_at,
MAX(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS last_finalized_at
FROM usage
WHERE wallet_id IS NOT NULL
AND billing_status = 'settled'
AND total_cost_usd > 0
AND finalized_at >= $1
AND finalized_at < $2
GROUP BY wallet_id
JOIN usage_settlement_snapshots
ON usage_settlement_snapshots.request_id = usage.request_id
WHERE usage_settlement_snapshots.wallet_id IS NOT NULL
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
AND usage.total_cost_usd > 0
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $1
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $2
GROUP BY usage_settlement_snapshots.wallet_id
"#;
const UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL: &str = r#"
INSERT INTO wallet_daily_usage_ledgers (
@@ -171,11 +173,13 @@ WHERE ledgers.billing_date = $1
AND NOT EXISTS (
SELECT 1
FROM usage
WHERE usage.wallet_id = ledgers.wallet_id
AND usage.billing_status = 'settled'
JOIN usage_settlement_snapshots
ON usage_settlement_snapshots.request_id = usage.request_id
WHERE usage_settlement_snapshots.wallet_id = ledgers.wallet_id
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
AND usage.total_cost_usd > 0
AND usage.finalized_at >= $3
AND usage.finalized_at < $4
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $3
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $4
)
"#;
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
@@ -258,7 +262,7 @@ USING doomed
WHERE usage_rows.id = doomed.id
"#;
const SELECT_USAGE_HEADER_BATCH_SQL: &str = r#"
SELECT id
SELECT id, request_id
FROM usage
WHERE created_at < $1
AND ($2::timestamptz IS NULL OR created_at >= $2)
@@ -267,6 +271,17 @@ WHERE created_at < $1
OR response_headers IS NOT NULL
OR provider_request_headers IS NOT NULL
OR client_response_headers IS NOT NULL
OR EXISTS (
SELECT 1
FROM usage_http_audits
WHERE usage_http_audits.request_id = usage.request_id
AND (
usage_http_audits.request_headers IS NOT NULL
OR usage_http_audits.response_headers IS NOT NULL
OR usage_http_audits.provider_request_headers IS NOT NULL
OR usage_http_audits.client_response_headers IS NOT NULL
)
)
)
ORDER BY created_at ASC, id ASC
LIMIT $3
@@ -279,8 +294,17 @@ SET request_headers = NULL,
client_response_headers = NULL
WHERE id = ANY($1)
"#;
const CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL: &str = r#"
UPDATE usage_http_audits
SET request_headers = NULL,
response_headers = NULL,
provider_request_headers = NULL,
client_response_headers = NULL,
updated_at = NOW()
WHERE request_id = ANY($1)
"#;
const SELECT_USAGE_STALE_BODY_BATCH_SQL: &str = r#"
SELECT id
SELECT id, request_id
FROM usage
WHERE created_at < $1
AND ($2::timestamptz IS NULL OR created_at >= $2)
@@ -293,6 +317,22 @@ WHERE created_at < $1
OR response_body_compressed IS NOT NULL
OR provider_request_body_compressed IS NOT NULL
OR client_response_body_compressed IS NOT NULL
OR EXISTS (
SELECT 1
FROM usage_body_blobs
WHERE usage_body_blobs.request_id = usage.request_id
)
OR EXISTS (
SELECT 1
FROM usage_http_audits
WHERE usage_http_audits.request_id = usage.request_id
AND (
usage_http_audits.request_body_ref IS NOT NULL
OR usage_http_audits.provider_request_body_ref IS NOT NULL
OR usage_http_audits.response_body_ref IS NOT NULL
OR usage_http_audits.client_response_body_ref IS NOT NULL
)
)
)
ORDER BY created_at ASC, id ASC
LIMIT $3
@@ -309,35 +349,150 @@ SET request_body = NULL,
client_response_body_compressed = NULL
WHERE id = ANY($1)
"#;
const DELETE_USAGE_BODY_BLOBS_SQL: &str = r#"
DELETE FROM usage_body_blobs
WHERE request_id = ANY($1)
"#;
const CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL: &str = r#"
UPDATE usage_http_audits
SET request_body_ref = NULL,
provider_request_body_ref = NULL,
response_body_ref = NULL,
client_response_body_ref = NULL,
body_capture_mode = 'none',
updated_at = NOW()
WHERE request_id = ANY($1)
"#;
const DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL: &str = r#"
DELETE FROM usage_http_audits
WHERE request_id = ANY($1)
AND request_headers IS NULL
AND response_headers IS NULL
AND provider_request_headers IS NULL
AND client_response_headers IS NULL
AND request_body_ref IS NULL
AND provider_request_body_ref IS NULL
AND response_body_ref IS NULL
AND client_response_body_ref IS NULL
"#;
const SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL: &str = r#"
SELECT
id,
request_id,
request_body,
request_body_compressed,
response_body,
response_body_compressed,
provider_request_body,
provider_request_body_compressed,
client_response_body
,
client_response_body_compressed
FROM usage
WHERE created_at < $1
AND ($2::timestamptz IS NULL OR created_at >= $2)
AND (
request_body IS NOT NULL
OR request_body_compressed IS NOT NULL
OR response_body IS NOT NULL
OR response_body_compressed IS NOT NULL
OR provider_request_body IS NOT NULL
OR provider_request_body_compressed IS NOT NULL
OR client_response_body IS NOT NULL
OR client_response_body_compressed IS NOT NULL
)
ORDER BY created_at ASC, id ASC
LIMIT $3
"#;
const SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL: &str = r#"
SELECT
id,
request_id,
request_metadata
FROM usage
WHERE created_at < $1
AND ($2::timestamptz IS NULL OR created_at >= $2)
AND request_metadata IS NOT NULL
AND (
request_metadata::jsonb ? 'request_body_ref'
OR request_metadata::jsonb ? 'provider_request_body_ref'
OR request_metadata::jsonb ? 'response_body_ref'
OR request_metadata::jsonb ? 'client_response_body_ref'
)
ORDER BY created_at ASC, id ASC
LIMIT $3
"#;
const UPSERT_USAGE_BODY_BLOB_SQL: &str = r#"
INSERT INTO usage_body_blobs (
body_ref,
request_id,
body_field,
payload_gzip
) VALUES (
$1,
$2,
$3,
$4
)
ON CONFLICT (body_ref)
DO UPDATE SET
payload_gzip = EXCLUDED.payload_gzip,
updated_at = NOW()
"#;
const UPDATE_USAGE_REQUEST_METADATA_SQL: &str = r#"
UPDATE usage
SET request_metadata = $2::json,
updated_at = NOW()
WHERE id = $1
"#;
const UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL: &str = r#"
INSERT INTO usage_http_audits (
request_id,
request_body_ref,
provider_request_body_ref,
response_body_ref,
client_response_body_ref,
body_capture_mode
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6
)
ON CONFLICT (request_id)
DO UPDATE SET
request_body_ref = COALESCE(EXCLUDED.request_body_ref, usage_http_audits.request_body_ref),
provider_request_body_ref = COALESCE(
EXCLUDED.provider_request_body_ref,
usage_http_audits.provider_request_body_ref
),
response_body_ref = COALESCE(EXCLUDED.response_body_ref, usage_http_audits.response_body_ref),
client_response_body_ref = COALESCE(
EXCLUDED.client_response_body_ref,
usage_http_audits.client_response_body_ref
),
body_capture_mode = CASE
WHEN EXCLUDED.request_body_ref IS NOT NULL
OR EXCLUDED.provider_request_body_ref IS NOT NULL
OR EXCLUDED.response_body_ref IS NOT NULL
OR EXCLUDED.client_response_body_ref IS NOT NULL
THEN EXCLUDED.body_capture_mode
ELSE usage_http_audits.body_capture_mode
END,
updated_at = NOW()
"#;
const UPDATE_USAGE_BODY_COMPRESSION_SQL: &str = r#"
UPDATE usage
SET request_body = NULL,
response_body = NULL,
provider_request_body = NULL,
client_response_body = NULL,
request_body_compressed = $2,
response_body_compressed = $3,
provider_request_body_compressed = $4,
client_response_body_compressed = $5
request_body_compressed = NULL,
response_body_compressed = NULL,
provider_request_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE id = $1
"#;
const SELECT_EXPIRED_ACTIVE_API_KEYS_SQL: &str = r#"
@@ -988,7 +1143,8 @@ struct PercentileSummary {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
struct UsageCleanupSummary {
body_compressed: usize,
body_externalized: usize,
legacy_body_refs_migrated: usize,
body_cleaned: usize,
header_cleaned: usize,
keys_cleaned: usize,
@@ -1016,10 +1172,21 @@ struct UsageCleanupWindow {
#[derive(Debug, Clone, PartialEq)]
struct UsageBodyCompressionRow {
id: String,
request_id: String,
request_body: Option<Value>,
request_body_compressed: Option<Vec<u8>>,
response_body: Option<Value>,
response_body_compressed: Option<Vec<u8>>,
provider_request_body: Option<Value>,
provider_request_body_compressed: Option<Vec<u8>>,
client_response_body: Option<Value>,
client_response_body_compressed: Option<Vec<u8>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct UsageBodyCleanupRow {
id: String,
request_id: String,
}
#[derive(Debug, Clone, Copy, PartialEq)]

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

@@ -188,7 +188,8 @@ pub(super) async fn run_stats_aggregation_once(
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
if summary.body_externalized > 0
|| summary.legacy_body_refs_migrated > 0
|| summary.body_cleaned > 0
|| summary.header_cleaned > 0
|| summary.keys_cleaned > 0
@@ -198,7 +199,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
event_name = "usage_cleanup_completed",
log_type = "ops",
worker = "usage_cleanup",
body_compressed = summary.body_compressed,
body_externalized = summary.body_externalized,
legacy_body_refs_migrated = summary.legacy_body_refs_migrated,
body_cleaned = summary.body_cleaned,
header_cleaned = summary.header_cleaned,
keys_cleaned = summary.keys_cleaned,

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

@@ -28,7 +28,8 @@ use super::{
stats_hourly_aggregation_target_hour, summarize_postgres_pool, usage_cleanup_settings,
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
UsageCleanupSettings, DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL,
SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
@@ -77,6 +78,18 @@ async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
assert!(spawn_pool_monitor_worker(Arc::new(GatewayDataState::disabled())).is_none());
}
#[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")
);
assert!(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_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"));
}
fn sample_connected_proxy_node(
node_id: &str,
heartbeat_interval: i32,

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

@@ -1,9 +1,12 @@
use std::io::Write;
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, UsageBodyField,
};
use aether_data_contracts::DataLayerError;
use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use serde_json::Value;
use serde_json::{Map, Value};
use sqlx::Row;
use tracing::warn;
@@ -11,13 +14,17 @@ use crate::data::GatewayDataState;
use super::{
system_config_bool, usage_cleanup_settings, usage_cleanup_window, ExpiredApiKeyRow,
UsageBodyCompressionRow, UsageCleanupSummary, CLEAR_USAGE_BODY_FIELDS_SQL,
CLEAR_USAGE_HEADER_FIELDS_SQL, DELETE_EXPIRED_API_KEY_SQL, DELETE_OLD_USAGE_RECORDS_SQL,
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, 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_STALE_BODY_BATCH_SQL,
UPDATE_USAGE_BODY_COMPRESSION_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,
};
pub(super) async fn perform_usage_cleanup_once(
@@ -41,6 +48,13 @@ pub(super) async fn perform_usage_cleanup_once(
Some(window.log_cutoff),
)
.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,
@@ -48,7 +62,7 @@ pub(super) async fn perform_usage_cleanup_once(
Some(window.log_cutoff),
)
.await?;
let body_compressed = compress_usage_body_fields(
let body_externalized = compress_usage_body_fields(
&pool,
window.detail_cutoff,
settings.batch_size,
@@ -65,7 +79,8 @@ pub(super) async fn perform_usage_cleanup_once(
};
Ok(UsageCleanupSummary {
body_compressed,
body_externalized,
legacy_body_refs_migrated,
body_cleaned,
header_cleaned,
keys_cleaned,
@@ -73,6 +88,89 @@ pub(super) async fn perform_usage_cleanup_once(
})
}
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>,
@@ -113,7 +211,7 @@ async fn cleanup_usage_header_fields(
let mut total_cleaned = 0usize;
loop {
let ids = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
let rows = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
@@ -121,11 +219,23 @@ async fn cleanup_usage_header_fields(
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
.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>>()?;
if ids.is_empty() {
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)
@@ -133,6 +243,16 @@ async fn cleanup_usage_header_fields(
.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 {
@@ -159,7 +279,7 @@ async fn cleanup_usage_stale_body_fields(
let mut total_cleaned = 0usize;
loop {
let ids = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
let rows = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
@@ -167,11 +287,23 @@ async fn cleanup_usage_stale_body_fields(
.await
.map_err(postgres_error)?
.into_iter()
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
.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>>()?;
if ids.is_empty() {
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)
@@ -179,6 +311,21 @@ async fn cleanup_usage_stale_body_fields(
.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 {
@@ -218,18 +365,33 @@ async fn compress_usage_body_fields(
.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>>()?;
@@ -239,18 +401,44 @@ async fn compress_usage_body_fields(
let mut batch_success = 0usize;
for row in rows {
let compressed = (
compress_usage_json_value(row.request_body.as_ref()),
compress_usage_json_value(row.response_body.as_ref()),
compress_usage_json_value(row.provider_request_body.as_ref()),
compress_usage_json_value(row.client_response_body.as_ref()),
);
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)
.bind(compressed.0)
.bind(compressed.1)
.bind(compressed.2)
.bind(compressed.3)
.bind(&row.id)
.execute(pool)
.await
.map_err(postgres_error)?
@@ -276,12 +464,174 @@ async fn compress_usage_body_fields(
Ok(total_compressed)
}
fn compress_usage_json_value(value: Option<&Value>) -> Option<Vec<u8>> {
let value = value?;
let bytes = serde_json::to_vec(value).ok()?;
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).ok()?;
encoder.finish().ok()
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(
@@ -375,3 +725,137 @@ async fn nullify_expired_api_key_candidate_refs(
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
}))
);
}
}