i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 01:17:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(security): harden gateway boundaries and usage policies

Browse Source 
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/usage/cleanup.rs
+141 -462
View File
@@ -1,11 +1,8 @@
use std::io::Write;
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, UsageBodyField, UsageCleanupExecutionMode,
UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
UsageCleanupExecutionMode, UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets,
UsageCleanupWindow,
};
use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use futures_util::TryStreamExt;
use serde_json::Value;
use sqlx::Row;
@@ -214,8 +211,7 @@ SET request_body_ref = NULL,
WHERE request_id = ANY($1)
"#;
const SELECT_USAGE_BODY_COMPRESSION_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)
@@ -228,26 +224,26 @@ WHERE created_at < $1
OR provider_request_body_compressed IS NOT NULL
OR client_response_body 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
"#;
const SELECT_USAGE_BODY_COMPRESSION_ROW_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 id = $1
LIMIT 1
"#;
const SELECT_EXPIRED_ACTIVE_API_KEYS_SQL: &str = r#"
SELECT id, auto_delete_on_expiry
FROM api_keys
@@ -274,30 +270,6 @@ WHERE id = $1
AND is_active IS TRUE
"#;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageDetachedBodyBlobWrite {
pub body_ref: String,
pub body_field: &'static str,
pub payload_gzip: Vec<u8>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UsageDetachedBodyRefs {
pub request_body_ref: Option<String>,
pub provider_request_body_ref: Option<String>,
pub response_body_ref: Option<String>,
pub client_response_body_ref: Option<String>,
}
impl UsageDetachedBodyRefs {
pub 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, PartialEq)]
pub struct UsageLegacyBodyRefMetadataRow {
pub id: String,
@@ -305,30 +277,9 @@ pub struct UsageLegacyBodyRefMetadataRow {
pub request_metadata: Option<Value>,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct UsageLegacyBodyRefMigrationPlan {
pub refs: UsageDetachedBodyRefs,
pub request_metadata: Option<Value>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct UsageBodyCompressionRow {
pub id: String,
pub request_id: String,
pub request_body: Option<Value>,
pub request_body_compressed: Option<Vec<u8>>,
pub response_body: Option<Value>,
pub response_body_compressed: Option<Vec<u8>>,
pub provider_request_body: Option<Value>,
pub provider_request_body_compressed: Option<Vec<u8>>,
pub client_response_body: Option<Value>,
pub client_response_body_compressed: Option<Vec<u8>>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UsageBodyExternalizationPlan {
pub blobs: Vec<UsageDetachedBodyBlobWrite>,
pub refs: UsageDetachedBodyRefs,
pub struct UsageLegacyBodyRefPurgePlan {
pub request_metadata: Option<Value>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -343,57 +294,23 @@ struct ExpiredApiKeyRow<'a> {
auto_delete_on_expiry: Option<bool>,
}
pub fn compress_usage_json_value(value: &Value) -> Result<Vec<u8>, DataLayerError> {
let bytes = serde_json::to_vec(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to serialize usage json for gzip: {err}"))
})?;
let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
encoder.write_all(&bytes).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to gzip usage json: {err}"))
})?;
encoder.finish().map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to finish gzipped usage json: {err}"))
})
}
pub fn migrate_legacy_body_ref_metadata_plan(
request_id: &str,
pub fn purge_legacy_body_ref_metadata_plan(
request_metadata: Option<Value>,
) -> Option<UsageLegacyBodyRefMigrationPlan> {
) -> Option<UsageLegacyBodyRefPurgePlan> {
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,
for key in [
"request_body_ref",
"provider_request_body_ref",
"response_body_ref",
"client_response_body_ref",
] {
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 metadata.remove(key).is_some() {
removed_any = true;
}
}
@@ -401,47 +318,11 @@ pub fn migrate_legacy_body_ref_metadata_plan(
return None;
}
Some(UsageLegacyBodyRefMigrationPlan {
refs,
Some(UsageLegacyBodyRefPurgePlan {
request_metadata: (!metadata.is_empty()).then_some(Value::Object(metadata)),
})
}
pub 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)
}
impl SqlxUsageReadRepository {
pub async fn cleanup_usage(
&self,
@@ -477,6 +358,8 @@ impl SqlxUsageReadRepository {
header_cleaned: 0,
keys_cleaned: 0,
records_deleted: 0,
cost_reservations_deleted: 0,
request_admissions_deleted: 0,
});
}
@@ -509,7 +392,7 @@ impl SqlxUsageReadRepository {
};
let detail_body_newer_than = detail_body_newer_than(window, targets);
let legacy_body_refs_migrated = if targets.detail_body {
migrate_legacy_usage_body_ref_metadata(
purge_legacy_usage_body_ref_metadata(
&self.pool,
window.detail_cutoff,
batch_size,
@@ -520,7 +403,7 @@ impl SqlxUsageReadRepository {
0
};
let body_externalized = if targets.detail_body {
compress_usage_body_fields(
purge_usage_detail_body_fields(
&self.pool,
window.detail_cutoff,
batch_size,
@@ -549,6 +432,8 @@ impl SqlxUsageReadRepository {
header_cleaned,
keys_cleaned,
records_deleted,
cost_reservations_deleted: 0,
request_admissions_deleted: 0,
})
}
}
@@ -646,12 +531,33 @@ async fn cleanup_usage_raw_body_fields(
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 mut tx = pool.begin().await.map_err(postgres_error)?;
let cleaned = sqlx::query(CLEAR_USAGE_RAW_BODY_FIELDS_SQL)
.bind(ids)
.execute(pool)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
sqlx::query(DELETE_USAGE_BODY_BLOBS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
.bind(request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
total_cleaned += cleaned;
if rows.len() < batch_size {
@@ -973,7 +879,7 @@ async fn delete_old_usage_records(
Ok(total_deleted)
}
async fn migrate_legacy_usage_body_ref_metadata(
async fn purge_legacy_usage_body_ref_metadata(
pool: &PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
@@ -983,7 +889,7 @@ async fn migrate_legacy_usage_body_ref_metadata(
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"usage cleanup legacy body-ref migration skipped due to invalid window"
"usage cleanup legacy body-ref purge skipped due to invalid window"
);
return Ok(0);
}
@@ -1014,26 +920,12 @@ async fn migrate_legacy_usage_body_ref_metadata(
break;
}
let mut batch_migrated = 0usize;
let mut batch_purged = 0usize;
for row in rows {
let Some(plan) =
migrate_legacy_body_ref_metadata_plan(&row.request_id, row.request_metadata)
else {
let Some(plan) = purge_legacy_body_ref_metadata_plan(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)
@@ -1041,14 +933,30 @@ async fn migrate_legacy_usage_body_ref_metadata(
.await
.map_err(postgres_error)?
.rows_affected();
let request_ids = vec![row.request_id];
sqlx::query(DELETE_USAGE_BODY_BLOBS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
.bind(request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
if updated > 0 {
batch_migrated += 1;
batch_purged += 1;
}
}
total_migrated += batch_migrated;
if batch_migrated == 0 || batch_migrated < batch_size {
total_migrated += batch_purged;
if batch_purged == 0 || batch_purged < batch_size {
break;
}
}
@@ -1191,7 +1099,7 @@ async fn cleanup_usage_stale_body_fields(
Ok(total_cleaned)
}
async fn compress_usage_body_fields(
async fn purge_usage_detail_body_fields(
pool: &PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
@@ -1201,129 +1109,65 @@ async fn compress_usage_body_fields(
warn!(
cutoff_time = %cutoff_time,
newer_than = ?newer_than,
"usage cleanup body compression skipped due to invalid window"
"usage cleanup detail body purge skipped due to invalid window"
);
return Ok(0);
}
let mut total_compressed = 0usize;
let mut no_progress_count = 0usize;
let batch_size = batch_size.clamp(1, 25);
let mut total_purged = 0usize;
loop {
let mut stream = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
.bind(cutoff_time)
.bind(newer_than)
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
.fetch(pool);
let mut ids = Vec::new();
let mut rows = Vec::new();
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
ids.push(row.try_get::<String, _>("id").map_err(postgres_error)?);
}
if ids.is_empty() {
break;
}
let mut batch_success = 0usize;
for id in ids {
let row = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_ROW_SQL)
.bind(&id)
.fetch_optional(pool)
.await
.map_err(postgres_error)?;
let Some(row) = row else {
continue;
};
let row = UsageBodyCompressionRow {
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)?,
request_body: row
.try_get::<Option<Value>, _>("request_body")
.map_err(postgres_error)?,
request_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("request_body_compressed")
.map_err(postgres_error)?,
response_body: row
.try_get::<Option<Value>, _>("response_body")
.map_err(postgres_error)?,
response_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("response_body_compressed")
.map_err(postgres_error)?,
provider_request_body: row
.try_get::<Option<Value>, _>("provider_request_body")
.map_err(postgres_error)?,
provider_request_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
.map_err(postgres_error)?,
client_response_body: row
.try_get::<Option<Value>, _>("client_response_body")
.map_err(postgres_error)?,
client_response_body_compressed: row
.try_get::<Option<Vec<u8>>, _>("client_response_body_compressed")
.map_err(postgres_error)?,
};
let detached = build_usage_body_externalization(&row)?;
if detached.refs.any_present() {
let mut tx = pool.begin().await.map_err(postgres_error)?;
for blob in &detached.blobs {
sqlx::query(super::UPSERT_USAGE_BODY_BLOB_SQL)
.bind(&blob.body_ref)
.bind(&row.request_id)
.bind(blob.body_field)
.bind(&blob.payload_gzip)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
}
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&row.request_id)
.bind(detached.refs.request_body_ref.as_deref())
.bind(detached.refs.provider_request_body_ref.as_deref())
.bind(detached.refs.response_body_ref.as_deref())
.bind(detached.refs.client_response_body_ref.as_deref())
.bind("ref_backed")
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
.bind(&row.id)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
tx.commit().await.map_err(postgres_error)?;
if updated > 0 {
batch_success += 1;
}
continue;
}
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
.bind(&row.id)
.execute(pool)
.await
.map_err(postgres_error)?
.rows_affected();
if updated > 0 {
batch_success += 1;
}
});
}
if batch_success == 0 {
no_progress_count += 1;
if no_progress_count >= 3 {
warn!(
"usage cleanup body compression stopped after repeated zero-progress batches"
);
break;
}
} else {
no_progress_count = 0;
if rows.is_empty() {
break;
}
let row_count = rows.len();
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 mut tx = pool.begin().await.map_err(postgres_error)?;
let updated = sqlx::query(CLEAR_USAGE_BODY_FIELDS_SQL)
.bind(ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?
.rows_affected();
sqlx::query(DELETE_USAGE_BODY_BLOBS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
.bind(&request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
.bind(request_ids)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
total_purged = total_purged.saturating_add(usize::try_from(updated).unwrap_or(usize::MAX));
if row_count < batch_size {
break;
}
total_compressed += batch_success;
}
Ok(total_compressed)
Ok(total_purged)
}
async fn cleanup_expired_api_keys(
@@ -1374,74 +1218,6 @@ async fn cleanup_expired_api_keys(
Ok(cleaned)
}
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(())
}
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_REQUEST_METADATA_SQL: &str = r#"
UPDATE usage
SET request_metadata = $2::json,
@@ -1449,41 +1225,15 @@ SET request_metadata = $2::json,
WHERE id = $1
"#;
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 = NULL,
response_body_compressed = NULL,
provider_request_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE id = $1
"#;
#[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,
purge_legacy_body_ref_metadata_plan, SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL,
SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL,
};
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 legacy_body_ref_cleanup_index_is_embedded_and_matches_batch_predicate() {
const MIGRATION_VERSION: i64 = 20_260_715_000_000;
@@ -1532,88 +1282,21 @@ mod tests {
}
#[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})
);
fn detail_body_cleanup_selects_detached_capture_for_deletion() {
assert!(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL.contains("usage_body_blobs"));
assert!(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL.contains("usage_http_audits"));
assert!(!SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL.contains("payload_gzip"));
}
#[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"
})),
)
fn legacy_body_ref_metadata_purge_strips_all_ref_keys() {
let plan = purge_legacy_body_ref_metadata_plan(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!({
@@ -1623,18 +1306,14 @@ mod tests {
}
#[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
})),
)
fn legacy_body_ref_metadata_purge_does_not_preserve_untrusted_refs() {
let plan = purge_legacy_body_ref_metadata_plan(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!({
crates/aether-data/adapters/postgres/src/usage/mod.rs
+296 -212
View File
@@ -1,8 +1,9 @@
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, ApiKeyLastUsedDelta, ManagementTokenCounterDelta,
ProxyNodeCounterDelta, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
canonical_usage_body_ref_for, parse_usage_body_ref, read_decompressed_usage_json,
usage_body_ref, ApiKeyLastUsedDelta, ManagementTokenCounterDelta, ProxyNodeCounterDelta,
StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow,
StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow,
StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, StoredUsageErrorDistributionRow,
StoredUsageLeaderboardSummary, StoredUsagePerformancePercentilesRow,
@@ -33,7 +34,7 @@ use sqlx::{
PgPool, Postgres, QueryBuilder, Row,
};
use std::collections::{BTreeMap, BTreeSet};
use std::io::{Read, Write};
use std::io::Write;
use uuid::Uuid;
use crate::{
@@ -41,11 +42,12 @@ use crate::{
PostgresTransactionRunner,
};
use aether_data_contracts::repository::usage::{
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
model_usage_contribution, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, ApiKeyUsageDelta, ModelUsageDelta,
PendingUsageCleanupSummary, ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
api_key_usage_contribution, model_usage_contribution, provider_api_key_usage_contribution,
sanitize_usage_capture_controls_for_persistence, sanitize_usage_for_persistence,
sanitize_usage_request_metadata, usage_can_recover_terminal_failure,
usage_error_category_for_status_code, usage_lifecycle_update_allowed, ApiKeyUsageDelta,
ModelUsageDelta, PendingUsageCleanupSummary, ProviderApiKeyUsageContribution,
ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageCounterFlushSummary,
UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery,
@@ -61,9 +63,7 @@ pub mod cleanup;
// newly captured bodies always spill to usage_body_blobs and resolve through usage_http_audits.
const MAX_INLINE_USAGE_BODY_BYTES: usize = 0;
const MAX_SUPPORTED_UNIX_SECS: u64 = 253_402_300_799;
const FIND_USAGE_BODY_BLOB_BY_REF_SQL: &str =
r#"SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = $1 LIMIT 1"#;
const UPSERT_USAGE_BODY_BLOB_SQL: &str = include_str!("queries/upsert_usage_body_blob_sql.sql");
const FIND_USAGE_BODY_BLOB_BY_REF_SQL: &str = r#"SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = $1 AND request_id = $2 AND body_field = $3 LIMIT 1"#;
const DELETE_USAGE_BODY_BLOB_SQL: &str = include_str!("queries/delete_usage_body_blob_sql.sql");
#[derive(Debug, Clone, PartialEq, Eq, Default)]
@@ -1372,9 +1372,10 @@ SELECT
FROM "usage" AS u
WHERE u.request_id = ANY($1)
"#;
const UPSERT_USAGE_HTTP_AUDIT_SQL: &str = include_str!("queries/upsert_usage_http_audit_sql.sql");
const UPSERT_USAGE_ROUTING_SNAPSHOT_SQL: &str =
include_str!("queries/upsert_usage_routing_snapshot_sql.sql");
#[cfg(test)]
const UPSERT_USAGE_HTTP_AUDIT_SQL: &str = include_str!("queries/upsert_usage_http_audit_sql.sql");
const UPSERT_USAGE_SETTLEMENT_PRICING_SNAPSHOT_SQL: &str =
include_str!("queries/upsert_usage_settlement_pricing_snapshot_sql.sql");
@@ -1418,6 +1419,7 @@ INSERT INTO usage_counter_deltas (
request_id,
kind,
target_id,
target_tunnel_generation,
request_count_delta,
total_requests_delta,
success_count_delta,
@@ -1433,7 +1435,7 @@ INSERT INTO usage_counter_deltas (
removed_last_used_at_unix_secs,
usage_created_at_unix_secs
) VALUES (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19
)
"#;
const INSERT_USAGE_COUNTER_DELTAS_PREFIX_SQL: &str = r#"
@@ -1442,6 +1444,7 @@ INSERT INTO usage_counter_deltas (
request_id,
kind,
target_id,
target_tunnel_generation,
request_count_delta,
total_requests_delta,
success_count_delta,
@@ -1457,7 +1460,7 @@ INSERT INTO usage_counter_deltas (
removed_last_used_at_unix_secs,
usage_created_at_unix_secs
) "#;
const USAGE_COUNTER_DELTA_INSERT_BINDS_PER_ROW: usize = 18;
const USAGE_COUNTER_DELTA_INSERT_BINDS_PER_ROW: usize = 19;
const USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE: usize =
u16::MAX as usize / USAGE_COUNTER_DELTA_INSERT_BINDS_PER_ROW;
@@ -1474,6 +1477,7 @@ SELECT
delta.id,
delta.kind,
delta.target_id,
delta.target_tunnel_generation,
delta.request_count_delta,
delta.total_requests_delta,
delta.success_count_delta,
@@ -1586,6 +1590,7 @@ SET
stream_errors = stream_errors + GREATEST($5::bigint, 0),
updated_at = NOW()
WHERE id = $1
AND tunnel_generation = $6
"#;
const APPLY_MANAGEMENT_TOKEN_COUNTER_DELTA_SQL: &str = r#"
@@ -1937,7 +1942,6 @@ INSERT INTO "usage" (
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
SELECT
usage.request_id,
usage.status,
COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) AS billing_status
FROM usage
LEFT JOIN usage_settlement_snapshots
@@ -1966,15 +1970,15 @@ const UPDATE_RECOVERED_STALE_USAGE_SQL: &str = r#"
UPDATE usage
SET status = 'completed',
status_code = 200,
error_message = NULL
error_message = NULL,
error_category = NULL
WHERE request_id = $1
"#;
const SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL: &str = r#"
SELECT DISTINCT ON (request_id)
request_id,
status_code,
error_message
status_code
FROM request_candidates
WHERE request_id = ANY($1)
AND status IN ('failed', 'cancelled')
@@ -1988,8 +1992,9 @@ ORDER BY request_id,
const UPDATE_FAILED_STALE_USAGE_SQL: &str = r#"
UPDATE usage
SET status = 'failed',
status_code = $3,
error_message = $2
status_code = $2,
error_message = NULL,
error_category = $3
WHERE request_id = $1
"#;
@@ -1997,10 +2002,11 @@ const UPDATE_FAILED_VOID_STALE_USAGE_SQL: &str = r#"
WITH updated_usage AS (
UPDATE usage
SET status = 'failed',
status_code = $4,
error_message = $2,
status_code = $3,
error_message = NULL,
error_category = $4,
billing_status = 'void',
finalized_at = $3,
finalized_at = $2,
total_cost_usd = 0,
request_cost_usd = 0,
actual_total_cost_usd = 0,
@@ -2013,7 +2019,7 @@ INSERT INTO usage_settlement_snapshots (
billing_status,
finalized_at
)
SELECT request_id, 'void', $3
SELECT request_id, 'void', $2
FROM updated_usage
ON CONFLICT (request_id)
DO UPDATE SET
@@ -2037,7 +2043,8 @@ const UPDATE_FAILED_PENDING_CANDIDATES_SQL: &str = r#"
UPDATE request_candidates
SET status = 'failed',
finished_at = $2,
error_message = '请求超时（服务器可能已重启）'
error_type = 'internal',
error_message = NULL
WHERE request_id = $1
AND status IN ('pending', 'streaming')
"#;
@@ -2097,8 +2104,12 @@ impl PreparedPendingUsage {
));
}
let usage = strip_deprecated_usage_display_fields(usage);
let prepared = prepare_usage_upsert_context(&usage)?;
// Keep the capture input separate from the accounting row. The persistence sanitizer
// intentionally removes HTTP bodies/headers/states, but the pending batch still needs
// those values to populate the canonical audit/blob tables.
let capture_usage = usage.clone();
let usage = sanitize_usage_for_persistence(usage);
let prepared = prepare_usage_upsert_context(&capture_usage)?;
let input_tokens = usage
.input_tokens
.map(to_i32)
@@ -2215,7 +2226,7 @@ impl PreparedFirstByteUsage {
));
}
let usage = strip_deprecated_usage_display_fields(usage);
let usage = sanitize_usage_for_persistence(usage);
let request_metadata_json = json_bind_text(usage.request_metadata.as_ref())?;
let response_time_ms = usage.response_time_ms.map(to_i32).transpose()?;
let first_byte_time_ms = usage.first_byte_time_ms.map(to_i32).transpose()?;
@@ -2251,6 +2262,7 @@ fn partition_first_byte_usages(
let mut batch_rows = Vec::new();
let mut fallback_rows = Vec::new();
for (sequence, usage) in usages.into_iter().enumerate() {
let original_usage = usage.clone();
let prepared = PreparedFirstByteUsage::try_from_usage(usage)?;
if request_id_counts
.get(&prepared.usage.request_id)
@@ -2260,7 +2272,10 @@ fn partition_first_byte_usages(
{
batch_rows.push(prepared);
} else {
fallback_rows.push((sequence, prepared.usage));
// Duplicate rows are replayed through the canonical upsert, which owns the
// persistence sanitizer. Keep the original metadata here so replay semantics remain
// lossless up to that boundary.
fallback_rows.push((sequence, original_usage));
}
}
Ok((batch_rows, fallback_rows))
@@ -2848,8 +2863,14 @@ ORDER BY request_count DESC, "usage".provider_name ASC
}
pub async fn resolve_body_ref(&self, body_ref: &str) -> Result<Option<Value>, DataLayerError> {
let Some((request_id, field)) = parse_usage_body_ref(body_ref) else {
return Ok(None);
};
let canonical_ref = usage_body_ref(&request_id, field);
let blob_row = sqlx::query(FIND_USAGE_BODY_BLOB_BY_REF_SQL)
.bind(body_ref)
.bind(&canonical_ref)
.bind(&request_id)
.bind(field.as_storage_field())
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -2859,9 +2880,6 @@ ORDER BY request_count DESC, "usage".provider_name ASC
.map_postgres_err()?;
return inflate_usage_json_value(&payload_gzip).map(Some);
}
let Some((request_id, field)) = parse_usage_body_ref(body_ref) else {
return Ok(None);
};
let (inline_column, compressed_column) = usage_body_sql_columns(field);
let row = sqlx::query(&format!(
"SELECT {inline_column} AS inline_body, {compressed_column} AS compressed_body FROM \"usage\" WHERE request_id = $1 LIMIT 1"
@@ -2908,8 +2926,10 @@ ORDER BY request_count DESC, "usage".provider_name ASC
usage: &StoredRequestUsageAudit,
field: UsageBodyField,
) -> Result<Option<Value>, DataLayerError> {
let body_ref = usage.body_ref(field);
match body_ref {
let body_ref = usage
.body_ref(field)
.and_then(|body_ref| canonical_usage_body_ref_for(body_ref, &usage.request_id, field));
match body_ref.as_deref() {
Some(body_ref) => self.resolve_body_ref(body_ref).await,
None => Ok(None),
}
@@ -8367,40 +8387,41 @@ ORDER BY "usage".user_id ASC
usage: UpsertUsageRecord,
) -> Result<StoredRequestUsageAudit, DataLayerError> {
usage.validate()?;
let usage = strip_deprecated_usage_display_fields(usage);
let prepared = prepare_usage_upsert_context(&usage)?;
// `usage` is the sanitized accounting projection; prepare the auxiliary capture and
// snapshots from the original event so typed `none` markers can clear prior facts.
let capture_usage = usage.clone();
let usage = sanitize_usage_for_persistence(usage);
self.tx_runner
.run_read_write(|tx| {
let PreparedUsageUpsert {
request_headers_json,
provider_request_headers_json,
response_headers_json,
client_response_headers_json,
request_body_storage,
provider_request_body_storage,
response_body_storage,
client_response_body_storage,
http_audit_refs,
http_audit_states,
http_audit_capture_mode,
routing_snapshot,
settlement_pricing_snapshot,
mut request_metadata_value,
mut request_metadata_json,
replace_client_request_body_facts,
replace_provider_request_body_facts,
clear_request_body,
clear_provider_request_body,
clear_response_body,
clear_client_response_body,
} = prepared;
Box::pin(async move {
lock_usage_request_id_in_tx(tx, &usage.request_id).await?;
if incoming_usage_can_recover_terminal_failure(
usage.status.as_str(),
usage.billing_status.as_str(),
) {
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
if let Some(previous) = previous_usage.as_ref() {
if !usage_lifecycle_update_allowed(
&previous.status,
&previous.billing_status,
previous.updated_at_unix_secs,
previous.finalized_at_unix_secs,
&usage.status,
&usage.billing_status,
usage.updated_at_unix_secs,
usage.finalized_at_unix_secs,
) {
return Ok(previous.clone());
}
}
let recovers_terminal_failure =
previous_usage.as_ref().is_some_and(|previous| {
usage_can_recover_terminal_failure(
&previous.status,
&previous.billing_status,
&usage.status,
&usage.billing_status,
)
});
if recovers_terminal_failure {
sqlx::query(RESET_STALE_VOID_USAGE_SQL)
.bind(&usage.request_id)
.execute(&mut **tx)
@@ -8413,14 +8434,36 @@ ORDER BY "usage".user_id ASC
.map_postgres_err()?;
}
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
let capture_update_allowed = usage_capture_update_allowed(
previous_usage
.as_ref()
.map(|stored| (stored.status.as_str(), stored.billing_status.as_str())),
usage.status.as_str(),
);
let PreparedUsageUpsert {
request_headers_json,
provider_request_headers_json,
response_headers_json,
client_response_headers_json,
request_body_storage,
provider_request_body_storage,
response_body_storage,
client_response_body_storage,
http_audit_refs,
http_audit_states,
http_audit_capture_mode,
routing_snapshot,
settlement_pricing_snapshot,
mut request_metadata_value,
mut request_metadata_json,
replace_client_request_body_facts,
replace_provider_request_body_facts,
clear_request_body,
clear_provider_request_body,
clear_response_body,
clear_client_response_body,
} = prepare_usage_upsert_context(&capture_usage)?;
let capture_update_allowed = recovers_terminal_failure
|| usage_capture_update_allowed(
previous_usage.as_ref().map(|stored| {
(stored.status.as_str(), stored.billing_status.as_str())
}),
usage.status.as_str(),
);
let replace_terminal_snapshots =
matches!(usage.status.as_str(), "completed" | "failed" | "cancelled");
if capture_update_allowed
@@ -8432,7 +8475,10 @@ ORDER BY "usage".user_id ASC
let previous_metadata = previous_usage
.as_ref()
.and_then(|stored| stored.request_metadata.as_ref());
request_metadata_value = Some(if replace_terminal_snapshots {
let preserve_empty_tombstone = !replace_terminal_snapshots
&& (replace_client_request_body_facts
|| replace_provider_request_body_facts);
let previous_metadata = if replace_terminal_snapshots {
retain_previous_request_audit_metadata(
previous_metadata,
!replace_client_request_body_facts,
@@ -8443,7 +8489,14 @@ ORDER BY "usage".user_id ASC
replace_client_request_body_facts,
replace_provider_request_body_facts,
)
});
};
request_metadata_value = Some(
project_usage_request_metadata(
Some(previous_metadata),
preserve_empty_tombstone,
)
.unwrap_or_else(|| Value::Object(Map::new())),
);
request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
}
let _row = sqlx::query(UPSERT_SQL)
@@ -8834,6 +8887,7 @@ ORDER BY "usage".user_id ASC
let mut batch_rows = Vec::<(usize, PreparedPendingUsage)>::new();
let mut fallback_rows = Vec::<(usize, UpsertUsageRecord)>::new();
for (sequence, usage) in usages.into_iter().enumerate() {
let original_usage = usage.clone();
let prepared = PreparedPendingUsage::try_from_usage(usage)?;
if request_id_counts
.get(&prepared.usage.request_id)
@@ -8843,7 +8897,9 @@ ORDER BY "usage".user_id ASC
{
batch_rows.push((sequence, prepared));
} else {
fallback_rows.push((sequence, prepared.usage));
// Preserve capture markers for the canonical fallback; that path performs the
// sanitized bind only after preparing the auxiliary audit/blob state.
fallback_rows.push((sequence, original_usage));
}
}
@@ -9448,7 +9504,7 @@ removed_last_used_at_unix_secs, usage_created_at_unix_secs
));
}
let usage = strip_deprecated_usage_display_fields(usage);
let usage = sanitize_usage_for_persistence(usage);
let request_metadata_json = json_bind_text(usage.request_metadata.as_ref())?;
let response_time_ms = usage.response_time_ms.map(to_i32).transpose()?;
let first_byte_time_ms = usage.first_byte_time_ms.map(to_i32).transpose()?;
@@ -9675,9 +9731,13 @@ DO UPDATE SET
COALESCE(NULLIF(EXCLUDED.updated_at_unix_secs, 0), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
WHERE "usage".billing_status = 'pending'
WHERE "usage".billing_status = 'pending'
AND "usage".status IN ('pending', 'streaming')
AND "usage".finalized_at IS NULL
AND EXCLUDED.updated_at_unix_secs >= COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
RETURNING
request_id,
provider_api_key_id,
@@ -9736,8 +9796,14 @@ RETURNING
apply_provider_monthly_usage_delta_in_tx(tx, provider_id.as_str(), *delta)
.await?;
}
for (node_id, delta) in &aggregates.proxy_nodes {
apply_proxy_node_counter_delta_in_tx(tx, node_id.as_str(), delta).await?;
for ((node_id, tunnel_generation), delta) in &aggregates.proxy_nodes {
apply_proxy_node_counter_delta_in_tx(
tx,
node_id.as_str(),
tunnel_generation.as_str(),
delta,
)
.await?;
}
for (token_id, delta) in &aggregates.management_tokens {
apply_management_token_counter_delta_in_tx(tx, token_id.as_str(), delta)
@@ -9772,17 +9838,50 @@ RETURNING
if delta.is_noop() {
return Ok(false);
}
let Some(expected_tunnel_generation) = delta
.expected_tunnel_generation
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| value.len() <= 64)
.map(ToOwned::to_owned)
else {
// Do not turn a bare node id into an implicit current-incarnation
// binding. Missing fences are rejected so stale plans fail closed.
return Ok(false);
};
let node_id = delta.node_id.trim().to_string();
let request_id = format!("proxy_node:{node_id}:{}", Uuid::new_v4());
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
// Keep the parent lookup lock-free because flush claims
// outbox rows before updating proxy_nodes. The generation
// is persisted in the outbox row and checked again by the
// flush UPDATE, so id reuse can only retire this delta.
let tunnel_generation: Option<String> = sqlx::query_scalar(
"SELECT tunnel_generation FROM proxy_nodes WHERE id = $1 AND tunnel_generation = $2 LIMIT 1",
)
.bind(&node_id)
.bind(&expected_tunnel_generation)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?;
let Some(_tunnel_generation) = tunnel_generation
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
else {
return Ok(false);
};
insert_usage_counter_delta_in_tx(
tx,
UsageCounterDeltaInsert {
request_id: &request_id,
kind: USAGE_COUNTER_KIND_PROXY_NODE,
target_id: &node_id,
target_tunnel_generation: Some(&expected_tunnel_generation),
request_count_delta: 0,
total_requests_delta: delta.total_requests_delta,
success_count_delta: 0,
@@ -9833,6 +9932,7 @@ RETURNING
request_id: &request_id,
kind: USAGE_COUNTER_KIND_MANAGEMENT_TOKEN,
target_id: &token_id,
target_tunnel_generation: None,
request_count_delta: delta.usage_count_delta,
total_requests_delta: 0,
success_count_delta: 0,
@@ -9874,6 +9974,7 @@ RETURNING
request_id: &request_id,
kind: USAGE_COUNTER_KIND_API_KEY_LAST_USED,
target_id: &api_key_id,
target_tunnel_generation: None,
request_count_delta: 0,
total_requests_delta: 0,
success_count_delta: 0,
@@ -9910,7 +10011,7 @@ RETURNING
&self,
cutoff_unix_secs: u64,
now_unix_secs: u64,
timeout_minutes: u64,
_timeout_minutes: u64,
batch_size: usize,
) -> Result<PendingUsageCleanupSummary, DataLayerError> {
if batch_size == 0 {
@@ -9963,7 +10064,6 @@ RETURNING
.map(|row| {
Ok(StalePendingUsageRow {
request_id: row.try_get("request_id").map_postgres_err()?,
status: row.try_get("status").map_postgres_err()?,
billing_status: row.try_get("billing_status").map_postgres_err()?,
})
})
@@ -9996,18 +10096,7 @@ RETURNING
.try_get::<Option<i32>, _>("status_code")
.map_postgres_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_postgres_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
failed_map.insert(
request_id,
FailedCandidateCleanupInfo {
status_code,
error_message,
},
);
failed_map.insert(request_id, FailedCandidateCleanupInfo { status_code });
}
(completed, failed_map)
};
@@ -10030,23 +10119,23 @@ RETURNING
}
let candidate_info = failed_candidate_info.get(&row.request_id);
let (status_code, error_message) =
resolve_stale_pending_failure(candidate_info, &row.status, timeout_minutes);
let status_code = resolve_stale_pending_status_code(candidate_info);
let error_category = usage_error_category_for_status_code(status_code);
let status_code_i32 = i32::from(status_code);
if row.billing_status == "pending" {
sqlx::query(UPDATE_FAILED_VOID_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(now)
.bind(status_code_i32)
.bind(error_category)
.execute(&mut *tx)
.await
.map_postgres_err()?;
} else {
sqlx::query(UPDATE_FAILED_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(status_code_i32)
.bind(error_category)
.execute(&mut *tx)
.await
.map_postgres_err()?;
@@ -10572,33 +10661,17 @@ impl UsageWriteRepository for SqlxUsageReadRepository {
struct StalePendingUsageRow {
request_id: String,
status: String,
billing_status: String,
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
fn resolve_stale_pending_status_code(candidate: Option<&FailedCandidateCleanupInfo>) -> u16 {
candidate
.and_then(|info| info.status_code)
.unwrap_or(if candidate.is_some() { 502 } else { 504 })
}
async fn find_usage_by_request_id_in_tx(
@@ -10768,6 +10841,7 @@ fn prepare_first_byte_provider_contribution_transitions(
request_id: &transition.request_id,
kind: USAGE_COUNTER_KIND_PROVIDER_API_KEY,
target_id: &transition.key_id,
target_tunnel_generation: None,
request_count_delta: transition.delta.request_count,
total_requests_delta: 0,
success_count_delta: transition.delta.success_count,
@@ -10830,6 +10904,7 @@ struct UsageCounterDeltaRow {
id: String,
kind: String,
target_id: String,
target_tunnel_generation: Option<String>,
request_count_delta: i64,
total_requests_delta: i64,
success_count_delta: i64,
@@ -10852,7 +10927,7 @@ struct UsageCounterDeltaAggregates {
provider_api_keys: BTreeMap<String, ProviderApiKeyUsageDelta>,
models: BTreeMap<String, ModelUsageDelta>,
provider_monthly: BTreeMap<String, f64>,
proxy_nodes: BTreeMap<String, ProxyNodeCounterDelta>,
proxy_nodes: BTreeMap<(String, String), ProxyNodeCounterDelta>,
management_tokens: BTreeMap<String, ManagementTokenCounterDelta>,
api_key_last_used: BTreeMap<String, ApiKeyLastUsedDelta>,
}
@@ -10921,16 +10996,28 @@ impl UsageCounterDeltaAggregates {
*entry += row.total_cost_usd_delta;
}
USAGE_COUNTER_KIND_PROXY_NODE => {
let entry = aggregates
.proxy_nodes
.entry(row.target_id.clone())
.or_insert(ProxyNodeCounterDelta {
let Some(tunnel_generation) = row
.target_tunnel_generation
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
else {
// Legacy rows without a generation fence are retired
// without applying them to a reused node id.
continue;
};
let aggregate_key = (row.target_id.clone(), tunnel_generation.clone());
let entry = aggregates.proxy_nodes.entry(aggregate_key).or_insert(
ProxyNodeCounterDelta {
node_id: row.target_id.clone(),
expected_tunnel_generation: Some(tunnel_generation),
total_requests_delta: 0,
failed_requests_delta: 0,
dns_failures_delta: 0,
stream_errors_delta: 0,
});
},
);
entry.total_requests_delta += row.total_requests_delta;
entry.failed_requests_delta += row.error_count_delta;
entry.dns_failures_delta += row.dns_failures_delta;
@@ -11055,6 +11142,7 @@ async fn enqueue_api_key_usage_delta_in_tx(
request_id,
kind: USAGE_COUNTER_KIND_API_KEY,
target_id: api_key_id,
target_tunnel_generation: None,
request_count_delta: 0,
total_requests_delta: delta.total_requests,
success_count_delta: 0,
@@ -11089,6 +11177,7 @@ async fn enqueue_model_usage_delta_in_tx(
request_id,
kind: USAGE_COUNTER_KIND_MODEL,
target_id: model,
target_tunnel_generation: None,
request_count_delta: delta.request_count,
total_requests_delta: 0,
success_count_delta: 0,
@@ -11128,6 +11217,7 @@ async fn enqueue_provider_api_key_usage_delta_in_tx(
request_id,
kind: USAGE_COUNTER_KIND_PROVIDER_API_KEY,
target_id: key_id,
target_tunnel_generation: None,
request_count_delta: delta.request_count,
total_requests_delta: 0,
success_count_delta: delta.success_count,
@@ -11151,6 +11241,7 @@ struct UsageCounterDeltaInsert<'a> {
request_id: &'a str,
kind: &'a str,
target_id: &'a str,
target_tunnel_generation: Option<&'a str>,
request_count_delta: i64,
total_requests_delta: i64,
success_count_delta: i64,
@@ -11173,6 +11264,7 @@ struct PreparedUsageCounterDeltaInsert {
request_id: String,
kind: String,
target_id: String,
target_tunnel_generation: Option<String>,
request_count_delta: i64,
total_requests_delta: i64,
success_count_delta: i64,
@@ -11224,6 +11316,11 @@ fn prepare_usage_counter_delta_insert(
request_id: request_id.to_string(),
kind: input.kind.to_string(),
target_id: target_id.to_string(),
target_tunnel_generation: input
.target_tunnel_generation
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
request_count_delta: input.request_count_delta,
total_requests_delta: input.total_requests_delta,
success_count_delta: input.success_count_delta,
@@ -11258,6 +11355,7 @@ async fn insert_usage_counter_delta_in_tx(
.bind(input.request_id)
.bind(input.kind)
.bind(input.target_id)
.bind(input.target_tunnel_generation)
.bind(input.request_count_delta)
.bind(input.total_requests_delta)
.bind(input.success_count_delta)
@@ -11290,6 +11388,7 @@ async fn insert_usage_counter_deltas_batch_in_tx(
.push_bind(input.request_id.clone())
.push_bind(input.kind.clone())
.push_bind(input.target_id.clone())
.push_bind(input.target_tunnel_generation.clone())
.push_bind(input.request_count_delta)
.push_bind(input.total_requests_delta)
.push_bind(input.success_count_delta)
@@ -11340,6 +11439,9 @@ fn map_usage_counter_delta_row(row: &PgRow) -> Result<UsageCounterDeltaRow, Data
id: row.try_get::<String, _>("id").map_postgres_err()?,
kind: row.try_get::<String, _>("kind").map_postgres_err()?,
target_id: row.try_get::<String, _>("target_id").map_postgres_err()?,
target_tunnel_generation: row
.try_get::<Option<String>, _>("target_tunnel_generation")
.map_postgres_err()?,
request_count_delta: row
.try_get::<i64, _>("request_count_delta")
.map_postgres_err()?,
@@ -11565,9 +11667,10 @@ async fn apply_provider_monthly_usage_delta_in_tx(
async fn apply_proxy_node_counter_delta_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
node_id: &str,
tunnel_generation: &str,
delta: &ProxyNodeCounterDelta,
) -> Result<(), DataLayerError> {
if delta.is_noop() || node_id.trim().is_empty() {
if delta.is_noop() || node_id.trim().is_empty() || tunnel_generation.trim().is_empty() {
return Ok(());
}
@@ -11577,6 +11680,7 @@ async fn apply_proxy_node_counter_delta_in_tx(
.bind(delta.failed_requests_delta)
.bind(delta.dns_failures_delta)
.bind(delta.stream_errors_delta)
.bind(tunnel_generation)
.execute(&mut **tx)
.await
.map_postgres_err()?;
@@ -12243,9 +12347,23 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
.transpose()
}
fn project_usage_request_metadata(
value: Option<Value>,
preserve_empty_tombstone: bool,
) -> Option<Value> {
let projected = sanitize_usage_request_metadata(value);
if preserve_empty_tombstone && projected.is_none() {
Some(Value::Object(Map::new()))
} else {
projected
}
}
fn prepare_usage_upsert_context(
usage: &UpsertUsageRecord,
) -> Result<PreparedUsageUpsert, DataLayerError> {
let usage = sanitize_usage_capture_controls_for_persistence(usage.clone());
let usage = &usage;
let replace_client_request_body_facts = request_body_capture_replaces_derived_facts(
usage.request_body.as_ref(),
usage.request_body_state,
@@ -12382,6 +12500,14 @@ fn prepare_usage_upsert_context(
clear_provider_request_body,
));
}
// The raw event is used above to build the capture and billing snapshots, but only the
// allow-listed metadata may reach the accounting row. Keep an explicit empty object when a
// body `none` marker cleared the last derived fact: PostgreSQL's sparse upsert uses COALESCE
// and would otherwise resurrect the previous candidate's provider metadata from NULL.
let preserve_empty_metadata_tombstone =
(clear_request_body || clear_provider_request_body) && request_metadata_value.is_some();
request_metadata_value =
project_usage_request_metadata(request_metadata_value, preserve_empty_metadata_tombstone);
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
@@ -12460,14 +12586,10 @@ fn resolved_read_usage_body_ref(
http_audit_ref: Option<&str>,
) -> Option<String> {
explicit_ref
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field))
.or_else(|| {
http_audit_ref
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field))
})
.or_else(|| has_compressed_storage.then(|| usage_body_ref(request_id, field)))
.or_else(|| metadata_usage_body_ref_value(metadata, request_id, field))
@@ -12481,15 +12603,11 @@ fn resolved_write_usage_body_ref(
http_audit_ref: Option<&str>,
) -> Option<String> {
explicit_ref
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field))
.or_else(|| has_compressed_storage.then(|| usage_body_ref(request_id, field)))
.or_else(|| {
http_audit_ref
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field))
})
}
@@ -12511,11 +12629,7 @@ fn metadata_usage_body_ref_value(
field: UsageBodyField,
) -> Option<String> {
metadata_ref_value(metadata, field.as_ref_key())
.and_then(|value| parse_usage_body_ref(&value))
.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))
.and_then(|body_ref| canonical_usage_body_ref_for(&body_ref, request_id, field))
}
fn metadata_number_value(
@@ -13076,11 +13190,7 @@ fn usage_json_column(
}
fn inflate_usage_json_value(bytes: &[u8]) -> Result<Value, DataLayerError> {
let mut decoder = GzDecoder::new(bytes);
let mut json_bytes = Vec::new();
decoder.read_to_end(&mut json_bytes).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to decompress usage json: {err}"))
})?;
let json_bytes = read_decompressed_usage_json(GzDecoder::new(bytes))?;
serde_json::from_slice(&json_bytes).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to parse decompressed usage json: {err}"))
})
@@ -13252,42 +13362,19 @@ async fn sync_usage_body_blob_storage<'e, E>(
executor: E,
request_id: &str,
field: UsageBodyField,
value: Option<&Value>,
storage: &UsageBodyStorage,
clear_existing: bool,
_value: Option<&Value>,
_storage: &UsageBodyStorage,
_clear_existing: bool,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
let body_ref = usage_body_ref(request_id, field);
if clear_existing {
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
return Ok(());
}
if let Some(payload_gzip) = storage.detached_blob_bytes.as_ref() {
sqlx::query(UPSERT_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.bind(request_id)
.bind(field.as_storage_field())
.bind(payload_gzip)
.execute(executor)
.await
.map_postgres_err()?;
return Ok(());
}
if value.is_some() {
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
}
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
@@ -13296,46 +13383,43 @@ async fn sync_usage_http_audit_storage<'e, E>(
request_id: &str,
headers: &UsageHttpAuditHeaders<'_>,
refs: &UsageHttpAuditRefs,
states: &UsageHttpAuditStates,
_states: &UsageHttpAuditStates,
body_capture_mode: &str,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if !headers.any_present()
&& !refs.any_present()
&& !states.any_present()
&& body_capture_mode == "none"
{
return Ok(());
if headers.any_present() || refs.any_present() || body_capture_mode != "none" {
return Err(DataLayerError::InvalidInput(
"usage HTTP capture persistence is disabled".to_string(),
));
}
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_SQL)
.bind(request_id)
.bind(headers.request_headers_json)
.bind(headers.provider_request_headers_json)
.bind(headers.response_headers_json)
.bind(headers.client_response_headers_json)
.bind(refs.request_body_ref.as_deref())
.bind(refs.provider_request_body_ref.as_deref())
.bind(refs.response_body_ref.as_deref())
.bind(refs.client_response_body_ref.as_deref())
.bind(usage_body_capture_state_bind_text(
states.request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.provider_request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.response_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.client_response_body_state,
))
.bind(body_capture_mode)
.execute(executor)
.await
.map_postgres_err()?;
sqlx::query(
r#"
WITH deleted_audit AS (
DELETE FROM usage_http_audits WHERE request_id = $1
)
UPDATE usage
SET request_headers = NULL,
request_body = NULL,
provider_request_headers = NULL,
provider_request_body = NULL,
response_headers = NULL,
response_body = NULL,
client_response_headers = NULL,
client_response_body = NULL,
request_body_compressed = NULL,
provider_request_body_compressed = NULL,
response_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE request_id = $1
"#,
)
.bind(request_id)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
crates/aether-data/adapters/postgres/src/usage/queries/find_by_id_sql.sql
+7 -5
View File
@@ -221,11 +221,13 @@ SELECT
usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id,
usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version,
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT) AS created_at_unix_ms,
GREATEST(
COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
GREATEST(
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/adapters/postgres/src/usage/queries/find_by_request_id_sql.sql
+7 -5
View File
@@ -225,11 +225,13 @@ SELECT
usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id,
usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version,
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT) AS created_at_unix_ms,
GREATEST(
COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
GREATEST(
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql
+7 -5
View File
@@ -327,11 +327,13 @@ SELECT
usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id,
usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version,
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT) AS created_at_unix_ms,
GREATEST(
COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
GREATEST(
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql
+7 -5
View File
@@ -327,11 +327,13 @@ SELECT
usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id,
usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version,
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT) AS created_at_unix_ms,
GREATEST(
COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
GREATEST(
COALESCE(CAST(EXTRACT(EPOCH FROM usage_settlement_snapshots.finalized_at) AS BIGINT), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM "usage".finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/adapters/postgres/src/usage/queries/upsert_first_byte_sql.sql
+4
View File
@@ -119,6 +119,10 @@ DO UPDATE SET
WHERE "usage".billing_status = 'pending'
AND "usage".status IN ('pending', 'streaming')
AND "usage".finalized_at IS NULL
AND EXCLUDED.updated_at_unix_secs >= COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
RETURNING
request_id,
provider_api_key_id,
crates/aether-data/adapters/postgres/src/usage/tests.rs
+233 -28
View File
@@ -187,6 +187,156 @@ async fn pending_batch_is_opt_in_and_rejects_non_pending_before_connecting() {
.contains("pending usage batch requires pending status"));
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_stale_terminal_event_is_a_full_transaction_noop() {
let database_url = std::env::var("AETHER_TEST_DATABASE_URL")
.expect("AETHER_TEST_DATABASE_URL must point at the test database");
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url,
min_connections: 1,
max_connections: 2,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let repository =
SqlxUsageReadRepository::new(factory.connect_lazy().expect("lazy pool should build"));
crate::run_migrations(repository.pool())
.await
.expect("test database migrations should succeed");
let suffix = uuid::Uuid::new_v4().simple().to_string();
let request_id = format!("req-stale-terminal-{suffix}");
let provider_name = format!("stale-provider-{suffix}");
let now_unix_secs = Utc::now().timestamp().max(2) as u64;
let mut newer = fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs,
true,
UsageBodyCaptureState::None,
None,
);
newer.candidate_id = Some("candidate-new".to_string());
newer.route_kind = Some("route-new".to_string());
newer.total_cost_usd = Some(0.5);
newer.actual_total_cost_usd = Some(0.4);
repository
.upsert(newer)
.await
.expect("newer terminal usage should upsert");
let counter_rows_before: i64 = sqlx::query_scalar(
"SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("counter rows should count");
let routing_before = sqlx::query(
"SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("routing snapshot should load");
let routing_before = (
routing_before
.try_get::<Option<String>, _>("candidate_id")
.unwrap(),
routing_before
.try_get::<Option<String>, _>("route_kind")
.unwrap(),
);
let settlement_before = sqlx::query(
"SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("settlement snapshot should load");
let settlement_before = (
settlement_before
.try_get::<String, _>("billing_status")
.unwrap(),
settlement_before
.try_get::<Option<f64>, _>("billing_total_cost_usd")
.unwrap(),
);
let mut stale = fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs - 2,
true,
UsageBodyCaptureState::None,
None,
);
stale.status = "failed".to_string();
stale.billing_status = "void".to_string();
stale.status_code = Some(503);
stale.total_cost_usd = Some(99.0);
stale.actual_total_cost_usd = Some(98.0);
stale.candidate_id = Some("candidate-stale".to_string());
stale.route_kind = Some("route-stale".to_string());
let stored = repository
.upsert(stale)
.await
.expect("stale terminal usage should be ignored");
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.status_code, Some(200));
assert_eq!(stored.total_cost_usd, 0.5);
assert_eq!(stored.routing_candidate_id(), Some("candidate-new"));
assert_eq!(stored.routing_route_kind(), Some("route-new"));
let counter_rows_after: i64 = sqlx::query_scalar(
"SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("counter rows should count");
let routing_after = sqlx::query(
"SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("routing snapshot should load");
let routing_after = (
routing_after
.try_get::<Option<String>, _>("candidate_id")
.unwrap(),
routing_after
.try_get::<Option<String>, _>("route_kind")
.unwrap(),
);
let settlement_after = sqlx::query(
"SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("settlement snapshot should load");
let settlement_after = (
settlement_after
.try_get::<String, _>("billing_status")
.unwrap(),
settlement_after
.try_get::<Option<f64>, _>("billing_total_cost_usd")
.unwrap(),
);
assert_eq!(counter_rows_after, counter_rows_before);
assert_eq!(routing_after, routing_before);
assert_eq!(settlement_after, settlement_before);
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_pending_batch_persists_auxiliary_state_and_preserves_terminal_conflicts() {
@@ -385,37 +535,14 @@ async fn live_pending_batch_persists_auxiliary_state_and_preserves_terminal_conf
"late pending must not clear the first-byte observation"
);
let http = sqlx::query(
"SELECT request_headers, provider_request_headers, response_headers, client_response_headers, request_body_ref, provider_request_body_ref, response_body_ref, client_response_body_ref, request_body_state, provider_request_body_state, response_body_state, client_response_body_state FROM usage_http_audits WHERE request_id = $1",
let http_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*)::BIGINT FROM usage_http_audits WHERE request_id = $1",
)
.bind(&rich_request_id)
.fetch_one(repository.pool())
.await
.expect("rich HTTP audit should exist");
assert_eq!(
http.try_get::<serde_json::Value, _>("request_headers")
.unwrap(),
json!({"x-request": "request-value"})
);
for field in [
"request_body_ref",
"provider_request_body_ref",
"response_body_ref",
"client_response_body_ref",
] {
assert!(http.try_get::<Option<String>, _>(field).unwrap().is_some());
}
for field in [
"request_body_state",
"provider_request_body_state",
"response_body_state",
"client_response_body_state",
] {
assert_eq!(
http.try_get::<Option<String>, _>(field).unwrap().as_deref(),
Some("reference")
);
}
.expect("HTTP audit count should be readable");
assert_eq!(http_count, 0);
let blob_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*)::BIGINT FROM usage_body_blobs WHERE request_id = $1",
)
@@ -423,7 +550,7 @@ async fn live_pending_batch_persists_auxiliary_state_and_preserves_terminal_conf
.fetch_one(repository.pool())
.await
.expect("body blob count should be readable");
assert_eq!(blob_count, 4);
assert_eq!(blob_count, 0);
let routing = sqlx::query(
"SELECT candidate_id, candidate_index, selected_provider_api_key_id FROM usage_routing_snapshots WHERE request_id = $1",
@@ -2264,6 +2391,22 @@ fn usage_sql_does_not_require_updated_at_column() {
assert!(!super::UPSERT_SQL.contains("updated_at = CASE"));
}
#[test]
fn usage_sql_preserves_nonzero_lifecycle_updated_revision() {
for sql in [
super::FIND_BY_REQUEST_ID_SQL,
super::FIND_BY_ID_SQL,
super::LIST_USAGE_AUDITS_PREFIX,
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
] {
assert!(sql
.contains("COALESCE(\n NULLIF(\"usage\".updated_at_unix_secs, 0),\n GREATEST("));
assert!(
!sql.contains("GREATEST(\n COALESCE(NULLIF(\"usage\".updated_at_unix_secs, 0), 0),")
);
}
}
#[test]
fn usage_sql_summarizes_tokens_by_api_key_ids_in_database() {
let sql = super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL;
@@ -3586,6 +3729,7 @@ fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() {
fn stale_cleanup_failed_candidate_sql_orders_by_effective_timestamp() {
let sql = super::SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL;
assert!(sql.contains("COALESCE(finished_at, started_at, created_at) DESC"));
assert!(!sql.contains("error_message"));
assert!(!sql.contains("finished_at DESC NULLS LAST"));
assert!(!sql.contains("started_at DESC NULLS LAST"));
}
@@ -3609,6 +3753,10 @@ fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() {
#[test]
fn first_byte_upsert_sql_is_single_row_guarded_and_preserves_existing_metadata() {
let sql = normalize_newlines(super::UPSERT_FIRST_BYTE_SQL);
let revision_guard = r#"AND EXCLUDED.updated_at_unix_secs >= COALESCE(
NULLIF("usage".updated_at_unix_secs, 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)"#;
assert_eq!(sql.matches("INSERT INTO").count(), 1);
assert!(!sql.contains("usage_http_audits"));
assert!(!sql.contains("usage_routing_snapshots"));
@@ -3619,6 +3767,8 @@ fn first_byte_upsert_sql_is_single_row_guarded_and_preserves_existing_metadata()
assert!(sql.contains("WHERE \"usage\".billing_status = 'pending'"));
assert!(sql.contains("\"usage\".status IN ('pending', 'streaming')"));
assert!(sql.contains("\"usage\".finalized_at IS NULL"));
assert!(sql.contains(revision_guard));
assert!(normalize_newlines(include_str!("mod.rs")).contains(revision_guard));
assert!(sql.contains("$22::json->>'upstream_is_stream'"));
assert!(sql.contains("\"usage\".upstream_is_stream"));
@@ -3788,6 +3938,7 @@ fn first_byte_provider_counter_batch_prepares_all_columns_before_query_building(
request_id: " req-counter-prepared ",
kind: "provider_api_key",
target_id: " key-counter-prepared ",
target_tunnel_generation: None,
request_count_delta: 1,
total_requests_delta: 2,
success_count_delta: 3,
@@ -3841,6 +3992,7 @@ fn first_byte_provider_counter_batch_prepares_all_columns_before_query_building(
request_id: "req-counter-out-of-range",
kind: "provider_api_key",
target_id: "key-counter-out-of-range",
target_tunnel_generation: None,
request_count_delta: 0,
total_requests_delta: 0,
success_count_delta: 0,
@@ -4284,6 +4436,39 @@ fn resolved_read_usage_body_ref_prefers_typed_then_http_audit_then_compressed_th
),
Some("usage://request/req-123/client_response_body".to_string())
);
assert_eq!(
resolved_read_usage_body_ref(
Some("usage://request/req-other/request_body"),
None,
"req-123",
UsageBodyField::RequestBody,
false,
Some("usage://request/req-123/request_body"),
),
Some(usage_body_ref("req-123", UsageBodyField::RequestBody))
);
assert_eq!(
resolved_read_usage_body_ref(
None,
None,
"req-123",
UsageBodyField::RequestBody,
false,
Some("usage://request/req-other/request_body"),
),
None
);
assert_eq!(
resolved_read_usage_body_ref(
None,
None,
"req-123",
UsageBodyField::RequestBody,
false,
Some("usage://request/req-123/response_body"),
),
None
);
}
#[test]
@@ -4322,6 +4507,26 @@ fn resolved_write_usage_body_ref_ignores_metadata_compatibility_keys() {
),
Some("usage://request/req-123/client_response_body".to_string())
);
assert_eq!(
resolved_write_usage_body_ref(
Some("usage://request/req-other/request_body"),
"req-123",
UsageBodyField::RequestBody,
false,
Some("usage://request/req-123/request_body"),
),
Some(usage_body_ref("req-123", UsageBodyField::RequestBody))
);
assert_eq!(
resolved_write_usage_body_ref(
Some("usage://request/req-123/response_body"),
"req-123",
UsageBodyField::RequestBody,
false,
Some("usage://request/req-other/request_body"),
),
None
);
}
#[test]