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

Refactor usage body capture and stream terminal reporting

Browse Source
This commit is contained in:
fawney19
2026-04-18 17:48:21 +08:00
parent 569242d72f
commit 3363592751
36 changed files with 2673 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/aether-data/bootstrap/20260413020000_baseline_v2.sql
View File

@@ -4623,6 +4623,10 @@ CREATE TABLE IF NOT EXISTS public.usage_http_audits (
provider_request_body_ref character varying(160),
response_body_ref character varying(160),
client_response_body_ref character varying(160),
request_body_state character varying(32),
provider_request_body_state character varying(32),
response_body_state character varying(32),
client_response_body_state character varying(32),
body_capture_mode character varying(32) DEFAULT 'none' NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,

4
crates/aether-data/migrations/20260413020000_squash_usage_schema_split.sql
View File

@@ -35,6 +35,10 @@ CREATE TABLE IF NOT EXISTS public.usage_http_audits (
provider_request_body_ref character varying(160),
response_body_ref character varying(160),
client_response_body_ref character varying(160),
request_body_state character varying(32),
provider_request_body_state character varying(32),
response_body_state character varying(32),
client_response_body_state character varying(32),
body_capture_mode character varying(32) DEFAULT 'none' NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,

5
crates/aether-data/migrations/20260418000000_add_usage_http_audit_body_states.sql Normal file
View File

@@ -0,0 +1,5 @@
ALTER TABLE public.usage_http_audits
ADD COLUMN IF NOT EXISTS request_body_state character varying(32),
ADD COLUMN IF NOT EXISTS provider_request_body_state character varying(32),
ADD COLUMN IF NOT EXISTS response_body_state character varying(32),
ADD COLUMN IF NOT EXISTS client_response_body_state character varying(32);

10
crates/aether-data/src/migrate.rs
View File

@@ -8,7 +8,7 @@ use tracing::{error, info, warn};
static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
static BASELINE_V2_SQL: &str = include_str!("../bootstrap/20260413020000_baseline_v2.sql");
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260415000000;
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260418000000;
const MIGRATIONS_TABLE_EXISTS_SQL: &str =
"SELECT to_regclass('public._sqlx_migrations') IS NOT NULL";
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
@@ -623,6 +623,7 @@ mod tests {
20260413020000,
20260413030000,
20260415000000,
20260418000000,
]
);
}
@@ -632,6 +633,10 @@ mod tests {
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_body_blobs"));
assert!(BASELINE_V2_SQL.contains("ix_usage_body_blobs_request_id"));
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_http_audits"));
assert!(BASELINE_V2_SQL.contains("request_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("provider_request_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("response_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("client_response_body_state character varying(32)"));
assert!(
BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_routing_snapshots")
);
@@ -717,7 +722,8 @@ mod tests {
20260410000000,
20260413020000,
20260413030000,
20260415000000
20260415000000,
20260418000000,
]
);
}

68
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -2194,6 +2194,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.request_body
.or_else(|| existing.and_then(|existing| existing.request_body.clone())),
request_body_ref,
request_body_state: usage
.request_body_state
.or_else(|| existing.and_then(|existing| existing.request_body_state)),
provider_request_headers: usage.provider_request_headers.or_else(|| {
existing.and_then(|existing| existing.provider_request_headers.clone())
}),
@@ -2201,6 +2204,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.provider_request_body
.or_else(|| existing.and_then(|existing| existing.provider_request_body.clone())),
provider_request_body_ref,
provider_request_body_state: usage
.provider_request_body_state
.or_else(|| existing.and_then(|existing| existing.provider_request_body_state)),
response_headers: usage
.response_headers
.or_else(|| existing.and_then(|existing| existing.response_headers.clone())),
@@ -2208,6 +2214,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.response_body
.or_else(|| existing.and_then(|existing| existing.response_body.clone())),
response_body_ref,
response_body_state: usage
.response_body_state
.or_else(|| existing.and_then(|existing| existing.response_body_state)),
client_response_headers: usage
.client_response_headers
.or_else(|| existing.and_then(|existing| existing.client_response_headers.clone())),
@@ -2215,6 +2224,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.client_response_body
.or_else(|| existing.and_then(|existing| existing.client_response_body.clone())),
client_response_body_ref,
client_response_body_state: usage
.client_response_body_state
.or_else(|| existing.and_then(|existing| existing.client_response_body_state)),
candidate_id: usage.candidate_id.or_else(|| {
existing.and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned))
}),
@@ -2382,6 +2394,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2451,6 +2467,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2534,6 +2554,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2603,6 +2627,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(1),
key_name: Some("primary".to_string()),
@@ -2694,6 +2722,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(1),
key_name: Some("primary".to_string()),
@@ -2765,6 +2797,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2965,6 +3001,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3052,6 +3092,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3127,6 +3171,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3205,6 +3253,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3288,6 +3340,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-typed".to_string()),
candidate_index: Some(2),
key_name: Some("primary".to_string()),
@@ -3382,6 +3438,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3457,6 +3517,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3523,6 +3587,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,

