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

feat(usage): expose request timing details

Browse Source
This commit is contained in:
elky
2026-06-10 09:16:15 +08:00
parent 84f41dae77
commit 8edcbdcb29
41 changed files with 1884 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/repository/candidates/mysql.rs
+17
View File
@@ -86,6 +86,23 @@ impl RequestCandidateReadRepository for MysqlRequestCandidateRepository {
rows.iter().map(map_candidate_row).collect()
}
async fn list_attempted_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let rows = sqlx::query(&format!(
"{CANDIDATE_COLUMNS} WHERE request_id = ? \
AND (status IN ('streaming', 'success', 'failed', 'cancelled') \
OR (status = 'pending' AND started_at IS NOT NULL)) \
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC"
))
.bind(request_id)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_candidate_row).collect()
}
async fn list_recent(
&self,
limit: usize,
crates/aether-data/src/repository/candidates/postgres.rs
+27
View File
@@ -229,6 +229,26 @@ impl SqlxRequestCandidateReadRepository {
collect_query_rows(builder.build().fetch(&self.pool), map_request_candidate_row).await
}
pub async fn list_attempted_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(candidate_columns());
let mut where_clause = WhereClause::new();
push_eq(
&mut builder,
&mut where_clause,
"request_id",
request_id.to_string(),
);
builder.push(
" AND (status IN ('streaming', 'success', 'failed', 'cancelled') \
OR (status = 'pending' AND started_at IS NOT NULL)) \
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC",
);
collect_query_rows(builder.build().fetch(&self.pool), map_request_candidate_row).await
}
pub async fn list_recent(
&self,
limit: usize,
@@ -513,6 +533,13 @@ impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
Self::list_by_request_id(self, request_id).await
}
async fn list_attempted_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_attempted_by_request_id(self, request_id).await
}
async fn list_recent(
&self,
limit: usize,
crates/aether-data/src/repository/candidates/sqlite.rs
+17
View File
@@ -87,6 +87,23 @@ impl RequestCandidateReadRepository for SqliteRequestCandidateRepository {
rows.iter().map(map_candidate_row).collect()
}
async fn list_attempted_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let rows = sqlx::query(&format!(
"{CANDIDATE_COLUMNS} WHERE request_id = ? \
AND (status IN ('streaming', 'success', 'failed', 'cancelled') \
OR (status = 'pending' AND started_at IS NOT NULL)) \
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC"
))
.bind(request_id)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_candidate_row).collect()
}
async fn list_recent(
&self,
limit: usize,
crates/aether-data/src/repository/usage/memory.rs
+49
View File
@@ -225,6 +225,25 @@ fn accumulate_api_key_usage_contribution(
};
}
fn usage_audit_client_family(item: &StoredRequestUsageAudit) -> Option<&str> {
item.client_family
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| usage_request_metadata_client_family(item.request_metadata.as_ref()))
}
fn usage_admin_unknown_label(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"unknown" | "unknow"
)
}
fn usage_has_admin_unknown_model_or_provider(item: &StoredRequestUsageAudit) -> bool {
usage_admin_unknown_label(&item.model) || usage_admin_unknown_label(&item.provider_name)
}
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
// The field is historically named `created_at_unix_ms`, but usage audit rows
// across gateway handlers, SQL repositories and tests are stored as epoch seconds.
@@ -258,6 +277,21 @@ fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditLi
return false;
}
}
if let Some(client_family) = query
.client_family
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if !usage_audit_client_family(item)
.is_some_and(|value| value.eq_ignore_ascii_case(client_family))
{
return false;
}
}
if query.exclude_unknown_model_or_provider && usage_has_admin_unknown_model_or_provider(item) {
return false;
}
if let Some(statuses) = query.statuses.as_ref() {
if !statuses.iter().any(|status| status == &item.status) {
return false;
@@ -318,6 +352,21 @@ fn usage_matches_keyword_search_query(
return false;
}
}
if let Some(client_family) = query
.client_family
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if !usage_audit_client_family(item)
.is_some_and(|value| value.eq_ignore_ascii_case(client_family))
{
return false;
}
}
if query.exclude_unknown_model_or_provider && usage_has_admin_unknown_model_or_provider(item) {
return false;
}
if let Some(statuses) = query.statuses.as_ref() {
if !statuses.iter().any(|status| status == &item.status) {
return false;
crates/aether-data/src/repository/usage/postgres/mod.rs
+146
View File
@@ -1520,6 +1520,45 @@ const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
fn push_postgres_usage_where(builder: &mut QueryBuilder<'_, Postgres>, has_where: &mut bool) {
builder.push(if *has_where { " AND " } else { " WHERE " });
*has_where = true;
}
fn push_postgres_usage_client_family_filter(
builder: &mut QueryBuilder<'_, Postgres>,
has_where: &mut bool,
client_family: Option<&str>,
) {
let Some(client_family) = client_family
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return;
};
push_postgres_usage_where(builder, has_where);
builder
.push("LOWER(COALESCE(NULLIF(BTRIM(\"usage\".request_metadata->'client_session_affinity'->>'client_family'), ''), NULLIF(BTRIM(\"usage\".request_metadata->>'client_family'), ''))) = ")
.push_bind(client_family.to_ascii_lowercase());
}
fn push_postgres_usage_exclude_unknown_filter(
builder: &mut QueryBuilder<'_, Postgres>,
has_where: &mut bool,
exclude_unknown_model_or_provider: bool,
) {
if !exclude_unknown_model_or_provider {
return;
}
push_postgres_usage_where(builder, has_where);
builder.push(
"(lower(BTRIM(COALESCE(\"usage\".model, ''))) NOT IN ('unknown', 'unknow') \
AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow'))",
);
}
const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = r#" AND (
(
BTRIM(COALESCE("usage".provider_id, '')) <> ''
@@ -2370,6 +2409,21 @@ ORDER BY request_count DESC, "usage".provider_name ASC
}
}
pub async fn find_by_request_id_shallow(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
let sql = shallow_usage_body_projection_sql(FIND_BY_REQUEST_ID_SQL);
let row = sqlx::query(&sql)
.bind(request_id)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref()
.map(|row| map_usage_row(row, false))
.transpose()
}
pub async fn find_by_id(
&self,
id: &str,
@@ -2529,6 +2583,7 @@ ORDER BY request_count DESC, "usage".provider_name ASC
}
if let Some(user_id) = query.user_id.as_deref() {
builder.push(if has_where { " AND " } else { " WHERE " });
has_where = true;
builder
.push("\"usage\".user_id = ")
.push_bind(user_id.to_string());
@@ -2554,6 +2609,16 @@ ORDER BY request_count DESC, "usage".provider_name ASC
.push("\"usage\".api_format = ")
.push_bind(api_format.to_string());
}
push_postgres_usage_client_family_filter(
&mut builder,
&mut has_where,
query.client_family.as_deref(),
);
push_postgres_usage_exclude_unknown_filter(
&mut builder,
&mut has_where,
query.exclude_unknown_model_or_provider,
);
if let Some(statuses) = query.statuses.as_deref() {
if !statuses.is_empty() {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -2625,6 +2690,7 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> ''
}
if let Some(user_id) = query.user_id.as_deref() {
builder.push(if has_where { " AND " } else { " WHERE " });
has_where = true;
builder
.push("\"usage\".user_id = ")
.push_bind(user_id.to_string());
@@ -2650,6 +2716,16 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> ''
.push("\"usage\".api_format = ")
.push_bind(api_format.to_string());
}
push_postgres_usage_client_family_filter(
&mut builder,
&mut has_where,
query.client_family.as_deref(),
);
push_postgres_usage_exclude_unknown_filter(
&mut builder,
&mut has_where,
query.exclude_unknown_model_or_provider,
);
if let Some(statuses) = query.statuses.as_deref() {
if !statuses.is_empty() {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -2828,6 +2904,16 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> ''
.push("\"usage\".api_format = ")
.push_bind(api_format.to_string());
}
push_postgres_usage_client_family_filter(
&mut builder,
&mut has_where,
query.client_family.as_deref(),
);
push_postgres_usage_exclude_unknown_filter(
&mut builder,
&mut has_where,
query.exclude_unknown_model_or_provider,
);
if let Some(statuses) = query.statuses.as_deref() {
if !statuses.is_empty() {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -2914,6 +3000,16 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> ''
.push("\"usage\".api_format = ")
.push_bind(api_format.to_string());
}
push_postgres_usage_client_family_filter(
&mut builder,
&mut has_where,
query.client_family.as_deref(),
);
push_postgres_usage_exclude_unknown_filter(
&mut builder,
&mut has_where,
query.exclude_unknown_model_or_provider,
);
if let Some(statuses) = query.statuses.as_deref() {
if !statuses.is_empty() {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -8035,6 +8131,12 @@ ORDER BY "usage".user_id ASC
.bind(&request_metadata_json)
.bind(usage.finalized_at_unix_secs.map(|value| value as f64))
.bind(usage.created_at_unix_ms.map(|value| value as f64))
.bind(i64::try_from(usage.updated_at_unix_secs).map_err(|_| {
DataLayerError::InvalidInput(format!(
"usage.updated_at_unix_secs out of range: {}",
usage.updated_at_unix_secs
))
})?)
.bind(request_body_storage.has_detached_blob())
.bind(provider_request_body_storage.has_detached_blob())
.bind(response_body_storage.has_detached_blob())
@@ -8785,6 +8887,13 @@ impl UsageReadRepository for SqlxUsageReadRepository {
Self::find_by_request_id(self, request_id).await
}
async fn find_by_request_id_shallow(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
Self::find_by_request_id_shallow(self, request_id).await
}
async fn resolve_body_ref(&self, body_ref: &str) -> Result<Option<Value>, DataLayerError> {
Self::resolve_body_ref(self, body_ref).await
}
@@ -10075,6 +10184,43 @@ fn map_usage_row(
Ok(usage)
}
fn shallow_usage_body_projection_sql(sql: &str) -> String {
let replacements = [
("\"usage\".request_body,", "NULL::json AS request_body,"),
(
"\"usage\".request_body_compressed,",
"CASE WHEN \"usage\".request_body_compressed IS NULL THEN NULL ELSE ''::bytea END AS request_body_compressed,",
),
(
"\"usage\".provider_request_body,",
"NULL::json AS provider_request_body,",
),
(
"\"usage\".provider_request_body_compressed,",
"CASE WHEN \"usage\".provider_request_body_compressed IS NULL THEN NULL ELSE ''::bytea END AS provider_request_body_compressed,",
),
("\"usage\".response_body,", "NULL::json AS response_body,"),
(
"\"usage\".response_body_compressed,",
"CASE WHEN \"usage\".response_body_compressed IS NULL THEN NULL ELSE ''::bytea END AS response_body_compressed,",
),
(
"\"usage\".client_response_body,",
"NULL::json AS client_response_body,",
),
(
"\"usage\".client_response_body_compressed,",
"CASE WHEN \"usage\".client_response_body_compressed IS NULL THEN NULL ELSE ''::bytea END AS client_response_body_compressed,",
),
];
replacements
.into_iter()
.fold(sql.to_string(), |current, (from, to)| {
current.replace(from, to)
})
}
fn to_i32(value: u64) -> Result<i32, DataLayerError> {
i32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!("invalid usage integer value: {value}"))
crates/aether-data/src/repository/usage/postgres/queries/find_by_id_sql.sql
+5 -8
View File
@@ -156,14 +156,11 @@ 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,
CAST(
EXTRACT(
EPOCH FROM COALESCE(
usage_settlement_snapshots.finalized_at,
"usage".finalized_at,
"usage".created_at
)
) AS BIGINT
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)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/src/repository/usage/postgres/queries/find_by_request_id_sql.sql
+5 -8
View File
@@ -160,14 +160,11 @@ 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,
CAST(
EXTRACT(
EPOCH FROM COALESCE(
usage_settlement_snapshots.finalized_at,
"usage".finalized_at,
"usage".created_at
)
) AS BIGINT
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)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql
+9 -32
View File
@@ -114,18 +114,8 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'output_config'->>'effort'), '')
)
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
END IS NOT NULL
OR CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
END IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
THEN jsonb_strip_nulls(jsonb_build_object(
@@ -138,19 +128,9 @@ SELECT
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'provider_reasoning_effort',
CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'output_config'->>'effort'), '')
)
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
END,
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
'provider_service_tier',
CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
END,
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
@@ -238,14 +218,11 @@ 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,
CAST(
EXTRACT(
EPOCH FROM COALESCE(
usage_settlement_snapshots.finalized_at,
"usage".finalized_at,
"usage".created_at
)
) AS BIGINT
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)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql
+9 -32
View File
@@ -114,18 +114,8 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'output_config'->>'effort'), '')
)
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
END IS NOT NULL
OR CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
END IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
THEN jsonb_strip_nulls(jsonb_build_object(
@@ -138,19 +128,9 @@ SELECT
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'provider_reasoning_effort',
CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), ''),
NULLIF(BTRIM("usage".provider_request_body->'output_config'->>'effort'), '')
)
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
END,
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
'provider_service_tier',
CASE
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
END,
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
@@ -238,14 +218,11 @@ 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,
CAST(
EXTRACT(
EPOCH FROM COALESCE(
usage_settlement_snapshots.finalized_at,
"usage".finalized_at,
"usage".created_at
)
) AS BIGINT
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)
) AS updated_at_unix_secs,
CAST(
EXTRACT(
crates/aether-data/src/repository/usage/postgres/queries/upsert_sql.sql
+28 -12
View File
@@ -56,7 +56,8 @@ INSERT INTO "usage" (
client_response_body_compressed,
request_metadata,
finalized_at,
created_at
created_at,
updated_at_unix_secs
) VALUES (
$1,
$2,
@@ -144,7 +145,11 @@ INSERT INTO "usage" (
WHEN $54 IS NULL THEN NULL
ELSE TO_TIMESTAMP($54::double precision)
END,
COALESCE(TO_TIMESTAMP($55::double precision), NOW())
COALESCE(TO_TIMESTAMP($55::double precision), NOW()),
COALESCE(
NULLIF($56::bigint, 0),
CAST(EXTRACT(EPOCH FROM COALESCE(TO_TIMESTAMP($55::double precision), NOW())) AS BIGINT)
)
)
ON CONFLICT (request_id)
DO UPDATE SET
@@ -218,42 +223,49 @@ DO UPDATE SET
billing_status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.billing_status ELSE "usage".billing_status END,
request_headers = NULL,
request_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.request_body_compressed IS NOT NULL OR $56 THEN NULL
WHEN EXCLUDED.request_body_compressed IS NOT NULL OR $57 THEN NULL
ELSE COALESCE(EXCLUDED.request_body, "usage".request_body)
END ELSE "usage".request_body END,
request_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.request_body IS NOT NULL OR $56 THEN NULL
WHEN EXCLUDED.request_body IS NOT NULL OR $57 THEN NULL
ELSE COALESCE(EXCLUDED.request_body_compressed, "usage".request_body_compressed)
END ELSE "usage".request_body_compressed END,
provider_request_headers = NULL,
provider_request_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.provider_request_body_compressed IS NOT NULL OR $57 THEN NULL
WHEN EXCLUDED.provider_request_body_compressed IS NOT NULL OR $58 THEN NULL
ELSE COALESCE(EXCLUDED.provider_request_body, "usage".provider_request_body)
END ELSE "usage".provider_request_body END,
provider_request_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.provider_request_body IS NOT NULL OR $57 THEN NULL
WHEN EXCLUDED.provider_request_body IS NOT NULL OR $58 THEN NULL
ELSE COALESCE(EXCLUDED.provider_request_body_compressed, "usage".provider_request_body_compressed)
END ELSE "usage".provider_request_body_compressed END,
response_headers = NULL,
response_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.response_body_compressed IS NOT NULL OR $58 THEN NULL
WHEN EXCLUDED.response_body_compressed IS NOT NULL OR $59 THEN NULL
ELSE COALESCE(EXCLUDED.response_body, "usage".response_body)
END ELSE "usage".response_body END,
response_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.response_body IS NOT NULL OR $58 THEN NULL
WHEN EXCLUDED.response_body IS NOT NULL OR $59 THEN NULL
ELSE COALESCE(EXCLUDED.response_body_compressed, "usage".response_body_compressed)
END ELSE "usage".response_body_compressed END,
client_response_headers = NULL,
client_response_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $59 THEN NULL
WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $60 THEN NULL
ELSE COALESCE(EXCLUDED.client_response_body, "usage".client_response_body)
END ELSE "usage".client_response_body END,
client_response_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN EXCLUDED.client_response_body IS NOT NULL OR $59 THEN NULL
WHEN EXCLUDED.client_response_body IS NOT NULL OR $60 THEN NULL
ELSE COALESCE(EXCLUDED.client_response_body_compressed, "usage".client_response_body_compressed)
END ELSE "usage".client_response_body_compressed END,
request_metadata = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_metadata, "usage".request_metadata) ELSE "usage".request_metadata END,
finalized_at = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.finalized_at, "usage".finalized_at) ELSE "usage".finalized_at END
finalized_at = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.finalized_at, "usage".finalized_at) ELSE "usage".finalized_at END,
updated_at_unix_secs = CASE WHEN "usage".billing_status = 'pending' THEN
GREATEST(
COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0),
COALESCE(NULLIF(EXCLUDED.updated_at_unix_secs, 0), 0),
CAST(EXTRACT(EPOCH FROM "usage".created_at) AS BIGINT)
)
ELSE "usage".updated_at_unix_secs END
RETURNING
id,
request_id,
@@ -352,5 +364,9 @@ RETURNING
NULL::varchar AS settlement_billing_rule_id,
NULL::varchar AS settlement_billing_rule_version,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms,
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
GREATEST(
COALESCE(NULLIF(updated_at_unix_secs, 0), 0),
COALESCE(CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT), 0),
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT)
) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
crates/aether-data/src/repository/usage/postgres/tests.rs
+33 -5
View File
@@ -725,7 +725,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
] {
assert!(sql.contains("jsonb_strip_nulls(jsonb_build_object("));
assert!(sql.contains("jsonb_typeof(\"usage\".provider_request_body::jsonb)"));
assert!(!sql.contains("jsonb_typeof(\"usage\".provider_request_body::jsonb)"));
assert!(!sql.contains("\"usage\".provider_request_body->"));
assert!(!sql.contains("jsonb_typeof(\"usage\".provider_request_body)"));
assert!(sql.contains("'client_ip'"));
assert!(sql.contains("request_metadata->>'client_ip'"));
@@ -747,6 +748,33 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
assert!(!super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("NULL::jsonb"));
}
#[test]
fn usage_sql_admin_record_filters_are_pushed_into_postgres_queries() {
let source = include_str!("mod.rs");
assert!(source.contains("push_postgres_usage_client_family_filter"));
assert!(source.contains("request_metadata->'client_session_affinity'->>'client_family'"));
assert!(source.contains("request_metadata->>'client_family'"));
assert!(source.contains("exclude_unknown_model_or_provider"));
assert!(source.contains("NOT IN ('unknown', 'unknow')"));
}
#[test]
fn usage_sql_keyword_search_error_filter_keeps_where_state_before_keywords() {
let source = include_str!("mod.rs");
let function = source
.split("pub async fn list_usage_audits_by_keyword_search")
.nth(1)
.and_then(|tail| tail.split("pub async fn count_usage_audits").next())
.expect("keyword search function should be present");
let error_filter = function
.split("if query.error_only")
.nth(1)
.and_then(|tail| tail.split("for (index, keyword)").next())
.expect("error filter should precede keyword loop");
assert!(error_filter.contains("has_where = true;"));
}
#[test]
fn usage_sql_reads_list_output_price_from_settlement_snapshots_before_legacy_usage_column() {
assert!(super::LIST_USAGE_AUDITS_PREFIX
@@ -880,13 +908,13 @@ fn usage_sql_clears_legacy_header_columns_on_upsert() {
#[test]
fn usage_sql_detached_body_flags_clear_inline_and_compressed_columns() {
assert!(super::UPSERT_SQL
.contains("WHEN EXCLUDED.request_body_compressed IS NOT NULL OR $56 THEN NULL"));
.contains("WHEN EXCLUDED.request_body_compressed IS NOT NULL OR $57 THEN NULL"));
assert!(super::UPSERT_SQL
.contains("WHEN EXCLUDED.provider_request_body_compressed IS NOT NULL OR $57 THEN NULL"));
.contains("WHEN EXCLUDED.provider_request_body_compressed IS NOT NULL OR $58 THEN NULL"));
assert!(super::UPSERT_SQL
.contains("WHEN EXCLUDED.response_body_compressed IS NOT NULL OR $58 THEN NULL"));
.contains("WHEN EXCLUDED.response_body_compressed IS NOT NULL OR $59 THEN NULL"));
assert!(super::UPSERT_SQL
.contains("WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $59 THEN NULL"));
.contains("WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $60 THEN NULL"));
}
#[test]
crates/aether-data/src/repository/usage/sqlite.rs
+17
View File
@@ -455,6 +455,21 @@ fn push_sqlite_usage_list_filters(
.push("api_format = ")
.push_bind(api_format.to_string());
}
if let Some(client_family) = query.client_family.as_deref().map(str::trim) {
if !client_family.is_empty() {
push_sqlite_usage_where(builder, has_where);
builder
.push("LOWER(COALESCE(NULLIF(TRIM(CAST(json_extract(request_metadata, '$.client_session_affinity.client_family') AS TEXT)), ''), NULLIF(TRIM(CAST(json_extract(request_metadata, '$.client_family') AS TEXT)), ''))) = ")
.push_bind(client_family.to_ascii_lowercase());
}
}
if query.exclude_unknown_model_or_provider {
push_sqlite_usage_where(builder, has_where);
builder.push(
"(LOWER(TRIM(COALESCE(model, ''))) NOT IN ('unknown', 'unknow') \
AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow'))",
);
}
if let Some(statuses) = query.statuses.as_deref() {
if !statuses.is_empty() {
push_sqlite_usage_where(builder, has_where);
@@ -496,6 +511,8 @@ fn push_sqlite_usage_keyword_filters(
provider_name: query.provider_name.clone(),
model: query.model.clone(),
api_format: query.api_format.clone(),
client_family: query.client_family.clone(),
exclude_unknown_model_or_provider: query.exclude_unknown_model_or_provider,
statuses: query.statuses.clone(),
is_stream: query.is_stream,
error_only: query.error_only,