feat(stats): 新增聚合读路径与回填机制并重构 dashboard/usage 读取链路

- 新增 stats_user_summary 及 user_daily_provider/api_format/cost_savings 等聚合表
- 扩展 stats_daily/hourly 有效 token 与响应时间等字段,maintenance runtime 同步写入
- 新增 backfill 模块与 --apply-backfills 命令补齐历史聚合数据
- 重写 dashboard_filters、usage_heatmap、user_rollups 查询改走聚合表
- 同步更新 baseline_v2.sql 与 migration 集,README/dev.sh 补充回填用法
This commit is contained in:
fawney19
2026-04-22 17:29:39 +08:00
parent 063ef02306
commit a5e6bd3b62
45 changed files with 15323 additions and 1039 deletions

View File

@@ -2765,6 +2765,7 @@ mod tests {
Some(200),
false,
9,
0,
1.75,
false,
)
@@ -2785,6 +2786,7 @@ mod tests {
None,
true,
2,
0,
0.5,
true,
)

View File

@@ -166,6 +166,7 @@ pub(super) async fn build_admin_create_api_key_response(
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
},
)

View File

@@ -222,6 +222,7 @@ pub(super) fn sample_monitoring_export_api_key(
None,
false,
0,
0,
0.0,
false,
)

View File

@@ -1,10 +1,13 @@
use crate::handlers::admin::request::AdminAppState;
use crate::query::usage_heatmap::{
list_usage_heatmap_aggregate_rows, read_stats_daily_cutoff_date,
};
use crate::GatewayError;
use aether_admin::observability::stats::round_to;
use aether_admin::observability::usage::{
admin_usage_data_unavailable_response, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data_contracts::repository::usage::UsageDailyHeatmapQuery;
use aether_data_contracts::repository::usage::{StoredUsageDailySummary, UsageDailyHeatmapQuery};
use axum::{
body::Body,
response::{IntoResponse, Response},
@@ -34,17 +37,13 @@ pub(super) async fn build_admin_usage_heatmap_response(
)
.unwrap_or_default();
let summaries = state
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
created_from_unix_secs,
user_id: None,
admin_mode: true,
})
.await?;
let summaries =
build_admin_heatmap_summaries(state, created_from_unix_secs, start_date, today).await?;
let grouped: BTreeMap<String, _> = summaries.into_iter().map(|s| (s.date.clone(), s)).collect();
let mut max_requests = 0_u64;
let mut active_days = 0_u64;
let mut cursor = start_date;
let mut days = Vec::new();
while cursor <= today {
@@ -61,6 +60,9 @@ pub(super) async fn build_admin_usage_heatmap_response(
(0, 0, 0.0, 0.0)
};
max_requests = max_requests.max(requests);
if requests > 0 {
active_days = active_days.saturating_add(1);
}
days.push(json!({
"date": date_str,
"requests": requests,
@@ -77,8 +79,59 @@ pub(super) async fn build_admin_usage_heatmap_response(
"start_date": start_date.to_string(),
"end_date": today.to_string(),
"total_days": days.len(),
"active_days": active_days,
"max_requests": max_requests,
"days": days,
}))
.into_response())
}
async fn build_admin_heatmap_summaries(
state: &AdminAppState<'_>,
created_from_unix_secs: u64,
start_date: chrono::NaiveDate,
today: chrono::NaiveDate,
) -> Result<Vec<StoredUsageDailySummary>, GatewayError> {
let query = UsageDailyHeatmapQuery {
created_from_unix_secs,
user_id: None,
admin_mode: true,
};
let Some(pool) = state.app().postgres_pool() else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let Some(cutoff_date) = read_stats_daily_cutoff_date(&pool).await? else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let cutoff_day = cutoff_date.date_naive().min(today);
let mut summaries =
list_usage_heatmap_aggregate_rows(&pool, start_date, cutoff_day, None).await?;
let raw_start_date = start_date.max(cutoff_day);
if raw_start_date <= today {
let raw_start_of_day = raw_start_date
.and_hms_opt(0, 0, 0)
.expect("heatmap day start should be valid");
let raw_created_from_unix_secs = u64::try_from(
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
raw_start_of_day,
chrono::Utc,
)
.timestamp(),
)
.unwrap_or_default();
summaries.extend(
state
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
created_from_unix_secs: raw_created_from_unix_secs,
user_id: None,
admin_mode: true,
})
.await?,
);
}
summaries.sort_by(|left, right| left.date.cmp(&right.date));
Ok(summaries)
}

View File

