mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
feat: revamp analytics dashboards and harden database migrations
Add dashboard and overview analytics, health monitoring, provider expense tracking, and announcement updates across the gateway and frontend. Keep schema migrations free of historical backfills while preserving automatic backfill execution. Bound migration deadlines, run schema preparation before Compose replacement, and anonymize deleted dashboard users. Include the current documentation cleanup and regression coverage.
This commit is contained in:
elky
@@ -168,6 +168,26 @@ impl DataBackends {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn rebuild_overview_buckets(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.rebuild_overview_buckets(input).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn drain_overview_dirty_events(
|
||||
&self,
|
||||
now: chrono::DateTime<chrono::Utc>,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.drain_overview_dirty_events(now).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
@@ -417,6 +437,30 @@ impl<'a> SqlBackendRef<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn rebuild_overview_buckets(
|
||||
self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(postgres) => postgres.rebuild_overview_buckets(input).await,
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn drain_overview_dirty_events(
|
||||
self,
|
||||
now: chrono::DateTime<chrono::Utc>,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Postgres(postgres) => postgres.drain_overview_dirty_events(now).await,
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn aggregate_stats_hourly(
|
||||
self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
|
||||
@@ -13,6 +13,32 @@ mod sql;
|
||||
use self::sql::*;
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn drain_overview_dirty_events(
|
||||
&self,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let repository = aether_data_postgres::SqlxUsageReadRepository::new(self.pool().clone());
|
||||
let merged = repository.merge_overview_dirty_events().await?;
|
||||
let retained = repository.maintain_dashboard_projection(now, 1_000).await?;
|
||||
Ok(merged + u64::from(retained))
|
||||
}
|
||||
|
||||
pub async fn rebuild_overview_buckets(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let repository = aether_data_postgres::SqlxUsageReadRepository::new(self.pool().clone());
|
||||
let cleaned = repository
|
||||
.maintain_dashboard_projection(input.aggregated_at, 1_000)
|
||||
.await?;
|
||||
let rebuilt = repository
|
||||
.rebuild_overview_buckets(input.target_hour_utc + chrono::Duration::hours(1), 8)
|
||||
.await?;
|
||||
// Retention work uses the worker's existing bounded catch-up loop too,
|
||||
// so a busy installation can retire more than one batch per hour.
|
||||
Ok(rebuilt + usize::from(cleaned))
|
||||
}
|
||||
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
|
||||
@@ -15,6 +15,8 @@ use crate::{DataLayerError, DatabaseDriver, SqlDatabaseConfig};
|
||||
#[cfg(feature = "postgres")]
|
||||
mod postgres;
|
||||
|
||||
#[cfg(all(test, feature = "postgres"))]
|
||||
mod dashboard_snapshot_tests;
|
||||
#[cfg(all(test, feature = "postgres"))]
|
||||
mod tests;
|
||||
|
||||
@@ -97,6 +99,38 @@ struct AuxiliaryTable {
|
||||
}
|
||||
|
||||
const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_stats_state",
|
||||
primary_key: &["singleton"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_stats_total",
|
||||
primary_key: &["shard"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_stats_minute",
|
||||
primary_key: &["bucket_start", "shard"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_activity_hour",
|
||||
primary_key: &["bucket_start", "shard"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_activity_minute",
|
||||
primary_key: &["bucket_start", "shard"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_actor_minute",
|
||||
primary_key: &["bucket_start", "shard", "actor_user_id"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_user_events_minute",
|
||||
primary_key: &["bucket_start", "shard"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "dashboard_request_contributions",
|
||||
primary_key: &["request_id"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "audit_logs",
|
||||
primary_key: &["id"],
|
||||
@@ -181,6 +215,10 @@ const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
|
||||
name: "payment_gateway_configs",
|
||||
primary_key: &["provider"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "provider_expenses",
|
||||
primary_key: &["id"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "billing_plans",
|
||||
primary_key: &["id"],
|
||||
@@ -213,6 +251,10 @@ const AUXILIARY_TABLES: &[AuxiliaryTable] = &[
|
||||
name: "usage_routing_snapshots",
|
||||
primary_key: &["request_id"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "usage_attribution_snapshots",
|
||||
primary_key: &["request_id"],
|
||||
},
|
||||
AuxiliaryTable {
|
||||
name: "usage_counter_deltas",
|
||||
primary_key: &["id"],
|
||||
@@ -389,6 +431,22 @@ pub struct DataExportManifest {
|
||||
pub created_at_unix_secs: u64,
|
||||
pub source_driver: Option<DatabaseDriver>,
|
||||
pub domains: Vec<ExportDomain>,
|
||||
/// Complete dashboard projection, restored atomically rather than merged by row.
|
||||
/// Older exports omit this field and retain their ordinary import behavior.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub dashboard_snapshot: Option<DashboardSnapshotManifest>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DashboardSnapshotManifest {
|
||||
pub version: u32,
|
||||
pub tables: BTreeMap<String, DashboardSnapshotTable>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DashboardSnapshotTable {
|
||||
pub rows: usize,
|
||||
pub sha256: String,
|
||||
}
|
||||
|
||||
impl DataExportManifest {
|
||||
@@ -405,6 +463,7 @@ impl DataExportManifest {
|
||||
created_at_unix_secs,
|
||||
source_driver,
|
||||
domains,
|
||||
dashboard_snapshot: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
use super::*;
|
||||
use crate::lifecycle::postgres_test_support::ManagedPostgresServer;
|
||||
use sqlx::PgPool;
|
||||
|
||||
async fn migrate(pool: &PgPool) {
|
||||
crate::lifecycle::migrate::prepare_database_for_startup(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
crate::lifecycle::migrate::run_migrations(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn request(pool: &PgPool, id: &str, input: i32) {
|
||||
sqlx::query("INSERT INTO usage(id,request_id,user_id,api_key_id,provider_name,model,status,billing_status,created_at,input_tokens,output_tokens,total_tokens) VALUES ($1,$1,'backup-user','backup-key','test','test','completed','settled',clock_timestamp(),$2,3,$2+3)")
|
||||
.bind(id).bind(input).execute(pool).await.unwrap();
|
||||
}
|
||||
|
||||
async fn dashboard_rows(pool: &PgPool) -> BTreeMap<String, Vec<Value>> {
|
||||
let mut result = BTreeMap::new();
|
||||
for table in AUXILIARY_TABLES
|
||||
.iter()
|
||||
.filter(|table| table.name.starts_with("dashboard_"))
|
||||
{
|
||||
let mut rows =
|
||||
sqlx::query_scalar::<_, Value>(&format!("SELECT to_jsonb(t) FROM {} t", table.name))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
rows.sort_by_key(Value::to_string);
|
||||
result.insert(table.name.to_owned(), rows);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn total(pool: &PgPool) -> i64 {
|
||||
sqlx::query_scalar("SELECT COALESCE(sum((metrics->>'request_count')::bigint),0)::bigint FROM dashboard_stats_total").fetch_one(pool).await.unwrap()
|
||||
}
|
||||
|
||||
async fn tokens(pool: &PgPool) -> i64 {
|
||||
sqlx::query_scalar("SELECT COALESCE(sum((metrics->>'total_tokens')::bigint),0)::bigint FROM dashboard_stats_total")
|
||||
.fetch_one(pool).await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_dashboard_snapshot_roundtrip_preserves_purged_totals_and_future_updates() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let source = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.connect(server.database_url())
|
||||
.await
|
||||
.unwrap();
|
||||
migrate(&source).await;
|
||||
sqlx::raw_sql("INSERT INTO users(id,username,email_verified) VALUES('backup-user','backup-user',false); INSERT INTO api_keys(id,user_id,key_hash) VALUES('backup-key','backup-user',repeat('e',64));").execute(&source).await.unwrap();
|
||||
request(&source, "backup-kept", 11).await;
|
||||
request(&source, "backup-purged", 23).await;
|
||||
sqlx::query("DELETE FROM usage WHERE request_id='backup-purged'")
|
||||
.execute(&source)
|
||||
.await
|
||||
.unwrap();
|
||||
// Simulate the retention worker dropping a contribution whose raw request is gone.
|
||||
sqlx::query("DELETE FROM dashboard_request_contributions WHERE request_id='backup-purged'")
|
||||
.execute(&source)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(total(&source).await, 2);
|
||||
assert_eq!(tokens(&source).await, 40);
|
||||
let expected = dashboard_rows(&source).await;
|
||||
let export = export_postgres_core_jsonl(&source, 1_800_000_000)
|
||||
.await
|
||||
.unwrap();
|
||||
let plan = build_import_plan(&export).unwrap();
|
||||
assert!(plan.manifest.dashboard_snapshot.is_some());
|
||||
assert!(!plan
|
||||
.rows(ExportDomain::Auxiliary)
|
||||
.iter()
|
||||
.any(|row| row.payload["__table"] == "dashboard_stats_pending"));
|
||||
sqlx::query("CREATE DATABASE dashboard_restore_test")
|
||||
.execute(&source)
|
||||
.await
|
||||
.unwrap();
|
||||
let target_url = server
|
||||
.database_url()
|
||||
.strip_suffix("/postgres")
|
||||
.unwrap()
|
||||
.to_owned()
|
||||
+ "/dashboard_restore_test";
|
||||
let target = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&target_url)
|
||||
.await
|
||||
.unwrap();
|
||||
migrate(&target).await;
|
||||
// The gateway initializes one verified local admin and a zero-value unlimited
|
||||
// wallet before the operator can run restore. Preserve that account while
|
||||
// replacing its one initialization event with the source statistics.
|
||||
sqlx::raw_sql("INSERT INTO users(id,username,email_verified,role,auth_source,is_active) VALUES('bootstrap-admin','custom-root-name',true,'admin','local',true); INSERT INTO wallets(id,user_id,limit_mode,currency,status,created_at,updated_at) VALUES('bootstrap-wallet','bootstrap-admin','unlimited','USD','active',clock_timestamp(),clock_timestamp());")
|
||||
.execute(&target).await.unwrap();
|
||||
sqlx::query("UPDATE wallets SET balance=1,total_recharged=1 WHERE id='bootstrap-wallet'")
|
||||
.execute(&target)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = import_postgres_jsonl(&target, &export).await.unwrap_err();
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("conflicts with existing statistics"),
|
||||
"{error}"
|
||||
);
|
||||
sqlx::query("UPDATE wallets SET balance=0,total_recharged=0 WHERE id='bootstrap-wallet'")
|
||||
.execute(&target)
|
||||
.await
|
||||
.unwrap();
|
||||
import_postgres_jsonl(&target, &export).await.unwrap();
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM users WHERE id='bootstrap-admin'")
|
||||
.fetch_one(&target)
|
||||
.await
|
||||
.unwrap(),
|
||||
1,
|
||||
"the initialized admin account is not removed"
|
||||
);
|
||||
assert_eq!(dashboard_rows(&target).await,expected,"restore includes exact activation timestamp, narrow history, user events and purged request counts");
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM usage")
|
||||
.fetch_one(&target)
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
import_postgres_jsonl(&target, &export).await.unwrap();
|
||||
assert_eq!(
|
||||
dashboard_rows(&target).await,
|
||||
expected,
|
||||
"repeated import must not count users or requests twice"
|
||||
);
|
||||
let mut invalid_source = decode_jsonl(&export).unwrap();
|
||||
for record in &mut invalid_source {
|
||||
if let DataExportRecord::Row {
|
||||
domain: ExportDomain::Users,
|
||||
payload,
|
||||
..
|
||||
} = record
|
||||
{
|
||||
payload["auth_source"] = Value::String("invalid-auth-source".into());
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
import_postgres_jsonl(&target, &encode_jsonl(&invalid_source).unwrap())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard_rows(&target).await,
|
||||
expected,
|
||||
"failure after the restore guard and deletes rolls the entire transaction back"
|
||||
);
|
||||
let guard: Option<String> =
|
||||
sqlx::query_scalar("SELECT current_setting('aether.dashboard_restore',true)")
|
||||
.fetch_one(&target)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(guard.as_deref(), Some("on"));
|
||||
sqlx::query("UPDATE usage SET input_tokens=17,total_tokens=20 WHERE request_id='backup-kept'")
|
||||
.execute(&target)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
total(&target).await,
|
||||
2,
|
||||
"existing restored contribution updates by delta"
|
||||
);
|
||||
assert_eq!(tokens(&target).await, 46);
|
||||
request(&target, "backup-next", 31).await;
|
||||
assert_eq!(
|
||||
total(&target).await,
|
||||
3,
|
||||
"future writes resume ordinary aggregation"
|
||||
);
|
||||
assert_eq!(tokens(&target).await, 80);
|
||||
let before_conflict = dashboard_rows(&target).await;
|
||||
let error = import_postgres_jsonl(&target, &export).await.unwrap_err();
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("conflicts with existing statistics"),
|
||||
"{error}"
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard_rows(&target).await,
|
||||
before_conflict,
|
||||
"conflicting snapshots must roll back without overwriting accumulated data"
|
||||
);
|
||||
let mut truncated = decode_jsonl(&export).unwrap();
|
||||
truncated.retain(|row| !matches!(row,DataExportRecord::Row { payload,.. } if payload["__table"]=="dashboard_stats_total" && payload["shard"]==serde_json::json!(1)));
|
||||
let error = import_postgres_jsonl(&target, &encode_jsonl(&truncated).unwrap())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("missing, truncated"), "{error}");
|
||||
// A legacy file has neither aggregate rows nor the optional snapshot manifest.
|
||||
// Its ordinary source import must still trigger aggregation.
|
||||
let mut legacy = decode_jsonl(&export).unwrap();
|
||||
if let DataExportRecord::Manifest { manifest } = &mut legacy[0] {
|
||||
manifest.dashboard_snapshot = None;
|
||||
}
|
||||
legacy.retain(|row| !matches!(row,DataExportRecord::Row { payload,.. } if payload["__table"].as_str().is_some_and(|name| name.starts_with("dashboard_"))));
|
||||
sqlx::query("CREATE DATABASE dashboard_legacy_restore_test")
|
||||
.execute(&source)
|
||||
.await
|
||||
.unwrap();
|
||||
let legacy_url = server
|
||||
.database_url()
|
||||
.strip_suffix("/postgres")
|
||||
.unwrap()
|
||||
.to_owned()
|
||||
+ "/dashboard_legacy_restore_test";
|
||||
let legacy_target = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.connect(&legacy_url)
|
||||
.await
|
||||
.unwrap();
|
||||
migrate(&legacy_target).await;
|
||||
sqlx::query("UPDATE dashboard_stats_state SET stats_since='2000-01-01'")
|
||||
.execute(&legacy_target)
|
||||
.await
|
||||
.unwrap();
|
||||
let legacy_export = encode_jsonl(&legacy).unwrap();
|
||||
import_postgres_jsonl(&legacy_target, &legacy_export)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(total(&legacy_target).await, 1);
|
||||
import_postgres_jsonl(&legacy_target, &legacy_export)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(total(&legacy_target).await, 1);
|
||||
legacy_target.close().await;
|
||||
target.close().await;
|
||||
source.close().await;
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
use super::*;
|
||||
|
||||
mod dashboard_snapshot;
|
||||
|
||||
pub async fn export_postgres_core_jsonl(
|
||||
pool: &crate::driver::postgres::PostgresPool,
|
||||
created_at_unix_secs: u64,
|
||||
@@ -17,6 +19,10 @@ pub async fn export_postgres_jsonl(
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("SET LOCAL TIME ZONE 'UTC'")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let manifest = DataExportManifest::new(
|
||||
created_at_unix_secs,
|
||||
Some(DatabaseDriver::Postgres),
|
||||
@@ -51,6 +57,7 @@ pub async fn export_postgres_jsonl(
|
||||
}
|
||||
}
|
||||
|
||||
dashboard_snapshot::attach_manifest(&mut records)?;
|
||||
tx.commit().await.map_sql_err()?;
|
||||
encode_jsonl(&records)
|
||||
}
|
||||
@@ -85,6 +92,11 @@ async fn import_postgres_plan_with_options(
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let identity_scope = IdentityImportScope::from_plan(plan)?;
|
||||
let mut tx = pool.begin().await.map_sql_err()?;
|
||||
sqlx::query("SET LOCAL TIME ZONE 'UTC'")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
dashboard_snapshot::prepare_restore(&mut tx, plan).await?;
|
||||
let identity_state = capture_postgres_identity_import_state(&mut tx, &identity_scope).await?;
|
||||
let mut imported = 0usize;
|
||||
let mut column_cache = BTreeMap::<String, PostgresImportColumns>::new();
|
||||
|
||||
@@ -0,0 +1,289 @@
|
||||
use super::*;
|
||||
|
||||
fn tables() -> impl Iterator<Item = &'static AuxiliaryTable> {
|
||||
AUXILIARY_TABLES
|
||||
.iter()
|
||||
.filter(|table| table.name.starts_with("dashboard_"))
|
||||
}
|
||||
|
||||
fn fingerprint(rows: &[Value]) -> DashboardSnapshotTable {
|
||||
// Sorting canonical JSON makes fingerprints independent of file row order.
|
||||
let mut encoded = rows
|
||||
.iter()
|
||||
.map(|row| {
|
||||
let mut canonical = row.clone();
|
||||
canonical.sort_all_objects();
|
||||
canonical.to_string()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
encoded.sort();
|
||||
let mut digest = Sha256::new();
|
||||
for row in encoded {
|
||||
digest.update(row.as_bytes());
|
||||
digest.update(b"\n");
|
||||
}
|
||||
DashboardSnapshotTable {
|
||||
rows: rows.len(),
|
||||
sha256: format!("{:x}", digest.finalize()),
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot_rows<'a>(
|
||||
rows: impl Iterator<Item = &'a ExportRow>,
|
||||
) -> Result<BTreeMap<String, Vec<Value>>, DataLayerError> {
|
||||
let mut result = tables()
|
||||
.map(|table| (table.name.to_owned(), Vec::new()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for row in rows {
|
||||
let (name, payload) = domain_payload_table(row, "auxiliary", None)?;
|
||||
if let Some(entries) = result.get_mut(&name) {
|
||||
entries.push(payload);
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(super) fn attach_manifest(records: &mut [DataExportRecord]) -> Result<(), DataLayerError> {
|
||||
let Some(DataExportRecord::Manifest { manifest }) = records.first() else {
|
||||
return Ok(());
|
||||
};
|
||||
if !manifest.domains.contains(&ExportDomain::Auxiliary) {
|
||||
return Ok(());
|
||||
}
|
||||
let rows = records
|
||||
.iter()
|
||||
.filter_map(|record| match record {
|
||||
DataExportRecord::Row {
|
||||
domain: ExportDomain::Auxiliary,
|
||||
id,
|
||||
payload,
|
||||
} => Some(ExportRow {
|
||||
id: id.clone(),
|
||||
payload: payload.clone(),
|
||||
}),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let values = snapshot_rows(rows.iter())?;
|
||||
let tables = values
|
||||
.into_iter()
|
||||
.map(|(name, values)| (name, fingerprint(&values)))
|
||||
.collect();
|
||||
if let Some(DataExportRecord::Manifest { manifest }) = records.first_mut() {
|
||||
manifest.dashboard_snapshot = Some(DashboardSnapshotManifest { version: 1, tables });
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn invalid(detail: &str) -> DataLayerError {
|
||||
DataLayerError::InvalidInput(format!("dashboard snapshot {detail}"))
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_restore(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
plan: &DataImportPlan,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let values = snapshot_rows(plan.rows(ExportDomain::Auxiliary).iter())?;
|
||||
let Some(manifest) = &plan.manifest.dashboard_snapshot else {
|
||||
if values.values().any(|rows| !rows.is_empty()) {
|
||||
return Err(invalid(
|
||||
"requires its complete manifest; partial aggregate imports cannot be merged",
|
||||
));
|
||||
}
|
||||
return Ok(()); // Legacy backups deliberately keep normal trigger behavior.
|
||||
};
|
||||
if manifest.version != 1
|
||||
|| !plan.imports_domain(ExportDomain::Auxiliary)
|
||||
|| manifest.tables.len() != values.len()
|
||||
{
|
||||
return Err(invalid(
|
||||
"has an unsupported version or incomplete table inventory",
|
||||
));
|
||||
}
|
||||
for (name, rows) in &values {
|
||||
if manifest.tables.get(name) != Some(&fingerprint(rows)) {
|
||||
return Err(invalid(&format!(
|
||||
"table '{name}' is missing, truncated, or has changed"
|
||||
)));
|
||||
}
|
||||
}
|
||||
let state = &values["dashboard_stats_state"];
|
||||
let totals = &values["dashboard_stats_total"];
|
||||
let shards = totals
|
||||
.iter()
|
||||
.filter_map(|row| row.get("shard").and_then(Value::as_u64))
|
||||
.collect::<BTreeSet<_>>();
|
||||
if state.len() != 1
|
||||
|| state[0].get("singleton") != Some(&Value::Bool(true))
|
||||
|| state[0]
|
||||
.get("stats_since")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(parse_imported_datetime)
|
||||
.is_none()
|
||||
|| totals.len() != 16
|
||||
|| shards != (0..16).collect()
|
||||
{
|
||||
return Err(invalid(
|
||||
"must include one activation state and all 16 total shards",
|
||||
));
|
||||
}
|
||||
validate_request_counts(&values)?;
|
||||
|
||||
// Exclude concurrent source writes before examining or replacing projections.
|
||||
// Their triggers acquire projection locks in this same source-first order.
|
||||
sqlx::query("LOCK TABLE public.users, public.usage, public.usage_settlement_snapshots, public.usage_attribution_snapshots IN SHARE ROW EXCLUSIVE MODE")
|
||||
.execute(&mut **tx).await.map_sql_err()?;
|
||||
// Match retention's shard -> minute -> activity -> actor -> event -> state -> ledger
|
||||
// order, otherwise a maintenance pass could deadlock against the restore.
|
||||
for table in [
|
||||
"dashboard_stats_total",
|
||||
"dashboard_stats_minute",
|
||||
"dashboard_activity_minute",
|
||||
"dashboard_actor_minute",
|
||||
"dashboard_user_events_minute",
|
||||
"dashboard_stats_state",
|
||||
"dashboard_request_contributions",
|
||||
"dashboard_activity_hour",
|
||||
] {
|
||||
sqlx::query(&format!(
|
||||
"LOCK TABLE public.{table} IN SHARE ROW EXCLUSIVE MODE"
|
||||
))
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
}
|
||||
let bootstrap_only = bootstrap_admin_only(tx).await?;
|
||||
let mut identical = true;
|
||||
let mut empty = true;
|
||||
for table in tables() {
|
||||
let mut current = sqlx::query_scalar::<_, Value>(&format!(
|
||||
"SELECT to_jsonb(t) FROM public.{} t",
|
||||
table.name
|
||||
))
|
||||
.fetch_all(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
// The bounded cleanup cursor is operational progress, not a change to
|
||||
// statistics; moving it alone must not make a repeated restore conflict.
|
||||
if table.name == "dashboard_stats_state" {
|
||||
for row in &mut current {
|
||||
if let Some(object) = row.as_object_mut() {
|
||||
if let Some(cursor) = state[0].get("contributions_cleanup_cursor") {
|
||||
object.insert("contributions_cleanup_cursor".into(), cursor.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
identical &= manifest.tables.get(table.name) == Some(&fingerprint(¤t));
|
||||
empty &= match table.name {
|
||||
"dashboard_stats_state" => true, // A freshly migrated database already has an activation timestamp.
|
||||
"dashboard_stats_total" => current.iter().all(|row| {
|
||||
row.get("metrics")
|
||||
.and_then(Value::as_object)
|
||||
.is_some_and(|metrics| metrics.values().all(|n| n.as_f64() == Some(0.0)))
|
||||
}),
|
||||
"dashboard_user_events_minute" => current.is_empty() || bootstrap_only,
|
||||
_ => current.is_empty(),
|
||||
};
|
||||
}
|
||||
if !empty && !identical {
|
||||
return Err(invalid("conflicts with existing statistics; restore into an empty database. Complete statistics cannot be incrementally merged"));
|
||||
}
|
||||
if empty && !identical {
|
||||
let since =
|
||||
parse_imported_datetime(state[0]["stats_since"].as_str().expect("validated state"))
|
||||
.expect("validated timestamp");
|
||||
if sqlx::query_scalar::<_, bool>(
|
||||
"SELECT EXISTS(SELECT 1 FROM public.usage WHERE created_at >= $1)",
|
||||
)
|
||||
.bind(since)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
{
|
||||
return Err(invalid("conflicts with existing requests within the restored activation period; restore into an empty database"));
|
||||
}
|
||||
}
|
||||
// A local custom setting affects only dashboard triggers, not integrity or
|
||||
// billing triggers, and is automatically reverted on both commit and rollback.
|
||||
sqlx::query("SET LOCAL aether.dashboard_restore = 'on'")
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
for table in tables() {
|
||||
sqlx::query(&format!("DELETE FROM public.{}", table.name))
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn bootstrap_admin_only(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
// bootstrap_admin_from_env has no persistent marker and permits a configured
|
||||
// username. Recognize only its otherwise untouched single-admin/zero-wallet
|
||||
// state, never an installation with financial or request history.
|
||||
sqlx::query("LOCK TABLE public.api_keys, public.wallets, public.payment_orders, public.wallet_transactions, public.provider_expenses IN SHARE ROW EXCLUSIVE MODE")
|
||||
.execute(&mut **tx).await.map_sql_err()?;
|
||||
sqlx::query_scalar(r#"
|
||||
SELECT (SELECT count(*) FROM users)=1
|
||||
AND (SELECT count(*) FROM wallets)=1
|
||||
AND (SELECT count(*) FROM dashboard_user_events_minute)=1
|
||||
AND NOT EXISTS(SELECT 1 FROM usage)
|
||||
AND NOT EXISTS(SELECT 1 FROM api_keys)
|
||||
AND NOT EXISTS(SELECT 1 FROM payment_orders)
|
||||
AND NOT EXISTS(SELECT 1 FROM wallet_transactions)
|
||||
AND NOT EXISTS(SELECT 1 FROM provider_expenses)
|
||||
AND EXISTS(
|
||||
SELECT 1 FROM users u JOIN wallets w ON w.user_id=u.id
|
||||
JOIN dashboard_user_events_minute e
|
||||
ON e.bucket_start=date_trunc('minute',u.created_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'
|
||||
AND e.shard=(hashtextextended(u.id,0)&15)::smallint
|
||||
JOIN dashboard_stats_state s ON s.singleton
|
||||
WHERE u.role='admin' AND u.auth_source='local' AND u.is_active
|
||||
AND NOT u.is_deleted AND u.email_verified AND u.created_at>=s.stats_since
|
||||
AND w.api_key_id IS NULL AND w.limit_mode='unlimited'
|
||||
AND w.currency='USD' AND w.status='active'
|
||||
AND w.balance=0 AND w.gift_balance=0 AND w.total_recharged=0
|
||||
AND w.total_consumed=0 AND w.total_refunded=0 AND w.total_adjusted=0
|
||||
AND e.created_count=1 AND e.deleted_count=0
|
||||
)
|
||||
"#).fetch_one(&mut **tx).await.map_sql_err()
|
||||
}
|
||||
|
||||
fn validate_request_counts(values: &BTreeMap<String, Vec<Value>>) -> Result<(), DataLayerError> {
|
||||
let mut totals = [0u64; 16];
|
||||
for row in &values["dashboard_stats_total"] {
|
||||
let shard = row["shard"]
|
||||
.as_u64()
|
||||
.ok_or_else(|| invalid("has an invalid shard"))? as usize;
|
||||
totals[shard] = request_count(&row["metrics"]["request_count"])?;
|
||||
}
|
||||
let mut hours = [0u64; 16];
|
||||
for row in &values["dashboard_activity_hour"] {
|
||||
let shard = row["shard"]
|
||||
.as_u64()
|
||||
.filter(|n| *n < 16)
|
||||
.ok_or_else(|| invalid("has an invalid activity shard"))? as usize;
|
||||
hours[shard] = hours[shard]
|
||||
.checked_add(request_count(&row["request_count"])?)
|
||||
.ok_or_else(|| invalid("activity counts overflow"))?;
|
||||
}
|
||||
if hours != totals {
|
||||
return Err(invalid(
|
||||
"hourly activity and cumulative request counts disagree",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn request_count(value: &Value) -> Result<u64, DataLayerError> {
|
||||
if value.is_null() {
|
||||
return Ok(0);
|
||||
}
|
||||
value
|
||||
.as_u64()
|
||||
.ok_or_else(|| invalid("contains an invalid request count"))
|
||||
}
|
||||
@@ -28,6 +28,12 @@ use crate::lifecycle::bootstrap::postgres::{
|
||||
};
|
||||
|
||||
mod policy_nulls;
|
||||
mod overview_dirty_events;
|
||||
mod provider_expenses;
|
||||
mod migration_deadlines;
|
||||
mod overview_migration_safety;
|
||||
mod legacy_overview_upgrade;
|
||||
mod dashboard_user_anonymization;
|
||||
|
||||
/// A clean PostgreSQL database is bootstrapped from the schema snapshot first;
|
||||
/// migrations after the privacy/security frontier are intentionally left
|
||||
@@ -1575,6 +1581,18 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260901000000,
|
||||
20260903000000,
|
||||
20260908000000,
|
||||
20260911000000,
|
||||
20260917000000,
|
||||
20260917000100,
|
||||
20260918000000,
|
||||
20260918000100,
|
||||
20260919000000,
|
||||
20260920000000,
|
||||
20260920120000,
|
||||
20260921010000,
|
||||
20260921020000,
|
||||
20260921020100,
|
||||
20261001000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -1875,6 +1893,169 @@ WHERE id = 'metadata-migration-key'
|
||||
.expect("provider migration fixture should clean up");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_migrations_support_fresh_and_legacy_install() {
|
||||
for legacy in [false, true] {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
.await
|
||||
.expect("overview PostgreSQL should start or skip")
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if legacy {
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
connection.ensure_migrations_table().await.unwrap();
|
||||
for migration in POSTGRES_MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.version < 20260911000000)
|
||||
{
|
||||
connection.apply(migration).await.unwrap();
|
||||
}
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
INSERT INTO users(id, username, email_verified)
|
||||
VALUES ('overview-legacy-owner', 'overview-legacy-owner', false);
|
||||
INSERT INTO api_keys(id, user_id, key_hash)
|
||||
VALUES ('overview-legacy-key', 'overview-legacy-owner', repeat('e', 64));
|
||||
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
|
||||
status, billing_status, created_at)
|
||||
VALUES ('overview-legacy-request', 'overview-legacy-request',
|
||||
'overview-legacy-owner', 'overview-legacy-key', 'test', 'test',
|
||||
'completed', 'settled', '2020-01-01 12:34:00+00');
|
||||
"#,
|
||||
)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
prepare_and_apply_clean_postgres_database(&pool).await;
|
||||
for table in [
|
||||
"usage_attribution_snapshots",
|
||||
"stats_bucket_state",
|
||||
"stats_overview_hourly",
|
||||
"stats_overview_daily",
|
||||
] {
|
||||
assert!(table_exists(&pool, table).await.unwrap());
|
||||
}
|
||||
assert!(!column_exists(&pool, "api_keys", "credential_kind")
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(
|
||||
!column_exists(&pool, "usage_attribution_snapshots", "credential_kind")
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert!(column_exists(&pool, "stats_bucket_state", "last_failed_at")
|
||||
.await
|
||||
.unwrap());
|
||||
let precision:(i32,i32)=sqlx::query_as("SELECT numeric_precision::integer,numeric_scale::integer FROM information_schema.columns WHERE table_schema='public' AND table_name='usage_settlement_snapshots' AND column_name='wallet_debit_amount_usd'").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(precision, (20, 8));
|
||||
if legacy {
|
||||
let facts:(Option<String>,Option<String>,String,Option<String>)=sqlx::query_as("SELECT actor_user_id,credential_owner_id,attribution_source,wallet_debit_amount::text FROM usage_analytics_facts_v1 WHERE request_id='overview-legacy-request'").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(
|
||||
facts,
|
||||
(
|
||||
Some("overview-legacy-owner".into()),
|
||||
Some("overview-legacy-owner".into()),
|
||||
"user_account".into(),
|
||||
None
|
||||
)
|
||||
);
|
||||
}
|
||||
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
|
||||
let target = chrono::DateTime::parse_from_rfc3339("2026-09-17T00:00:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&chrono::Utc);
|
||||
for _ in 0..2 {
|
||||
let counts: (i64, i64, i64, i64) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT (SELECT COUNT(*) FROM usage_attribution_snapshots),
|
||||
(SELECT COUNT(*) FROM stats_bucket_state),
|
||||
(SELECT COUNT(*) FROM stats_overview_hourly),
|
||||
(SELECT COUNT(*) FROM stats_overview_daily)
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts, (0, 0, 0, 0));
|
||||
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 0);
|
||||
super::run_migrations(&pool).await.unwrap();
|
||||
}
|
||||
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
INSERT INTO users(id, username, email_verified)
|
||||
VALUES ('overview-new-owner', 'overview-new-owner', false);
|
||||
INSERT INTO api_keys(id, user_id, key_hash)
|
||||
VALUES ('overview-new-key', 'overview-new-owner', repeat('f', 64));
|
||||
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
|
||||
status, billing_status, created_at, request_metadata)
|
||||
VALUES ('overview-new-request', 'overview-new-request',
|
||||
'overview-new-owner', 'overview-new-key', 'test', 'test',
|
||||
'completed', 'settled', '2026-09-15 12:34:00+00',
|
||||
'{"analytics_attribution":{"is_standalone":false}}'::json);
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let identity: (String, String, String, String) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT actor_user_id, credential_owner_id, attribution_kind,
|
||||
attribution_source
|
||||
FROM usage_attribution_snapshots WHERE request_id = 'overview-new-request'
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
identity,
|
||||
(
|
||||
"overview-new-owner".into(),
|
||||
"overview-new-owner".into(),
|
||||
"employee".into(),
|
||||
"user_account".into(),
|
||||
)
|
||||
);
|
||||
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 2);
|
||||
assert_eq!(repo.rebuild_overview_buckets(target, 8).await.unwrap(), 0);
|
||||
let counts: (i64, i64, i64, i64, i64) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT (SELECT COUNT(*) FROM usage_attribution_snapshots),
|
||||
(SELECT COUNT(*) FROM usage_attribution_snapshots
|
||||
WHERE request_id = 'overview-legacy-request'),
|
||||
(SELECT COUNT(*) FROM stats_bucket_state),
|
||||
(SELECT COUNT(*) FROM stats_overview_hourly),
|
||||
(SELECT COUNT(*) FROM stats_overview_daily)
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts, (1, 0, 2, 1, 1));
|
||||
let unexpected_bucket_count: i64 = query_scalar(
|
||||
r#"
|
||||
SELECT COUNT(*) FROM stats_bucket_state
|
||||
WHERE projection_version <> 'overview-v2'
|
||||
OR (granularity, bucket_start) NOT IN (
|
||||
('day', '2026-09-15 00:00:00+00'::timestamptz),
|
||||
('hour', '2026-09-15 12:00:00+00'::timestamptz)
|
||||
)
|
||||
OR coverage_status <> 'complete'
|
||||
OR source_revision <> built_revision
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(unexpected_bucket_count, 0);
|
||||
pool.close().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_database_for_startup_bootstraps_clean_database() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
@@ -2810,9 +2991,7 @@ WHERE request_id = 'billing-facts-cache-create'
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
|
||||
const MIGRATION_VERSION: i64 = 20260715000000;
|
||||
|
||||
async fn postgres_migrations_repair_invalid_concurrent_indexes() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
.await
|
||||
.expect("postgres migration retry test should start or skip")
|
||||
@@ -2824,11 +3003,6 @@ async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
|
||||
.await
|
||||
.expect("pool should connect");
|
||||
prepare_and_apply_clean_postgres_database(&pool).await;
|
||||
|
||||
query("DROP INDEX CONCURRENTLY public.idx_usage_legacy_body_ref_cleanup_created_at")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("snapshot cleanup index should exist");
|
||||
query("CREATE TABLE public.concurrent_index_failure_fixture (value integer NOT NULL)")
|
||||
.execute(&pool)
|
||||
.await
|
||||
@@ -2838,62 +3012,89 @@ async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
|
||||
.await
|
||||
.expect("duplicate failure fixtures should be inserted");
|
||||
|
||||
query(
|
||||
"CREATE UNIQUE INDEX CONCURRENTLY idx_usage_legacy_body_ref_cleanup_created_at ON public.concurrent_index_failure_fixture (value)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect_err("duplicate values should leave a failed concurrent index build");
|
||||
|
||||
let invalid_index_exists: bool = query_scalar(
|
||||
r#"
|
||||
for (migration_version, index_name, table_name) in [
|
||||
(
|
||||
20260715000000_i64,
|
||||
"idx_usage_legacy_body_ref_cleanup_created_at",
|
||||
"public.usage",
|
||||
),
|
||||
(
|
||||
20260918000000,
|
||||
"idx_usage_settlement_dashboard_cover_v2",
|
||||
"public.usage_settlement_snapshots",
|
||||
),
|
||||
(
|
||||
20260920000000,
|
||||
"idx_payment_orders_status_credited_user",
|
||||
"public.payment_orders",
|
||||
),
|
||||
(
|
||||
20260921020100,
|
||||
"ix_usage_analytics_actor_metadata",
|
||||
"public.usage",
|
||||
),
|
||||
] {
|
||||
query(&format!("DROP INDEX CONCURRENTLY public.{index_name}"))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("migrated index should exist");
|
||||
query(&format!("CREATE UNIQUE INDEX CONCURRENTLY {index_name} ON public.concurrent_index_failure_fixture (value)"))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect_err("duplicate values should leave a failed concurrent index build");
|
||||
let invalid_index_exists: bool = query_scalar(
|
||||
r#"
|
||||
SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM pg_catalog.pg_class AS index_relation
|
||||
JOIN pg_catalog.pg_namespace AS index_namespace
|
||||
ON index_namespace.oid = index_relation.relnamespace
|
||||
JOIN pg_catalog.pg_index AS index_state
|
||||
ON index_state.indexrelid = index_relation.oid
|
||||
WHERE index_namespace.nspname = 'public'
|
||||
AND index_relation.relname = 'idx_usage_legacy_body_ref_cleanup_created_at'
|
||||
AND NOT index_state.indisvalid
|
||||
)
|
||||
"#,
|
||||
SELECT 1 FROM pg_catalog.pg_index
|
||||
WHERE indexrelid = to_regclass($1) AND NOT indisvalid
|
||||
)"#,
|
||||
)
|
||||
.bind(format!("public.{index_name}"))
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("failed concurrent index state should be readable");
|
||||
assert!(invalid_index_exists);
|
||||
|
||||
query("DELETE FROM public._sqlx_migrations WHERE version = $1")
|
||||
.bind(migration_version)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("index migration stamp should be reset");
|
||||
super::run_migrations(&pool)
|
||||
.await
|
||||
.expect("migration retry should replace the invalid index");
|
||||
let valid_index_exists: bool = query_scalar(
|
||||
r#"
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM pg_catalog.pg_index
|
||||
WHERE indexrelid = to_regclass($1) AND indrelid = $2::regclass AND indisvalid
|
||||
)"#,
|
||||
)
|
||||
.bind(format!("public.{index_name}"))
|
||||
.bind(table_name)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("rebuilt index state should be readable");
|
||||
assert!(valid_index_exists);
|
||||
}
|
||||
let index_definition: String = query_scalar(
|
||||
"SELECT pg_get_indexdef('public.idx_usage_settlement_dashboard_cover_v2'::regclass)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("failed concurrent index state should be readable");
|
||||
assert!(invalid_index_exists);
|
||||
|
||||
query("DELETE FROM public._sqlx_migrations WHERE version = $1")
|
||||
.bind(MIGRATION_VERSION)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("cleanup index migration stamp should be reset");
|
||||
super::run_migrations(&pool)
|
||||
.await
|
||||
.expect("migration retry should replace the invalid index");
|
||||
|
||||
let valid_usage_index_exists: bool = query_scalar(
|
||||
r#"
|
||||
SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM pg_catalog.pg_class AS index_relation
|
||||
JOIN pg_catalog.pg_namespace AS index_namespace
|
||||
ON index_namespace.oid = index_relation.relnamespace
|
||||
JOIN pg_catalog.pg_index AS index_state
|
||||
ON index_state.indexrelid = index_relation.oid
|
||||
WHERE index_namespace.nspname = 'public'
|
||||
AND index_relation.relname = 'idx_usage_legacy_body_ref_cleanup_created_at'
|
||||
AND index_state.indrelid = 'public.usage'::regclass
|
||||
AND index_state.indisvalid
|
||||
)
|
||||
"#,
|
||||
.expect("replacement index should be installed");
|
||||
assert!(index_definition.contains("billing_status"));
|
||||
assert!(index_definition.contains("allocation_status"));
|
||||
let old_cover_exists: bool = query_scalar(
|
||||
"SELECT to_regclass('public.idx_usage_settlement_dashboard_cover') IS NOT NULL",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("rebuilt cleanup index state should be readable");
|
||||
assert!(valid_usage_index_exists);
|
||||
.expect("old covering index state should be readable");
|
||||
assert!(
|
||||
!old_cover_exists,
|
||||
"successful migration must retire the redundant old cover"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
+231
@@ -0,0 +1,231 @@
|
||||
use super::*;
|
||||
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
|
||||
|
||||
const ANONYMIZATION_VERSION: i64 = 20261001000000;
|
||||
const INSERT_USER: &str = "INSERT INTO users(id,username,email_verified) VALUES($1,$1,false)";
|
||||
const INSERT_USAGE: &str = "INSERT INTO usage(id,request_id,user_id,model,provider_name,status,billing_status,created_at) VALUES($1,$1,$2,'anonymization','test','completed','settled',clock_timestamp())";
|
||||
|
||||
async fn assert_anonymous(pool: &PgPool, user: &str) {
|
||||
let counts: (i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT count(*) FROM dashboard_actor_minute WHERE actor_user_id=$1), (SELECT count(*) FROM dashboard_request_contributions WHERE actor_user_id=$1)",
|
||||
)
|
||||
.bind(user)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts, (0, 0), "retained identity for {user}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_user_anonymization_preserves_totals_without_migration_backfill() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
connection.ensure_migrations_table().await.unwrap();
|
||||
for migration in POSTGRES_MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.version < ANONYMIZATION_VERSION)
|
||||
{
|
||||
connection.apply(migration).await.unwrap();
|
||||
}
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
query(INSERT_USER)
|
||||
.bind("legacy-deleted")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind("legacy-request")
|
||||
.bind("legacy-deleted")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
query("DELETE FROM usage WHERE request_id='legacy-request'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
query("DELETE FROM users WHERE id='legacy-deleted'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let actor_before: String =
|
||||
query_scalar("SELECT jsonb_agg(to_jsonb(a))::text FROM dashboard_actor_minute a")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// An upgrade must succeed even while historical projection tables are
|
||||
// inaccessible: install definitions without reading or rewriting their rows.
|
||||
let mut blocked_history = pool.begin().await.unwrap();
|
||||
query("LOCK TABLE dashboard_actor_minute, dashboard_request_contributions IN ACCESS EXCLUSIVE MODE")
|
||||
.execute(&mut *blocked_history).await.unwrap();
|
||||
query("SET lock_timeout='500ms'")
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
let migration = POSTGRES_MIGRATOR
|
||||
.iter()
|
||||
.find(|migration| migration.version == ANONYMIZATION_VERSION)
|
||||
.unwrap();
|
||||
connection.apply(migration).await.unwrap();
|
||||
blocked_history.rollback().await.unwrap();
|
||||
let actor_after: String =
|
||||
query_scalar("SELECT jsonb_agg(to_jsonb(a))::text FROM dashboard_actor_minute a")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
actor_before, actor_after,
|
||||
"migration must not rewrite old actors"
|
||||
);
|
||||
|
||||
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
|
||||
let summary_query = UsageDashboardAnalyticsQuery {
|
||||
timezone: "UTC".into(),
|
||||
};
|
||||
let legacy = repo.query_dashboard_summary(&summary_query).await.unwrap();
|
||||
assert_eq!(
|
||||
legacy.today.active_users, 0,
|
||||
"old orphan actors must be excluded without backfill"
|
||||
);
|
||||
assert_eq!(legacy.total.request_count, 1);
|
||||
|
||||
for (user, purge, soft) in [
|
||||
("hard-live", false, false),
|
||||
("hard-purged", true, false),
|
||||
("soft-live", false, true),
|
||||
("soft-purged", true, true),
|
||||
] {
|
||||
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind(user)
|
||||
.bind(user)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
if purge {
|
||||
query("DELETE FROM usage WHERE request_id=$1")
|
||||
.bind(user)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let before = repo.query_dashboard_summary(&summary_query).await.unwrap();
|
||||
let sql = if soft {
|
||||
"UPDATE users SET is_deleted=true WHERE id=$1"
|
||||
} else {
|
||||
"DELETE FROM users WHERE id=$1"
|
||||
};
|
||||
query(sql).bind(user).execute(&pool).await.unwrap();
|
||||
assert_anonymous(&pool, user).await;
|
||||
let after = repo.query_dashboard_summary(&summary_query).await.unwrap();
|
||||
assert_eq!(
|
||||
after.total, before.total,
|
||||
"deletion must preserve request totals"
|
||||
);
|
||||
assert_eq!(after.today.active_users, 0);
|
||||
if !purge {
|
||||
query("UPDATE usage SET response_time_ms=200 WHERE request_id=$1")
|
||||
.bind(user)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_anonymous(&pool, user).await;
|
||||
}
|
||||
if soft {
|
||||
query("DELETE FROM users WHERE id=$1")
|
||||
.bind(user)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_anonymous(&pool, user).await;
|
||||
}
|
||||
}
|
||||
|
||||
// Exercise both deferred-event orders, plus explicitly immediate constraint
|
||||
// triggers. A captured deleted_fact must never recreate a deleted actor.
|
||||
for (user, user_first, immediate) in [
|
||||
("deferred-request-first", false, false),
|
||||
("deferred-user-first", true, false),
|
||||
("immediate-user", true, true),
|
||||
] {
|
||||
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
|
||||
let mut tx = pool.begin().await.unwrap();
|
||||
if immediate {
|
||||
query("SET CONSTRAINTS ALL IMMEDIATE")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
query(INSERT_USAGE)
|
||||
.bind(user)
|
||||
.bind(user)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
if user_first {
|
||||
query("DELETE FROM users WHERE id=$1")
|
||||
.bind(user)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
query("DELETE FROM usage WHERE request_id=$1")
|
||||
.bind(user)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
if !user_first {
|
||||
query("DELETE FROM users WHERE id=$1")
|
||||
.bind(user)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
tx.commit().await.unwrap();
|
||||
assert_anonymous(&pool, user).await;
|
||||
}
|
||||
|
||||
// A concurrent writer can still see a user after its deletion statement
|
||||
// but before that transaction commits. Check both commit orders: the new
|
||||
// actor must be rejected after the shard lock, or removed by the deleter.
|
||||
for (user, writer_first) in [
|
||||
("concurrent-delete-first", false),
|
||||
("concurrent-writer-first", true),
|
||||
] {
|
||||
query(INSERT_USER).bind(user).execute(&pool).await.unwrap();
|
||||
let mut deletion = pool.begin().await.unwrap();
|
||||
query("DELETE FROM users WHERE id=$1")
|
||||
.bind(user)
|
||||
.execute(&mut *deletion)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut writer = pool.begin().await.unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind(user)
|
||||
.bind(user)
|
||||
.execute(&mut *writer)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
if writer_first {
|
||||
writer.commit().await.unwrap();
|
||||
deletion.commit().await.unwrap();
|
||||
} else {
|
||||
deletion.commit().await.unwrap();
|
||||
writer.commit().await.unwrap();
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("concurrent deletion and usage must not deadlock");
|
||||
assert_anonymous(&pool, user).await;
|
||||
}
|
||||
let final_summary = repo.query_dashboard_summary(&summary_query).await.unwrap();
|
||||
assert_eq!(final_summary.total.request_count, 10);
|
||||
assert_eq!(final_summary.today.active_users, 0);
|
||||
let invalid: (i64, i64, i64) = sqlx::query_as("SELECT (SELECT count(*) FROM dashboard_actor_minute WHERE request_count < 0), (SELECT count(*) FROM dashboard_stats_pending), (SELECT count(*) FROM dashboard_user_anonymization_pending)")
|
||||
.fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(invalid, (0, 0, 0));
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
use super::*;
|
||||
|
||||
const ACCOUNT_ATTRIBUTION: i64 = 20260917000000;
|
||||
const DIRTY_EVENTS: i64 = 20260917000100;
|
||||
|
||||
async fn rows_snapshot(pool: &PgPool, table: &str) -> String {
|
||||
query_scalar(&format!(
|
||||
"SELECT COALESCE(jsonb_agg(to_jsonb(t) ORDER BY to_jsonb(t)::text), '[]')::text FROM {table} t"
|
||||
))
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_overview_upgrade_preserves_applied_history_and_existing_statistics() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
connection.ensure_migrations_table().await.unwrap();
|
||||
for migration in POSTGRES_MIGRATOR.iter().filter(|migration| {
|
||||
migration.version <= 20260920120000 && migration.version != DIRTY_EVENTS
|
||||
}) {
|
||||
connection.apply(migration).await.unwrap();
|
||||
}
|
||||
connection.close().await.unwrap();
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
|
||||
// The restored September 17 migration installs the historical direct-write
|
||||
// trigger. Remove schema additions folded into the September 11 baseline
|
||||
// so this exercises an already-running database that never received them.
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
DROP TABLE public.stats_overview_dirty_events;
|
||||
DROP INDEX public.ix_usage_attribution_owner_request;
|
||||
UPDATE public.dashboard_stats_state SET stats_since = clock_timestamp() - INTERVAL '1 day';
|
||||
INSERT INTO users(id, username, email_verified)
|
||||
VALUES ('legacy-upgrade-owner', 'legacy-upgrade-owner', false);
|
||||
INSERT INTO api_keys(id, user_id, key_hash)
|
||||
VALUES ('legacy-upgrade-key', 'legacy-upgrade-owner', repeat('a', 64));
|
||||
INSERT INTO usage(id, request_id, user_id, api_key_id, provider_name, model,
|
||||
status, billing_status, created_at, response_time_ms)
|
||||
VALUES ('legacy-upgrade-request', 'legacy-upgrade-request',
|
||||
'legacy-upgrade-owner', 'legacy-upgrade-key', 'test', 'test',
|
||||
'completed', 'settled', clock_timestamp() - INTERVAL '1 minute', 100);
|
||||
INSERT INTO stats_overview_hourly(projection_version, bucket_start, dimensions, metrics)
|
||||
SELECT projection_version, bucket_start, '{}'::jsonb, '{"request_count":1}'::jsonb
|
||||
FROM stats_bucket_state WHERE granularity = 'hour';
|
||||
INSERT INTO stats_overview_daily(projection_version, bucket_start, dimensions, metrics)
|
||||
SELECT projection_version, bucket_start, '{}'::jsonb, '{"request_count":1}'::jsonb
|
||||
FROM stats_bucket_state WHERE granularity = 'day';
|
||||
UPDATE stats_bucket_state SET built_revision=source_revision, coverage_status='complete';
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Model the old September 19 schema after creating real dashboard totals.
|
||||
// No fact is changed while its later retention helpers are absent.
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
DROP FUNCTION public.dashboard_ensure_activity_minute(timestamptz, smallint);
|
||||
DROP TABLE public.dashboard_activity_minute;
|
||||
ALTER TABLE public.dashboard_stats_state DROP COLUMN contributions_cleanup_cursor;
|
||||
UPDATE _sqlx_migrations SET checksum=decode(
|
||||
'1dd622827b22ae0540f5e43232e7419727aa02d0742649af7e045185e9a13f68655ef7b31007bd0bf6e9e63177022815', 'hex')
|
||||
WHERE version=20260911000000;
|
||||
UPDATE _sqlx_migrations SET checksum=decode(
|
||||
'c64293c5d95fba6c3c43fff764a89da0225b386cfbef14743ab585e1db012fea35fe2cb2eee132e0fb5dac834c0410f4', 'hex')
|
||||
WHERE version=20260919000000;
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let historical_records = rows_snapshot(&pool, "_sqlx_migrations").await;
|
||||
let preserved_tables = [
|
||||
"users",
|
||||
"api_keys",
|
||||
"usage",
|
||||
"usage_attribution_snapshots",
|
||||
"stats_bucket_state",
|
||||
"stats_overview_hourly",
|
||||
"stats_overview_daily",
|
||||
"dashboard_request_contributions",
|
||||
"dashboard_stats_total",
|
||||
"dashboard_stats_minute",
|
||||
"dashboard_activity_hour",
|
||||
"dashboard_actor_minute",
|
||||
];
|
||||
let mut original_rows = Vec::new();
|
||||
for table in preserved_tables {
|
||||
original_rows.push((table, rows_snapshot(&pool, table).await));
|
||||
}
|
||||
let account_record: String =
|
||||
query_scalar("SELECT to_jsonb(m)::text FROM _sqlx_migrations m WHERE version=$1")
|
||||
.bind(ACCOUNT_ATTRIBUTION)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let pending = prepare_database_for_startup(&pool).await.unwrap();
|
||||
assert_eq!(
|
||||
pending
|
||||
.iter()
|
||||
.map(|migration| migration.version)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![DIRTY_EVENTS, 20260921010000, 20260921020000, 20260921020100, 20261001000000]
|
||||
);
|
||||
assert_eq!(
|
||||
rows_snapshot(&pool, "_sqlx_migrations").await,
|
||||
historical_records
|
||||
);
|
||||
|
||||
for _ in 0..2 {
|
||||
super::super::run_migrations(&pool).await.unwrap();
|
||||
assert!(super::super::pending_migrations(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
for (table, expected) in &original_rows {
|
||||
assert_eq!(&rows_snapshot(&pool, table).await, expected, "{table}");
|
||||
}
|
||||
let after: String =
|
||||
query_scalar("SELECT to_jsonb(m)::text FROM _sqlx_migrations m WHERE version=$1")
|
||||
.bind(ACCOUNT_ATTRIBUTION)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
after, account_record,
|
||||
"the historical migration must not be restamped"
|
||||
);
|
||||
}
|
||||
assert!(table_exists(&pool, "stats_overview_dirty_events")
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(table_exists(&pool, "dashboard_activity_minute")
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(column_exists(
|
||||
&pool,
|
||||
"dashboard_stats_state",
|
||||
"contributions_cleanup_cursor"
|
||||
)
|
||||
.await
|
||||
.unwrap());
|
||||
for index in [
|
||||
"ix_usage_attribution_owner_request",
|
||||
"ix_usage_analytics_actor_metadata",
|
||||
] {
|
||||
let valid: bool =
|
||||
query_scalar("SELECT indisvalid FROM pg_index WHERE indexrelid=to_regclass($1)")
|
||||
.bind(index)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(valid, "{index}");
|
||||
}
|
||||
let empty_projections: (i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT count(*) FROM stats_overview_dirty_events), (SELECT count(*) FROM dashboard_activity_minute)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(empty_projections, (0, 0), "upgrade must not backfill usage");
|
||||
|
||||
// Correcting an existing request must seed the retained activity counter
|
||||
// from its old detailed minute, without counting that request twice.
|
||||
query("UPDATE usage SET response_time_ms=200 WHERE request_id='legacy-upgrade-request'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let corrected: (i64, i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM dashboard_stats_total), (SELECT sum(request_count)::bigint FROM dashboard_activity_minute), (SELECT sum((metrics->>'response_sum_ms')::bigint)::bigint FROM dashboard_stats_total)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(corrected, (1, 1, 200));
|
||||
query(
|
||||
"INSERT INTO usage(id,request_id,user_id,api_key_id,provider_name,model,status,billing_status,created_at,response_time_ms) SELECT 'after-upgrade','after-upgrade',user_id,api_key_id,provider_name,model,status,billing_status,created_at,300 FROM usage WHERE request_id='legacy-upgrade-request'",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let written: (i64, i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM dashboard_stats_total), (SELECT sum(request_count)::bigint FROM dashboard_activity_minute), (SELECT count(*) FROM stats_overview_dirty_events)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(written, (2, 2, 4));
|
||||
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(chrono::Utc::now() + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let rebuilt: (i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT sum((metrics->>'request_count')::bigint)::bigint FROM stats_overview_hourly), (SELECT count(*) FROM stats_overview_dirty_events)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rebuilt, (2, 0));
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -0,0 +1,318 @@
|
||||
use super::*;
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
const OVERVIEW_MIGRATION: i64 = 20260911000000;
|
||||
|
||||
async fn legacy_connection(server: &ManagedPostgresServer) -> PgConnection {
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
connection.ensure_migrations_table().await.unwrap();
|
||||
for migration in POSTGRES_MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.version < OVERVIEW_MIGRATION)
|
||||
{
|
||||
connection.apply(migration).await.unwrap();
|
||||
}
|
||||
connection
|
||||
}
|
||||
|
||||
async fn wait_for_settlement_ddl(connection: &mut PgConnection) {
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
let waiting: bool = query_scalar(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_locks WHERE relation='public.usage_settlement_snapshots'::regclass AND mode='AccessExclusiveLock' AND NOT granted)",
|
||||
)
|
||||
.fetch_one(&mut *connection)
|
||||
.await
|
||||
.unwrap();
|
||||
if waiting {
|
||||
return;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("migration should reach the blocked settlement ALTER TABLE");
|
||||
}
|
||||
|
||||
async fn assert_overview_migration_rolled_back(pool: &PgPool) {
|
||||
assert!(
|
||||
!column_exists(pool, "usage", "failure_origin")
|
||||
.await
|
||||
.unwrap(),
|
||||
"the earlier ALTER TABLE must roll back with the blocked statement"
|
||||
);
|
||||
let stamped: bool =
|
||||
query_scalar("SELECT EXISTS (SELECT 1 FROM public._sqlx_migrations WHERE version=$1)")
|
||||
.bind(OVERVIEW_MIGRATION)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
!stamped,
|
||||
"a failed migration must not receive a success stamp"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn migration_deadlines_release_queued_usage_work_and_roll_back_before_retry() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut observer = legacy_connection(&server).await;
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.after_connect(|connection, _| {
|
||||
Box::pin(async move {
|
||||
query("SET statement_timeout='30s'")
|
||||
.execute(&mut *connection)
|
||||
.await?;
|
||||
query("SET lock_timeout='3s'")
|
||||
.execute(&mut *connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.connect(server.database_url())
|
||||
.await
|
||||
.unwrap();
|
||||
let original_pid: i32 = query_scalar("SELECT pg_backend_pid()")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut business = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
let mut business_transaction = business.begin().await.unwrap();
|
||||
query("LOCK TABLE public.usage_settlement_snapshots IN ACCESS SHARE MODE")
|
||||
.execute(&mut *business_transaction)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut queued_business = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
|
||||
// The usage ALTERs take their table exclusively, then settlement ALTER
|
||||
// waits behind an existing reader. Business writes wait on the usage lock
|
||||
// already held by this transaction and must resume when it rolls back.
|
||||
let started = Instant::now();
|
||||
let (migration_result, ()) = tokio::join!(super::super::run_migrations(&pool), async {
|
||||
wait_for_settlement_ddl(&mut observer).await;
|
||||
let acquired_first_lock: bool = query_scalar(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_locks WHERE relation='public.usage'::regclass AND mode='AccessExclusiveLock' AND granted)",
|
||||
)
|
||||
.fetch_one(&mut observer)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
acquired_first_lock,
|
||||
"the migration must have already changed the first table"
|
||||
);
|
||||
let queued_write = query("INSERT INTO public.usage(id,request_id,model,provider_name,status,billing_status,created_at) VALUES ('migration-live-request','migration-live-request','test','test','completed','settled',NOW())")
|
||||
.execute(&mut queued_business);
|
||||
tokio::pin!(queued_write);
|
||||
tokio::select! {
|
||||
result = &mut queued_write => panic!("the business write should initially queue behind the DDL: {result:?}"),
|
||||
() = tokio::time::sleep(Duration::from_millis(100)) => {}
|
||||
}
|
||||
let written = tokio::time::timeout(Duration::from_secs(4), queued_write)
|
||||
.await
|
||||
.expect("queued business work must resume after the migration's lock timeout")
|
||||
.unwrap();
|
||||
assert_eq!(written.rows_affected(), 1);
|
||||
});
|
||||
let error =
|
||||
migration_result.expect_err("busy settlement table must defer this upgrade attempt");
|
||||
assert!(
|
||||
error.to_string().contains("lock timeout"),
|
||||
"the failure should identify the bounded lock wait: {error}"
|
||||
);
|
||||
assert!(started.elapsed() < Duration::from_secs(5));
|
||||
assert_overview_migration_rolled_back(&pool).await;
|
||||
|
||||
let (new_pid, statement_timeout, lock_timeout): (i32, String, String) = sqlx::query_as(
|
||||
"SELECT pg_backend_pid(), current_setting('statement_timeout'), current_setting('lock_timeout')",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(
|
||||
new_pid, original_pid,
|
||||
"failed migration connection must be discarded"
|
||||
);
|
||||
assert_eq!(statement_timeout, "30s");
|
||||
assert_eq!(lock_timeout, "3s");
|
||||
|
||||
business_transaction.rollback().await.unwrap();
|
||||
super::super::run_migrations(&pool).await.unwrap();
|
||||
assert!(super::super::pending_migrations(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
let success: bool =
|
||||
query_scalar("SELECT success FROM public._sqlx_migrations WHERE version=$1")
|
||||
.bind(OVERVIEW_MIGRATION)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
success,
|
||||
"a later quiet retry must succeed without manual stamp repair"
|
||||
);
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>(
|
||||
"SELECT count(*) FROM public.usage WHERE request_id='migration-live-request'"
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
1,
|
||||
"the resumed business write must survive the upgrade retry"
|
||||
);
|
||||
let settings: (String, String) = sqlx::query_as(
|
||||
"SELECT current_setting('statement_timeout'), current_setting('lock_timeout')",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(settings, ("30s".into(), "3s".into()));
|
||||
|
||||
// Startup preparation also takes SQLx's advisory migration lock. Another
|
||||
// upgrade process must not leave startup waiting indefinitely for it.
|
||||
observer.lock().await.unwrap();
|
||||
let preparation =
|
||||
tokio::time::timeout(Duration::from_secs(4), prepare_database_for_startup(&pool))
|
||||
.await
|
||||
.expect("startup preparation must bound its advisory-lock wait");
|
||||
let preparation_error = preparation.expect_err("another migration owns the advisory lock");
|
||||
assert!(preparation_error.to_string().contains("lock timeout"));
|
||||
observer.unlock().await.unwrap();
|
||||
assert!(prepare_database_for_startup(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
pool.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn migration_deadlines_caller_cancellation_releases_ddl_locks_and_rolls_back() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = legacy_connection(&server).await;
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
CREATE FUNCTION public.test_cancelled_migration_ddl() RETURNS event_trigger LANGUAGE plpgsql AS $$
|
||||
BEGIN PERFORM pg_sleep(8); END $$;
|
||||
CREATE EVENT TRIGGER test_cancelled_migration_ddl ON ddl_command_end
|
||||
WHEN TAG IN ('ALTER TABLE') EXECUTE FUNCTION public.test_cancelled_migration_ddl();
|
||||
"#,
|
||||
)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(server.database_url())
|
||||
.await
|
||||
.unwrap();
|
||||
// Keep this future in an inner scope: leaving it actually drops the entire
|
||||
// migration operation, rather than merely dropping a pinned reference.
|
||||
{
|
||||
let migrate = super::super::run_migrations(&pool);
|
||||
tokio::pin!(migrate);
|
||||
tokio::select! {
|
||||
result = &mut migrate => panic!("the caller must cancel before the slow migration finishes: {result:?}"),
|
||||
() = async {
|
||||
tokio::time::timeout(Duration::from_secs(4), async {
|
||||
loop {
|
||||
let running: bool = query_scalar(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_locks l JOIN pg_stat_activity a USING(pid) WHERE l.relation='public.usage'::regclass AND l.mode='AccessExclusiveLock' AND l.granted AND a.wait_event='PgSleep')",
|
||||
)
|
||||
.fetch_one(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
if running {
|
||||
return;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
}).await.expect("migration must acquire its first DDL lock before caller cancellation");
|
||||
} => {}
|
||||
}
|
||||
}
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(4),
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM public.usage").fetch_one(&pool),
|
||||
)
|
||||
.await
|
||||
.expect("dropping the caller future must stop the server-side statement and release DDL locks")
|
||||
.unwrap();
|
||||
assert_overview_migration_rolled_back(&pool).await;
|
||||
sqlx::raw_sql(
|
||||
"DROP EVENT TRIGGER test_cancelled_migration_ddl; DROP FUNCTION public.test_cancelled_migration_ddl()",
|
||||
)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
super::super::run_migrations(&pool).await.unwrap();
|
||||
assert!(super::super::pending_migrations(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
pool.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn migration_deadlines_bound_the_whole_transaction_not_only_each_statement() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = legacy_connection(&server).await;
|
||||
// Each ALTER finishes within the ten-second statement limit, but the
|
||||
// entire migration exceeds it. A per-statement timeout alone is insufficient.
|
||||
sqlx::raw_sql(
|
||||
r#"
|
||||
CREATE FUNCTION public.test_slow_migration_ddl() RETURNS event_trigger LANGUAGE plpgsql AS $$
|
||||
BEGIN PERFORM pg_sleep(3); END $$;
|
||||
CREATE EVENT TRIGGER test_slow_migration_ddl ON ddl_command_end
|
||||
WHEN TAG IN ('ALTER TABLE') EXECUTE FUNCTION public.test_slow_migration_ddl();
|
||||
"#,
|
||||
)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(server.database_url())
|
||||
.await
|
||||
.unwrap();
|
||||
let started = Instant::now();
|
||||
let result = tokio::time::timeout(Duration::from_secs(15), super::super::run_migrations(&pool))
|
||||
.await
|
||||
.expect("the whole migration deadline must fire before all slow statements finish");
|
||||
assert!(result.is_err(), "the over-budget migration must fail");
|
||||
assert!(
|
||||
started.elapsed() >= Duration::from_secs(9),
|
||||
"upgrade failed before its deadline: {result:?}"
|
||||
);
|
||||
assert!(started.elapsed() < Duration::from_secs(15));
|
||||
// The backend can still be unwinding the statement when its socket closes.
|
||||
// Reading the first locked table proves cancellation released the DDL lock.
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(4),
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM public.usage").fetch_one(&pool),
|
||||
)
|
||||
.await
|
||||
.expect("DDL locks must be released when the migration connection closes")
|
||||
.unwrap();
|
||||
assert_overview_migration_rolled_back(&pool).await;
|
||||
sqlx::raw_sql(
|
||||
"DROP EVENT TRIGGER test_slow_migration_ddl; DROP FUNCTION public.test_slow_migration_ddl()",
|
||||
)
|
||||
.execute(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
super::super::run_migrations(&pool).await.unwrap();
|
||||
assert!(super::super::pending_migrations(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
use super::*;
|
||||
use aether_data_contracts::repository::usage::{UsageAnalyticsQuery, UsageAnalyticsView};
|
||||
use chrono::{TimeZone, Utc};
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_dirty_events_keep_independent_usage_writes_concurrent_and_reads_fresh() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
.await
|
||||
.expect("overview PostgreSQL should start or skip")
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
prepare_and_apply_clean_postgres_database(&pool).await;
|
||||
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
|
||||
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
|
||||
let insert = "INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,created_at,response_time_ms) VALUES($1,$1,'dirty-event-test','test','completed','settled',$2,100)";
|
||||
let second_id: String = query_scalar(
|
||||
"SELECT 'second-' || n FROM generate_series(1,100) n WHERE (hashtextextended('second-' || n,0) & 15) <> (hashtextextended('first',0) & 15) LIMIT 1",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// A stays open while B writes the same hour/day. Before the migration,
|
||||
// B times out on the shared stats_bucket_state day row despite a distinct request.
|
||||
let mut first = pool.begin().await.unwrap();
|
||||
sqlx::query(insert)
|
||||
.bind("first")
|
||||
.bind(at)
|
||||
.execute(&mut *first)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut second = pool.begin().await.unwrap();
|
||||
sqlx::query("SET LOCAL lock_timeout='500ms'")
|
||||
.execute(&mut *second)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(insert)
|
||||
.bind(&second_id)
|
||||
.bind(at)
|
||||
.execute(&mut *second)
|
||||
.await
|
||||
.unwrap();
|
||||
second.commit().await.unwrap();
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let query = UsageAnalyticsQuery {
|
||||
from_unix_ms: at.timestamp_millis() as u64,
|
||||
to_unix_ms: (at + chrono::Duration::hours(1)).timestamp_millis() as u64,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Summary,
|
||||
limit: 100,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
repo.query_usage_analytics(&query)
|
||||
.await
|
||||
.unwrap()
|
||||
.summary
|
||||
.request_count,
|
||||
1
|
||||
);
|
||||
|
||||
// Commit order differs from transaction-ID order. Consuming B must never
|
||||
// acknowledge A, and A must invalidate the already-published clean projection.
|
||||
first.commit().await.unwrap();
|
||||
let pending = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(pending.summary.request_count, 2);
|
||||
assert_eq!(pending.coverage.dirty_bucket_count, 1);
|
||||
let queued: i64 = query_scalar("SELECT count(*) FROM stats_overview_dirty_events")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(queued, 2);
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let rebuilt = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(rebuilt.summary, pending.summary);
|
||||
assert_eq!(rebuilt.coverage.dirty_bucket_count, 0);
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
|
||||
let mut rolled_back = pool.begin().await.unwrap();
|
||||
sqlx::query(insert)
|
||||
.bind("rollback")
|
||||
.bind(at)
|
||||
.execute(&mut *rolled_back)
|
||||
.await
|
||||
.unwrap();
|
||||
rolled_back.rollback().await.unwrap();
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
|
||||
// Repeated mutations in one transaction deduplicate, but deleting facts
|
||||
// retains the irreversible-loss marker until it is merged into state.
|
||||
let mut deleted = pool.begin().await.unwrap();
|
||||
sqlx::query("UPDATE usage SET response_time_ms=200 WHERE request_id='first'")
|
||||
.execute(&mut *deleted)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("DELETE FROM usage WHERE request_id='first'")
|
||||
.execute(&mut *deleted)
|
||||
.await
|
||||
.unwrap();
|
||||
deleted.commit().await.unwrap();
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>(
|
||||
"SELECT count(*) FROM stats_overview_dirty_events WHERE unrecoverable"
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let lost = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(lost.summary.request_count, 1);
|
||||
assert_eq!(lost.unrecoverable_bucket_count, 1);
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
let merged = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(merged.unrecoverable_bucket_count, 1);
|
||||
assert_eq!(merged.summary, lost.summary);
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
use super::*;
|
||||
use aether_data_contracts::repository::usage::{UsageAnalyticsQuery, UsageAnalyticsView};
|
||||
use chrono::{TimeZone, Utc};
|
||||
|
||||
const OVERVIEW_START: i64 = 20260911000000;
|
||||
const BILLING_INDEX: i64 = 20260918000000;
|
||||
const INSERT_USAGE: &str = "INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,created_at) VALUES($1,$1,'overview-migration-test','test','completed','settled',$2)";
|
||||
|
||||
async fn apply_through_overview(connection: &mut PgConnection) {
|
||||
connection.ensure_migrations_table().await.unwrap();
|
||||
for migration in POSTGRES_MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.version <= OVERVIEW_START)
|
||||
{
|
||||
connection.apply(migration).await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
async fn assert_independent_same_bucket_writes(pool: &PgPool, prefix: &str) {
|
||||
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
|
||||
let mut first = pool.begin().await.unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind(format!("{prefix}-first"))
|
||||
.bind(at)
|
||||
.execute(&mut *first)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut second = pool.begin().await.unwrap();
|
||||
query("SET LOCAL lock_timeout='500ms'")
|
||||
.execute(&mut *second)
|
||||
.await
|
||||
.unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind(format!("{prefix}-second"))
|
||||
.bind(at)
|
||||
.execute(&mut *second)
|
||||
.await
|
||||
.expect("another request in the same hour/day must not wait for the first transaction");
|
||||
second.commit().await.unwrap();
|
||||
first.commit().await.unwrap();
|
||||
}
|
||||
|
||||
async fn is_stamped(pool: &PgPool, version: i64) -> bool {
|
||||
query_scalar("SELECT EXISTS(SELECT 1 FROM _sqlx_migrations WHERE version=$1 AND success)")
|
||||
.bind(version)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn queue_and_state(pool: &PgPool) -> (String, String) {
|
||||
sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
(SELECT COALESCE(jsonb_agg(to_jsonb(e) ORDER BY transaction_id,projection_version,granularity,bucket_start),'[]')::text FROM stats_overview_dirty_events e),
|
||||
(SELECT COALESCE(jsonb_agg(to_jsonb(s) ORDER BY projection_version,granularity,bucket_start),'[]')::text FROM stats_bucket_state s)
|
||||
"#,
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_queue_is_safe_from_its_first_installation() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
apply_through_overview(&mut connection).await;
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
|
||||
// The initial installation must be safe before any subsequent migration,
|
||||
// including the potentially long concurrent index build, has completed.
|
||||
assert_independent_same_bucket_writes(&pool, "first-install").await;
|
||||
let (queued, obsolete, direct): (i64, i64, i64) = sqlx::query_as(
|
||||
"SELECT (SELECT count(*) FROM stats_overview_dirty_events), (SELECT count(*) FROM stats_overview_dirty_events WHERE projection_version <> 'overview-v2'), (SELECT count(*) FROM stats_bucket_state)",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!((queued, obsolete, direct), (4, 0, 0));
|
||||
assert!(!is_stamped(&pool, BILLING_INDEX).await);
|
||||
pool.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_queue_survives_later_index_failure_and_retry_without_losing_data() {
|
||||
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
|
||||
return;
|
||||
};
|
||||
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
|
||||
apply_through_overview(&mut connection).await;
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
let at = Utc.with_ymd_and_hms(2026, 1, 2, 3, 0, 0).unwrap();
|
||||
query(INSERT_USAGE)
|
||||
.bind("before-index-failure")
|
||||
.bind(at)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let repo = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let original_state = queue_and_state(&pool).await.1;
|
||||
|
||||
let mut blocker = pool.begin().await.unwrap();
|
||||
query("LOCK TABLE usage_settlement_snapshots IN SHARE MODE")
|
||||
.execute(&mut *blocker)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = super::super::run_migrations(&pool)
|
||||
.await
|
||||
.expect_err("the concurrent index must encounter the held relation lock");
|
||||
assert!(error.to_string().contains("lock timeout"), "{error}");
|
||||
assert!(is_stamped(&pool, OVERVIEW_START).await);
|
||||
assert!(!is_stamped(&pool, BILLING_INDEX).await);
|
||||
|
||||
// A failed later migration must leave the safe trigger committed and able
|
||||
// to accept independent business writes until the upgrade can be retried.
|
||||
assert_independent_same_bucket_writes(&pool, "failed-upgrade").await;
|
||||
let after_failure = queue_and_state(&pool).await;
|
||||
assert_eq!(after_failure.1, original_state);
|
||||
assert_eq!(
|
||||
query_scalar::<_, i64>("SELECT count(*) FROM stats_overview_dirty_events")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
4
|
||||
);
|
||||
blocker.rollback().await.unwrap();
|
||||
super::super::run_migrations(&pool).await.unwrap();
|
||||
assert_eq!(queue_and_state(&pool).await, after_failure);
|
||||
assert!(super::super::pending_migrations(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
|
||||
assert_eq!(
|
||||
repo.rebuild_overview_buckets(at + chrono::Duration::days(1), 8)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
let result = repo
|
||||
.query_usage_analytics(&UsageAnalyticsQuery {
|
||||
from_unix_ms: at.timestamp_millis() as u64,
|
||||
to_unix_ms: (at + chrono::Duration::hours(1)).timestamp_millis() as u64,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Summary,
|
||||
limit: 100,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.summary.request_count, 3);
|
||||
assert_eq!(result.coverage.dirty_bucket_count, 0);
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
use super::*;
|
||||
use aether_data_contracts::repository::billing::{
|
||||
AdminBillingMutationOutcome, BillingReadRepository, ProviderExpenseInput, ProviderExpenseQuery,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_provider_expense_ledger_migration_preserves_exact_sums_idempotency_and_voids() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
.await
|
||||
.expect("local postgres should start or skip")
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let pool = PgPool::connect(server.database_url()).await.unwrap();
|
||||
// Exercise only this self-contained migration against an isolated database.
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../schema/bootstrap/postgres/210_provider_expenses.sql"
|
||||
))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let repository = crate::repository::billing::SqlxBillingReadRepository::new(pool.clone());
|
||||
let input = ProviderExpenseInput {
|
||||
client_request_id: uuid::Uuid::new_v4().to_string(),
|
||||
provider_id: "deleted-later".into(),
|
||||
provider_name: "Supplier".into(),
|
||||
kind: "recharge".into(),
|
||||
amount: "0.10000000".into(),
|
||||
currency: "USD".into(),
|
||||
paid_at_unix_ms: 1000,
|
||||
period_start_unix_ms: None,
|
||||
period_end_unix_ms: None,
|
||||
note: None,
|
||||
external_reference: None,
|
||||
created_by: Some("admin".into()),
|
||||
};
|
||||
let (a, b) = tokio::join!(
|
||||
repository.create_provider_expense(&input),
|
||||
repository.create_provider_expense(&input)
|
||||
);
|
||||
let (AdminBillingMutationOutcome::Applied(a), AdminBillingMutationOutcome::Applied(b)) =
|
||||
(a.unwrap(), b.unwrap())
|
||||
else {
|
||||
panic!("expected applied")
|
||||
};
|
||||
assert_eq!(a.id, b.id);
|
||||
let mut subscription = input.clone();
|
||||
subscription.client_request_id = uuid::Uuid::new_v4().to_string();
|
||||
subscription.amount = "0.20000000".into();
|
||||
subscription.kind = "subscription".into();
|
||||
let AdminBillingMutationOutcome::Applied(second) = repository
|
||||
.create_provider_expense(&subscription)
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected applied")
|
||||
};
|
||||
let mut cny = input.clone();
|
||||
cny.client_request_id = uuid::Uuid::new_v4().to_string();
|
||||
cny.currency = "CNY".into();
|
||||
cny.amount = "7.00000000".into();
|
||||
repository.create_provider_expense(&cny).await.unwrap();
|
||||
let query = ProviderExpenseQuery {
|
||||
from_unix_ms: 0,
|
||||
to_unix_ms: 2000,
|
||||
limit: 1,
|
||||
offset: 1,
|
||||
};
|
||||
let page = repository
|
||||
.list_provider_expenses(&query)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(page.total, 3);
|
||||
assert_eq!(page.items.len(), 1);
|
||||
assert_eq!(page.totals[0].currency, "CNY");
|
||||
assert_eq!(page.totals[1].amount, "0.30000000");
|
||||
assert_eq!(page.totals[1].subscription_amount, "0.20000000");
|
||||
let mut conflict = input.clone();
|
||||
conflict.amount = "5.00000000".into();
|
||||
assert!(matches!(
|
||||
repository.create_provider_expense(&conflict).await.unwrap(),
|
||||
AdminBillingMutationOutcome::Invalid(_)
|
||||
));
|
||||
let AdminBillingMutationOutcome::Applied(voided) = repository
|
||||
.void_provider_expense(&second.id, Some("operator-1"))
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected void")
|
||||
};
|
||||
let AdminBillingMutationOutcome::Applied(again) = repository
|
||||
.void_provider_expense(&second.id, Some("operator-2"))
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected repeat void")
|
||||
};
|
||||
assert_eq!(voided, again);
|
||||
let page = repository
|
||||
.list_provider_expenses(&query)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(page.total, 2);
|
||||
assert_eq!(page.totals[1].amount, "0.10000000");
|
||||
let raw_count: i64 = sqlx::query_scalar("SELECT count(*) FROM provider_expenses")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(raw_count, 3);
|
||||
let invalid=sqlx::query("INSERT INTO provider_expenses(id,client_request_id,provider_id,provider_name,kind,amount,currency,paid_at) VALUES('bad','bad','p','P','recharge',-1,'USD',NOW())").execute(&pool).await;
|
||||
assert!(invalid.is_err());
|
||||
pool.close().await;
|
||||
}
|
||||
@@ -8,7 +8,8 @@ use uuid::Uuid;
|
||||
use crate::DataLayerError;
|
||||
use aether_data_contracts::repository::announcements::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, StoredUserAnnouncement,
|
||||
StoredUserAnnouncementPage, UpdateAnnouncementRecord, UserAnnouncementListQuery,
|
||||
};
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -105,6 +106,65 @@ impl AnnouncementReadRepository for InMemoryAnnouncementReadRepository {
|
||||
Ok(StoredAnnouncementPage { items, total })
|
||||
}
|
||||
|
||||
async fn list_user_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
query: &UserAnnouncementListQuery,
|
||||
) -> Result<StoredUserAnnouncementPage, DataLayerError> {
|
||||
query.validate()?;
|
||||
let announcements = self
|
||||
.announcements
|
||||
.read()
|
||||
.expect("announcement repository lock");
|
||||
let reads = self
|
||||
.announcement_reads
|
||||
.read()
|
||||
.expect("announcement reads repository lock");
|
||||
let mut unread_count = 0;
|
||||
let mut items = announcements
|
||||
.iter()
|
||||
.filter(|announcement| {
|
||||
announcement.is_active
|
||||
&& announcement
|
||||
.start_time_unix_secs
|
||||
.is_none_or(|value| value <= query.now_unix_secs)
|
||||
&& announcement
|
||||
.end_time_unix_secs
|
||||
.is_none_or(|value| value >= query.now_unix_secs)
|
||||
})
|
||||
.filter_map(|announcement| {
|
||||
let is_read = reads.contains(&(user_id.to_string(), announcement.id.clone()));
|
||||
if !is_read {
|
||||
unread_count += 1;
|
||||
}
|
||||
(!query.unread_only || !is_read).then_some((announcement, is_read))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by(|(left, _), (right, _)| {
|
||||
right
|
||||
.is_pinned
|
||||
.cmp(&left.is_pinned)
|
||||
.then_with(|| right.priority.cmp(&left.priority))
|
||||
.then_with(|| right.created_at_unix_ms.cmp(&left.created_at_unix_ms))
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
let total = items.len() as u64;
|
||||
let items = items
|
||||
.into_iter()
|
||||
.skip(query.offset)
|
||||
.take(query.limit)
|
||||
.map(|(announcement, is_read)| StoredUserAnnouncement {
|
||||
announcement: announcement.clone(),
|
||||
is_read,
|
||||
})
|
||||
.collect();
|
||||
Ok(StoredUserAnnouncementPage {
|
||||
items,
|
||||
total,
|
||||
unread_count,
|
||||
})
|
||||
}
|
||||
|
||||
async fn count_unread_active_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -294,9 +354,95 @@ mod tests {
|
||||
use super::InMemoryAnnouncementReadRepository;
|
||||
use crate::repository::announcements::{
|
||||
AnnouncementReadRepository, AnnouncementWriteRepository, CreateAnnouncementRecord,
|
||||
StoredAnnouncement, UpdateAnnouncementRecord,
|
||||
StoredAnnouncement, UpdateAnnouncementRecord, UserAnnouncementListQuery,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn personal_announcement_page_preserves_global_unread_and_visibility() {
|
||||
let now = 1_800_000_000;
|
||||
let announcements = [
|
||||
("pinned", true, true, 0, None, None),
|
||||
("normal-a", true, false, 20, Some(now), Some(now)),
|
||||
("normal-b", true, false, 20, None, None),
|
||||
("draft", false, false, 99, None, None),
|
||||
("future", true, false, 99, Some(now + 1), None),
|
||||
("expired", true, false, 99, None, Some(now - 1)),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(id, active, pinned, priority, start, end)| {
|
||||
StoredAnnouncement::new(
|
||||
id.into(),
|
||||
id.into(),
|
||||
"content".into(),
|
||||
"info".into(),
|
||||
priority,
|
||||
active,
|
||||
pinned,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
start,
|
||||
end,
|
||||
now,
|
||||
now,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let repository = InMemoryAnnouncementReadRepository::seed_with_reads(
|
||||
announcements,
|
||||
[("reader".into(), "pinned".into())],
|
||||
);
|
||||
let mut query = UserAnnouncementListQuery {
|
||||
unread_only: false,
|
||||
offset: 0,
|
||||
limit: 1,
|
||||
now_unix_secs: now as u64,
|
||||
};
|
||||
let first = repository
|
||||
.list_user_announcements("reader", &query)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!((first.total, first.unread_count), (3, 2));
|
||||
assert_eq!(first.items[0].announcement.id, "pinned");
|
||||
assert!(first.items[0].is_read);
|
||||
query.unread_only = true;
|
||||
query.offset = 1;
|
||||
let second = repository
|
||||
.list_user_announcements("reader", &query)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!((second.total, second.unread_count), (2, 2));
|
||||
assert_eq!(second.items[0].announcement.id, "normal-b");
|
||||
assert!(!second.items[0].is_read);
|
||||
query.offset = i64::MAX as usize;
|
||||
let empty = repository
|
||||
.list_user_announcements("reader", &query)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(empty.items.is_empty());
|
||||
assert_eq!((empty.total, empty.unread_count), (2, 2));
|
||||
let other = repository
|
||||
.list_user_announcements("other", &query)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!((other.total, other.unread_count), (3, 3));
|
||||
repository
|
||||
.mark_announcement_as_read("reader", "normal-a", now as u64)
|
||||
.await
|
||||
.unwrap();
|
||||
query.offset = 0;
|
||||
let after_read = repository
|
||||
.list_user_announcements("reader", &query)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!((after_read.total, after_read.unread_count), (1, 1));
|
||||
query.limit = 101;
|
||||
assert!(repository
|
||||
.list_user_announcements("reader", &query)
|
||||
.await
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_seeded_announcements() {
|
||||
let repository = InMemoryAnnouncementReadRepository::seed(vec![StoredAnnouncement::new(
|
||||
|
||||
@@ -2,7 +2,8 @@ mod memory;
|
||||
|
||||
pub use aether_data_contracts::repository::announcements::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, StoredUserAnnouncement,
|
||||
StoredUserAnnouncementPage, UpdateAnnouncementRecord, UserAnnouncementListQuery,
|
||||
};
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use aether_data_postgres::SqlxAnnouncementReadRepository;
|
||||
|
||||
@@ -229,6 +229,24 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn standalone_flags(&self) -> BTreeMap<String, bool> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
index
|
||||
.export_by_api_key_id
|
||||
.iter()
|
||||
.map(|(id, record)| (id.clone(), record.is_standalone))
|
||||
.chain(
|
||||
index
|
||||
.by_api_key_id
|
||||
.iter()
|
||||
.map(|(id, snapshot)| (id.clone(), snapshot.api_key_is_standalone)),
|
||||
)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn snapshot_lookup_count(&self, api_key_id: &str) -> usize {
|
||||
self.index
|
||||
.read()
|
||||
@@ -2101,6 +2119,43 @@ mod tests {
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_flags_cover_export_only_keys_and_remove_deleted_keys() {
|
||||
let mut standalone = sample_snapshot("standalone-key", "user-1");
|
||||
standalone.api_key_is_standalone = true;
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed([
|
||||
(None, sample_snapshot("member-key", "user-1")),
|
||||
(None, standalone),
|
||||
]);
|
||||
let mut export = repository
|
||||
.list_export_api_keys_by_ids(&["standalone-key".into()])
|
||||
.await
|
||||
.unwrap()
|
||||
.remove(0);
|
||||
export.api_key_id = "export-only-key".into();
|
||||
let repository = repository.with_export_records([export]);
|
||||
|
||||
assert_eq!(
|
||||
repository.standalone_flags(),
|
||||
std::collections::BTreeMap::from([
|
||||
("member-key".into(), false),
|
||||
("standalone-key".into(), true),
|
||||
("export-only-key".into(), true),
|
||||
])
|
||||
);
|
||||
assert!(repository
|
||||
.delete_standalone_api_key("standalone-key")
|
||||
.await
|
||||
.unwrap());
|
||||
assert_eq!(
|
||||
repository.standalone_flags(),
|
||||
std::collections::BTreeMap::from([
|
||||
("member-key".into(), false),
|
||||
("export-only-key".into(), true),
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_auth_snapshot_by_all_supported_keys() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
|
||||
@@ -16,6 +16,7 @@ type BillingContextMap = BTreeMap<BillingContextKey, StoredBillingModelContext>;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryBillingReadRepository {
|
||||
provider_expenses: RwLock<BTreeMap<String, super::ProviderExpenseRecord>>,
|
||||
by_key: RwLock<BillingContextMap>,
|
||||
gateway_configs_by_provider: RwLock<BTreeMap<String, PaymentGatewayConfigRecord>>,
|
||||
billing_plans_by_id: RwLock<BTreeMap<String, BillingPlanRecord>>,
|
||||
@@ -23,6 +24,20 @@ pub struct InMemoryBillingReadRepository {
|
||||
}
|
||||
|
||||
impl InMemoryBillingReadRepository {
|
||||
pub fn seed_user_plan_entitlements(
|
||||
items: impl IntoIterator<Item = UserPlanEntitlementRecord>,
|
||||
) -> Self {
|
||||
Self {
|
||||
entitlements_by_id: RwLock::new(
|
||||
items
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect(),
|
||||
),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredBillingModelContext>,
|
||||
@@ -39,6 +54,7 @@ impl InMemoryBillingReadRepository {
|
||||
);
|
||||
}
|
||||
Self {
|
||||
provider_expenses: RwLock::default(),
|
||||
by_key: RwLock::new(by_key),
|
||||
gateway_configs_by_provider: RwLock::new(BTreeMap::new()),
|
||||
billing_plans_by_id: RwLock::new(BTreeMap::new()),
|
||||
@@ -325,6 +341,68 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
Ok(AdminBillingMutationOutcome::Applied(record))
|
||||
}
|
||||
|
||||
async fn list_provider_expenses(
|
||||
&self,
|
||||
query: &super::ProviderExpenseQuery,
|
||||
) -> Result<Option<super::ProviderExpensePage>, DataLayerError> {
|
||||
let guard = self
|
||||
.provider_expenses
|
||||
.read()
|
||||
.expect("provider expense store should lock");
|
||||
super::provider_expense_memory_page(guard.values().cloned(), query).map(Some)
|
||||
}
|
||||
async fn create_provider_expense(
|
||||
&self,
|
||||
input: &super::ProviderExpenseInput,
|
||||
) -> Result<AdminBillingMutationOutcome<super::ProviderExpenseRecord>, DataLayerError> {
|
||||
if let Err(detail) = input.validate() {
|
||||
return Ok(AdminBillingMutationOutcome::Invalid(detail));
|
||||
}
|
||||
let mut guard = self
|
||||
.provider_expenses
|
||||
.write()
|
||||
.expect("provider expense store should lock");
|
||||
if let Some(existing) = guard
|
||||
.values()
|
||||
.find(|r| r.entry.client_request_id == input.client_request_id)
|
||||
{
|
||||
return Ok(if existing.entry.same_request_as(input) {
|
||||
AdminBillingMutationOutcome::Applied(existing.clone())
|
||||
} else {
|
||||
AdminBillingMutationOutcome::Invalid(
|
||||
"client_request_id was already used for another expense".into(),
|
||||
)
|
||||
});
|
||||
}
|
||||
let record = super::ProviderExpenseRecord {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
entry: input.clone(),
|
||||
created_at_unix_ms: chrono::Utc::now().timestamp_millis().max(0) as u64,
|
||||
voided_at_unix_ms: None,
|
||||
voided_by: None,
|
||||
};
|
||||
guard.insert(record.id.clone(), record.clone());
|
||||
Ok(AdminBillingMutationOutcome::Applied(record))
|
||||
}
|
||||
async fn void_provider_expense(
|
||||
&self,
|
||||
id: &str,
|
||||
operator: Option<&str>,
|
||||
) -> Result<AdminBillingMutationOutcome<super::ProviderExpenseRecord>, DataLayerError> {
|
||||
let mut guard = self
|
||||
.provider_expenses
|
||||
.write()
|
||||
.expect("provider expense store should lock");
|
||||
let Some(record) = guard.get_mut(id) else {
|
||||
return Ok(AdminBillingMutationOutcome::NotFound);
|
||||
};
|
||||
if record.voided_at_unix_ms.is_none() {
|
||||
record.voided_at_unix_ms = Some(chrono::Utc::now().timestamp_millis().max(0) as u64);
|
||||
record.voided_by = operator.map(str::to_owned);
|
||||
}
|
||||
Ok(AdminBillingMutationOutcome::Applied(record.clone()))
|
||||
}
|
||||
|
||||
async fn list_billing_plans(
|
||||
&self,
|
||||
include_disabled: bool,
|
||||
@@ -435,6 +513,15 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
async fn list_user_plan_entitlements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
|
||||
self.list_user_plan_entitlements_with_history(user_id, false)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_user_plan_entitlements_with_history(
|
||||
&self,
|
||||
user_id: &str,
|
||||
include_inactive: bool,
|
||||
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
|
||||
let now = current_unix_secs();
|
||||
let mut items = self
|
||||
@@ -444,12 +531,12 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
|
||||
.values()
|
||||
.filter(|item| {
|
||||
item.user_id == user_id
|
||||
&& item.status == "active"
|
||||
&& item.expires_at_unix_secs > now
|
||||
&& (include_inactive
|
||||
|| (item.status == "active" && item.expires_at_unix_secs > now))
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
items.sort_by_key(|item| item.expires_at_unix_secs);
|
||||
items.sort_by_key(|item| (item.expires_at_unix_secs, item.created_at_unix_secs));
|
||||
Ok(Some(items))
|
||||
}
|
||||
|
||||
@@ -592,6 +679,59 @@ mod tests {
|
||||
.expect("billing context should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_user_plan_entitlements_history_includes_inactive_only_for_selected_user() {
|
||||
let repository = InMemoryBillingReadRepository::default();
|
||||
let now = super::current_unix_secs();
|
||||
for (id, user_id, status, expires_at) in [
|
||||
("active", "user-1", "active", now + 3600),
|
||||
("expired", "user-1", "active", now - 60),
|
||||
("revoked", "user-1", "revoked", now + 3600),
|
||||
("replaced", "user-1", "replaced", now + 3600),
|
||||
("another-user", "user-2", "revoked", now + 3600),
|
||||
] {
|
||||
repository.entitlements_by_id.write().unwrap().insert(
|
||||
id.to_string(),
|
||||
super::UserPlanEntitlementRecord {
|
||||
id: id.to_string(),
|
||||
user_id: user_id.to_string(),
|
||||
plan_id: "plan-1".to_string(),
|
||||
payment_order_id: format!("order-{id}"),
|
||||
status: status.to_string(),
|
||||
starts_at_unix_secs: now - 120,
|
||||
expires_at_unix_secs: expires_at,
|
||||
entitlements_snapshot: json!([]),
|
||||
created_at_unix_secs: now - 120,
|
||||
updated_at_unix_secs: now - 60,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let active = repository
|
||||
.list_user_plan_entitlements("user-1")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(active.len(), 1);
|
||||
assert_eq!(active[0].id, "active");
|
||||
let explicit_active = repository
|
||||
.list_user_plan_entitlements_with_history("user-1", false)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(active, explicit_active);
|
||||
let history = repository
|
||||
.list_user_plan_entitlements_with_history("user-1", true)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(history.len(), 4);
|
||||
assert!(history.iter().all(|item| item.user_id == "user-1"));
|
||||
for expected in ["active", "expired", "revoked", "replaced"] {
|
||||
assert!(history.iter().any(|item| item.id == expected));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn falls_back_to_provider_without_key_scope() {
|
||||
let repository = InMemoryBillingReadRepository::seed(vec![sample_context()]);
|
||||
@@ -779,3 +919,108 @@ mod tests {
|
||||
assert_eq!(after, expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod provider_expense_tests {
|
||||
use super::*;
|
||||
use crate::repository::billing::{ProviderExpenseInput, ProviderExpenseQuery};
|
||||
fn input() -> ProviderExpenseInput {
|
||||
ProviderExpenseInput {
|
||||
client_request_id: uuid::Uuid::new_v4().to_string(),
|
||||
provider_id: "provider-1".into(),
|
||||
provider_name: "Supplier".into(),
|
||||
kind: "recharge".into(),
|
||||
amount: "12.34000000".into(),
|
||||
currency: "USD".into(),
|
||||
paid_at_unix_ms: 1000,
|
||||
period_start_unix_ms: None,
|
||||
period_end_unix_ms: None,
|
||||
note: Some("test".into()),
|
||||
external_reference: None,
|
||||
created_by: Some("admin-1".into()),
|
||||
}
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn provider_expense_retries_and_void_are_idempotent_without_erasing_audit() {
|
||||
let repo = InMemoryBillingReadRepository::default();
|
||||
let input = input();
|
||||
let AdminBillingMutationOutcome::Applied(first) =
|
||||
repo.create_provider_expense(&input).await.unwrap()
|
||||
else {
|
||||
panic!("expected record")
|
||||
};
|
||||
let mut retried = input.clone();
|
||||
retried.provider_name = "Renamed supplier".into();
|
||||
let AdminBillingMutationOutcome::Applied(second) =
|
||||
repo.create_provider_expense(&retried).await.unwrap()
|
||||
else {
|
||||
panic!("expected retry")
|
||||
};
|
||||
assert_eq!(first, second);
|
||||
retried.amount = "99".into();
|
||||
assert!(matches!(
|
||||
repo.create_provider_expense(&retried).await.unwrap(),
|
||||
AdminBillingMutationOutcome::Invalid(_)
|
||||
));
|
||||
let query = ProviderExpenseQuery {
|
||||
from_unix_ms: 0,
|
||||
to_unix_ms: 2000,
|
||||
limit: 20,
|
||||
offset: 0,
|
||||
};
|
||||
assert_eq!(
|
||||
repo.list_provider_expenses(&query)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.total,
|
||||
1
|
||||
);
|
||||
let AdminBillingMutationOutcome::Applied(voided) = repo
|
||||
.void_provider_expense(&first.id, Some("admin-2"))
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected void")
|
||||
};
|
||||
let AdminBillingMutationOutcome::Applied(again) = repo
|
||||
.void_provider_expense(&first.id, Some("admin-3"))
|
||||
.await
|
||||
.unwrap()
|
||||
else {
|
||||
panic!("expected retry void")
|
||||
};
|
||||
assert_eq!(voided, again);
|
||||
assert_eq!(again.voided_by.as_deref(), Some("admin-2"));
|
||||
let page = repo.list_provider_expenses(&query).await.unwrap().unwrap();
|
||||
assert_eq!(page.total, 0);
|
||||
assert!(page.totals.is_empty());
|
||||
let AdminBillingMutationOutcome::Applied(after_void) =
|
||||
repo.create_provider_expense(&input).await.unwrap()
|
||||
else {
|
||||
panic!("expected original tombstone")
|
||||
};
|
||||
assert_eq!(after_void, again);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn provider_expense_duplicate_submissions_record_once() {
|
||||
let repo = InMemoryBillingReadRepository::default();
|
||||
let input = input();
|
||||
let (a, b) = tokio::join!(
|
||||
repo.create_provider_expense(&input),
|
||||
repo.create_provider_expense(&input)
|
||||
);
|
||||
let (AdminBillingMutationOutcome::Applied(a), AdminBillingMutationOutcome::Applied(b)) =
|
||||
(a.unwrap(), b.unwrap())
|
||||
else {
|
||||
panic!("expected records")
|
||||
};
|
||||
assert_eq!(a.id, b.id);
|
||||
let mut invalid = input.clone();
|
||||
invalid.period_start_unix_ms = Some(100);
|
||||
assert!(matches!(
|
||||
repo.create_provider_expense(&invalid).await.unwrap(),
|
||||
AdminBillingMutationOutcome::Invalid(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use super::{
|
||||
ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput,
|
||||
ReserveUsagePolicyRequestOutcome, SettlementWriteRepository, StoredUsagePolicyCostReservation,
|
||||
StoredUsagePolicyRequestAdmission, StoredUsageSettlement, UsagePolicyCostReservationState,
|
||||
UsagePolicyRequestAdmissionState, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
UsagePolicyRequestAdmissionState, UsageSettlementInput,
|
||||
};
|
||||
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use crate::DataLayerError;
|
||||
@@ -511,9 +511,7 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
|
||||
settlement.wallet_recharge_balance_after = Some(wallet.balance);
|
||||
settlement.wallet_gift_balance_after = Some(wallet.gift_balance);
|
||||
settlement.wallet_balance_after = Some(wallet.balance + wallet.gift_balance);
|
||||
} else if final_billing_status == "settled"
|
||||
&& billable_cost_usd > SETTLEMENT_EPSILON_USD
|
||||
{
|
||||
} else if final_billing_status == "settled" && billable_cost_usd > 0.0 {
|
||||
final_billing_status = "insufficient_quota".to_string();
|
||||
settlement.billing_status = final_billing_status.clone();
|
||||
}
|
||||
|
||||
@@ -44,17 +44,33 @@ use super::{
|
||||
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
|
||||
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use crate::DataLayerError;
|
||||
mod analytics;
|
||||
mod dashboard_summary;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUsageReadRepository {
|
||||
dashboard_projection: RwLock<dashboard_summary::DashboardProjection>,
|
||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||
detached_bodies: RwLock<BTreeMap<String, Value>>,
|
||||
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
|
||||
auth_api_keys: Option<Arc<InMemoryAuthApiKeySnapshotRepository>>,
|
||||
provider_catalog: Option<Arc<InMemoryProviderCatalogReadRepository>>,
|
||||
analytics_users: RwLock<Vec<aether_data_contracts::repository::users::StoredUserSummary>>,
|
||||
analytics_candidates:
|
||||
RwLock<Vec<aether_data_contracts::repository::candidates::StoredRequestCandidate>>,
|
||||
analytics_allocations: RwLock<
|
||||
BTreeMap<String, aether_data_contracts::repository::usage::UsageAnalyticsAllocation>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl InMemoryUsageReadRepository {
|
||||
fn analytics_key_flags(&self) -> BTreeMap<String, bool> {
|
||||
self.auth_api_keys
|
||||
.as_ref()
|
||||
.map(|repository| repository.standalone_flags())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRequestUsageAudit>,
|
||||
@@ -66,11 +82,15 @@ impl InMemoryUsageReadRepository {
|
||||
by_request_id.insert(item.request_id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
dashboard_projection: Default::default(),
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
detached_bodies: RwLock::new(BTreeMap::new()),
|
||||
provider_usage_windows: RwLock::new(Vec::new()),
|
||||
auth_api_keys: None,
|
||||
provider_catalog: None,
|
||||
analytics_users: Default::default(),
|
||||
analytics_candidates: Default::default(),
|
||||
analytics_allocations: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -119,11 +139,15 @@ impl InMemoryUsageReadRepository {
|
||||
by_request_id.insert(request_id, item);
|
||||
}
|
||||
Self {
|
||||
dashboard_projection: Default::default(),
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
detached_bodies: RwLock::new(detached_bodies),
|
||||
provider_usage_windows: RwLock::new(Vec::new()),
|
||||
auth_api_keys: None,
|
||||
provider_catalog: None,
|
||||
analytics_users: Default::default(),
|
||||
analytics_candidates: Default::default(),
|
||||
analytics_allocations: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -132,11 +156,15 @@ impl InMemoryUsageReadRepository {
|
||||
I: IntoIterator<Item = StoredProviderUsageWindow>,
|
||||
{
|
||||
Self {
|
||||
dashboard_projection: self.dashboard_projection,
|
||||
by_request_id: self.by_request_id,
|
||||
detached_bodies: self.detached_bodies,
|
||||
provider_usage_windows: RwLock::new(items.into_iter().collect()),
|
||||
auth_api_keys: self.auth_api_keys,
|
||||
provider_catalog: self.provider_catalog,
|
||||
analytics_users: self.analytics_users,
|
||||
analytics_candidates: self.analytics_candidates,
|
||||
analytics_allocations: self.analytics_allocations,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -244,7 +272,48 @@ fn usage_has_admin_unknown_model_or_provider(item: &StoredRequestUsageAudit) ->
|
||||
usage_admin_unknown_label(&item.model) || usage_admin_unknown_label(&item.provider_name)
|
||||
}
|
||||
|
||||
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
|
||||
fn usage_matches_list_query(
|
||||
item: &StoredRequestUsageAudit,
|
||||
query: &UsageAuditListQuery,
|
||||
keys: &BTreeMap<String, bool>,
|
||||
) -> bool {
|
||||
if query
|
||||
.provider_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.provider_id.as_ref() != Some(id))
|
||||
|| query
|
||||
.endpoint_kind
|
||||
.as_ref()
|
||||
.is_some_and(|value| item.endpoint_kind.as_ref() != Some(value))
|
||||
|| query
|
||||
.request_type
|
||||
.as_ref()
|
||||
.is_some_and(|value| item.request_type.as_ref() != Some(value))
|
||||
|| query
|
||||
.slow_threshold_ms
|
||||
.is_some_and(|value| item.response_time_ms.is_none_or(|latency| latency < value))
|
||||
|| query
|
||||
.has_format_conversion
|
||||
.is_some_and(|value| item.has_format_conversion != value)
|
||||
|| query
|
||||
.api_key_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.api_key_id.as_ref() != Some(id))
|
||||
|| query
|
||||
.request_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.request_id != *id)
|
||||
|| query
|
||||
.attribution_kind
|
||||
.as_deref()
|
||||
.is_some_and(|kind| analytics::attribution(item, keys) != kind)
|
||||
|| query
|
||||
.actor_user_id
|
||||
.as_deref()
|
||||
.is_some_and(|id| analytics::actor(item, keys) != Some(id))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
// 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.
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
@@ -332,7 +401,45 @@ fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditLi
|
||||
fn usage_matches_keyword_search_query(
|
||||
item: &StoredRequestUsageAudit,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
keys: &BTreeMap<String, bool>,
|
||||
) -> bool {
|
||||
if query
|
||||
.provider_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.provider_id.as_ref() != Some(id))
|
||||
|| query
|
||||
.endpoint_kind
|
||||
.as_ref()
|
||||
.is_some_and(|value| item.endpoint_kind.as_ref() != Some(value))
|
||||
|| query
|
||||
.request_type
|
||||
.as_ref()
|
||||
.is_some_and(|value| item.request_type.as_ref() != Some(value))
|
||||
|| query
|
||||
.slow_threshold_ms
|
||||
.is_some_and(|value| item.response_time_ms.is_none_or(|latency| latency < value))
|
||||
|| query
|
||||
.has_format_conversion
|
||||
.is_some_and(|value| item.has_format_conversion != value)
|
||||
|| query
|
||||
.api_key_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.api_key_id.as_ref() != Some(id))
|
||||
|| query
|
||||
.request_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| item.request_id != *id)
|
||||
|| query
|
||||
.attribution_kind
|
||||
.as_deref()
|
||||
.is_some_and(|kind| analytics::attribution(item, keys) != kind)
|
||||
|| query
|
||||
.actor_user_id
|
||||
.as_deref()
|
||||
.is_some_and(|id| analytics::actor(item, keys) != Some(id))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
if item.created_at_unix_ms < created_from_unix_secs {
|
||||
return false;
|
||||
@@ -1127,6 +1234,38 @@ fn usage_provider_aggregation_identity(
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn query_dashboard_summary(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
|
||||
) -> Result<aether_data_contracts::repository::usage::StoredDashboardSummary, DataLayerError> {
|
||||
self.dashboard_summary_query(query)
|
||||
}
|
||||
|
||||
async fn query_dashboard_analytics(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
|
||||
) -> Result<
|
||||
aether_data_contracts::repository::usage::StoredUsageDashboardAnalytics,
|
||||
DataLayerError,
|
||||
> {
|
||||
self.dashboard_analytics_query(query)
|
||||
}
|
||||
|
||||
async fn query_usage_analytics(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::usage::UsageAnalyticsQuery,
|
||||
) -> Result<aether_data_contracts::repository::usage::StoredUsageAnalytics, DataLayerError>
|
||||
{
|
||||
self.analytics_query(query)
|
||||
}
|
||||
|
||||
async fn summarize_health_observations(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::usage::HealthObservationQuery,
|
||||
) -> Result<aether_data_contracts::repository::usage::HealthObservationSummary, DataLayerError>
|
||||
{
|
||||
self.health_observations(query)
|
||||
}
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
@@ -1205,12 +1344,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let keys = self.analytics_key_flags();
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| usage_matches_list_query(item, query))
|
||||
.filter(|item| usage_matches_list_query(item, query, &keys))
|
||||
.cloned()
|
||||
.collect();
|
||||
sort_usage_items(&mut items, query.newest_first);
|
||||
@@ -1231,12 +1371,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
&self,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let keys = self.analytics_key_flags();
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| usage_matches_keyword_search_query(item, query))
|
||||
.filter(|item| usage_matches_keyword_search_query(item, query, &keys))
|
||||
.cloned()
|
||||
.collect();
|
||||
sort_usage_items(&mut items, query.newest_first);
|
||||
@@ -1254,12 +1395,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
}
|
||||
|
||||
async fn count_usage_audits(&self, query: &UsageAuditListQuery) -> Result<u64, DataLayerError> {
|
||||
let keys = self.analytics_key_flags();
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| usage_matches_list_query(item, query))
|
||||
.filter(|item| usage_matches_list_query(item, query, &keys))
|
||||
.count() as u64)
|
||||
}
|
||||
|
||||
@@ -1267,12 +1409,13 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
&self,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let keys = self.analytics_key_flags();
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| usage_matches_keyword_search_query(item, query))
|
||||
.filter(|item| usage_matches_keyword_search_query(item, query, &keys))
|
||||
.count() as u64)
|
||||
}
|
||||
|
||||
@@ -3281,6 +3424,8 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
|
||||
finalized_at_unix_secs: usage.finalized_at_unix_secs,
|
||||
};
|
||||
|
||||
self.dashboard_projection.write().expect("dashboard projection lock")
|
||||
.record(&stored, &self.analytics_key_flags());
|
||||
by_request_id.insert(stored.request_id.clone(), stored.clone());
|
||||
if let Some(auth_api_keys) = self.auth_api_keys.as_ref() {
|
||||
let before_contribution = existing.as_ref().and_then(api_key_usage_contribution);
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,228 @@
|
||||
use super::{
|
||||
analytics, usage_cache_creation_tokens, usage_total_input_context, usage_total_tokens,
|
||||
InMemoryUsageReadRepository, StoredRequestUsageAudit,
|
||||
};
|
||||
use aether_data_contracts::{repository::usage::*, DataLayerError};
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) struct DashboardProjection {
|
||||
pub since: DateTime<Utc>,
|
||||
entries: BTreeMap<String, Contribution>,
|
||||
}
|
||||
impl Default for DashboardProjection {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
since: Utc::now(),
|
||||
entries: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug)]
|
||||
struct Contribution {
|
||||
at: DateTime<Utc>,
|
||||
actor: Option<String>,
|
||||
metrics: DashboardSummaryMetrics,
|
||||
billable_units: Option<i128>,
|
||||
}
|
||||
impl DashboardProjection {
|
||||
pub fn record(&mut self, row: &StoredRequestUsageAudit, keys: &BTreeMap<String, bool>) {
|
||||
let Some(at) = DateTime::from_timestamp(row.created_at_unix_ms as i64, 0) else {
|
||||
return;
|
||||
};
|
||||
if at < self.since {
|
||||
return;
|
||||
}
|
||||
if row
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|m| m.pointer("/analytics_attribution/record_kind"))
|
||||
.and_then(|v| v.as_str())
|
||||
== Some("session")
|
||||
{
|
||||
self.entries.remove(&row.request_id);
|
||||
return;
|
||||
}
|
||||
let available = |key| {
|
||||
row.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|m| m.get(key))
|
||||
.and_then(|v| v.as_bool())
|
||||
!= Some(false)
|
||||
};
|
||||
let usage = available(USAGE_AVAILABLE_METADATA_KEY);
|
||||
let priced =
|
||||
available(USAGE_PRICING_AVAILABLE_METADATA_KEY) && row.billing_status == "settled";
|
||||
let stream = row
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|m| m.get("upstream_is_stream"))
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(row.is_stream);
|
||||
let metrics = DashboardSummaryMetrics {
|
||||
request_count: 1,
|
||||
input_tokens: if usage { row.input_tokens } else { 0 },
|
||||
output_tokens: if usage { row.output_tokens } else { 0 },
|
||||
total_tokens: if usage {
|
||||
if row.total_tokens > 0 {
|
||||
row.total_tokens
|
||||
} else {
|
||||
usage_total_tokens(row)
|
||||
}
|
||||
} else {
|
||||
0
|
||||
},
|
||||
usage_available_count: u64::from(usage),
|
||||
pricing_available_count: u64::from(priced),
|
||||
cache_read_tokens: if usage {
|
||||
row.cache_read_input_tokens
|
||||
} else {
|
||||
0
|
||||
},
|
||||
cache_creation_tokens: if usage {
|
||||
usage_cache_creation_tokens(row)
|
||||
} else {
|
||||
0
|
||||
},
|
||||
cache_input_tokens: if usage {
|
||||
usage_total_input_context(row)
|
||||
} else {
|
||||
0
|
||||
},
|
||||
first_byte_sum_ms: row.first_byte_time_ms.unwrap_or(0) as f64,
|
||||
first_byte_sample_count: u64::from(row.first_byte_time_ms.is_some()),
|
||||
response_sum_ms: row.response_time_ms.unwrap_or(0) as f64,
|
||||
response_sample_count: u64::from(row.response_time_ms.is_some()),
|
||||
stream_requests: u64::from(stream),
|
||||
standard_requests: u64::from(!stream),
|
||||
..Default::default()
|
||||
};
|
||||
self.entries.insert(
|
||||
row.request_id.clone(),
|
||||
Contribution {
|
||||
at,
|
||||
actor: analytics::actor(row, keys).map(str::to_owned),
|
||||
metrics,
|
||||
billable_units: priced
|
||||
.then(|| (row.actual_total_cost_usd * 100_000_000.0).round() as i128),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
fn sum_metrics<'a>(rows: impl Iterator<Item = &'a Contribution>) -> DashboardSummaryMetrics {
|
||||
let mut sum = DashboardSummaryMetrics::default();
|
||||
let mut users = BTreeSet::new();
|
||||
let mut units = 0_i128;
|
||||
for row in rows {
|
||||
let m = &row.metrics;
|
||||
sum.request_count += m.request_count;
|
||||
sum.input_tokens += m.input_tokens;
|
||||
sum.output_tokens += m.output_tokens;
|
||||
sum.total_tokens += m.total_tokens;
|
||||
sum.usage_available_count += m.usage_available_count;
|
||||
sum.pricing_available_count += m.pricing_available_count;
|
||||
sum.cache_read_tokens += m.cache_read_tokens;
|
||||
sum.cache_creation_tokens += m.cache_creation_tokens;
|
||||
sum.cache_input_tokens += m.cache_input_tokens;
|
||||
sum.first_byte_sum_ms += m.first_byte_sum_ms;
|
||||
sum.first_byte_sample_count += m.first_byte_sample_count;
|
||||
sum.response_sum_ms += m.response_sum_ms;
|
||||
sum.response_sample_count += m.response_sample_count;
|
||||
sum.stream_requests += m.stream_requests;
|
||||
sum.standard_requests += m.standard_requests;
|
||||
units += row.billable_units.unwrap_or(0);
|
||||
if let Some(actor) = row.actor.as_deref() {
|
||||
users.insert(actor);
|
||||
}
|
||||
}
|
||||
sum.active_users = users.len() as u64;
|
||||
if sum.request_count == 0 || sum.pricing_available_count > 0 {
|
||||
sum.billable_amount = Some(format!(
|
||||
"{}{}.{:08}",
|
||||
if units < 0 { "-" } else { "" },
|
||||
units.abs() / 100_000_000,
|
||||
units.abs() % 100_000_000
|
||||
));
|
||||
}
|
||||
sum
|
||||
}
|
||||
impl InMemoryUsageReadRepository {
|
||||
/// Test/embedded initialization boundary; seeded older audit rows stay excluded.
|
||||
pub fn with_dashboard_stats_since(self, since: DateTime<Utc>) -> Self {
|
||||
let keys = self.analytics_key_flags();
|
||||
let mut projection = self
|
||||
.dashboard_projection
|
||||
.write()
|
||||
.expect("dashboard projection lock");
|
||||
projection.since = since;
|
||||
projection.entries.clear();
|
||||
for row in self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
{
|
||||
projection.record(row, &keys);
|
||||
}
|
||||
drop(projection);
|
||||
self
|
||||
}
|
||||
pub(super) fn dashboard_summary_query(
|
||||
&self,
|
||||
query: &UsageDashboardAnalyticsQuery,
|
||||
) -> Result<StoredDashboardSummary, DataLayerError> {
|
||||
query.validate()?;
|
||||
let projection = self.dashboard_projection.read().map_err(|_| {
|
||||
DataLayerError::UnexpectedValue("dashboard projection lock poisoned".into())
|
||||
})?;
|
||||
let now = Utc::now();
|
||||
let today_from = query.today_start(now)?.max(projection.since);
|
||||
let tz = query
|
||||
.timezone
|
||||
.parse::<chrono_tz::Tz>()
|
||||
.map_err(|_| DataLayerError::InvalidInput("invalid timezone".into()))?;
|
||||
let rows = || projection.entries.values().filter(|row| row.at < now);
|
||||
let today = sum_metrics(rows().filter(|row| row.at >= today_from));
|
||||
let total = sum_metrics(rows());
|
||||
let mut days = BTreeMap::<NaiveDate, u64>::new();
|
||||
for row in rows() {
|
||||
*days
|
||||
.entry(row.at.with_timezone(&tz).date_naive())
|
||||
.or_default() += 1;
|
||||
}
|
||||
let active_days = days.len() as u64;
|
||||
let local_today = now.with_timezone(&tz).date_naive();
|
||||
let consecutive_active_days =
|
||||
dashboard_consecutive_active_days(days.keys().copied(), local_today);
|
||||
let earliest_day = local_today - chrono::Duration::days(364);
|
||||
let activity_days = days
|
||||
.into_iter()
|
||||
.filter(|(day, _)| day >= &earliest_day)
|
||||
.map(|(date, requests)| DashboardActivityDay {
|
||||
date: date.to_string(),
|
||||
requests,
|
||||
})
|
||||
.collect();
|
||||
let users = self
|
||||
.analytics_users
|
||||
.read()
|
||||
.map_err(|_| DataLayerError::UnexpectedValue("users lock poisoned".into()))?;
|
||||
Ok(StoredDashboardSummary {
|
||||
stats_since: projection.since.to_rfc3339(),
|
||||
generated_at: now.to_rfc3339(),
|
||||
timezone: query.timezone.clone(),
|
||||
today_from: today_from.to_rfc3339(),
|
||||
window_seconds: (now - today_from).num_milliseconds().max(0) as f64 / 1000.0,
|
||||
today,
|
||||
total,
|
||||
users: DashboardUserCounts {
|
||||
total: users.iter().filter(|user| !user.is_deleted).count() as u64,
|
||||
..Default::default()
|
||||
},
|
||||
active_days,
|
||||
consecutive_active_days,
|
||||
activity_days,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -22,6 +22,486 @@ use aether_data_contracts::repository::usage::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_model_performance_merges_provider_samples_without_pagination() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let at = chrono::DateTime::parse_from_rfc3339("2026-09-12T10:05:00Z").unwrap();
|
||||
let mut records = Vec::new();
|
||||
for (index, provider, first_byte, response_time, output_tokens) in [
|
||||
(0, "provider-a", 100, 1100, 100),
|
||||
(1, "provider-a", 300, 1300, 200),
|
||||
(2, "provider-b", 500, 2500, 400),
|
||||
] {
|
||||
let mut row = sample_usage(&format!("model-sample-{index}"), at.timestamp());
|
||||
row.provider_id = Some(provider.into());
|
||||
row.model = "shared-model".into();
|
||||
row.target_model = Some(format!("{provider}-deployment"));
|
||||
row.first_byte_time_ms = Some(first_byte);
|
||||
row.response_time_ms = Some(response_time);
|
||||
row.output_tokens = output_tokens;
|
||||
row.is_stream = true;
|
||||
records.push(row);
|
||||
}
|
||||
let mut failed = sample_usage("model-failed", at.timestamp());
|
||||
failed.model = "shared-model".into();
|
||||
failed.status = "failed".into();
|
||||
failed.first_byte_time_ms = None;
|
||||
failed.response_time_ms = None;
|
||||
records.push(failed);
|
||||
let mut pending = sample_usage("model-pending", at.timestamp());
|
||||
pending.model = "pending-model".into();
|
||||
pending.status = "pending".into();
|
||||
pending.first_byte_time_ms = None;
|
||||
pending.response_time_ms = None;
|
||||
records.push(pending);
|
||||
|
||||
let repo = InMemoryUsageReadRepository::seed(records);
|
||||
let query = UsageAnalyticsQuery {
|
||||
from_unix_ms: (at - chrono::Duration::minutes(5)).timestamp_millis() as u64,
|
||||
to_unix_ms: (at + chrono::Duration::minutes(55)).timestamp_millis() as u64,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Performance,
|
||||
limit: 1,
|
||||
offset: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let result = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(result.model_rows.len(), 2);
|
||||
assert_eq!(result.model_rows[0].id.as_deref(), Some("shared-model"));
|
||||
assert!(result
|
||||
.model_rows
|
||||
.iter()
|
||||
.all(|row| row.bucket_start.is_none()));
|
||||
let metrics = &result.model_rows[0].metrics;
|
||||
assert_eq!(metrics.request_count, 4);
|
||||
assert_eq!(metrics.successful_request_count, 3);
|
||||
assert_eq!(metrics.failed_request_count, 1);
|
||||
assert_eq!(metrics.first_byte_sample_count, 3);
|
||||
assert_eq!(metrics.first_byte_sum_ms, 900.0);
|
||||
assert_eq!(metrics.latency_sample_count, 3);
|
||||
assert_eq!(metrics.latency_sum_ms, 4900.0);
|
||||
assert_eq!(metrics.output_tps_sample_count, 3);
|
||||
assert_eq!(metrics.output_tps_sum, 500.0);
|
||||
assert_eq!(result.model_rows[1].metrics.first_byte_sample_count, 0);
|
||||
assert_eq!(result.model_rows[1].metrics.in_flight_request_count, 1);
|
||||
assert_eq!(
|
||||
result
|
||||
.model_rows
|
||||
.iter()
|
||||
.map(|row| row.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
result.summary.request_count
|
||||
);
|
||||
|
||||
let filtered = repo
|
||||
.query_usage_analytics(&UsageAnalyticsQuery {
|
||||
model: Some("shared-model".into()),
|
||||
..query
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(filtered.model_rows, vec![result.model_rows[0].clone()]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_chart_hour_buckets_are_utc_in_half_hour_zones() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let at = chrono::DateTime::parse_from_rfc3339("2026-09-12T10:05:00Z").unwrap();
|
||||
let repo = InMemoryUsageReadRepository::seed([sample_usage("hour-zone", at.timestamp())]);
|
||||
let query = UsageAnalyticsQuery {
|
||||
from_unix_ms: (at - chrono::Duration::minutes(5)).timestamp_millis() as u64,
|
||||
to_unix_ms: (at + chrono::Duration::minutes(55)).timestamp_millis() as u64,
|
||||
timezone: "Asia/Kolkata".into(),
|
||||
view: UsageAnalyticsView::DashboardCharts,
|
||||
granularity: UsageAnalyticsGranularity::Hour,
|
||||
limit: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let charts = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(charts.summary.request_count, 1);
|
||||
assert_eq!(charts.rows.len(), 1);
|
||||
assert_eq!(charts.rows[0].metrics.request_count, 1);
|
||||
assert_eq!(
|
||||
charts.rows[0].bucket_start,
|
||||
charts.model_rows[0].bucket_start
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_chart_days_survive_skipped_midnight() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let moments = [
|
||||
"2026-09-05T12:00:00Z",
|
||||
"2026-09-06T12:00:00Z",
|
||||
"2026-09-07T12:00:00Z",
|
||||
];
|
||||
let repo = InMemoryUsageReadRepository::seed(moments.map(|moment| {
|
||||
sample_usage(
|
||||
moment,
|
||||
chrono::DateTime::parse_from_rfc3339(moment)
|
||||
.unwrap()
|
||||
.timestamp(),
|
||||
)
|
||||
}));
|
||||
let query = UsageAnalyticsQuery {
|
||||
from_unix_ms: chrono::DateTime::parse_from_rfc3339("2026-09-05T04:00:00Z")
|
||||
.unwrap()
|
||||
.timestamp_millis() as u64,
|
||||
to_unix_ms: chrono::DateTime::parse_from_rfc3339("2026-09-08T03:00:00Z")
|
||||
.unwrap()
|
||||
.timestamp_millis() as u64,
|
||||
timezone: "America/Santiago".into(),
|
||||
view: UsageAnalyticsView::DashboardCharts,
|
||||
limit: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let charts = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(charts.summary.request_count, 3);
|
||||
assert_eq!(charts.rows.len(), 3);
|
||||
assert_eq!(charts.model_rows.len(), 3);
|
||||
for (series, model) in charts.rows.iter().zip(&charts.model_rows) {
|
||||
assert_eq!(series.metrics.request_count, 1);
|
||||
assert_eq!(series.bucket_start, model.bucket_start);
|
||||
assert_eq!(
|
||||
series.metrics.billable_amount,
|
||||
model.metrics.billable_amount
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_dashboard_keeps_all_history_and_chart_dimensions() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let now = chrono::Utc::now();
|
||||
let mut old = sample_usage(
|
||||
"old-dashboard",
|
||||
(now - chrono::Duration::days(800)).timestamp(),
|
||||
);
|
||||
old.actual_total_cost_usd = 2.0;
|
||||
old.model = "old-model".into();
|
||||
let mut current = sample_usage("current-dashboard", now.timestamp());
|
||||
current.actual_total_cost_usd = 0.5;
|
||||
current.model = "new-model".into();
|
||||
let repo = InMemoryUsageReadRepository::seed([old, current]);
|
||||
let dashboard = repo
|
||||
.query_dashboard_analytics(&UsageDashboardAnalyticsQuery {
|
||||
timezone: "Asia/Shanghai".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(dashboard.today.summary.request_count, 1);
|
||||
assert_eq!(dashboard.total.summary.request_count, 2);
|
||||
assert_eq!(
|
||||
dashboard.today.summary.billable_amount.as_deref(),
|
||||
Some("0.50000000")
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard.total.summary.billable_amount.as_deref(),
|
||||
Some("2.50000000")
|
||||
);
|
||||
assert_eq!(dashboard.today.read_revision, dashboard.total.read_revision);
|
||||
assert_eq!(dashboard.history_complete, None);
|
||||
let charts = repo
|
||||
.query_usage_analytics(&UsageAnalyticsQuery {
|
||||
from_unix_ms: (now - chrono::Duration::days(1)).timestamp_millis() as u64,
|
||||
to_unix_ms: (now + chrono::Duration::seconds(1)).timestamp_millis() as u64,
|
||||
timezone: "Asia/Shanghai".into(),
|
||||
view: UsageAnalyticsView::DashboardCharts,
|
||||
limit: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(charts.model_rows.len(), 1);
|
||||
assert_eq!(charts.model_rows[0].id.as_deref(), Some("new-model"));
|
||||
assert_eq!(charts.provider_rows.len(), 1);
|
||||
assert_eq!(
|
||||
charts.model_rows[0].metrics.billable_amount,
|
||||
charts.summary.billable_amount
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_dashboard_total_keeps_card_coverage_without_historical_diagnostics() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let now = chrono::Utc::now();
|
||||
let mut known = sample_usage("dashboard-total-known", now.timestamp());
|
||||
known.request_metadata = Some(json!({"analytics_attribution":{"is_standalone":false}}));
|
||||
let mut missing = sample_usage(
|
||||
"dashboard-total-missing",
|
||||
(now - chrono::Duration::days(800)).timestamp(),
|
||||
);
|
||||
missing.billing_status = "pending".into();
|
||||
missing.request_metadata =
|
||||
Some(json!({"usage_available":false,"usage_pricing_available":false}));
|
||||
let mut session = sample_usage("dashboard-total-session", now.timestamp());
|
||||
session.request_metadata = Some(json!({"analytics_attribution":{"record_kind":"session"}}));
|
||||
let future = sample_usage(
|
||||
"dashboard-total-future",
|
||||
(now + chrono::Duration::days(1)).timestamp(),
|
||||
);
|
||||
let repo = InMemoryUsageReadRepository::seed([known, missing, session, future])
|
||||
.with_analytics_allocations([UsageAnalyticsAllocation {
|
||||
request_id: "dashboard-total-known".into(),
|
||||
complete: true,
|
||||
wallet_debit_amount: Some("0.18000000".into()),
|
||||
..Default::default()
|
||||
}]);
|
||||
let dashboard = repo
|
||||
.query_dashboard_analytics(&UsageDashboardAnalyticsQuery {
|
||||
timezone: "UTC".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let total = dashboard.total.summary;
|
||||
assert_eq!(total.request_count, 2);
|
||||
assert_eq!(total.total_tokens, 150);
|
||||
assert_eq!(total.billable_amount.as_deref(), Some("0.18000000"));
|
||||
assert_eq!(total.usage_available_count, 1);
|
||||
assert_eq!(total.pricing_available_count, 1);
|
||||
assert_eq!(total.settled_count, 1);
|
||||
assert_eq!(total.allocation_available_count, 1);
|
||||
assert!(total.latency_p95_ms.is_none());
|
||||
assert!(total.wallet_debit_amount.is_none());
|
||||
assert_eq!(dashboard.today.summary.latency_p95_ms, Some(420.0));
|
||||
assert_eq!(dashboard.today.summary.successful_request_count, 1);
|
||||
assert_eq!(
|
||||
dashboard.today.summary.wallet_debit_amount.as_deref(),
|
||||
Some("0.18000000")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_employee_roster_and_allocations_are_global() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
use aether_data_contracts::repository::users::StoredUserSummary;
|
||||
let mut usage = sample_usage("overview-request", 1_700_000_000);
|
||||
usage.user_id = Some("alice".into());
|
||||
usage.request_metadata = Some(json!({"analytics_attribution":{"is_standalone":false}}));
|
||||
usage.billing_status = "settled".into();
|
||||
usage.actual_total_cost_usd = 0.00000001;
|
||||
let repo = InMemoryUsageReadRepository::seed([usage])
|
||||
.with_analytics_users([
|
||||
StoredUserSummary::new(
|
||||
"alice".into(),
|
||||
"Alice".into(),
|
||||
None,
|
||||
"user".into(),
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.unwrap(),
|
||||
StoredUserSummary::new(
|
||||
"zero".into(),
|
||||
"Zero".into(),
|
||||
None,
|
||||
"user".into(),
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.unwrap(),
|
||||
])
|
||||
.with_analytics_allocations([UsageAnalyticsAllocation {
|
||||
request_id: "overview-request".into(),
|
||||
quota_covered_amount: Some("0.00000000".into()),
|
||||
wallet_consumed_amount: Some("0.00000001".into()),
|
||||
wallet_debit_amount: Some("0.00000000".into()),
|
||||
complete: true,
|
||||
..Default::default()
|
||||
}]);
|
||||
let mut query = UsageAnalyticsQuery {
|
||||
from_unix_ms: 1_700_000_000_000,
|
||||
to_unix_ms: 1_700_000_060_000,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Users,
|
||||
limit: 1,
|
||||
descending: true,
|
||||
..Default::default()
|
||||
};
|
||||
let first = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(first.total, 2);
|
||||
assert_eq!(first.users[0].user_id, "alice");
|
||||
assert_eq!(first.user_summary.as_ref().unwrap().user_count, 2);
|
||||
assert_eq!(first.user_summary.as_ref().unwrap().active_user_count, 1);
|
||||
assert!(first.user_finance_summary.is_none());
|
||||
assert!(first.user_payments.is_none());
|
||||
assert!(first.users[0].finance.is_none());
|
||||
assert_eq!(
|
||||
first.summary.wallet_consumed_amount.as_deref(),
|
||||
Some("0.00000001")
|
||||
);
|
||||
assert_eq!(
|
||||
first.summary.wallet_debit_amount.as_deref(),
|
||||
Some("0.00000000")
|
||||
);
|
||||
query.offset = 1;
|
||||
let second = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(second.users[0].user_id, "zero");
|
||||
assert_eq!(second.users[0].metrics.request_count, 0);
|
||||
assert_eq!(second.user_summary, first.user_summary);
|
||||
let searched = repo
|
||||
.query_usage_analytics(&UsageAnalyticsQuery {
|
||||
search: Some("Zero".into()),
|
||||
offset: 0,
|
||||
..query.clone()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(searched.user_summary.as_ref().unwrap().user_count, 1);
|
||||
assert_eq!(searched.user_summary.as_ref().unwrap().active_user_count, 0);
|
||||
assert_eq!(searched.summary.request_count, 0);
|
||||
query.from_unix_ms = query.to_unix_ms;
|
||||
query.to_unix_ms += 60_000;
|
||||
let outside = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(outside.summary.request_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_employee_is_grouped_by_account_owner() {
|
||||
use aether_data_contracts::repository::{usage::*, users::StoredUserSummary};
|
||||
let mut usage = sample_usage("trusted-request", 1_700_000_000);
|
||||
usage.user_id = Some("owner".into());
|
||||
usage.request_metadata =
|
||||
Some(json!({"analytics_attribution":{"is_standalone":false,"actor_user_id":"actor"}}));
|
||||
let repo = InMemoryUsageReadRepository::seed([usage]).with_analytics_users(
|
||||
["owner", "actor"].map(|id| {
|
||||
StoredUserSummary::new(id.into(), id.into(), None, "user".into(), true, false).unwrap()
|
||||
}),
|
||||
);
|
||||
let query = UsageAnalyticsQuery {
|
||||
from_unix_ms: 1_700_000_000_000,
|
||||
to_unix_ms: 1_700_000_060_000,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Users,
|
||||
attribution_kind: Some("employee".into()),
|
||||
has_usage: Some(true),
|
||||
limit: 100,
|
||||
..Default::default()
|
||||
};
|
||||
let result = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(result.total, 1);
|
||||
assert_eq!(result.users[0].user_id, "owner");
|
||||
let detail = repo
|
||||
.query_usage_analytics(&UsageAnalyticsQuery {
|
||||
actor_user_id: Some("owner".into()),
|
||||
..query
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(detail.users[0], result.users[0]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_legacy_requests_use_existing_key_account_flags() {
|
||||
use aether_data_contracts::repository::{usage::*, users::StoredUserSummary};
|
||||
let snapshots = [("member-key", false), ("standalone-key", true)].map(|(id, standalone)| {
|
||||
(
|
||||
None,
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
"user-1".into(),
|
||||
"alice".into(),
|
||||
None,
|
||||
"user".into(),
|
||||
"local".into(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id.into(),
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
standalone,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
)
|
||||
});
|
||||
let keys = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(snapshots));
|
||||
let rows = ["member-key", "standalone-key", "deleted-key"].map(|id| {
|
||||
let mut usage = sample_usage(id, 1_700_000_000);
|
||||
usage.api_key_id = Some(id.into());
|
||||
usage.request_metadata = None;
|
||||
usage
|
||||
});
|
||||
let repo = InMemoryUsageReadRepository::seed(rows)
|
||||
.with_auth_api_key_repository(keys)
|
||||
.with_analytics_users([StoredUserSummary::new(
|
||||
"user-1".into(),
|
||||
"alice".into(),
|
||||
None,
|
||||
"user".into(),
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.unwrap()]);
|
||||
let mut query = UsageAnalyticsQuery {
|
||||
from_unix_ms: 1_700_000_000_000,
|
||||
to_unix_ms: 1_700_000_060_000,
|
||||
timezone: "UTC".into(),
|
||||
view: UsageAnalyticsView::Consumption,
|
||||
limit: 100,
|
||||
..Default::default()
|
||||
};
|
||||
let result = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(result.summary.request_count, 3);
|
||||
assert_eq!(result.summary.usage_active_users, 1);
|
||||
assert_eq!(result.summary.trusted_attribution_count, 1);
|
||||
for (key, kind, source, user) in [
|
||||
("member-key", "employee", "user_account", Some("user-1")),
|
||||
("standalone-key", "standalone", "standalone_key", None),
|
||||
("deleted-key", "unknown", "unknown", None),
|
||||
] {
|
||||
let row = result
|
||||
.consumption
|
||||
.iter()
|
||||
.find(|row| row.request_id == key)
|
||||
.unwrap();
|
||||
assert_eq!(row.attribution_kind, kind);
|
||||
assert_eq!(row.attribution_source, source);
|
||||
assert_eq!(row.user_id.as_deref(), user);
|
||||
assert_eq!(row.credential_owner_id.as_deref(), Some("user-1"));
|
||||
}
|
||||
query.view = UsageAnalyticsView::Users;
|
||||
query.attribution_kind = Some("employee".into());
|
||||
let employees = repo.query_usage_analytics(&query).await.unwrap();
|
||||
assert_eq!(employees.users[0].metrics.request_count, 1);
|
||||
assert_eq!(employees.summary.request_count, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overview_memory_health_does_not_treat_upstream_cancellation_as_client() {
|
||||
use aether_data_contracts::repository::usage::*;
|
||||
let mut upstream = sample_usage("upstream-cancel", 1_700_000_000);
|
||||
upstream.status = "cancelled".into();
|
||||
upstream.request_metadata =
|
||||
Some(json!({"analytics_failure":{"origin":"upstream","reason":"provider_cancelled"}}));
|
||||
let mut client = upstream.clone();
|
||||
client.request_id = "client-cancel".into();
|
||||
client.request_metadata =
|
||||
Some(json!({"analytics_failure":{"origin":"client","reason":"downstream_disconnect"}}));
|
||||
let repo = InMemoryUsageReadRepository::seed([upstream, client]);
|
||||
let summary = repo
|
||||
.summarize_health_observations(&HealthObservationQuery {
|
||||
from_unix_ms: 1_700_000_000_000,
|
||||
to_unix_ms: 1_700_000_060_000,
|
||||
object_kind: HealthObservationObjectKind::Model,
|
||||
object_values: None,
|
||||
segments: 4,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(summary.overall.request_count, 2);
|
||||
assert_eq!(summary.overall.service_failed_count, 1);
|
||||
assert_eq!(summary.overall.excluded_count, 1);
|
||||
}
|
||||
|
||||
fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
@@ -2763,3 +3243,99 @@ async fn summarize_usage_provider_performance_computes_tps_and_top_provider_time
|
||||
assert_eq!(without_timeline.providers, summary.providers);
|
||||
assert!(without_timeline.timeline.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn future_dashboard_summary_excludes_old_rows_and_preserves_canonical_cache_samples() {
|
||||
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
|
||||
let now = chrono::Utc::now();
|
||||
let since = now - chrono::Duration::minutes(5);
|
||||
let mut included = sample_usage(
|
||||
"future-summary",
|
||||
(now - chrono::Duration::seconds(1)).timestamp(),
|
||||
);
|
||||
included.api_key_id = None;
|
||||
included.api_format = Some("claude:messages".into());
|
||||
included.endpoint_api_format = Some("claude:messages".into());
|
||||
included.input_tokens = 80;
|
||||
included.output_tokens = 20;
|
||||
included.total_tokens = 150;
|
||||
included.cache_read_input_tokens = 40;
|
||||
included.cache_creation_input_tokens = 10;
|
||||
included.response_time_ms = Some(800);
|
||||
included.first_byte_time_ms = Some(100);
|
||||
included.is_stream = false;
|
||||
included.request_metadata = Some(json!({"upstream_is_stream": true}));
|
||||
included.actual_total_cost_usd = 0.12345678;
|
||||
included.billing_status = "settled".into();
|
||||
let mut excluded = included.clone();
|
||||
excluded.request_id = "before-activation".into();
|
||||
excluded.created_at_unix_ms = (since - chrono::Duration::days(500)).timestamp() as u64;
|
||||
let repo =
|
||||
InMemoryUsageReadRepository::seed([included, excluded]).with_dashboard_stats_since(since);
|
||||
let first = repo
|
||||
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
|
||||
timezone: "Asia/Kathmandu".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(first.total.request_count, 1);
|
||||
assert_eq!(first.total.total_tokens, 150);
|
||||
assert_eq!(first.total.billable_amount.as_deref(), Some("0.12345678"));
|
||||
assert_eq!(first.today.cache_input_tokens, 130);
|
||||
assert_eq!(first.today.cache_read_tokens, 40);
|
||||
assert_eq!(first.today.first_byte_sample_count, 1);
|
||||
assert_eq!(first.today.response_sample_count, 1);
|
||||
assert_eq!(first.today.stream_requests, 1);
|
||||
assert_eq!(first.today.standard_requests, 0);
|
||||
assert_eq!(first.active_days, 1);
|
||||
assert_eq!(first.consecutive_active_days, 1);
|
||||
assert_eq!(
|
||||
first.activity_days.iter().map(|d| d.requests).sum::<u64>(),
|
||||
1
|
||||
);
|
||||
// Audit retention does not own the additive projection.
|
||||
repo.by_request_id.write().unwrap().clear();
|
||||
let retained = repo
|
||||
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
|
||||
timezone: "Asia/Kathmandu".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(retained.total, first.total);
|
||||
assert_eq!(retained.stats_since, since.to_rfc3339());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn future_dashboard_summary_streak_uses_local_days_before_heatmap_truncation() {
|
||||
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
|
||||
use chrono::TimeZone;
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let tz = chrono_tz::Asia::Kathmandu;
|
||||
let today = now.with_timezone(&tz).date_naive();
|
||||
let rows = (1..=400).map(|offset| {
|
||||
let day = today - chrono::Duration::days(offset);
|
||||
let at = tz
|
||||
.from_local_datetime(&day.and_hms_opt(12, 0, 0).unwrap())
|
||||
.single()
|
||||
.unwrap();
|
||||
sample_usage(&format!("streak-{offset}"), at.timestamp())
|
||||
});
|
||||
let repo = InMemoryUsageReadRepository::seed(rows)
|
||||
.with_dashboard_stats_since(now - chrono::Duration::days(401));
|
||||
let summary = repo
|
||||
.query_dashboard_summary(&UsageDashboardAnalyticsQuery {
|
||||
timezone: tz.to_string(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(summary.today.request_count, 0);
|
||||
assert_eq!(summary.active_days, 400);
|
||||
assert_eq!(summary.consecutive_active_days, 400);
|
||||
assert_eq!(summary.activity_days.len(), 364);
|
||||
assert_eq!(
|
||||
summary.activity_days.last().unwrap().date,
|
||||
today.pred_opt().unwrap().to_string()
|
||||
);
|
||||
}
|
||||
|
||||
@@ -831,6 +831,12 @@ impl WalletReadRepository for InMemoryWalletRepository {
|
||||
.as_deref()
|
||||
.is_none_or(|expected| wallet.status == expected)
|
||||
})
|
||||
.filter(|wallet| {
|
||||
query
|
||||
.user_id
|
||||
.as_deref()
|
||||
.is_none_or(|expected| wallet.user_id.as_deref() == Some(expected))
|
||||
})
|
||||
.filter(|wallet| match query.owner_type.as_deref() {
|
||||
Some("user") => wallet.user_id.is_some(),
|
||||
Some("api_key") => wallet.api_key_id.is_some(),
|
||||
@@ -3418,6 +3424,7 @@ mod tests {
|
||||
|
||||
let page = repository
|
||||
.list_admin_wallets(&AdminWalletListQuery {
|
||||
user_id: None,
|
||||
status: Some("active".to_string()),
|
||||
owner_type: Some("api_key".to_string()),
|
||||
limit: 1,
|
||||
@@ -3430,6 +3437,31 @@ mod tests {
|
||||
assert_eq!(page.items.len(), 1);
|
||||
assert_eq!(page.items[0].id, "wallet-3");
|
||||
assert_eq!(page.items[0].updated_at_unix_secs, Some(110));
|
||||
|
||||
let query = AdminWalletListQuery {
|
||||
user_id: Some("user-2".to_string()),
|
||||
owner_type: Some("user".to_string()),
|
||||
limit: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let selected = repository.list_admin_wallets(&query).await.unwrap();
|
||||
assert_eq!(selected.total, 1);
|
||||
assert_eq!(selected.items[0].id, "wallet-2");
|
||||
let missing = repository
|
||||
.list_admin_wallets(&AdminWalletListQuery {
|
||||
user_id: Some("missing-user".to_string()),
|
||||
..query.clone()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(missing.total, 0);
|
||||
assert!(missing.items.is_empty());
|
||||
let beyond_page = repository
|
||||
.list_admin_wallets(&AdminWalletListQuery { offset: 1, ..query })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(beyond_page.total, 1);
|
||||
assert!(beyond_page.items.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
Reference in New Issue
Block a user