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

chore: update gateway pressure observability

Browse Source
This commit is contained in:
elky
2026-06-30 17:01:39 +08:00
parent 974def5fef
commit f179ee72f9
69 changed files with 14843 additions and 1609 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/backend/maintenance.rs
+649 -1
View File
@@ -3,7 +3,8 @@ use super::{
};
use crate::error::{SqlResultExt, SqlxResultExt};
use crate::maintenance::{
DatabaseMaintenanceSummary, DatabasePoolSummary, StatsDailyAggregationInput,
DatabaseMaintenanceSummary, DatabasePoolSummary, DatabasePostgresActivityGroup,
DatabasePostgresObservabilitySnapshot, StatsDailyAggregationInput,
StatsDailyAggregationSummary, StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
@@ -14,6 +15,7 @@ use crate::repository::system::{
};
use crate::DataLayerError;
use sqlx::migrate::MigrateError;
use sqlx::Row;
fn maintenance_identifier(value: &str) -> Result<&str, DataLayerError> {
let valid = !value.is_empty()
@@ -109,6 +111,25 @@ impl DataBackends {
self.sql_backend().map(SqlBackendRef::database_pool_summary)
}
pub async fn postgres_observability_snapshot(
&self,
) -> Result<Option<DatabasePostgresObservabilitySnapshot>, DataLayerError> {
match self.postgres() {
Some(postgres) => postgres.postgres_observability_snapshot().await.map(Some),
None => Ok(None),
}
}
pub async fn postgres_activity_groups(
&self,
limit: i64,
) -> Result<Vec<DatabasePostgresActivityGroup>, DataLayerError> {
match self.postgres() {
Some(postgres) => postgres.postgres_activity_groups(limit).await,
None => Ok(Vec::new()),
}
}
pub async fn aggregate_wallet_daily_usage(
&self,
input: &WalletDailyUsageAggregationInput,
@@ -260,6 +281,633 @@ impl PostgresBackend {
}
Ok(summary)
}
pub async fn postgres_observability_snapshot(
&self,
) -> Result<DatabasePostgresObservabilitySnapshot, DataLayerError> {
const ACTIVITY_SQL: &str = r#"
SELECT
COUNT(*) FILTER (WHERE state = 'active')::BIGINT AS active_connections,
COUNT(*) FILTER (WHERE state = 'idle')::BIGINT AS idle_connections,
COUNT(*) FILTER (WHERE state = 'idle in transaction')::BIGINT AS idle_in_transaction_connections,
COUNT(*) FILTER (WHERE state = 'active' AND wait_event_type IS NOT NULL)::BIGINT AS waiting_connections,
COUNT(*) FILTER (WHERE state = 'active' AND wait_event_type = 'Lock')::BIGINT AS lock_waiting_connections,
COALESCE(MAX(EXTRACT(EPOCH FROM now() - query_start) * 1000) FILTER (WHERE state = 'active' AND query_start IS NOT NULL), 0)::BIGINT AS oldest_active_query_age_ms,
COALESCE(MAX(EXTRACT(EPOCH FROM now() - xact_start) * 1000) FILTER (WHERE xact_start IS NOT NULL), 0)::BIGINT AS oldest_transaction_age_ms
FROM pg_stat_activity
WHERE datname = current_database()
AND pid <> pg_backend_pid()
"#;
const DEADLOCKS_SQL: &str = r#"
SELECT
COALESCE(SUM(deadlocks), 0)::BIGINT AS deadlocks_total,
COALESCE(SUM(blks_read), 0)::BIGINT AS block_read_total,
COALESCE(SUM(blks_hit), 0)::BIGINT AS block_hit_total,
COALESCE(SUM(temp_files), 0)::BIGINT AS temp_files_total,
COALESCE(SUM(temp_bytes), 0)::BIGINT AS temp_bytes_total,
COALESCE(SUM(xact_commit), 0)::BIGINT AS xact_commit_total,
COALESCE(SUM(xact_rollback), 0)::BIGINT AS xact_rollback_total
FROM pg_stat_database
WHERE datname = current_database()
"#;
let activity = sqlx::query(ACTIVITY_SQL)
.fetch_one(self.pool())
.await
.map_postgres_err()?;
let database = sqlx::query(DEADLOCKS_SQL)
.fetch_one(self.pool())
.await
.map_postgres_err()?;
let wal = self.postgres_wal_observability_snapshot().await;
let checkpoint = self.postgres_checkpoint_observability_snapshot().await;
let statements = self.postgres_statement_observability_snapshot().await;
let block_read_total = row_u64(&database, "block_read_total")?;
let block_hit_total = row_u64(&database, "block_hit_total")?;
Ok(DatabasePostgresObservabilitySnapshot {
active_connections: row_u64(&activity, "active_connections")?,
idle_connections: row_u64(&activity, "idle_connections")?,
idle_in_transaction_connections: row_u64(&activity, "idle_in_transaction_connections")?,
waiting_connections: row_u64(&activity, "waiting_connections")?,
lock_waiting_connections: row_u64(&activity, "lock_waiting_connections")?,
oldest_active_query_age_ms: row_u64(&activity, "oldest_active_query_age_ms")?,
oldest_transaction_age_ms: row_u64(&activity, "oldest_transaction_age_ms")?,
deadlocks_total: row_u64(&database, "deadlocks_total")?,
block_read_total,
block_hit_total,
block_cache_hit_rate_basis_points: ratio_to_basis_points(
block_hit_total,
block_read_total.saturating_add(block_hit_total),
),
temp_files_total: row_u64(&database, "temp_files_total")?,
temp_bytes_total: row_u64(&database, "temp_bytes_total")?,
xact_commit_total: row_u64(&database, "xact_commit_total")?,
xact_rollback_total: row_u64(&database, "xact_rollback_total")?,
wal_observability_available: wal.available,
wal_observability_unavailable: wal.unavailable,
wal_records_total: wal.records_total,
wal_fpi_total: wal.fpi_total,
wal_bytes_total: wal.bytes_total,
wal_buffers_full_total: wal.buffers_full_total,
wal_write_total: wal.write_total,
wal_sync_total: wal.sync_total,
wal_write_time_ms_total: wal.write_time_ms_total,
wal_sync_time_ms_total: wal.sync_time_ms_total,
checkpoint_observability_available: checkpoint.available,
checkpoint_observability_unavailable: checkpoint.unavailable,
checkpoints_timed_total: checkpoint.timed_total,
checkpoints_requested_total: checkpoint.requested_total,
checkpoint_write_time_ms_total: checkpoint.write_time_ms_total,
checkpoint_sync_time_ms_total: checkpoint.sync_time_ms_total,
buffers_checkpoint_total: checkpoint.buffers_checkpoint_total,
buffers_backend_total: checkpoint.buffers_backend_total,
statement_observability_available: statements.available,
statement_observability_unavailable: statements.unavailable,
statement_top_calls_total: statements.top_calls_total,
statement_top_exec_time_ms_total: statements.top_exec_time_ms_total,
statement_top_max_mean_exec_time_ms: statements.top_max_mean_exec_time_ms,
statement_top_max_exec_time_ms: statements.top_max_exec_time_ms,
statement_top_shared_blks_read_total: statements.top_shared_blks_read_total,
statement_top_shared_blks_hit_total: statements.top_shared_blks_hit_total,
statement_top_temp_blks_total: statements.top_temp_blks_total,
})
}
pub async fn postgres_activity_groups(
&self,
limit: i64,
) -> Result<Vec<DatabasePostgresActivityGroup>, DataLayerError> {
const ACTIVITY_GROUP_SQL: &str = r#"
WITH normalized_activity AS (
SELECT
COALESCE(NULLIF(state, ''), 'unknown') AS state,
COALESCE(NULLIF(wait_event_type, ''), 'none') AS wait_event_type,
COALESCE(NULLIF(wait_event, ''), 'none') AS wait_event,
LEFT(
regexp_replace(
regexp_replace(
COALESCE(NULLIF(query, ''), '<empty>'),
'\s+',
' ',
'g'
),
'([0-9a-fA-F]{8,}|[0-9]+)',
'?',
'g'
),
160
) AS query_prefix,
COALESCE(EXTRACT(EPOCH FROM now() - query_start) * 1000, 0)::BIGINT AS query_age_ms,
COALESCE(EXTRACT(EPOCH FROM now() - xact_start) * 1000, 0)::BIGINT AS transaction_age_ms
FROM pg_stat_activity
WHERE datname = current_database()
AND pid <> pg_backend_pid()
)
SELECT
state,
wait_event_type,
wait_event,
query_prefix,
COUNT(*)::BIGINT AS connections,
COALESCE(MAX(query_age_ms), 0)::BIGINT AS max_query_age_ms,
COALESCE(MAX(transaction_age_ms), 0)::BIGINT AS max_transaction_age_ms
FROM normalized_activity
GROUP BY state, wait_event_type, wait_event, query_prefix
ORDER BY connections DESC, max_transaction_age_ms DESC, max_query_age_ms DESC
LIMIT $1
"#;
let rows = sqlx::query(ACTIVITY_GROUP_SQL)
.bind(limit.clamp(1, 20))
.fetch_all(self.pool())
.await
.map_postgres_err()?;
rows.into_iter()
.map(|row| {
Ok(DatabasePostgresActivityGroup {
state: row.try_get::<String, _>("state").map_postgres_err()?,
wait_event_type: row
.try_get::<String, _>("wait_event_type")
.map_postgres_err()?,
wait_event: row.try_get::<String, _>("wait_event").map_postgres_err()?,
query_prefix: row
.try_get::<String, _>("query_prefix")
.map_postgres_err()?,
connections: row_u64(&row, "connections")?,
max_query_age_ms: row_u64(&row, "max_query_age_ms")?,
max_transaction_age_ms: row_u64(&row, "max_transaction_age_ms")?,
})
})
.collect()
}
async fn postgres_wal_observability_snapshot(&self) -> PostgresWalObservabilitySnapshot {
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_wal",
&["wal_records", "wal_fpi", "wal_bytes", "wal_buffers_full"],
)
.await
{
return PostgresWalObservabilitySnapshot::default();
}
const WAL_SQL: &str = r#"
SELECT
COALESCE(SUM(wal_records), 0)::BIGINT AS records_total,
COALESCE(SUM(wal_fpi), 0)::BIGINT AS fpi_total,
COALESCE(SUM(wal_bytes), 0)::BIGINT AS bytes_total,
COALESCE(SUM(wal_buffers_full), 0)::BIGINT AS buffers_full_total
FROM pg_stat_wal
"#;
match sqlx::query(WAL_SQL).fetch_one(self.pool()).await {
Ok(row) => {
let io = self.postgres_wal_io_observability_snapshot().await;
PostgresWalObservabilitySnapshot {
available: 1,
records_total: row_u64(&row, "records_total").unwrap_or_default(),
fpi_total: row_u64(&row, "fpi_total").unwrap_or_default(),
bytes_total: row_u64(&row, "bytes_total").unwrap_or_default(),
buffers_full_total: row_u64(&row, "buffers_full_total").unwrap_or_default(),
write_total: io.write_total,
sync_total: io.sync_total,
write_time_ms_total: io.write_time_ms_total,
sync_time_ms_total: io.sync_time_ms_total,
..PostgresWalObservabilitySnapshot::default()
}
}
Err(_) => PostgresWalObservabilitySnapshot {
unavailable: 1,
..PostgresWalObservabilitySnapshot::default()
},
}
}
async fn postgres_wal_io_observability_snapshot(&self) -> PostgresWalIoObservabilitySnapshot {
if self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_wal",
&["wal_write", "wal_sync", "wal_write_time", "wal_sync_time"],
)
.await
{
return self.postgres_wal_legacy_io_observability_snapshot().await;
}
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_io",
&["object", "writes", "fsyncs", "write_time", "fsync_time"],
)
.await
{
return PostgresWalIoObservabilitySnapshot::default();
}
const WAL_IO_SQL: &str = r#"
SELECT
COALESCE(SUM(writes), 0)::BIGINT AS write_total,
COALESCE(SUM(fsyncs), 0)::BIGINT AS sync_total,
COALESCE(SUM(write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(fsync_time), 0)::BIGINT AS sync_time_ms_total
FROM pg_stat_io
WHERE object = 'wal'
"#;
match sqlx::query(WAL_IO_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresWalIoObservabilitySnapshot {
write_total: row_u64(&row, "write_total").unwrap_or_default(),
sync_total: row_u64(&row, "sync_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
},
Err(_) => PostgresWalIoObservabilitySnapshot::default(),
}
}
async fn postgres_wal_legacy_io_observability_snapshot(
&self,
) -> PostgresWalIoObservabilitySnapshot {
const WAL_IO_SQL: &str = r#"
SELECT
COALESCE(SUM(wal_write), 0)::BIGINT AS write_total,
COALESCE(SUM(wal_sync), 0)::BIGINT AS sync_total,
COALESCE(SUM(wal_write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(wal_sync_time), 0)::BIGINT AS sync_time_ms_total
FROM pg_stat_wal
"#;
match sqlx::query(WAL_IO_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresWalIoObservabilitySnapshot {
write_total: row_u64(&row, "write_total").unwrap_or_default(),
sync_total: row_u64(&row, "sync_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
},
Err(_) => PostgresWalIoObservabilitySnapshot::default(),
}
}
async fn postgres_checkpoint_observability_snapshot(
&self,
) -> PostgresCheckpointObservabilitySnapshot {
if self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_checkpointer",
&[
"num_timed",
"num_requested",
"write_time",
"sync_time",
"buffers_written",
],
)
.await
{
return self
.postgres_checkpoint_observability_snapshot_from_checkpointer()
.await;
}
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_bgwriter",
&[
"checkpoints_timed",
"checkpoints_req",
"checkpoint_write_time",
"checkpoint_sync_time",
"buffers_checkpoint",
"buffers_backend",
],
)
.await
{
return PostgresCheckpointObservabilitySnapshot::default();
}
const CHECKPOINT_SQL: &str = r#"
SELECT
COALESCE(SUM(checkpoints_timed), 0)::BIGINT AS timed_total,
COALESCE(SUM(checkpoints_req), 0)::BIGINT AS requested_total,
COALESCE(SUM(checkpoint_write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(checkpoint_sync_time), 0)::BIGINT AS sync_time_ms_total,
COALESCE(SUM(buffers_checkpoint), 0)::BIGINT AS buffers_checkpoint_total,
COALESCE(SUM(buffers_backend), 0)::BIGINT AS buffers_backend_total
FROM pg_stat_bgwriter
"#;
match sqlx::query(CHECKPOINT_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresCheckpointObservabilitySnapshot {
available: 1,
timed_total: row_u64(&row, "timed_total").unwrap_or_default(),
requested_total: row_u64(&row, "requested_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
buffers_checkpoint_total: row_u64(&row, "buffers_checkpoint_total")
.unwrap_or_default(),
buffers_backend_total: row_u64(&row, "buffers_backend_total").unwrap_or_default(),
..PostgresCheckpointObservabilitySnapshot::default()
},
Err(_) => PostgresCheckpointObservabilitySnapshot {
unavailable: 1,
..PostgresCheckpointObservabilitySnapshot::default()
},
}
}
async fn postgres_checkpoint_observability_snapshot_from_checkpointer(
&self,
) -> PostgresCheckpointObservabilitySnapshot {
const CHECKPOINT_SQL: &str = r#"
SELECT
COALESCE(SUM(num_timed), 0)::BIGINT AS timed_total,
COALESCE(SUM(num_requested), 0)::BIGINT AS requested_total,
COALESCE(SUM(write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(sync_time), 0)::BIGINT AS sync_time_ms_total,
COALESCE(SUM(buffers_written), 0)::BIGINT AS buffers_checkpoint_total
FROM pg_stat_checkpointer
"#;
match sqlx::query(CHECKPOINT_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresCheckpointObservabilitySnapshot {
available: 1,
timed_total: row_u64(&row, "timed_total").unwrap_or_default(),
requested_total: row_u64(&row, "requested_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
buffers_checkpoint_total: row_u64(&row, "buffers_checkpoint_total")
.unwrap_or_default(),
..PostgresCheckpointObservabilitySnapshot::default()
},
Err(_) => PostgresCheckpointObservabilitySnapshot {
unavailable: 1,
..PostgresCheckpointObservabilitySnapshot::default()
},
}
}
async fn postgres_statement_observability_snapshot(
&self,
) -> PostgresStatementObservabilitySnapshot {
let extension_installed = sqlx::query_scalar::<_, bool>(
"SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements')",
)
.fetch_one(self.pool())
.await
.unwrap_or(false);
if !extension_installed {
return PostgresStatementObservabilitySnapshot::default();
}
if self
.postgres_catalog_relation_has_columns(
"pg_stat_statements",
&[
"calls",
"total_exec_time",
"mean_exec_time",
"max_exec_time",
"shared_blks_read",
"shared_blks_hit",
"temp_blks_read",
"temp_blks_written",
"dbid",
],
)
.await
{
return self
.postgres_statement_observability_snapshot_with_exec_time()
.await;
}
if self
.postgres_catalog_relation_has_columns(
"pg_stat_statements",
&[
"calls",
"total_time",
"mean_time",
"max_time",
"shared_blks_read",
"shared_blks_hit",
"temp_blks_read",
"temp_blks_written",
"dbid",
],
)
.await
{
return self
.postgres_statement_observability_snapshot_with_total_time()
.await;
}
PostgresStatementObservabilitySnapshot::default()
}
async fn postgres_statement_observability_snapshot_with_exec_time(
&self,
) -> PostgresStatementObservabilitySnapshot {
const STATEMENTS_SQL: &str = r#"
SELECT
COALESCE(SUM(calls), 0)::BIGINT AS top_calls_total,
COALESCE(SUM(total_exec_time), 0)::BIGINT AS top_exec_time_ms_total,
COALESCE(MAX(mean_exec_time), 0)::BIGINT AS top_max_mean_exec_time_ms,
COALESCE(MAX(max_exec_time), 0)::BIGINT AS top_max_exec_time_ms,
COALESCE(SUM(shared_blks_read), 0)::BIGINT AS top_shared_blks_read_total,
COALESCE(SUM(shared_blks_hit), 0)::BIGINT AS top_shared_blks_hit_total,
COALESCE(SUM(temp_blks_read + temp_blks_written), 0)::BIGINT AS top_temp_blks_total
FROM (
SELECT
calls,
total_exec_time,
mean_exec_time,
max_exec_time,
shared_blks_read,
shared_blks_hit,
temp_blks_read,
temp_blks_written
FROM pg_stat_statements
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
ORDER BY total_exec_time DESC
LIMIT 20
) top_statements
"#;
match sqlx::query(STATEMENTS_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresStatementObservabilitySnapshot {
available: 1,
top_calls_total: row_u64(&row, "top_calls_total").unwrap_or_default(),
top_exec_time_ms_total: row_u64(&row, "top_exec_time_ms_total").unwrap_or_default(),
top_max_mean_exec_time_ms: row_u64(&row, "top_max_mean_exec_time_ms")
.unwrap_or_default(),
top_max_exec_time_ms: row_u64(&row, "top_max_exec_time_ms").unwrap_or_default(),
top_shared_blks_read_total: row_u64(&row, "top_shared_blks_read_total")
.unwrap_or_default(),
top_shared_blks_hit_total: row_u64(&row, "top_shared_blks_hit_total")
.unwrap_or_default(),
top_temp_blks_total: row_u64(&row, "top_temp_blks_total").unwrap_or_default(),
..PostgresStatementObservabilitySnapshot::default()
},
Err(_) => PostgresStatementObservabilitySnapshot {
unavailable: 1,
..PostgresStatementObservabilitySnapshot::default()
},
}
}
async fn postgres_statement_observability_snapshot_with_total_time(
&self,
) -> PostgresStatementObservabilitySnapshot {
const STATEMENTS_SQL: &str = r#"
SELECT
COALESCE(SUM(calls), 0)::BIGINT AS top_calls_total,
COALESCE(SUM(total_time), 0)::BIGINT AS top_exec_time_ms_total,
COALESCE(MAX(mean_time), 0)::BIGINT AS top_max_mean_exec_time_ms,
COALESCE(MAX(max_time), 0)::BIGINT AS top_max_exec_time_ms,
COALESCE(SUM(shared_blks_read), 0)::BIGINT AS top_shared_blks_read_total,
COALESCE(SUM(shared_blks_hit), 0)::BIGINT AS top_shared_blks_hit_total,
COALESCE(SUM(temp_blks_read + temp_blks_written), 0)::BIGINT AS top_temp_blks_total
FROM (
SELECT
calls,
total_time,
mean_time,
max_time,
shared_blks_read,
shared_blks_hit,
temp_blks_read,
temp_blks_written
FROM pg_stat_statements
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
ORDER BY total_time DESC
LIMIT 20
) top_statements
"#;
match sqlx::query(STATEMENTS_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresStatementObservabilitySnapshot {
available: 1,
top_calls_total: row_u64(&row, "top_calls_total").unwrap_or_default(),
top_exec_time_ms_total: row_u64(&row, "top_exec_time_ms_total").unwrap_or_default(),
top_max_mean_exec_time_ms: row_u64(&row, "top_max_mean_exec_time_ms")
.unwrap_or_default(),
top_max_exec_time_ms: row_u64(&row, "top_max_exec_time_ms").unwrap_or_default(),
top_shared_blks_read_total: row_u64(&row, "top_shared_blks_read_total")
.unwrap_or_default(),
top_shared_blks_hit_total: row_u64(&row, "top_shared_blks_hit_total")
.unwrap_or_default(),
top_temp_blks_total: row_u64(&row, "top_temp_blks_total").unwrap_or_default(),
..PostgresStatementObservabilitySnapshot::default()
},
Err(_) => PostgresStatementObservabilitySnapshot {
unavailable: 1,
..PostgresStatementObservabilitySnapshot::default()
},
}
}
async fn postgres_catalog_relation_has_columns(
&self,
relation: &str,
columns: &[&str],
) -> bool {
if !self.postgres_catalog_relation_exists(relation).await {
return false;
}
for column in columns {
if !self.postgres_catalog_column_exists(relation, column).await {
return false;
}
}
true
}
async fn postgres_catalog_column_exists(&self, relation: &str, column: &str) -> bool {
sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS (
SELECT 1
FROM pg_attribute
WHERE attrelid = to_regclass($1)
AND attname = $2
AND NOT attisdropped
)
"#,
)
.bind(relation)
.bind(column)
.fetch_one(self.pool())
.await
.unwrap_or(false)
}
async fn postgres_catalog_relation_exists(&self, relation: &str) -> bool {
sqlx::query_scalar::<_, Option<String>>("SELECT to_regclass($1)::TEXT")
.bind(relation)
.fetch_one(self.pool())
.await
.ok()
.flatten()
.is_some()
}
}
fn row_u64(row: &sqlx::postgres::PgRow, name: &str) -> Result<u64, DataLayerError> {
row.try_get::<i64, _>(name)
.map(u64_from_i64)
.map_postgres_err()
}
fn u64_from_i64(value: i64) -> u64 {
u64::try_from(value).unwrap_or_default()
}
fn ratio_to_basis_points(value: u64, total: u64) -> u64 {
value.saturating_mul(10_000).checked_div(total).unwrap_or(0)
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresWalObservabilitySnapshot {
available: u64,
unavailable: u64,
records_total: u64,
fpi_total: u64,
bytes_total: u64,
buffers_full_total: u64,
write_total: u64,
sync_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresWalIoObservabilitySnapshot {
write_total: u64,
sync_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresCheckpointObservabilitySnapshot {
available: u64,
unavailable: u64,
timed_total: u64,
requested_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
buffers_checkpoint_total: u64,
buffers_backend_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresStatementObservabilitySnapshot {
available: u64,
unavailable: u64,
top_calls_total: u64,
top_exec_time_ms_total: u64,
top_max_mean_exec_time_ms: u64,
top_max_exec_time_ms: u64,
top_shared_blks_read_total: u64,
top_shared_blks_hit_total: u64,
top_temp_blks_total: u64,
}
impl MysqlBackend {
crates/aether-data/src/lib.rs
+2 -1
View File
@@ -24,7 +24,8 @@ pub use config::DataLayerConfig;
pub use database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL};
pub use error::DataLayerError;
pub use maintenance::{
DatabaseMaintenanceSummary, DatabasePoolSummary, StatsDailyAggregationInput,
DatabaseMaintenanceSummary, DatabasePoolSummary, DatabasePostgresActivityGroup,
DatabasePostgresObservabilitySnapshot, StatsDailyAggregationInput,
StatsDailyAggregationSummary, StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
crates/aether-data/src/maintenance.rs
+57
View File
@@ -21,6 +21,63 @@ pub struct DatabasePoolSummary {
pub usage_rate: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct DatabasePostgresObservabilitySnapshot {
pub active_connections: u64,
pub idle_connections: u64,
pub idle_in_transaction_connections: u64,
pub waiting_connections: u64,
pub lock_waiting_connections: u64,
pub oldest_active_query_age_ms: u64,
pub oldest_transaction_age_ms: u64,
pub deadlocks_total: u64,
pub block_read_total: u64,
pub block_hit_total: u64,
pub block_cache_hit_rate_basis_points: u64,
pub temp_files_total: u64,
pub temp_bytes_total: u64,
pub xact_commit_total: u64,
pub xact_rollback_total: u64,
pub wal_observability_available: u64,
pub wal_observability_unavailable: u64,
pub wal_records_total: u64,
pub wal_fpi_total: u64,
pub wal_bytes_total: u64,
pub wal_buffers_full_total: u64,
pub wal_write_total: u64,
pub wal_sync_total: u64,
pub wal_write_time_ms_total: u64,
pub wal_sync_time_ms_total: u64,
pub checkpoint_observability_available: u64,
pub checkpoint_observability_unavailable: u64,
pub checkpoints_timed_total: u64,
pub checkpoints_requested_total: u64,
pub checkpoint_write_time_ms_total: u64,
pub checkpoint_sync_time_ms_total: u64,
pub buffers_checkpoint_total: u64,
pub buffers_backend_total: u64,
pub statement_observability_available: u64,
pub statement_observability_unavailable: u64,
pub statement_top_calls_total: u64,
pub statement_top_exec_time_ms_total: u64,
pub statement_top_max_mean_exec_time_ms: u64,
pub statement_top_max_exec_time_ms: u64,
pub statement_top_shared_blks_read_total: u64,
pub statement_top_shared_blks_hit_total: u64,
pub statement_top_temp_blks_total: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct DatabasePostgresActivityGroup {
pub state: String,
pub wait_event_type: String,
pub wait_event: String,
pub query_prefix: String,
pub connections: u64,
pub max_query_age_ms: u64,
pub max_transaction_age_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WalletDailyUsageAggregationInput {
pub billing_date: String,
crates/aether-data/src/repository/candidates/mysql.rs
+148 -16
View File
@@ -374,6 +374,17 @@ fn merge_candidate(
candidate: UpsertRequestCandidateRecord,
existing: Option<StoredRequestCandidate>,
) -> Result<StoredRequestCandidate, DataLayerError> {
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| {
request_candidate_lifecycle_would_regress(value.status, candidate.status)
});
let merged_status = if preserve_existing_lifecycle {
existing
.as_ref()
.map(|value| value.status)
.unwrap_or(candidate.status)
} else {
candidate.status
};
let created_at_unix_ms = candidate
.created_at_unix_ms
.filter(|value| *value > 1000)
@@ -426,7 +437,7 @@ fn merge_candidate(
candidate
.key_id
.or_else(|| existing.as_ref().and_then(|value| value.key_id.clone())),
candidate.status,
merged_status,
candidate.skip_reason.or_else(|| {
existing
.as_ref()
@@ -435,25 +446,48 @@ fn merge_candidate(
candidate
.is_cached
.unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)),
candidate.status_code.map(i32::from).or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
}),
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())),
candidate.error_message.or_else(|| {
} else {
candidate.status_code.map(i32::from).or_else(|| {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
})
},
if preserve_existing_lifecycle {
existing.as_ref().and_then(|value| value.error_type.clone())
} else {
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone()))
},
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
}),
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
} else {
candidate.error_message.or_else(|| {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
})
},
if preserve_existing_lifecycle {
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
),
}
} else {
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
)
},
candidate.concurrent_requests.map(to_i32).transpose()?.or(
match existing
.as_ref()
@@ -475,18 +509,42 @@ fn merge_candidate(
.or_else(|| existing.as_ref().and_then(|value| value.started_at_unix_ms))
.map(|value| u64_to_i64(value, "request candidate started_at"))
.transpose()?,
candidate
.finished_at_unix_ms
.or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
} else {
candidate.finished_at_unix_ms.or_else(|| {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
})
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
}
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
)
}
fn request_candidate_lifecycle_would_regress(
existing: RequestCandidateStatus,
incoming: RequestCandidateStatus,
) -> bool {
matches!(
existing,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
) && matches!(
incoming,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
) || existing == RequestCandidateStatus::Streaming
&& incoming == RequestCandidateStatus::Pending
}
fn aggregate_timeline(
candidates: Vec<StoredRequestCandidate>,
since_unix_secs: u64,
@@ -681,6 +739,9 @@ fn optional_u64_to_i64(value: Option<u64>, name: &str) -> Result<Option<i64>, Da
#[cfg(test)]
mod tests {
use super::MysqlRequestCandidateRepository;
use crate::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
#[tokio::test]
async fn repository_builds_from_lazy_pool() {
@@ -692,4 +753,75 @@ mod tests {
let _repository = MysqlRequestCandidateRepository::new(pool);
}
#[test]
fn merge_candidate_keeps_terminal_status_when_streaming_arrives_late() {
let existing = StoredRequestCandidate::new(
"candidate-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(123),
None,
Some(serde_json::json!({"terminal": true})),
None,
1_000,
Some(1_001),
Some(1_123),
)
.expect("existing candidate should build");
let merged = super::merge_candidate(
UpsertRequestCandidateRecord {
id: "candidate-late".to_string(),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
status: RequestCandidateStatus::Streaming,
skip_reason: None,
is_cached: Some(false),
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(9_999),
concurrent_requests: None,
extra_data: Some(serde_json::json!({"late": true})),
required_capabilities: None,
created_at_unix_ms: Some(1_050),
started_at_unix_ms: Some(1_051),
finished_at_unix_ms: None,
},
Some(existing),
)
.expect("candidate should merge");
assert_eq!(merged.id, "candidate-1");
assert_eq!(merged.status, RequestCandidateStatus::Success);
assert_eq!(merged.latency_ms, Some(123));
assert_eq!(merged.finished_at_unix_ms, Some(1_123));
assert_eq!(
merged.extra_data,
Some(serde_json::json!({"terminal": true, "late": true}))
);
}
}
crates/aether-data/src/repository/candidates/postgres.rs
+169 -19
View File
@@ -132,13 +132,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = COALESCE($14, request_candidates.is_cached),
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -155,7 +190,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
RETURNING
id,
request_id,
@@ -193,13 +235,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = COALESCE(EXCLUDED.is_cached, request_candidates.is_cached),
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -216,7 +293,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
"#;
const UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL: &str = r#"
@@ -229,13 +313,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = request_candidates.is_cached,
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -252,7 +371,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
"#;
const UPSERT_MANY_PREFIX_SQL: &str = r#"
@@ -1015,7 +1141,10 @@ fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
#[cfg(test)]
mod tests {
use super::{SqlxRequestCandidateReadRepository, UPSERT_SQL};
use super::{
SqlxRequestCandidateReadRepository, UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
UPSERT_CONFLICT_SQL, UPSERT_SQL,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
@@ -1030,6 +1159,27 @@ mod tests {
assert!(UPSERT_SQL.contains("THEN EXCLUDED.created_at"));
}
#[test]
fn upsert_sql_keeps_terminal_candidate_state_when_lifecycle_events_arrive_late() {
for sql in [
UPSERT_SQL,
UPSERT_CONFLICT_SQL,
UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
] {
assert!(sql.contains(
"request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')"
));
assert!(
sql.contains("EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')")
);
assert!(sql.contains(
"request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'"
));
assert!(sql.contains("THEN request_candidates.status"));
assert!(sql.contains("THEN request_candidates.latency_ms"));
}
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
crates/aether-data/src/repository/candidates/sqlite.rs
+91 -16
View File
@@ -361,6 +361,17 @@ fn merge_candidate(
candidate: UpsertRequestCandidateRecord,
existing: Option<StoredRequestCandidate>,
) -> Result<StoredRequestCandidate, DataLayerError> {
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| {
request_candidate_lifecycle_would_regress(value.status, candidate.status)
});
let merged_status = if preserve_existing_lifecycle {
existing
.as_ref()
.map(|value| value.status)
.unwrap_or(candidate.status)
} else {
candidate.status
};
let created_at_unix_ms = candidate
.created_at_unix_ms
.filter(|value| *value > 1000)
@@ -413,7 +424,7 @@ fn merge_candidate(
candidate
.key_id
.or_else(|| existing.as_ref().and_then(|value| value.key_id.clone())),
candidate.status,
merged_status,
candidate.skip_reason.or_else(|| {
existing
.as_ref()
@@ -422,25 +433,48 @@ fn merge_candidate(
candidate
.is_cached
.unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)),
candidate.status_code.map(i32::from).or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
}),
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())),
candidate.error_message.or_else(|| {
} else {
candidate.status_code.map(i32::from).or_else(|| {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
})
},
if preserve_existing_lifecycle {
existing.as_ref().and_then(|value| value.error_type.clone())
} else {
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone()))
},
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
}),
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
} else {
candidate.error_message.or_else(|| {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
})
},
if preserve_existing_lifecycle {
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
),
}
} else {
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
)
},
candidate.concurrent_requests.map(to_i32).transpose()?.or(
match existing
.as_ref()
@@ -462,18 +496,42 @@ fn merge_candidate(
.or_else(|| existing.as_ref().and_then(|value| value.started_at_unix_ms))
.map(|value| u64_to_i64(value, "request candidate started_at"))
.transpose()?,
candidate
.finished_at_unix_ms
.or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
} else {
candidate.finished_at_unix_ms.or_else(|| {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
})
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
}
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
)
}
fn request_candidate_lifecycle_would_regress(
existing: RequestCandidateStatus,
incoming: RequestCandidateStatus,
) -> bool {
matches!(
existing,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
) && matches!(
incoming,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
) || existing == RequestCandidateStatus::Streaming
&& incoming == RequestCandidateStatus::Pending
}
fn aggregate_timeline(
candidates: Vec<StoredRequestCandidate>,
since_unix_secs: u64,
@@ -710,6 +768,23 @@ mod tests {
assert_eq!(updated.id, "candidate-1");
assert_eq!(updated.extra_data, Some(json!({"a": 1, "b": 2})));
let late_streaming = repository
.upsert(sample_upsert(
"candidate-late-streaming",
RequestCandidateStatus::Streaming,
Some(json!({"late": true})),
1_000_250,
))
.await
.expect("late streaming candidate should not regress terminal status");
assert_eq!(late_streaming.id, "candidate-1");
assert_eq!(late_streaming.status, RequestCandidateStatus::Success);
assert_eq!(late_streaming.finished_at_unix_ms, Some(1_000_502));
assert_eq!(
late_streaming.extra_data,
Some(json!({"a": 1, "b": 2, "late": true}))
);
assert_eq!(
repository
.list_by_request_id("request-1")
crates/aether-data/src/repository/usage/mysql.rs
+96 -19
View File
@@ -151,33 +151,85 @@ ON DUPLICATE KEY UPDATE
has_format_conversion = VALUES(has_format_conversion),
is_stream = VALUES(is_stream),
upstream_is_stream = VALUES(upstream_is_stream),
input_tokens = VALUES(input_tokens),
output_tokens = VALUES(output_tokens),
total_tokens = VALUES(total_tokens),
cache_creation_input_tokens = VALUES(cache_creation_input_tokens),
cache_creation_ephemeral_5m_input_tokens = VALUES(cache_creation_ephemeral_5m_input_tokens),
cache_creation_ephemeral_1h_input_tokens = VALUES(cache_creation_ephemeral_1h_input_tokens),
cache_read_input_tokens = VALUES(cache_read_input_tokens),
cache_creation_cost_usd = VALUES(cache_creation_cost_usd),
cache_read_cost_usd = VALUES(cache_read_cost_usd),
output_price_per_1m = VALUES(output_price_per_1m),
total_cost_usd = VALUES(total_cost_usd),
actual_total_cost_usd = VALUES(actual_total_cost_usd),
status_code = VALUES(status_code),
error_message = VALUES(error_message),
error_category = VALUES(error_category),
input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN input_tokens
ELSE VALUES(input_tokens)
END,
output_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_tokens
ELSE VALUES(output_tokens)
END,
total_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_tokens
ELSE VALUES(total_tokens)
END,
cache_creation_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_input_tokens
ELSE VALUES(cache_creation_input_tokens)
END,
cache_creation_ephemeral_5m_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_5m_input_tokens
ELSE VALUES(cache_creation_ephemeral_5m_input_tokens)
END,
cache_creation_ephemeral_1h_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_1h_input_tokens
ELSE VALUES(cache_creation_ephemeral_1h_input_tokens)
END,
cache_read_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_input_tokens
ELSE VALUES(cache_read_input_tokens)
END,
cache_creation_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_cost_usd
ELSE VALUES(cache_creation_cost_usd)
END,
cache_read_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_cost_usd
ELSE VALUES(cache_read_cost_usd)
END,
output_price_per_1m = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_price_per_1m
ELSE VALUES(output_price_per_1m)
END,
total_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_cost_usd
ELSE VALUES(total_cost_usd)
END,
actual_total_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN actual_total_cost_usd
ELSE VALUES(actual_total_cost_usd)
END,
status_code = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status_code
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status_code
ELSE VALUES(status_code)
END,
error_message = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_message
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_message
ELSE VALUES(error_message)
END,
error_category = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_category
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_category
ELSE VALUES(error_category)
END,
response_time_ms = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN response_time_ms
WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0
THEN COALESCE(response_time_ms, VALUES(response_time_ms))
ELSE VALUES(response_time_ms)
END,
first_byte_time_ms = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN first_byte_time_ms
WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0
THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms))
ELSE VALUES(first_byte_time_ms)
END,
status = VALUES(status),
billing_status = VALUES(billing_status),
billing_status = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN billing_status
ELSE VALUES(billing_status)
END,
request_metadata = VALUES(request_metadata),
candidate_id = VALUES(candidate_id),
candidate_index = VALUES(candidate_index),
@@ -187,8 +239,19 @@ ON DUPLICATE KEY UPDATE
route_kind = VALUES(route_kind),
execution_path = VALUES(execution_path),
local_execution_runtime_miss_reason = VALUES(local_execution_runtime_miss_reason),
finalized_at = VALUES(finalized_at),
updated_at_unix_secs = VALUES(updated_at_unix_secs)
finalized_at = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN finalized_at
ELSE VALUES(finalized_at)
END,
updated_at_unix_secs = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN updated_at_unix_secs
ELSE VALUES(updated_at_unix_secs)
END,
status = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status
ELSE VALUES(status)
END
"#;
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
@@ -1550,6 +1613,20 @@ mod tests {
assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests"));
}
#[test]
fn mysql_usage_upsert_keeps_terminal_state_when_streaming_arrives_late() {
assert!(super::UPSERT_USAGE_SQL.contains(
"status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')"
));
assert!(super::UPSERT_USAGE_SQL.contains("input_tokens = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("status_code = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("billing_status = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("finalized_at = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("updated_at_unix_secs = CASE"));
assert!(super::UPSERT_USAGE_SQL
.contains("WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status"));
}
#[tokio::test]
async fn mysql_usage_write_repository_upserts_when_url_is_set() {
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
crates/aether-data/src/repository/usage/postgres/mod.rs
+149 -96
View File
@@ -8029,8 +8029,26 @@ ORDER BY "usage".user_id ASC
) -> Result<StoredRequestUsageAudit, DataLayerError> {
usage.validate()?;
let usage = strip_deprecated_usage_display_fields(usage);
let prepared = prepare_usage_upsert_context(&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,
request_metadata_value,
request_metadata_json,
} = prepared;
Box::pin(async move {
lock_usage_request_id_in_tx(tx, &usage.request_id).await?;
@@ -8052,102 +8070,6 @@ ORDER BY "usage".user_id ASC
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage =
prepare_usage_body_storage(usage.request_body.as_ref())?;
let provider_request_headers_json =
json_bind_text(usage.provider_request_headers.as_ref())?;
let provider_request_body_storage =
prepare_usage_body_storage(usage.provider_request_body.as_ref())?;
let response_headers_json = json_bind_text(usage.response_headers.as_ref())?;
let response_body_storage =
prepare_usage_body_storage(usage.response_body.as_ref())?;
let client_response_headers_json =
json_bind_text(usage.client_response_headers.as_ref())?;
let client_response_body_storage =
prepare_usage_body_storage(usage.client_response_body.as_ref())?;
let http_audit_refs = UsageHttpAuditRefs {
request_body_ref: resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
None,
),
provider_request_body_ref: resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
None,
),
response_body_ref: resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
None,
),
client_response_body_ref: resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
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(),
[
(
UsageBodyField::RequestBody,
&request_body_storage,
usage.request_body.as_ref(),
usage.request_body_ref.as_deref(),
),
(
UsageBodyField::ProviderRequestBody,
&provider_request_body_storage,
usage.provider_request_body.as_ref(),
usage.provider_request_body_ref.as_deref(),
),
(
UsageBodyField::ResponseBody,
&response_body_storage,
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
),
(
UsageBodyField::ClientResponseBody,
&client_response_body_storage,
usage.client_response_body.as_ref(),
usage.client_response_body_ref.as_deref(),
),
],
);
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
usage.request_body.as_ref(),
usage.provider_request_body.as_ref(),
usage.response_body.as_ref(),
usage.client_response_body.as_ref(),
],
);
let routing_snapshot =
usage_routing_snapshot_from_usage(&usage, request_metadata_value.as_ref());
let settlement_pricing_snapshot = usage_settlement_pricing_snapshot_from_usage(
&usage,
request_metadata_value.as_ref(),
)?;
let request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
let _row = sqlx::query(UPSERT_SQL)
.bind(Uuid::new_v4().to_string())
.bind(&usage.request_id)
@@ -10486,6 +10408,25 @@ impl UsageSettlementPricingSnapshot {
}
}
#[derive(Debug)]
struct PreparedUsageUpsert {
request_headers_json: Option<String>,
provider_request_headers_json: Option<String>,
response_headers_json: Option<String>,
client_response_headers_json: Option<String>,
request_body_storage: UsageBodyStorage,
provider_request_body_storage: UsageBodyStorage,
response_body_storage: UsageBodyStorage,
client_response_body_storage: UsageBodyStorage,
http_audit_refs: UsageHttpAuditRefs,
http_audit_states: UsageHttpAuditStates,
http_audit_capture_mode: &'static str,
routing_snapshot: UsageRoutingSnapshot,
settlement_pricing_snapshot: UsageSettlementPricingSnapshot,
request_metadata_value: Option<Value>,
request_metadata_json: Option<String>,
}
fn prepare_usage_body_storage(value: Option<&Value>) -> Result<UsageBodyStorage, DataLayerError> {
let Some(value) = value else {
return Ok(UsageBodyStorage {
@@ -10530,6 +10471,118 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
.transpose()
}
fn prepare_usage_upsert_context(
usage: &UpsertUsageRecord,
) -> Result<PreparedUsageUpsert, DataLayerError> {
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage = prepare_usage_body_storage(usage.request_body.as_ref())?;
let provider_request_headers_json = json_bind_text(usage.provider_request_headers.as_ref())?;
let provider_request_body_storage =
prepare_usage_body_storage(usage.provider_request_body.as_ref())?;
let response_headers_json = json_bind_text(usage.response_headers.as_ref())?;
let response_body_storage = prepare_usage_body_storage(usage.response_body.as_ref())?;
let client_response_headers_json = json_bind_text(usage.client_response_headers.as_ref())?;
let client_response_body_storage =
prepare_usage_body_storage(usage.client_response_body.as_ref())?;
let http_audit_refs = UsageHttpAuditRefs {
request_body_ref: resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
None,
),
provider_request_body_ref: resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
None,
),
response_body_ref: resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
None,
),
client_response_body_ref: resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
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(),
[
(
UsageBodyField::RequestBody,
&request_body_storage,
usage.request_body.as_ref(),
usage.request_body_ref.as_deref(),
),
(
UsageBodyField::ProviderRequestBody,
&provider_request_body_storage,
usage.provider_request_body.as_ref(),
usage.provider_request_body_ref.as_deref(),
),
(
UsageBodyField::ResponseBody,
&response_body_storage,
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
),
(
UsageBodyField::ClientResponseBody,
&client_response_body_storage,
usage.client_response_body.as_ref(),
usage.client_response_body_ref.as_deref(),
),
],
);
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
usage.request_body.as_ref(),
usage.provider_request_body.as_ref(),
usage.response_body.as_ref(),
usage.client_response_body.as_ref(),
],
);
let routing_snapshot =
usage_routing_snapshot_from_usage(usage, request_metadata_value.as_ref());
let settlement_pricing_snapshot =
usage_settlement_pricing_snapshot_from_usage(usage, request_metadata_value.as_ref())?;
let request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
Ok(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,
request_metadata_value,
request_metadata_json,
})
}
fn usage_body_capture_state_bind_text(
value: Option<UsageBodyCaptureState>,
) -> Option<&'static str> {
crates/aether-data/src/repository/usage/sqlite.rs
+107 -15
View File
@@ -172,18 +172,54 @@ ON CONFLICT (request_id) DO UPDATE SET
has_format_conversion = excluded.has_format_conversion,
is_stream = excluded.is_stream,
upstream_is_stream = excluded.upstream_is_stream,
input_tokens = excluded.input_tokens,
output_tokens = excluded.output_tokens,
total_tokens = excluded.total_tokens,
cache_creation_input_tokens = excluded.cache_creation_input_tokens,
cache_creation_ephemeral_5m_input_tokens = excluded.cache_creation_ephemeral_5m_input_tokens,
cache_creation_ephemeral_1h_input_tokens = excluded.cache_creation_ephemeral_1h_input_tokens,
cache_read_input_tokens = excluded.cache_read_input_tokens,
cache_creation_cost_usd = excluded.cache_creation_cost_usd,
cache_read_cost_usd = excluded.cache_read_cost_usd,
output_price_per_1m = excluded.output_price_per_1m,
total_cost_usd = excluded.total_cost_usd,
actual_total_cost_usd = excluded.actual_total_cost_usd,
input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".input_tokens
ELSE excluded.input_tokens
END,
output_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_tokens
ELSE excluded.output_tokens
END,
total_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_tokens
ELSE excluded.total_tokens
END,
cache_creation_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_input_tokens
ELSE excluded.cache_creation_input_tokens
END,
cache_creation_ephemeral_5m_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_5m_input_tokens
ELSE excluded.cache_creation_ephemeral_5m_input_tokens
END,
cache_creation_ephemeral_1h_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_1h_input_tokens
ELSE excluded.cache_creation_ephemeral_1h_input_tokens
END,
cache_read_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_input_tokens
ELSE excluded.cache_read_input_tokens
END,
cache_creation_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_cost_usd
ELSE excluded.cache_creation_cost_usd
END,
cache_read_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_cost_usd
ELSE excluded.cache_read_cost_usd
END,
output_price_per_1m = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_price_per_1m
ELSE excluded.output_price_per_1m
END,
total_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_cost_usd
ELSE excluded.total_cost_usd
END,
actual_total_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".actual_total_cost_usd
ELSE excluded.actual_total_cost_usd
END,
status_code = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status_code
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status_code
@@ -214,7 +250,10 @@ ON CONFLICT (request_id) DO UPDATE SET
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status
ELSE excluded.status
END,
billing_status = excluded.billing_status,
billing_status = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".billing_status
ELSE excluded.billing_status
END,
request_metadata = excluded.request_metadata,
candidate_id = COALESCE(excluded.candidate_id, "usage".candidate_id),
candidate_index = COALESCE(excluded.candidate_index, "usage".candidate_index),
@@ -224,8 +263,14 @@ ON CONFLICT (request_id) DO UPDATE SET
route_kind = excluded.route_kind,
execution_path = excluded.execution_path,
local_execution_runtime_miss_reason = excluded.local_execution_runtime_miss_reason,
finalized_at = excluded.finalized_at,
updated_at_unix_secs = excluded.updated_at_unix_secs
finalized_at = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".finalized_at
ELSE excluded.finalized_at
END,
updated_at_unix_secs = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".updated_at_unix_secs
ELSE excluded.updated_at_unix_secs
END
"#;
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
@@ -4520,6 +4565,53 @@ mod tests {
assert_eq!(existing.updated_at_unix_secs, 1_000);
}
#[tokio::test]
async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_sqlite_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_stats_targets(&pool).await;
let repository = SqliteUsageWriteRepository::new(pool);
repository
.upsert(sample_usage("request-1", "completed", "pending", 1_000))
.await
.expect("terminal usage should upsert");
let mut late_streaming = sample_usage("request-1", "streaming", "pending", 1_001);
late_streaming.input_tokens = Some(0);
late_streaming.output_tokens = Some(0);
late_streaming.total_tokens = Some(0);
late_streaming.cache_read_input_tokens = Some(0);
late_streaming.cache_read_cost_usd = Some(0.0);
late_streaming.total_cost_usd = Some(0.0);
late_streaming.actual_total_cost_usd = Some(0.0);
late_streaming.response_time_ms = Some(9_999);
late_streaming.first_byte_time_ms = Some(9_999);
late_streaming.finalized_at_unix_secs = None;
let current = repository
.upsert(late_streaming)
.await
.expect("late streaming usage should not regress terminal usage");
assert_eq!(current.status, "completed");
assert_eq!(current.billing_status, "pending");
assert_eq!(current.total_tokens, 7);
assert_eq!(current.cache_read_input_tokens, 2);
assert_eq!(current.total_cost_usd, 0.5);
assert_eq!(current.actual_total_cost_usd, 0.4);
assert_eq!(current.response_time_ms, Some(42));
assert_eq!(current.first_byte_time_ms, Some(12));
assert_eq!(current.finalized_at_unix_secs, Some(1_000));
assert_eq!(current.updated_at_unix_secs, 1_000);
}
#[tokio::test]
async fn sqlite_usage_write_repository_cleans_stale_pending_requests() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()