@@ -297,6 +297,7 @@ impl<'a> AdminAppState<'a> {
json!(key.auto_delete_on_expiry),
),
("total_requests".to_string(), json!(key.total_requests)),
("total_tokens".to_string(), json!(key.total_tokens)),
("total_cost_usd".to_string(), json!(key.total_cost_usd)),
(
"wallet".to_string(),

View File

@@ -1982,6 +1982,11 @@ impl<'a> AdminAppState<'a> {
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
@@ -2047,6 +2052,7 @@ impl<'a> AdminAppState<'a> {
|| key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(format!(
@@ -2077,6 +2083,7 @@ impl<'a> AdminAppState<'a> {
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
})
.await?;
@@ -2169,6 +2176,11 @@ impl<'a> AdminAppState<'a> {
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
@@ -2225,6 +2237,7 @@ impl<'a> AdminAppState<'a> {
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("force_capabilities")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(
@@ -2263,6 +2276,7 @@ impl<'a> AdminAppState<'a> {
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
},
)

View File

@@ -142,6 +142,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
expires_at_unix_secs: None,
auto_delete_on_expiry: false,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
})
.await?

View File

@@ -4,6 +4,7 @@ use super::support::{
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{query_param_optional_bool, query_param_value};
use crate::query::user_rollups::list_user_usage_totals_from_stats_summary;
use crate::GatewayError;
use axum::{
body::Body,
@@ -42,10 +43,19 @@ pub(in super::super) async fn build_admin_list_users_response(
.iter()
.map(|row| row.id.clone())
.collect::<Vec<_>>();
let usage_totals_future = async {
let Some(pool) = state.app().postgres_pool() else {
return state.summarize_usage_totals_by_user_ids(&user_ids).await;
};
match list_user_usage_totals_from_stats_summary(&pool, &user_ids).await? {
Some(items) => Ok(items),
None => state.summarize_usage_totals_by_user_ids(&user_ids).await,
}
};
let (auth_rows_result, wallet_rows_result, usage_totals_result) = tokio::join!(
state.list_user_auth_by_ids(&user_ids),
state.list_wallet_snapshots_by_user_ids(&user_ids),
state.summarize_usage_totals_by_user_ids(&user_ids),
usage_totals_future,
);
let auth_by_user_id = auth_rows_result?
.into_iter()

View File

@@ -554,6 +554,7 @@ pub(super) async fn handle_users_me_api_key_create(
expires_at_unix_secs: None,
auto_delete_on_expiry: false,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
};
let Some(created) = (match state.create_user_api_key(record).await {

View File

@@ -4,9 +4,10 @@ use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, StoredUsageBreakdownSummaryRow, UsageAuditKeywordSearchQuery,
UsageAuditListQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageDashboardSummaryQuery,
StoredRequestUsageAudit, StoredUsageBreakdownSummaryRow, StoredUsageDailySummary,
UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageBreakdownGroupBy,
UsageBreakdownSummaryQuery, UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery,
UsageDashboardSummaryQuery,
};
use axum::{
body::Body,
@@ -17,6 +18,9 @@ use axum::{
use chrono::Utc;
use serde_json::json;
use crate::query::usage_heatmap::{
list_usage_heatmap_aggregate_rows, read_stats_daily_cutoff_date,
};
use crate::GatewayError;
use super::{
@@ -957,15 +961,14 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
)
.unwrap_or_default();
let summaries = match state
.summarize_usage_daily_heatmap(
&aether_data_contracts::repository::usage::UsageDailyHeatmapQuery {
created_from_unix_secs,
user_id: Some(auth.user.id.clone()),
admin_mode: false,
},
)
.await
let summaries = match build_usage_heatmap_summaries(
state,
created_from_unix_secs,
start_date,
today,
Some(auth.user.id.as_str()),
)
.await
{
Ok(value) => value,
Err(err) => {
@@ -982,6 +985,7 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
summaries.into_iter().map(|s| (s.date.clone(), s)).collect();
let mut max_requests = 0_u64;
let mut active_days = 0_u64;
let mut cursor = start_date;
let mut days = Vec::new();
while cursor <= today {
@@ -998,6 +1002,9 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
(0, 0, 0.0, 0.0)
};
max_requests = max_requests.max(requests);
if requests > 0 {
active_days = active_days.saturating_add(1);
}
let mut day = json!({
"date": date_str,
"requests": requests,
@@ -1017,12 +1024,66 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
"start_date": start_date.to_string(),
"end_date": today.to_string(),
"total_days": days.len(),
"active_days": active_days,
"max_requests": max_requests,
"days": days,
}))
.into_response()
}
async fn build_usage_heatmap_summaries(
state: &AppState,
created_from_unix_secs: u64,
start_date: chrono::NaiveDate,
today: chrono::NaiveDate,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, GatewayError> {
let query = aether_data_contracts::repository::usage::UsageDailyHeatmapQuery {
created_from_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
admin_mode: user_id.is_none(),
};
let Some(pool) = state.postgres_pool() else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let Some(cutoff_date) = read_stats_daily_cutoff_date(&pool).await? else {
return state.summarize_usage_daily_heatmap(&query).await;
};
let cutoff_day = cutoff_date.date_naive().min(today);
let mut summaries =
list_usage_heatmap_aggregate_rows(&pool, start_date, cutoff_day, user_id).await?;
let raw_start_date = start_date.max(cutoff_day);
if raw_start_date <= today {
let raw_start_of_day = raw_start_date
.and_hms_opt(0, 0, 0)
.expect("heatmap day start should be valid");
let raw_created_from_unix_secs = u64::try_from(
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
raw_start_of_day,
chrono::Utc,
)
.timestamp(),
)
.unwrap_or_default();
summaries.extend(
state
.summarize_usage_daily_heatmap(
&aether_data_contracts::repository::usage::UsageDailyHeatmapQuery {
created_from_unix_secs: raw_created_from_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
admin_mode: user_id.is_none(),
},
)
.await?,
);
}
summaries.sort_by(|left, right| left.date.cmp(&right.date));
Ok(summaries)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;

View File

@@ -532,6 +532,9 @@ struct Args {
#[arg(long, default_value_t = false)]
migrate: bool,
#[arg(long, default_value_t = false)]
apply_backfills: bool,
/// Path to frontend static files directory (SPA). When set, the gateway
/// serves the frontend directly without nginx.
#[arg(long, env = "AETHER_GATEWAY_STATIC_DIR")]
@@ -613,7 +616,7 @@ struct Args {
impl Args {
fn runtime_config(&self) -> Result<ServiceRuntimeConfig, std::io::Error> {
let default_log_filter = if self.migrate {
let default_log_filter = if self.migrate || self.apply_backfills {
"aether_gateway=info,aether_data=info"
} else {
"aether_gateway=info"
@@ -765,6 +768,10 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
init_service_runtime(args.runtime_config()?)?;
return run_explicit_migrations(&args).await;
}
if args.apply_backfills {
init_service_runtime(args.runtime_config()?)?;
return run_explicit_backfills(&args).await;
}
let app_port = validate_app_port(args.app_port)?;
let bind_addr = gateway_bind_addr(app_port)?;
set_gateway_frontdoor_app_port(app_port);
@@ -921,6 +928,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
"aether-gateway data layer configured"
);
ensure_postgres_schema_is_current(&state).await?;
ensure_postgres_backfills_are_current(&state).await?;
let reset_stale_proxy_nodes = state.reset_stale_proxy_node_tunnel_statuses().await?;
if reset_stale_proxy_nodes > 0 {
info!(
@@ -1026,6 +1034,44 @@ async fn run_explicit_migrations(args: &Args) -> Result<(), Box<dyn std::error::
Ok(())
}
async fn run_explicit_backfills(args: &Args) -> Result<(), Box<dyn std::error::Error>> {
args.data.effective_postgres_url().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL is required when running --apply-backfills",
)
})?;
let state = AppState::new()?.with_data_config(args.data.to_config())?;
ensure_postgres_schema_is_current(&state).await?;
let pending = state
.pending_postgres_backfills()
.await?
.unwrap_or_default();
if pending.is_empty() {
info!(
pending_backfills = 0,
"database backfills already up to date"
);
return Ok(());
}
let next = pending
.first()
.expect("pending backfills should have a first element");
info!(
pending_backfills = pending.len(),
next_version = next.version,
next_description = %next.description,
pending_versions = %format_pending_backfills(&pending),
"running database backfills by explicit request..."
);
if state.run_postgres_backfills().await? {
info!("database backfills complete");
}
Ok(())
}
fn format_pending_migrations(pending: &[aether_data::migrate::PendingMigrationInfo]) -> String {
pending
.iter()
@@ -1034,6 +1080,30 @@ fn format_pending_migrations(pending: &[aether_data::migrate::PendingMigrationIn
.join(", ")
}
fn format_pending_backfills(pending: &[aether_data::backfill::PendingBackfillInfo]) -> String {
pending
.iter()
.map(|backfill| format!("{} ({})", backfill.version, backfill.description))
.collect::<Vec<_>>()
.join(", ")
}
async fn ensure_postgres_backfills_are_current(
state: &AppState,
) -> Result<(), Box<dyn std::error::Error>> {
let Some(pending) = state.pending_postgres_backfills().await? else {
return Ok(());
};
if pending.is_empty() {
return Ok(());
}
let next = pending
.first()
.expect("pending backfills should have a first element");
Err(pending_backfills_error(pending.len(), next.version, &next.description).into())
}
async fn ensure_postgres_schema_is_current(
state: &AppState,
) -> Result<(), Box<dyn std::error::Error>> {
@@ -1063,10 +1133,24 @@ fn pending_schema_error(
))
}
fn pending_backfills_error(
pending_count: usize,
next_version: i64,
next_description: &str,
) -> std::io::Error {
std::io::Error::other(format!(
"database backfills are behind by {} backfill(s); next pending backfill is {} ({})\nrun `aether-gateway --apply-backfills` before starting the service",
pending_count,
next_version,
next_description
))
}
#[cfg(test)]
mod tests {
use super::{
ensure_postgres_schema_is_current, pending_schema_error, resolve_healthcheck_url, Args,
ensure_postgres_backfills_are_current, ensure_postgres_schema_is_current,
pending_backfills_error, pending_schema_error, resolve_healthcheck_url, Args,
DeploymentTopologyArg, GatewayDataArgs, GatewayFrontdoorArgs, GatewayLogDestinationArg,
GatewayLogFormatArg, GatewayLogRotationArg, GatewayLoggingArgs, GatewayRateLimitArgs,
GatewayUsageArgs, NodeRoleArg, VideoTaskTruthSourceArg,
@@ -1081,6 +1165,7 @@ mod tests {
deployment_topology: DeploymentTopologyArg::SingleNode,
node_role: NodeRoleArg::All,
migrate: false,
apply_backfills: false,
static_dir: None,
video_task_truth_source_mode: VideoTaskTruthSourceArg::PythonSyncReport,
video_task_poller_interval_ms: 5_000,
@@ -1171,6 +1256,17 @@ mod tests {
assert_eq!(config.default_log_filter, "aether_gateway=info");
}
#[test]
fn apply_backfills_runtime_config_enables_data_logs() {
let mut args = test_args();
args.apply_backfills = true;
let config = args.runtime_config().expect("runtime config should build");
assert_eq!(
config.default_log_filter,
"aether_gateway=info,aether_data=info"
);
}
#[test]
fn pending_schema_error_mentions_explicit_migrate_command() {
let error = pending_schema_error(2, 20260413020000, "squash usage schema split");
@@ -1181,6 +1277,21 @@ mod tests {
assert!(message.contains("aether-gateway --migrate"));
}
#[test]
fn pending_backfills_error_mentions_explicit_apply_backfills_command() {
let message = pending_backfills_error(
1,
20260422110000,
"backfill stats aggregate read path support",
)
.to_string();
assert!(message.contains("database backfills are behind by 1 backfill(s)"));
assert!(message.contains("20260422110000"));
assert!(message.contains("backfill stats aggregate read path support"));
assert!(message.contains("aether-gateway --apply-backfills"));
assert!(message.contains("before starting the service"));
}
#[tokio::test]
async fn ensure_postgres_schema_is_current_is_noop_without_postgres_pool() {
let state = AppState::new().expect("state should build");
@@ -1189,6 +1300,14 @@ mod tests {
.expect("disabled data backend should not block startup");
}
#[tokio::test]
async fn ensure_postgres_backfills_are_current_is_noop_without_postgres_pool() {
let state = AppState::new().expect("state should build");
ensure_postgres_backfills_are_current(&state)
.await
.expect("disabled data backend should not block startup");
}
#[tokio::test]
async fn explicit_migrate_requires_postgres_url() {
let args = test_args();
@@ -1212,4 +1331,15 @@ mod tests {
assert!(message.contains("AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL"));
assert!(!message.contains("APP_PORT"));
}
#[tokio::test]
async fn explicit_backfills_require_postgres_url() {
let args = test_args();
let error = super::run_explicit_backfills(&args)
.await
.expect_err("missing postgres URL should fail");
let message = error.to_string();
assert!(message.contains("AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL"));
assert!(message.contains("--apply-backfills"));
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -165,9 +165,9 @@ pub(super) async fn run_wallet_daily_usage_aggregation_once(
pub(super) async fn run_stats_aggregation_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
) -> Result<bool, DataLayerError> {
let Some(summary) = perform_stats_aggregation_once(data).await? else {
return Ok(());
return Ok(false);
};
info!(
@@ -183,7 +183,7 @@ pub(super) async fn run_stats_aggregation_once(
user_rows = summary.user_rows,
"gateway aggregated daily stats tables"
);
Ok(())
Ok(true)
}
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
@@ -248,9 +248,9 @@ pub(super) async fn run_pending_cleanup_once(
pub(super) async fn run_stats_hourly_aggregation_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
) -> Result<bool, DataLayerError> {
let Some(summary) = perform_stats_hourly_aggregation_once(data).await? else {
return Ok(());
return Ok(false);
};
info!(
@@ -260,11 +260,12 @@ pub(super) async fn run_stats_hourly_aggregation_once(
hour_utc = %summary.hour_utc,
total_requests = summary.total_requests,
user_rows = summary.user_rows,
user_model_rows = summary.user_model_rows,
model_rows = summary.model_rows,
provider_rows = summary.provider_rows,
"gateway aggregated stats hourly tables"
);
Ok(())
Ok(true)
}
pub(super) async fn run_provider_checkin_once(state: &AppState) -> Result<(), GatewayError> {

View File

@@ -9,12 +9,22 @@ use super::{
postgres_error, stats_aggregation_target_day, system_config_bool, PercentileSummary,
StatsAggregationSummary, DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL, INSERT_STATS_DAILY_ERROR_SQL,
INSERT_STATS_SUMMARY_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
SELECT_LATEST_STATS_DAILY_DATE_SQL, SELECT_NEXT_STATS_DAILY_BUCKET_SQL,
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL, SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL,
SELECT_STATS_SUMMARY_TOTALS_SQL, UPDATE_STATS_SUMMARY_SQL, UPSERT_STATS_DAILY_API_KEY_SQL,
UPSERT_STATS_DAILY_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
UPSERT_STATS_USER_DAILY_SQL,
UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL, UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_SQL,
UPSERT_STATS_DAILY_COST_SAVINGS_PROVIDER_SQL, UPSERT_STATS_DAILY_COST_SAVINGS_SQL,
UPSERT_STATS_DAILY_MODEL_PROVIDER_SQL, UPSERT_STATS_DAILY_MODEL_SQL,
UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
UPSERT_STATS_USER_DAILY_API_FORMAT_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_SQL,
UPSERT_STATS_USER_DAILY_COST_SAVINGS_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_COST_SAVINGS_SQL,
UPSERT_STATS_USER_DAILY_MODEL_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_MODEL_SQL,
UPSERT_STATS_USER_DAILY_PROVIDER_SQL, UPSERT_STATS_USER_DAILY_SQL,
UPSERT_STATS_USER_SUMMARY_SQL,
};
pub(super) async fn perform_stats_aggregation_once(
@@ -28,109 +38,125 @@ pub(super) async fn perform_stats_aggregation_once(
}
let now_utc = Utc::now();
let day_start_utc = stats_aggregation_target_day(now_utc);
let target_day_utc = stats_aggregation_target_day(now_utc);
let Some(day_start_utc) = next_stats_aggregation_day(&pool, target_day_utc)
.await
.map_err(postgres_error)?
else {
return Ok(None);
};
perform_stats_aggregation_for_day(&pool, day_start_utc, now_utc)
.await
.map(Some)
.map_err(postgres_error)
}
async fn next_stats_aggregation_day(
pool: &aether_data::postgres::PostgresPool,
target_day_utc: DateTime<Utc>,
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
let latest_row = sqlx::query(SELECT_LATEST_STATS_DAILY_DATE_SQL)
.fetch_one(pool)
.await?;
let latest_day = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_date")?;
let search_from = latest_day
.map(|value| value + chrono::Duration::days(1))
.unwrap_or_else(|| {
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
});
let search_until = target_day_utc + chrono::Duration::days(1);
if search_from >= search_until {
return Ok(None);
}
let next_row = sqlx::query(SELECT_NEXT_STATS_DAILY_BUCKET_SQL)
.bind(search_from)
.bind(search_until)
.fetch_one(pool)
.await?;
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
Ok(next_bucket.filter(|value| *value <= target_day_utc))
}
async fn perform_stats_aggregation_for_day(
pool: &aether_data::postgres::PostgresPool,
day_start_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<StatsAggregationSummary, sqlx::Error> {
let day_end_utc = day_start_utc + chrono::Duration::days(1);
let mut tx = pool.begin().await.map_err(postgres_error)?;
let mut tx = pool.begin().await?;
let aggregate_row = sqlx::query(SELECT_STATS_DAILY_AGGREGATE_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.fetch_one(&mut *tx)
.await
.map_err(postgres_error)?;
let total_requests = aggregate_row
.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?;
let error_requests = aggregate_row
.try_get::<i64, _>("error_requests")
.map_err(postgres_error)?;
.await?;
let total_requests = aggregate_row.try_get::<i64, _>("total_requests")?;
let error_requests = aggregate_row.try_get::<i64, _>("error_requests")?;
let success_requests = total_requests.saturating_sub(error_requests);
let fallback_count = sqlx::query(SELECT_STATS_DAILY_FALLBACK_COUNT_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(vec!["success", "failed"])
.fetch_one(&mut *tx)
.await
.map_err(postgres_error)?
.try_get::<i64, _>("fallback_count")
.map_err(postgres_error)?;
.await?
.try_get::<i64, _>("fallback_count")?;
let response_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await
.map_err(postgres_error)?;
.await?;
let first_byte_percentiles = fetch_stats_daily_percentiles(
&mut tx,
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
day_start_utc,
day_end_utc,
)
.await
.map_err(postgres_error)?;
.await?;
sqlx::query(UPSERT_STATS_DAILY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(day_start_utc)
.bind(total_requests)
.bind(aggregate_row.try_get::<i64, _>("cache_hit_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("cache_hit_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_hit_requests")?)
.bind(aggregate_row.try_get::<i64, _>("completed_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_cache_read_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("completed_total_input_context")?)
.bind(aggregate_row.try_get::<f64, _>("completed_cache_creation_cost")?)
.bind(aggregate_row.try_get::<f64, _>("completed_cache_read_cost")?)
.bind(aggregate_row.try_get::<f64, _>("settled_total_cost")?)
.bind(aggregate_row.try_get::<i64, _>("settled_total_requests")?)
.bind(aggregate_row.try_get::<i64, _>("settled_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_output_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("settled_cache_read_tokens")?)
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
.bind(success_requests)
.bind(error_requests)
.bind(
aggregate_row
.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("cache_creation_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("cache_read_tokens")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("actual_total_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("input_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("output_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("cache_creation_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("cache_read_cost")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<f64, _>("avg_response_time_ms")
.map_err(postgres_error)?,
)
.bind(aggregate_row.try_get::<i64, _>("input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("effective_input_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("output_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_5m_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_1h_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("cache_read_tokens")?)
.bind(aggregate_row.try_get::<i64, _>("total_input_context")?)
.bind(aggregate_row.try_get::<f64, _>("total_cost")?)
.bind(aggregate_row.try_get::<f64, _>("actual_total_cost")?)
.bind(aggregate_row.try_get::<f64, _>("input_cost")?)
.bind(aggregate_row.try_get::<f64, _>("output_cost")?)
.bind(aggregate_row.try_get::<f64, _>("cache_creation_cost")?)
.bind(aggregate_row.try_get::<f64, _>("cache_read_cost")?)
.bind(aggregate_row.try_get::<f64, _>("response_time_sum_ms")?)
.bind(aggregate_row.try_get::<i64, _>("response_time_samples")?)
.bind(aggregate_row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(response_percentiles.p50)
.bind(response_percentiles.p90)
.bind(response_percentiles.p99)
@@ -138,47 +164,65 @@ pub(super) async fn perform_stats_aggregation_once(
.bind(first_byte_percentiles.p90)
.bind(first_byte_percentiles.p99)
.bind(fallback_count)
.bind(
aggregate_row
.try_get::<i64, _>("unique_models")
.map_err(postgres_error)?,
)
.bind(
aggregate_row
.try_get::<i64, _>("unique_providers")
.map_err(postgres_error)?,
)
.bind(aggregate_row.try_get::<i64, _>("unique_models")?)
.bind(aggregate_row.try_get::<i64, _>("unique_providers")?)
.bind(true)
.bind(now_utc)
.bind(now_utc)
.bind(now_utc)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
.await?;
let model_rows = upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let model_rows =
upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let provider_rows =
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_daily_cost_savings_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_daily_cost_savings_model_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
let api_key_rows =
upsert_stats_daily_api_key_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let error_rows = refresh_stats_daily_error_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
let user_rows = upsert_stats_user_daily_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
upsert_stats_daily_api_key_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let error_rows =
refresh_stats_daily_error_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let user_rows =
upsert_stats_user_daily_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_user_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
upsert_stats_user_daily_cost_savings_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
upsert_stats_user_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
.await?;
upsert_stats_user_daily_cost_savings_model_provider_rows(
&mut tx,
day_start_utc,
day_end_utc,
now_utc,
)
.await?;
upsert_stats_user_daily_api_format_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc).await?;
refresh_stats_user_summary_rows(&mut tx, day_end_utc, now_utc).await?;
tx.commit().await?;
Ok(Some(StatsAggregationSummary {
Ok(StatsAggregationSummary {
day_start_utc,
total_requests,
model_rows,
@@ -186,7 +230,7 @@ pub(super) async fn perform_stats_aggregation_once(
api_key_rows,
error_rows,
user_rows,
}))
})
}
async fn fetch_stats_daily_percentiles(
@@ -250,6 +294,91 @@ async fn upsert_stats_daily_provider_rows(
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_cost_savings_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_daily_api_key_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
@@ -305,6 +434,142 @@ async fn upsert_stats_user_daily_rows(
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_cost_savings_model_provider_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_user_daily_api_format_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
day_start_utc: DateTime<Utc>,
day_end_utc: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_API_FORMAT_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn refresh_stats_summary_row(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
cutoff_date: DateTime<Utc>,
@@ -381,3 +646,17 @@ async fn refresh_stats_summary_row(
Ok(())
}
async fn refresh_stats_user_summary_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
cutoff_date: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<(), sqlx::Error> {
sqlx::query(UPSERT_STATS_USER_SUMMARY_SQL)
.bind(cutoff_date)
.bind(now_utc)
.execute(&mut **tx)
.await?;
Ok(())
}

View File

@@ -6,9 +6,10 @@ use crate::data::GatewayDataState;
use aether_data_contracts::DataLayerError;
use super::{
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_STATS_HOURLY_AGGREGATE_SQL,
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_LATEST_STATS_HOURLY_HOUR_SQL,
SELECT_NEXT_STATS_HOURLY_BUCKET_SQL, SELECT_STATS_HOURLY_AGGREGATE_SQL,
UPSERT_STATS_HOURLY_MODEL_SQL, UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL,
UPSERT_STATS_HOURLY_USER_SQL,
UPSERT_STATS_HOURLY_USER_MODEL_SQL, UPSERT_STATS_HOURLY_USER_SQL,
};
#[derive(Debug, Clone, PartialEq)]
@@ -16,6 +17,7 @@ pub(super) struct StatsHourlyAggregationSummary {
pub(super) hour_utc: DateTime<Utc>,
pub(super) total_requests: i64,
pub(super) user_rows: usize,
pub(super) user_model_rows: usize,
pub(super) model_rows: usize,
pub(super) provider_rows: usize,
}
@@ -31,81 +33,121 @@ pub(super) async fn perform_stats_hourly_aggregation_once(
}
let now_utc = Utc::now();
let hour_utc = stats_hourly_aggregation_target_hour(now_utc);
let target_hour_utc = stats_hourly_aggregation_target_hour(now_utc);
let Some(hour_utc) = next_stats_hourly_bucket(&pool, target_hour_utc)
.await
.map_err(postgres_error)?
else {
return Ok(None);
};
perform_stats_hourly_aggregation_for_hour(&pool, hour_utc, now_utc)
.await
.map(Some)
.map_err(postgres_error)
}
async fn next_stats_hourly_bucket(
pool: &aether_data::postgres::PostgresPool,
target_hour_utc: DateTime<Utc>,
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
let latest_row = sqlx::query(SELECT_LATEST_STATS_HOURLY_HOUR_SQL)
.fetch_one(pool)
.await?;
let latest_hour = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_hour")?;
let search_from = latest_hour
.map(|value| value + chrono::Duration::hours(1))
.unwrap_or_else(|| {
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
});
let search_until = target_hour_utc + chrono::Duration::hours(1);
if search_from >= search_until {
return Ok(None);
}
let next_row = sqlx::query(SELECT_NEXT_STATS_HOURLY_BUCKET_SQL)
.bind(search_from)
.bind(search_until)
.fetch_one(pool)
.await?;
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
Ok(next_bucket.filter(|value| *value <= target_hour_utc))
}
async fn perform_stats_hourly_aggregation_for_hour(
pool: &aether_data::postgres::PostgresPool,
hour_utc: DateTime<Utc>,
aggregated_at: DateTime<Utc>,
) -> Result<StatsHourlyAggregationSummary, sqlx::Error> {
let hour_end = hour_utc + chrono::Duration::hours(1);
let aggregated_at = now_utc;
let mut tx = pool.begin().await.map_err(postgres_error)?;
let mut tx = pool.begin().await?;
let row = sqlx::query(SELECT_STATS_HOURLY_AGGREGATE_SQL)
.bind(hour_utc)
.bind(hour_end)
.fetch_one(&mut *tx)
.await
.map_err(postgres_error)?;
let total_requests = row
.try_get::<i64, _>("total_requests")
.map_err(postgres_error)?;
let error_requests = row
.try_get::<i64, _>("error_requests")
.map_err(postgres_error)?;
.await?;
let total_requests = row.try_get::<i64, _>("total_requests")?;
let error_requests = row.try_get::<i64, _>("error_requests")?;
let success_requests = total_requests.saturating_sub(error_requests);
sqlx::query(UPSERT_STATS_HOURLY_SQL)
.bind(Uuid::new_v4().to_string())
.bind(hour_utc)
.bind(total_requests)
.bind(row.try_get::<i64, _>("cache_hit_total_requests")?)
.bind(row.try_get::<i64, _>("cache_hit_requests")?)
.bind(row.try_get::<i64, _>("completed_total_requests")?)
.bind(row.try_get::<i64, _>("completed_cache_hit_requests")?)
.bind(row.try_get::<i64, _>("completed_input_tokens")?)
.bind(row.try_get::<i64, _>("completed_cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("completed_cache_read_tokens")?)
.bind(row.try_get::<i64, _>("completed_total_input_context")?)
.bind(row.try_get::<f64, _>("completed_cache_creation_cost")?)
.bind(row.try_get::<f64, _>("completed_cache_read_cost")?)
.bind(row.try_get::<f64, _>("settled_total_cost")?)
.bind(row.try_get::<i64, _>("settled_total_requests")?)
.bind(row.try_get::<i64, _>("settled_input_tokens")?)
.bind(row.try_get::<i64, _>("settled_output_tokens")?)
.bind(row.try_get::<i64, _>("settled_cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("settled_cache_read_tokens")?)
.bind(row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
.bind(row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
.bind(success_requests)
.bind(error_requests)
.bind(
row.try_get::<i64, _>("input_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("output_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_creation_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<i64, _>("cache_read_tokens")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("total_cost")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("actual_total_cost")
.map_err(postgres_error)?,
)
.bind(
row.try_get::<f64, _>("avg_response_time_ms")
.map_err(postgres_error)?,
)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
.bind(row.try_get::<f64, _>("total_cost")?)
.bind(row.try_get::<f64, _>("actual_total_cost")?)
.bind(row.try_get::<f64, _>("response_time_sum_ms")?)
.bind(row.try_get::<i64, _>("response_time_samples")?)
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
.bind(true)
.bind(aggregated_at)
.bind(aggregated_at)
.bind(aggregated_at)
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
.await?;
let user_rows =
upsert_stats_hourly_user_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let user_model_rows =
upsert_stats_hourly_user_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let model_rows =
upsert_stats_hourly_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
let provider_rows =
upsert_stats_hourly_provider_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
tx.commit().await.map_err(postgres_error)?;
tx.commit().await?;
Ok(Some(StatsHourlyAggregationSummary {
Ok(StatsHourlyAggregationSummary {
hour_utc,
total_requests,
user_rows,
user_model_rows,
model_rows,
provider_rows,
}))
})
}
async fn upsert_stats_hourly_user_rows(
@@ -113,14 +155,13 @@ async fn upsert_stats_hourly_user_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
@@ -131,14 +172,30 @@ async fn upsert_stats_hourly_model_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
}
async fn upsert_stats_hourly_user_model_rows(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_MODEL_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
@@ -149,14 +206,13 @@ async fn upsert_stats_hourly_provider_rows(
hour_utc: DateTime<Utc>,
hour_end: DateTime<Utc>,
now_utc: DateTime<Utc>,
) -> Result<usize, DataLayerError> {
) -> Result<usize, sqlx::Error> {
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
.bind(hour_utc)
.bind(hour_end)
.bind(now_utc)
.execute(&mut **tx)
.await
.map_err(postgres_error)?
.await?
.rows_affected();
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))

View File

@@ -21,6 +21,9 @@ use super::{
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
const STATS_DAILY_CATCH_UP_BURST_LIMIT: usize = 14;
const STATS_HOURLY_CATCH_UP_BURST_LIMIT: usize = 72;
fn log_maintenance_worker_failure(
worker: &'static str,
phase: &'static str,
@@ -110,10 +113,23 @@ pub(crate) fn spawn_stats_aggregation_worker(
Some(tokio::spawn(async move {
loop {
tokio::time::sleep(duration_until_next_stats_aggregation_run(Utc::now())).await;
if let Err(err) = run_stats_aggregation_once(&data).await {
log_maintenance_worker_failure("stats_daily_aggregation", "tick", &err);
let mut processed = 0_usize;
while processed < STATS_DAILY_CATCH_UP_BURST_LIMIT {
match run_stats_aggregation_once(&data).await {
Ok(true) => processed += 1,
Ok(false) => break,
Err(err) => {
log_maintenance_worker_failure("stats_daily_aggregation", "tick", &err);
break;
}
}
}
if processed >= STATS_DAILY_CATCH_UP_BURST_LIMIT {
continue;
}
tokio::time::sleep(duration_until_next_stats_aggregation_run(Utc::now())).await;
}
}))
}
@@ -304,10 +320,23 @@ pub(crate) fn spawn_stats_hourly_aggregation_worker(
Some(tokio::spawn(async move {
loop {
tokio::time::sleep(duration_until_next_stats_hourly_aggregation_run(Utc::now())).await;
if let Err(err) = run_stats_hourly_aggregation_once(&data).await {
log_maintenance_worker_failure("stats_hourly_aggregation", "tick", &err);
let mut processed = 0_usize;
while processed < STATS_HOURLY_CATCH_UP_BURST_LIMIT {
match run_stats_hourly_aggregation_once(&data).await {
Ok(true) => processed += 1,
Ok(false) => break,
Err(err) => {
log_maintenance_worker_failure("stats_hourly_aggregation", "tick", &err);
break;
}
}
}
if processed >= STATS_HOURLY_CATCH_UP_BURST_LIMIT {
continue;
}
tokio::time::sleep(duration_until_next_stats_hourly_aggregation_run(Utc::now())).await;
}
}))
}

View File

@@ -0,0 +1,853 @@
use aether_data::postgres::PostgresPool;
use aether_data_contracts::repository::usage::StoredUsageDashboardSummary;
use chrono::{DateTime, Utc};
use futures_util::TryStreamExt;
use sqlx::Row;
use crate::GatewayError;
fn internal(err: impl ToString) -> GatewayError {
GatewayError::Internal(err.to_string())
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct DashboardDailyTotalsAggregateRow {
pub(crate) date: String,
pub(crate) requests: u64,
pub(crate) total_tokens: u64,
pub(crate) total_cost_usd: f64,
pub(crate) response_time_sum_ms: f64,
pub(crate) response_time_samples: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct DashboardDailyModelAggregateRow {
pub(crate) date: String,
pub(crate) model: String,
pub(crate) requests: u64,
pub(crate) total_tokens: u64,
pub(crate) total_cost_usd: f64,
pub(crate) response_time_sum_ms: f64,
pub(crate) response_time_samples: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct DashboardDailyProviderAggregateRow {
pub(crate) date: String,
pub(crate) provider: String,
pub(crate) requests: u64,
pub(crate) total_tokens: u64,
pub(crate) total_cost_usd: f64,
}
pub(crate) async fn summarize_dashboard_usage_from_daily_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: Option<&str>,
) -> Result<StoredUsageDashboardSummary, GatewayError> {
let row = if let Some(user_id) = user_id {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
CAST(COALESCE(SUM(cache_creation_cost), 0) AS DOUBLE PRECISION) AS cache_creation_cost_usd,
CAST(COALESCE(SUM(cache_read_cost), 0) AS DOUBLE PRECISION) AS cache_read_cost_usd,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd,
CAST(COALESCE(SUM(actual_total_cost), 0) AS DOUBLE PRECISION) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0)::BIGINT AS error_requests,
COALESCE(SUM(response_time_sum_ms), 0) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
"#,
)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch_one(pool)
.await
.map_err(|err| internal(format!("user daily aggregate summary lookup failed: {err}")))?
} else {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
CAST(COALESCE(SUM(cache_creation_cost), 0) AS DOUBLE PRECISION) AS cache_creation_cost_usd,
CAST(COALESCE(SUM(cache_read_cost), 0) AS DOUBLE PRECISION) AS cache_read_cost_usd,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd,
CAST(COALESCE(SUM(actual_total_cost), 0) AS DOUBLE PRECISION) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0)::BIGINT AS error_requests,
COALESCE(SUM(response_time_sum_ms), 0) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_daily
WHERE date >= $1
AND date < $2
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch_one(pool)
.await
.map_err(|err| internal(format!("daily aggregate summary lookup failed: {err}")))?
};
Ok(StoredUsageDashboardSummary {
total_requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
input_tokens: row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
effective_input_tokens: row
.try_get::<i64, _>("effective_input_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
output_tokens: row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
cache_creation_tokens: row
.try_get::<i64, _>("cache_creation_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
cache_read_tokens: row
.try_get::<i64, _>("cache_read_tokens")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
total_input_context: row
.try_get::<i64, _>("total_input_context")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
cache_creation_cost_usd: row
.try_get::<f64, _>("cache_creation_cost_usd")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?,
cache_read_cost_usd: row
.try_get::<f64, _>("cache_read_cost_usd")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?,
total_cost_usd: row
.try_get::<f64, _>("total_cost_usd")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?,
actual_total_cost_usd: row
.try_get::<f64, _>("actual_total_cost_usd")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?,
error_requests: row
.try_get::<i64, _>("error_requests")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| internal(format!("aggregate summary decode failed: {err}")))?
.max(0) as u64,
})
}
pub(crate) async fn read_stats_hourly_cutoff(
pool: &PostgresPool,
) -> Result<Option<DateTime<Utc>>, GatewayError> {
let row = sqlx::query(
r#"
SELECT MAX(hour_utc) AS latest_hour
FROM stats_hourly
WHERE is_complete IS TRUE
"#,
)
.fetch_one(pool)
.await
.map_err(|err| internal(format!("hourly aggregate cutoff lookup failed: {err}")))?;
let latest_hour = row
.try_get::<Option<DateTime<Utc>>, _>("latest_hour")
.map_err(|err| internal(format!("hourly aggregate cutoff decode failed: {err}")))?;
Ok(latest_hour.map(|value| value + chrono::Duration::hours(1)))
}
pub(crate) async fn list_admin_dashboard_daily_totals_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
) -> Result<Vec<DashboardDailyTotalsAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
date,
total_requests,
input_tokens,
output_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
response_time_sum_ms,
response_time_samples
FROM stats_daily
WHERE date >= $1
AND date < $2
ORDER BY date ASC
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("daily aggregate totals read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?;
items.push(DashboardDailyTotalsAggregateRow {
date: date.date_naive().to_string(),
requests: row
.try_get::<i32, _>("total_requests")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?
.max(0) as u64,
total_tokens: input_tokens.saturating_add(output_tokens).max(0) as u64,
total_cost_usd: row
.try_get::<f64, _>("total_cost")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| internal(format!("daily aggregate totals decode failed: {err}")))?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_admin_dashboard_hourly_totals_aggregates(
pool: &PostgresPool,
start_utc: DateTime<Utc>,
end_utc: DateTime<Utc>,
tz_offset_minutes: i32,
) -> Result<Vec<DashboardDailyTotalsAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
CAST(DATE(hour_utc + ($3::integer * INTERVAL '1 minute')) AS TEXT) AS date,
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(response_time_sum_ms), 0) AS DOUBLE PRECISION) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_hourly
WHERE hour_utc >= $1
AND hour_utc < $2
GROUP BY date
ORDER BY date ASC
"#,
)
.bind(start_utc)
.bind(end_utc)
.bind(tz_offset_minutes)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("hourly aggregate totals read failed: {err}")))?
{
items.push(DashboardDailyTotalsAggregateRow {
date: row
.try_get::<String, _>("date")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?,
requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?
.max(0) as u64,
total_cost_usd: row
.try_get::<f64, _>("total_cost")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| internal(format!("hourly aggregate totals decode failed: {err}")))?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_admin_dashboard_daily_model_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
) -> Result<Vec<DashboardDailyModelAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
date,
model,
total_requests,
input_tokens,
output_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
response_time_sum_ms,
response_time_samples
FROM stats_daily_model
WHERE date >= $1
AND date < $2
ORDER BY date ASC, total_cost DESC, model ASC
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("daily aggregate model read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?;
items.push(DashboardDailyModelAggregateRow {
date: date.date_naive().to_string(),
model: row
.try_get::<String, _>("model")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?,
requests: row
.try_get::<i32, _>("total_requests")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?
.max(0) as u64,
total_tokens: input_tokens.saturating_add(output_tokens).max(0) as u64,
total_cost_usd: row
.try_get::<f64, _>("total_cost")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| internal(format!("daily aggregate model decode failed: {err}")))?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_admin_dashboard_hourly_model_aggregates(
pool: &PostgresPool,
start_utc: DateTime<Utc>,
end_utc: DateTime<Utc>,
tz_offset_minutes: i32,
) -> Result<Vec<DashboardDailyModelAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
CAST(DATE(hour_utc + ($3::integer * INTERVAL '1 minute')) AS TEXT) AS date,
model,
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(response_time_sum_ms), 0) AS DOUBLE PRECISION) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_hourly_model
WHERE hour_utc >= $1
AND hour_utc < $2
GROUP BY date, model
ORDER BY date ASC, total_cost DESC, model ASC
"#,
)
.bind(start_utc)
.bind(end_utc)
.bind(tz_offset_minutes)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("hourly aggregate model read failed: {err}")))?
{
items.push(DashboardDailyModelAggregateRow {
date: row
.try_get::<String, _>("date")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?,
model: row
.try_get::<String, _>("model")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?,
requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?
.max(0) as u64,
total_cost_usd: row
.try_get::<f64, _>("total_cost")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| internal(format!("hourly aggregate model decode failed: {err}")))?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_admin_dashboard_daily_provider_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
) -> Result<Vec<DashboardDailyProviderAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
date,
provider_name,
total_requests,
input_tokens,
output_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost
FROM stats_daily_provider
WHERE date >= $1
AND date < $2
ORDER BY date ASC, total_cost DESC, provider_name ASC
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("daily aggregate provider read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("daily aggregate provider decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("daily aggregate provider decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("daily aggregate provider decode failed: {err}")))?;
items.push(DashboardDailyProviderAggregateRow {
date: date.date_naive().to_string(),
provider: row.try_get::<String, _>("provider_name").map_err(|err| {
internal(format!("daily aggregate provider decode failed: {err}"))
})?,
requests: row
.try_get::<i32, _>("total_requests")
.map_err(|err| internal(format!("daily aggregate provider decode failed: {err}")))?
.max(0) as u64,
total_tokens: input_tokens.saturating_add(output_tokens).max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("daily aggregate provider decode failed: {err}"))
})?,
});
}
Ok(items)
}
pub(crate) async fn list_admin_dashboard_hourly_provider_aggregates(
pool: &PostgresPool,
start_utc: DateTime<Utc>,
end_utc: DateTime<Utc>,
tz_offset_minutes: i32,
) -> Result<Vec<DashboardDailyProviderAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
CAST(DATE(hour_utc + ($3::integer * INTERVAL '1 minute')) AS TEXT) AS date,
provider_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost
FROM stats_hourly_provider
WHERE hour_utc >= $1
AND hour_utc < $2
GROUP BY date, provider_name
ORDER BY date ASC, total_cost DESC, provider_name ASC
"#,
)
.bind(start_utc)
.bind(end_utc)
.bind(tz_offset_minutes)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("hourly aggregate provider read failed: {err}")))?
{
items.push(DashboardDailyProviderAggregateRow {
date: row.try_get::<String, _>("date").map_err(|err| {
internal(format!("hourly aggregate provider decode failed: {err}"))
})?,
provider: row.try_get::<String, _>("provider_name").map_err(|err| {
internal(format!("hourly aggregate provider decode failed: {err}"))
})?,
requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| internal(format!("hourly aggregate provider decode failed: {err}")))?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("hourly aggregate provider decode failed: {err}")))?
.max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("hourly aggregate provider decode failed: {err}"))
})?,
});
}
Ok(items)
}
pub(crate) async fn list_user_dashboard_daily_totals_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: &str,
) -> Result<Vec<DashboardDailyTotalsAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
date,
total_requests,
input_tokens,
output_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
response_time_sum_ms,
response_time_samples
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
ORDER BY date ASC
"#,
)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("user daily aggregate totals read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("user daily aggregate totals decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("user daily aggregate totals decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("user daily aggregate totals decode failed: {err}")))?;
items.push(DashboardDailyTotalsAggregateRow {
date: date.date_naive().to_string(),
requests: row
.try_get::<i32, _>("total_requests")
.map_err(|err| {
internal(format!("user daily aggregate totals decode failed: {err}"))
})?
.max(0) as u64,
total_tokens: input_tokens.saturating_add(output_tokens).max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("user daily aggregate totals decode failed: {err}"))
})?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| {
internal(format!("user daily aggregate totals decode failed: {err}"))
})?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| {
internal(format!("user daily aggregate totals decode failed: {err}"))
})?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_user_dashboard_hourly_totals_aggregates(
pool: &PostgresPool,
start_utc: DateTime<Utc>,
end_utc: DateTime<Utc>,
tz_offset_minutes: i32,
user_id: &str,
) -> Result<Vec<DashboardDailyTotalsAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
CAST(DATE(hour_utc + ($4::integer * INTERVAL '1 minute')) AS TEXT) AS date,
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(response_time_sum_ms), 0) AS DOUBLE PRECISION) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_hourly_user
WHERE user_id = $1
AND hour_utc >= $2
AND hour_utc < $3
GROUP BY date
ORDER BY date ASC
"#,
)
.bind(user_id)
.bind(start_utc)
.bind(end_utc)
.bind(tz_offset_minutes)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("user hourly aggregate totals read failed: {err}")))?
{
items.push(DashboardDailyTotalsAggregateRow {
date: row.try_get::<String, _>("date").map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?,
requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?
.max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| {
internal(format!("user hourly aggregate totals decode failed: {err}"))
})?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_user_dashboard_daily_model_aggregates(
pool: &PostgresPool,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: &str,
) -> Result<Vec<DashboardDailyModelAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
date,
model,
total_requests,
input_tokens,
output_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
response_time_sum_ms,
response_time_samples
FROM stats_user_daily_model
WHERE user_id = $1
AND date >= $2
AND date < $3
ORDER BY date ASC, total_cost DESC, model ASC
"#,
)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("user daily aggregate model read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("user daily aggregate model decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("user daily aggregate model decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("user daily aggregate model decode failed: {err}")))?;
items.push(DashboardDailyModelAggregateRow {
date: date.date_naive().to_string(),
model: row.try_get::<String, _>("model").map_err(|err| {
internal(format!("user daily aggregate model decode failed: {err}"))
})?,
requests: row
.try_get::<i32, _>("total_requests")
.map_err(|err| {
internal(format!("user daily aggregate model decode failed: {err}"))
})?
.max(0) as u64,
total_tokens: input_tokens.saturating_add(output_tokens).max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("user daily aggregate model decode failed: {err}"))
})?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| {
internal(format!("user daily aggregate model decode failed: {err}"))
})?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| {
internal(format!("user daily aggregate model decode failed: {err}"))
})?
.max(0) as u64,
});
}
Ok(items)
}
pub(crate) async fn list_user_dashboard_hourly_model_aggregates(
pool: &PostgresPool,
start_utc: DateTime<Utc>,
end_utc: DateTime<Utc>,
tz_offset_minutes: i32,
user_id: &str,
) -> Result<Vec<DashboardDailyModelAggregateRow>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
CAST(DATE(hour_utc + ($4::integer * INTERVAL '1 minute')) AS TEXT) AS date,
model,
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(response_time_sum_ms), 0) AS DOUBLE PRECISION) AS response_time_sum_ms,
COALESCE(SUM(response_time_samples), 0)::BIGINT AS response_time_samples
FROM stats_hourly_user_model
WHERE user_id = $1
AND hour_utc >= $2
AND hour_utc < $3
GROUP BY date, model
ORDER BY date ASC, total_cost DESC, model ASC
"#,
)
.bind(user_id)
.bind(start_utc)
.bind(end_utc)
.bind(tz_offset_minutes)
.fetch(pool);
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("user hourly aggregate model read failed: {err}")))?
{
items.push(DashboardDailyModelAggregateRow {
date: row.try_get::<String, _>("date").map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?,
model: row.try_get::<String, _>("model").map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?,
requests: row
.try_get::<i64, _>("total_requests")
.map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?
.max(0) as u64,
total_tokens: row
.try_get::<i64, _>("total_tokens")
.map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?
.max(0) as u64,
total_cost_usd: row.try_get::<f64, _>("total_cost").map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?,
response_time_sum_ms: row
.try_get::<f64, _>("response_time_sum_ms")
.map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?,
response_time_samples: row
.try_get::<i64, _>("response_time_samples")
.map_err(|err| {
internal(format!("user hourly aggregate model decode failed: {err}"))
})?
.max(0) as u64,
});
}
Ok(items)
}

View File

@@ -1,2 +1,5 @@
pub(crate) mod billing;
pub(crate) mod dashboard_stats;
pub(crate) mod monitoring;
pub(crate) mod usage_heatmap;
pub(crate) mod user_rollups;

View File

@@ -0,0 +1,176 @@
use aether_data::postgres::PostgresPool;
use aether_data_contracts::repository::usage::StoredUsageDailySummary;
use chrono::{DateTime, NaiveDate, Utc};
use futures_util::TryStreamExt;
use sqlx::Row;
use crate::GatewayError;
const USER_HEATMAP_AGGREGATE_SQL: &str = r#"
SELECT
date,
total_requests,
input_tokens,
output_tokens,
cache_creation_tokens,
cache_read_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
COALESCE(actual_total_cost, 0)::DOUBLE PRECISION AS actual_total_cost
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
ORDER BY date ASC
"#;
const GLOBAL_HEATMAP_AGGREGATE_SQL: &str = r#"
SELECT
date,
total_requests,
input_tokens,
output_tokens,
cache_creation_tokens,
cache_read_tokens,
COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost,
COALESCE(actual_total_cost, 0)::DOUBLE PRECISION AS actual_total_cost
FROM stats_daily
WHERE date >= $1
AND date < $2
ORDER BY date ASC
"#;
fn internal(err: impl ToString) -> GatewayError {
GatewayError::Internal(err.to_string())
}
pub(crate) async fn read_stats_daily_cutoff_date(
pool: &PostgresPool,
) -> Result<Option<DateTime<Utc>>, GatewayError> {
let row = sqlx::query(
r#"
SELECT cutoff_date
FROM stats_summary
ORDER BY updated_at DESC, created_at DESC
LIMIT 1
"#,
)
.fetch_optional(pool)
.await
.map_err(|err| internal(format!("stats summary cutoff lookup failed: {err}")))?;
let Some(row) = row else {
return Ok(None);
};
row.try_get::<DateTime<Utc>, _>("cutoff_date")
.map(Some)
.map_err(|err| internal(format!("stats summary cutoff decode failed: {err}")))
}
pub(crate) async fn list_usage_heatmap_aggregate_rows(
pool: &PostgresPool,
start_date: NaiveDate,
end_date_exclusive: NaiveDate,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, GatewayError> {
if start_date >= end_date_exclusive {
return Ok(Vec::new());
}
let start_at = DateTime::<Utc>::from_naive_utc_and_offset(
start_date
.and_hms_opt(0, 0, 0)
.expect("midnight should be valid"),
Utc,
);
let end_at = DateTime::<Utc>::from_naive_utc_and_offset(
end_date_exclusive
.and_hms_opt(0, 0, 0)
.expect("midnight should be valid"),
Utc,
);
let mut rows = if let Some(user_id) = user_id {
sqlx::query(USER_HEATMAP_AGGREGATE_SQL)
.bind(user_id)
.bind(start_at)
.bind(end_at)
.fetch(pool)
} else {
sqlx::query(GLOBAL_HEATMAP_AGGREGATE_SQL)
.bind(start_at)
.bind(end_at)
.fetch(pool)
};
let mut items = Vec::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("aggregate heatmap read failed: {err}")))?
{
let date = row
.try_get::<DateTime<Utc>, _>("date")
.map_err(|err| internal(format!("aggregate heatmap date decode failed: {err}")))?;
let requests = row
.try_get::<i32, _>("total_requests")
.map_err(|err| internal(format!("aggregate heatmap request decode failed: {err}")))?;
let input_tokens = row
.try_get::<i64, _>("input_tokens")
.map_err(|err| internal(format!("aggregate heatmap token decode failed: {err}")))?;
let output_tokens = row
.try_get::<i64, _>("output_tokens")
.map_err(|err| internal(format!("aggregate heatmap token decode failed: {err}")))?;
let cache_creation_tokens = row
.try_get::<i64, _>("cache_creation_tokens")
.map_err(|err| internal(format!("aggregate heatmap token decode failed: {err}")))?;
let cache_read_tokens = row
.try_get::<i64, _>("cache_read_tokens")
.map_err(|err| internal(format!("aggregate heatmap token decode failed: {err}")))?;
let total_cost_usd = row
.try_get::<f64, _>("total_cost")
.map_err(|err| internal(format!("aggregate heatmap cost decode failed: {err}")))?;
let actual_total_cost_usd = row.try_get::<f64, _>("actual_total_cost").map_err(|err| {
internal(format!(
"aggregate heatmap actual cost decode failed: {err}"
))
})?;
items.push(StoredUsageDailySummary {
date: date.date_naive().to_string(),
requests: u64::try_from(requests.max(0)).unwrap_or_default(),
total_tokens: u64::try_from(
input_tokens
.saturating_add(output_tokens)
.saturating_add(cache_creation_tokens)
.saturating_add(cache_read_tokens)
.max(0),
)
.unwrap_or_default(),
total_cost_usd,
actual_total_cost_usd,
});
}
Ok(items)
}
#[cfg(test)]
mod tests {
use super::{GLOBAL_HEATMAP_AGGREGATE_SQL, USER_HEATMAP_AGGREGATE_SQL};
#[test]
fn user_heatmap_query_casts_cost_columns_to_double_precision() {
assert!(USER_HEATMAP_AGGREGATE_SQL
.contains("COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost"));
assert!(USER_HEATMAP_AGGREGATE_SQL
.contains("COALESCE(actual_total_cost, 0)::DOUBLE PRECISION AS actual_total_cost"));
}
#[test]
fn global_heatmap_query_casts_cost_columns_to_double_precision() {
assert!(GLOBAL_HEATMAP_AGGREGATE_SQL
.contains("COALESCE(total_cost, 0)::DOUBLE PRECISION AS total_cost"));
assert!(GLOBAL_HEATMAP_AGGREGATE_SQL
.contains("COALESCE(actual_total_cost, 0)::DOUBLE PRECISION AS actual_total_cost"));
}
}

View File

@@ -0,0 +1,154 @@
use std::collections::BTreeMap;
use aether_data::postgres::PostgresPool;
use aether_data_contracts::repository::usage::StoredUsageUserTotals;
use chrono::{DateTime, Utc};
use futures_util::TryStreamExt;
use sqlx::Row;
use crate::query::usage_heatmap::read_stats_daily_cutoff_date;
use crate::GatewayError;
fn internal(err: impl ToString) -> GatewayError {
GatewayError::Internal(err.to_string())
}
pub(crate) async fn list_user_usage_totals_from_stats_summary(
pool: &PostgresPool,
user_ids: &[String],
) -> Result<Option<Vec<StoredUsageUserTotals>>, GatewayError> {
if user_ids.is_empty() {
return Ok(Some(Vec::new()));
}
let Some(cutoff_date) = read_stats_daily_cutoff_date(pool).await? else {
return Ok(None);
};
let mut totals = load_stats_user_summary_rows(pool, user_ids).await?;
absorb_stats_user_summary_tail(pool, cutoff_date, user_ids, &mut totals).await?;
let mut items = user_ids
.iter()
.map(|user_id| {
totals
.remove(user_id)
.unwrap_or_else(|| StoredUsageUserTotals {
user_id: user_id.clone(),
request_count: 0,
total_tokens: 0,
})
})
.collect::<Vec<_>>();
items.sort_by(|left, right| left.user_id.cmp(&right.user_id));
Ok(Some(items))
}
async fn load_stats_user_summary_rows(
pool: &PostgresPool,
user_ids: &[String],
) -> Result<BTreeMap<String, StoredUsageUserTotals>, GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
user_id,
COALESCE(all_time_requests, 0)::BIGINT AS request_count,
COALESCE(
all_time_input_tokens
+ all_time_output_tokens
+ all_time_cache_creation_tokens
+ all_time_cache_read_tokens,
0
)::BIGINT AS total_tokens
FROM stats_user_summary
WHERE user_id = ANY($1::TEXT[])
ORDER BY user_id ASC
"#,
)
.bind(user_ids)
.fetch(pool);
let mut items = BTreeMap::new();
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("stats_user_summary lookup failed: {err}")))?
{
let user_id = row
.try_get::<String, _>("user_id")
.map_err(|err| internal(format!("stats_user_summary decode failed: {err}")))?;
let request_count = row
.try_get::<i64, _>("request_count")
.map_err(|err| internal(format!("stats_user_summary decode failed: {err}")))?
.max(0) as u64;
let total_tokens = row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("stats_user_summary decode failed: {err}")))?
.max(0) as u64;
items.insert(
user_id.clone(),
StoredUsageUserTotals {
user_id,
request_count,
total_tokens,
},
);
}
Ok(items)
}
async fn absorb_stats_user_summary_tail(
pool: &PostgresPool,
cutoff_date: DateTime<Utc>,
user_ids: &[String],
totals: &mut BTreeMap<String, StoredUsageUserTotals>,
) -> Result<(), GatewayError> {
let mut rows = sqlx::query(
r#"
SELECT
"usage".user_id,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)), 0)::BIGINT AS total_tokens
FROM "usage"
WHERE "usage".user_id = ANY($1::TEXT[])
AND "usage".created_at >= $2
AND "usage".status NOT IN ('pending', 'streaming')
AND "usage".provider_name NOT IN ('unknown', 'pending')
GROUP BY "usage".user_id
ORDER BY "usage".user_id ASC
"#,
)
.bind(user_ids)
.bind(cutoff_date)
.fetch(pool);
while let Some(row) = rows
.try_next()
.await
.map_err(|err| internal(format!("stats_user_summary tail lookup failed: {err}")))?
{
let user_id = row
.try_get::<String, _>("user_id")
.map_err(|err| internal(format!("stats_user_summary tail decode failed: {err}")))?;
let request_count = row
.try_get::<i64, _>("request_count")
.map_err(|err| internal(format!("stats_user_summary tail decode failed: {err}")))?
.max(0) as u64;
let total_tokens = row
.try_get::<i64, _>("total_tokens")
.map_err(|err| internal(format!("stats_user_summary tail decode failed: {err}")))?
.max(0) as u64;
let entry = totals
.entry(user_id.clone())
.or_insert_with(|| StoredUsageUserTotals {
user_id,
request_count: 0,
total_tokens: 0,
});
entry.request_count = entry.request_count.saturating_add(request_count);
entry.total_tokens = entry.total_tokens.saturating_add(total_tokens);
}
Ok(())
}

View File

@@ -936,6 +936,7 @@ mod tests {
None,
false,
0,
0,
0.0,
false,
)

View File

@@ -255,6 +255,14 @@ impl AppState {
Ok(true)
}
pub async fn run_postgres_backfills(&self) -> Result<bool, sqlx::migrate::MigrateError> {
let Some(pool) = self.postgres_pool() else {
return Ok(false);
};
aether_data::backfill::run_backfills(&pool).await?;
Ok(true)
}
pub async fn pending_postgres_migrations(
&self,
) -> Result<Option<Vec<aether_data::migrate::PendingMigrationInfo>>, sqlx::migrate::MigrateError>
@@ -277,6 +285,16 @@ impl AppState {
))
}
pub async fn pending_postgres_backfills(
&self,
) -> Result<Option<Vec<aether_data::backfill::PendingBackfillInfo>>, sqlx::migrate::MigrateError>
{
let Some(pool) = self.postgres_pool() else {
return Ok(None);
};
Ok(Some(aether_data::backfill::pending_backfills(&pool).await?))
}
pub fn with_video_task_poller_config(mut self, interval: Duration, batch_size: usize) -> Self {
self.video_task_poller = Some(VideoTaskPollerConfig {
interval,

View File

@@ -104,6 +104,7 @@ fn sample_standalone_export_record(
Some(4_102_444_800),
false,
7,
0,
1.25,
true,
)

View File

@@ -414,6 +414,7 @@ fn sample_monitoring_export_api_key(
None,
false,
0,
0,
0.0,
false,
)

View File

@@ -635,6 +635,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
Some(1_900_000_000),
false,
42,
420,
12.34,
false,
)
@@ -661,6 +662,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
None,
false,
7,
84,
3.21,
true,
)
@@ -749,10 +751,15 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["api_keys"][0]["is_standalone"],
json!(false)
);
assert_eq!(
payload["users"][0]["api_keys"][0]["total_tokens"],
json!(420)
);
assert_eq!(
payload["standalone_keys"][0]["key"],
json!("ak-standalone-live-1")
);
assert_eq!(payload["standalone_keys"][0]["total_tokens"], json!(84));
assert_eq!(
payload["standalone_keys"][0]["wallet"]["unlimited"],
json!(true)

View File

@@ -607,6 +607,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
"expires_at": "2099-01-01T00:00:00Z",
"auto_delete_on_expiry": false,
"total_requests": 12,
"total_tokens": 3456,
"total_cost_usd": 1.25
}]
}],
@@ -622,6 +623,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
"expires_at": "2099-02-01T00:00:00Z",
"auto_delete_on_expiry": false,
"total_requests": 3,
"total_tokens": 789,
"total_cost_usd": 0.75,
"wallet": {
"balance": 30.0,
@@ -704,6 +706,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.expect("user api keys should load");
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(
user_api_keys[0].allowed_api_formats,
Some(vec!["openai:chat".to_string()])
@@ -729,6 +732,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
standalone_keys[0].name.as_deref(),
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,

View File

@@ -674,6 +674,7 @@ async fn gateway_handles_admin_user_api_key_routes_locally_with_trusted_admin_pr
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -835,6 +836,7 @@ async fn gateway_returns_conflict_for_admin_create_user_api_key_when_writer_unav
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -963,6 +965,7 @@ async fn gateway_returns_conflict_for_admin_lock_user_api_key_when_writer_unavai
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -1088,6 +1091,7 @@ async fn gateway_returns_conflict_for_admin_update_user_api_key_when_writer_unav
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -1162,6 +1166,7 @@ async fn gateway_returns_conflict_for_admin_delete_user_api_key_when_writer_unav
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -1238,6 +1243,7 @@ async fn gateway_lists_admin_user_api_keys_locally_with_trusted_admin_principal(
Some(200),
false,
9,
0,
1.5,
false,
)
@@ -1428,6 +1434,7 @@ async fn gateway_reveals_admin_user_full_key_locally_with_trusted_admin_principa
Some(200),
false,
9,
0,
1.5,
false,
)

View File

@@ -5833,6 +5833,7 @@ async fn gateway_handles_users_me_api_keys_locally_without_proxying_upstream() {
Some(300),
false,
9,
0,
1.5,
false,
)
@@ -6143,6 +6144,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
Some(300),
false,
3,
0,
0.5,
false,
)
@@ -6411,6 +6413,7 @@ async fn gateway_returns_service_unavailable_for_users_me_api_key_writes_without
Some(300),
false,
3,
0,
0.5,
false,
)

View File

@@ -101,6 +101,7 @@ async fn gateway_handles_dashboard_stats_locally_without_proxying_upstream() {
None,
false,
5,
0,
1.5,
false,
)
@@ -121,6 +122,7 @@ async fn gateway_handles_dashboard_stats_locally_without_proxying_upstream() {
None,
false,
1,
0,
0.5,
false,
)
@@ -326,6 +328,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
None,
false,
5,
0,
1.5,
false,
)
@@ -346,6 +349,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
None,
false,
1,
0,
0.5,
false,
)
@@ -366,6 +370,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
None,
false,
3,
0,
0.75,
true,
)
@@ -564,6 +569,118 @@ async fn gateway_handles_dashboard_daily_stats_locally_without_proxying_upstream
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_user_dashboard_daily_stats_locally_without_proxying_upstream() {
let now = stable_dashboard_now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!(user.role)),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
(
"session_id".to_string(),
json!("session-dashboard-daily-stats-user"),
),
]),
chrono::Utc::now() + chrono::Duration::hours(1),
);
let session = sample_auth_session(
"user-auth-1",
"session-dashboard-daily-stats-user",
"device-dashboard-daily-stats-user",
"refresh-dashboard-daily-stats-user",
now,
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
"usage-dashboard-user-daily-1",
"req-dashboard-user-daily-1",
"user-auth-1",
"gpt-5",
"openai",
"completed",
now - chrono::Duration::hours(1),
),
sample_user_usage_audit(
"usage-dashboard-user-daily-2",
"req-dashboard-user-daily-2",
"user-auth-1",
"gpt-4.1",
"openai",
"completed",
now - chrono::Duration::days(1) - chrono::Duration::hours(2),
),
sample_user_usage_audit(
"usage-dashboard-user-daily-3",
"req-dashboard-user-daily-3",
"user-auth-2",
"claude-3-7",
"claude",
"completed",
now - chrono::Duration::hours(2),
),
]));
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
user.clone()
]));
let wallet_repository =
Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
"user-auth-1",
now,
)]));
let data_state = GatewayDataState::with_user_wallet_and_usage_for_tests(
user_repository,
wallet_repository,
usage_repository,
);
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([session])
})
.await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/dashboard/daily-stats?days=2"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-dashboard-daily-stats-user")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let daily_stats = payload["daily_stats"]
.as_array()
.expect("daily stats should be array");
assert_eq!(daily_stats.len(), 2);
assert_eq!(daily_stats[0]["requests"], 1);
assert_eq!(daily_stats[1]["requests"], 1);
assert_eq!(
daily_stats[1]["model_breakdown"].as_array().map(Vec::len),
Some(1)
);
let model_summary = payload["model_summary"]
.as_array()
.expect("model summary should exist");
assert_eq!(model_summary.len(), 2);
assert_eq!(payload.get("provider_summary"), None);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_dashboard_recent_requests_locally_without_proxying_upstream() {
let now = Utc::now();