8
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -45,6 +45,10 @@ pub(crate) fn usage_can_recover_terminal_failure(
&& incoming_usage_can_recover_terminal_failure(incoming_status, incoming_billing_status)
}
/// Clear legacy display-cache fields that still exist on `public.usage` for compatibility.
///
/// These values are no longer treated as authoritative read-model inputs; new writes should not
/// repopulate the deprecated mirror columns.
pub(crate) fn strip_deprecated_usage_display_fields(
mut usage: UpsertUsageRecord,
) -> UpsertUsageRecord {
@@ -114,6 +118,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,

171
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -6,8 +6,8 @@ use aether_data_contracts::repository::usage::{
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery,
UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, UsageBodyField, UsageBreakdownGroupBy,
UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery,
UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField,
UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery,
UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery,
UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery,
@@ -89,6 +89,10 @@ INSERT INTO usage_http_audits (
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,
body_capture_mode
) VALUES (
$1,
@@ -100,7 +104,11 @@ INSERT INTO usage_http_audits (
$7,
$8,
$9,
$10
$10,
$11,
$12,
$13,
$14
)
ON CONFLICT (request_id)
DO UPDATE SET
@@ -124,6 +132,22 @@ DO UPDATE SET
EXCLUDED.client_response_body_ref,
usage_http_audits.client_response_body_ref
),
request_body_state = COALESCE(
EXCLUDED.request_body_state,
usage_http_audits.request_body_state
),
provider_request_body_state = COALESCE(
EXCLUDED.provider_request_body_state,
usage_http_audits.provider_request_body_state
),
response_body_state = COALESCE(
EXCLUDED.response_body_state,
usage_http_audits.response_body_state
),
client_response_body_state = COALESCE(
EXCLUDED.client_response_body_state,
usage_http_audits.client_response_body_state
),
body_capture_mode = COALESCE(
NULLIF(EXCLUDED.body_capture_mode, 'none'),
usage_http_audits.body_capture_mode,
@@ -330,6 +354,10 @@ SELECT
usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref,
usage_http_audits.response_body_ref AS http_response_body_ref,
usage_http_audits.client_response_body_ref AS http_client_response_body_ref,
usage_http_audits.request_body_state AS http_request_body_state,
usage_http_audits.provider_request_body_state AS http_provider_request_body_state,
usage_http_audits.response_body_state AS http_response_body_state,
usage_http_audits.client_response_body_state AS http_client_response_body_state,
usage_routing_snapshots.candidate_id AS routing_candidate_id,
usage_routing_snapshots.candidate_index AS routing_candidate_index,
usage_routing_snapshots.key_name AS routing_key_name,
@@ -597,6 +625,10 @@ SELECT
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
COALESCE(
usage_routing_snapshots.candidate_id,
NULLIF(BTRIM("usage".request_metadata->>'candidate_id'), '')
@@ -813,6 +845,10 @@ SELECT
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
COALESCE(
usage_routing_snapshots.candidate_id,
NULLIF(BTRIM("usage".request_metadata->>'candidate_id'), '')
@@ -1148,6 +1184,10 @@ RETURNING
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
NULL::varchar AS routing_candidate_id,
NULL::integer AS routing_candidate_index,
NULL::varchar AS routing_key_name,
@@ -3991,6 +4031,12 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
None,
),
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: usage.request_body_state,
provider_request_body_state: usage.provider_request_body_state,
response_body_state: usage.response_body_state,
client_response_body_state: usage.client_response_body_state,
};
let request_metadata_value = prepare_request_metadata_for_body_storage(
usage.request_metadata.clone(),
[
@@ -4163,6 +4209,7 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
&usage.request_id,
&http_audit_headers,
&http_audit_refs,
&http_audit_states,
http_audit_capture_mode,
)
.await?;
@@ -4224,6 +4271,16 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
client_response_body_storage.has_detached_blob(),
http_audit_refs.client_response_body_ref.as_deref(),
);
stored.request_body_state =
usage.request_body_state.or(stored.request_body_state);
stored.provider_request_body_state = usage
.provider_request_body_state
.or(stored.provider_request_body_state);
stored.response_body_state =
usage.response_body_state.or(stored.response_body_state);
stored.client_response_body_state = usage
.client_response_body_state
.or(stored.client_response_body_state);
stored.candidate_id = routing_snapshot.candidate_id.clone();
stored.candidate_index = routing_snapshot.candidate_index;
stored.key_name = routing_snapshot.key_name.clone();
@@ -4462,6 +4519,10 @@ impl UsageWriteRepository for SqlxUsageReadRepository {
}
}
// Build the usage read model from the split storage layout.
//
// Query projections already prefer the newer audit/snapshot owners and only fall back to
// deprecated `public.usage` mirror columns for historical rows that predate the split schema.
fn map_usage_row(
row: &sqlx::postgres::PgRow,
resolve_compressed_bodies: bool,
@@ -4574,6 +4635,28 @@ fn map_usage_row(
.try_get("http_client_response_body_ref")
.map_postgres_err()?,
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: row
.try_get::<Option<String>, _>("http_request_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
provider_request_body_state: row
.try_get::<Option<String>, _>("http_provider_request_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
response_body_state: row
.try_get::<Option<String>, _>("http_response_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
client_response_body_state: row
.try_get::<Option<String>, _>("http_client_response_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
};
let routing_snapshot = usage_routing_snapshot_from_row(row)?;
let settlement_pricing_snapshot = usage_settlement_pricing_snapshot_from_row(row)?;
usage.request_body = request_body.value;
@@ -4613,6 +4696,10 @@ fn map_usage_row(
client_response_body.has_compressed_storage,
http_audit_refs.client_response_body_ref.as_deref(),
);
usage.request_body_state = http_audit_states.request_body_state;
usage.provider_request_body_state = http_audit_states.provider_request_body_state;
usage.response_body_state = http_audit_states.response_body_state;
usage.client_response_body_state = http_audit_states.client_response_body_state;
usage.candidate_id = routing_snapshot.candidate_id.clone();
usage.candidate_index = routing_snapshot.candidate_index;
usage.key_name = routing_snapshot.key_name.clone();
@@ -4675,6 +4762,23 @@ impl UsageHttpAuditRefs {
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageHttpAuditStates {
request_body_state: Option<UsageBodyCaptureState>,
provider_request_body_state: Option<UsageBodyCaptureState>,
response_body_state: Option<UsageBodyCaptureState>,
client_response_body_state: Option<UsageBodyCaptureState>,
}
impl UsageHttpAuditStates {
fn any_present(&self) -> bool {
self.request_body_state.is_some()
|| self.provider_request_body_state.is_some()
|| self.response_body_state.is_some()
|| self.client_response_body_state.is_some()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageHttpAuditHeaders<'a> {
request_headers_json: Option<&'a str>,
@@ -4802,6 +4906,24 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
.transpose()
}
fn usage_body_capture_state_bind_text(
value: Option<UsageBodyCaptureState>,
) -> Option<&'static str> {
value.map(UsageBodyCaptureState::as_str)
}
fn parse_usage_body_capture_state(value: &str) -> Option<UsageBodyCaptureState> {
match value.trim() {
"none" => Some(UsageBodyCaptureState::None),
"inline" => Some(UsageBodyCaptureState::Inline),
"reference" => Some(UsageBodyCaptureState::Reference),
"truncated" => Some(UsageBodyCaptureState::Truncated),
"disabled" => Some(UsageBodyCaptureState::Disabled),
"unavailable" => Some(UsageBodyCaptureState::Unavailable),
_ => None,
}
}
#[cfg(test)]
fn usage_http_audit_body_refs(metadata: Option<&Value>) -> UsageHttpAuditRefs {
let object = metadata.and_then(Value::as_object);
@@ -5042,6 +5164,10 @@ fn usage_settlement_pricing_snapshot_from_usage(
}
}
// Decode deprecated inline/compressed body columns from `public.usage`.
//
// New writes keep these columns empty by forcing body storage through `usage_body_blobs` and
// `usage_http_audits`; this helper exists only so older rows remain readable without backfill.
fn usage_json_column(
row: &sqlx::postgres::PgRow,
inline_column: &str,
@@ -5286,12 +5412,17 @@ async fn sync_usage_http_audit_storage<'e, E>(
request_id: &str,
headers: &UsageHttpAuditHeaders<'_>,
refs: &UsageHttpAuditRefs,
states: &UsageHttpAuditStates,
body_capture_mode: &str,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if !headers.any_present() && !refs.any_present() && body_capture_mode == "none" {
if !headers.any_present()
&& !refs.any_present()
&& !states.any_present()
&& body_capture_mode == "none"
{
return Ok(());
}
@@ -5305,6 +5436,18 @@ where
.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
@@ -5661,6 +5804,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6276,6 +6423,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6371,6 +6522,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6445,6 +6600,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-typed".to_string()),
candidate_index: Some(2),
key_name: Some("primary".to_string()),
@@ -6576,6 +6735,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,