From 778cfb1a5c5e67dba8cd0545d9ceaa0497c63e35 Mon Sep 17 00:00:00 2001 From: elky Date: Sat, 25 Jul 2026 21:28:21 +0800 Subject: [PATCH] feat(data): complete portable SQL backend parity Align MySQL and SQLite schemas, migrations, usage, stats, export, and backfill behavior with the shared data contracts. Extend gateway startup and maintenance support across all SQL drivers. --- Cargo.lock | 1 + apps/aether-gateway/src/data/state/core.rs | 30 +- apps/aether-gateway/src/main.rs | 132 +- .../src/state/bootstrap_admin.rs | 3 +- .../src/tests/architecture/sql_and_data.rs | 9 +- crates/aether-data/adapters/mysql/Cargo.toml | 1 + ...5010000_add_cross_driver_schema_parity.sql | 230 +++ ...260725020000_add_advanced_stats_parity.sql | 381 ++++ ...0000_add_usage_routing_snapshot_parity.sql | 69 + .../aether-data/adapters/mysql/src/billing.rs | 141 +- .../adapters/mysql/src/candidates.rs | 468 +++-- .../adapters/mysql/src/global_models.rs | 927 +++++++--- crates/aether-data/adapters/mysql/src/lib.rs | 2 +- .../adapters/mysql/src/migrations.rs | 184 +- crates/aether-data/adapters/mysql/src/pool.rs | 38 + .../adapters/mysql/src/provider_catalog.rs | 662 ++++++- .../adapters/mysql/src/settlement.rs | 199 ++- .../aether-data/adapters/mysql/src/usage.rs | 804 ++++++--- .../adapters/mysql/src/usage/cleanup.rs | 1177 +++++++++++++ .../adapters/mysql/src/usage/counters.rs | 1153 ++++++++++++ .../adapters/mysql/src/usage/http_capture.rs | 830 +++++++++ .../adapters/mysql/src/usage/read.rs | 811 +++++++++ .../adapters/mysql/src/usage/snapshots.rs | 1043 +++++++++++ .../adapters/mysql/src/usage/tests.rs | 858 ++++++++- .../adapters/mysql/src/video_tasks.rs | 226 ++- .../aether-data/adapters/mysql/src/wallet.rs | 1043 ++++++++--- .../adapters/mysql/src/wallet/tests.rs | 427 ++++- ...0_normalize_imported_timestamp_storage.sql | 784 +++++++++ ...5010000_add_cross_driver_schema_parity.sql | 270 +++ ...260725020000_add_advanced_stats_parity.sql | 447 +++++ ...0000_add_usage_routing_snapshot_parity.sql | 70 + ...e_remaining_imported_timestamp_storage.sql | 821 +++++++++ .../adapters/sqlite/src/billing.rs | 124 +- .../adapters/sqlite/src/candidates.rs | 435 +++-- .../adapters/sqlite/src/global_models.rs | 42 +- .../adapters/sqlite/src/migrations.rs | 501 +++++- .../adapters/sqlite/src/settlement.rs | 189 +- .../aether-data/adapters/sqlite/src/usage.rs | 1171 ++++++++++--- .../adapters/sqlite/src/usage/cleanup.rs | 1142 ++++++++++++ .../adapters/sqlite/src/usage/counters.rs | 1283 ++++++++++++++ .../adapters/sqlite/src/usage/http_capture.rs | 816 +++++++++ .../adapters/sqlite/src/usage/snapshots.rs | 1047 +++++++++++ .../adapters/sqlite/src/usage/tests.rs | 1384 ++++++++++++++- .../adapters/sqlite/src/video_tasks.rs | 167 +- ...422120000_backfill_api_key_usage_stats.sql | 36 + ...20000_rebuild_global_model_usage_count.sql | 15 + ...12000_sync_legacy_enabled_active_flags.sql | 11 + ...0_rebuild_canonical_usage_token_totals.sql | 229 +++ .../runtime/backfills/mysql/README.md | 3 +- ...422120000_backfill_api_key_usage_stats.sql | 28 + ...20000_rebuild_global_model_usage_count.sql | 9 + ...12000_sync_legacy_enabled_active_flags.sql | 11 + ...0_rebuild_canonical_usage_token_totals.sql | 225 +++ .../runtime/backfills/sqlite/README.md | 3 +- .../generated/mysql/baseline/006_usage.sql | 88 +- .../generated/mysql/baseline/007_stats.sql | 363 ++++ .../generated/postgres/baseline/006_usage.sql | 89 + .../generated/postgres/baseline/007_stats.sql | 391 ++++- .../generated/sqlite/baseline/006_usage.sql | 83 + .../generated/sqlite/baseline/007_stats.sql | 349 ++++ .../runtime/schema/logical/006_usage.toml | 402 +++++ .../runtime/schema/logical/007_stats.toml | 1542 +++++++++++++++++ .../runtime/src/backend/maintenance/mysql.rs | 45 +- .../aether-data/runtime/src/backend/mysql.rs | 79 +- .../aether-data/runtime/src/backend/sqlite.rs | 280 ++- .../runtime/src/backend/stats/mysql.rs | 588 ++++++- .../src/backend/stats/mysql/advanced.rs | 973 +++++++++++ .../runtime/src/backend/stats/sqlite.rs | 626 ++++++- .../src/backend/stats/sqlite/advanced.rs | 986 +++++++++++ .../runtime/src/backend/stats_common.rs | 6 - .../aether-data/runtime/src/backend/system.rs | 2 + .../runtime/src/backend/system/mysql.rs | 24 + .../runtime/src/backend/system/postgres.rs | 24 + .../runtime/src/backend/system/sqlite.rs | 153 ++ .../runtime/src/lifecycle/backfill/mysql.rs | 251 ++- .../runtime/src/lifecycle/backfill/sqlite.rs | 252 ++- .../runtime/src/lifecycle/backfill/tests.rs | 569 +++++- .../runtime/src/lifecycle/export.rs | 638 ++++++- .../runtime/src/lifecycle/export/mysql.rs | 350 +++- .../runtime/src/lifecycle/export/postgres.rs | 206 ++- .../runtime/src/lifecycle/export/sqlite.rs | 218 ++- .../runtime/src/lifecycle/export/tests.rs | 514 +++++- .../runtime/src/lifecycle/migrate/tests.rs | 459 ++++- .../runtime/src/repository/usage/mod.rs | 355 ---- .../runtime/src/repository/usage/mysql.rs | 407 ++++- 85 files changed, 32096 insertions(+), 2328 deletions(-) create mode 100644 crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql create mode 100644 crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql create mode 100644 crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql create mode 100644 crates/aether-data/adapters/mysql/src/usage/cleanup.rs create mode 100644 crates/aether-data/adapters/mysql/src/usage/counters.rs create mode 100644 crates/aether-data/adapters/mysql/src/usage/http_capture.rs create mode 100644 crates/aether-data/adapters/mysql/src/usage/read.rs create mode 100644 crates/aether-data/adapters/mysql/src/usage/snapshots.rs create mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql create mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql create mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql create mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql create mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql create mode 100644 crates/aether-data/adapters/sqlite/src/usage/cleanup.rs create mode 100644 crates/aether-data/adapters/sqlite/src/usage/counters.rs create mode 100644 crates/aether-data/adapters/sqlite/src/usage/http_capture.rs create mode 100644 crates/aether-data/adapters/sqlite/src/usage/snapshots.rs create mode 100644 crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql create mode 100644 crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql create mode 100644 crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql create mode 100644 crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql create mode 100644 crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql create mode 100644 crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql create mode 100644 crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql create mode 100644 crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql create mode 100644 crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs create mode 100644 crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs diff --git a/Cargo.lock b/Cargo.lock index 8a2b05fd0..6808176ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -208,6 +208,7 @@ dependencies = [ "async-trait", "chrono", "chrono-tz", + "flate2", "serde_json", "sha2", "sqlx", diff --git a/apps/aether-gateway/src/data/state/core.rs b/apps/aether-gateway/src/data/state/core.rs index e5458d6ab..992fc6a02 100644 --- a/apps/aether-gateway/src/data/state/core.rs +++ b/apps/aether-gateway/src/data/state/core.rs @@ -764,7 +764,8 @@ impl GatewayDataState { } pub(crate) fn has_usage_counter_flush_backend(&self) -> bool { - self.has_usage_writer() && self.database_driver() == Some(DatabaseDriver::Postgres) + self.has_usage_writer() + && database_driver_supports_usage_counter_flush(self.database_driver()) } pub(crate) fn has_usage_worker_queue(&self) -> bool { @@ -1075,3 +1076,30 @@ impl GatewayDataState { } } } + +fn database_driver_supports_usage_counter_flush(driver: Option) -> bool { + matches!( + driver, + Some(DatabaseDriver::Postgres | DatabaseDriver::Mysql | DatabaseDriver::Sqlite) + ) +} + +#[cfg(test)] +mod usage_counter_flush_backend_tests { + use super::database_driver_supports_usage_counter_flush; + use aether_data::DatabaseDriver; + + #[test] + fn every_sql_driver_supports_usage_counter_flush() { + assert!(database_driver_supports_usage_counter_flush(Some( + DatabaseDriver::Postgres + ))); + assert!(database_driver_supports_usage_counter_flush(Some( + DatabaseDriver::Mysql + ))); + assert!(database_driver_supports_usage_counter_flush(Some( + DatabaseDriver::Sqlite + ))); + assert!(!database_driver_supports_usage_counter_flush(None)); + } +} diff --git a/apps/aether-gateway/src/main.rs b/apps/aether-gateway/src/main.rs index 58b405aa6..13e95d533 100644 --- a/apps/aether-gateway/src/main.rs +++ b/apps/aether-gateway/src/main.rs @@ -85,10 +85,17 @@ enum ExportDomainArg { Endpoints, Models, GlobalModels, + AuthModules, + OAuthProviders, + UserOAuthLinks, + UserGroups, + UserGroupMembers, + ProxyNodes, SystemConfigs, Wallets, Usage, Billing, + Auxiliary, } impl From for ExportDomain { @@ -101,10 +108,17 @@ impl From for ExportDomain { ExportDomainArg::Endpoints => ExportDomain::Endpoints, ExportDomainArg::Models => ExportDomain::Models, ExportDomainArg::GlobalModels => ExportDomain::GlobalModels, + ExportDomainArg::AuthModules => ExportDomain::AuthModules, + ExportDomainArg::OAuthProviders => ExportDomain::OAuthProviders, + ExportDomainArg::UserOAuthLinks => ExportDomain::UserOAuthLinks, + ExportDomainArg::UserGroups => ExportDomain::UserGroups, + ExportDomainArg::UserGroupMembers => ExportDomain::UserGroupMembers, + ExportDomainArg::ProxyNodes => ExportDomain::ProxyNodes, ExportDomainArg::SystemConfigs => ExportDomain::SystemConfigs, ExportDomainArg::Wallets => ExportDomain::Wallets, ExportDomainArg::Usage => ExportDomain::Usage, ExportDomainArg::Billing => ExportDomain::Billing, + ExportDomainArg::Auxiliary => ExportDomain::Auxiliary, } } } @@ -586,14 +600,23 @@ impl GatewayDataArgs { .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned); + let legacy_postgres_url = self + .postgres_url + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let generic_database_url = std::env::var("DATABASE_URL") + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); - match (self.effective_database_driver(), configured_url) { - (Some(DatabaseDriver::Sqlite), None) => Some(DEFAULT_SQLITE_DATABASE_URL.to_string()), - (_, Some(url)) => Some(url), - (None, None) => self.effective_postgres_url(), - (Some(DatabaseDriver::Postgres), None) => self.effective_postgres_url(), - (Some(DatabaseDriver::Mysql), None) => None, - } + resolve_database_url( + self.effective_database_driver(), + configured_url, + legacy_postgres_url, + generic_database_url, + ) } fn effective_sql_database_config(&self) -> Option { @@ -649,20 +672,6 @@ impl GatewayDataArgs { } } - fn effective_postgres_url(&self) -> Option { - self.postgres_url - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - .or_else(|| { - std::env::var("DATABASE_URL") - .ok() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - }) - } - fn effective_redis_url(&self) -> Option { self.redis_url .as_deref() @@ -725,6 +734,25 @@ impl GatewayDataArgs { } } +fn resolve_database_url( + driver: Option, + configured_url: Option, + legacy_postgres_url: Option, + generic_database_url: Option, +) -> Option { + if configured_url.is_some() { + return configured_url; + } + + match driver { + Some(DatabaseDriver::Sqlite) => { + generic_database_url.or_else(|| Some(DEFAULT_SQLITE_DATABASE_URL.to_string())) + } + Some(DatabaseDriver::Mysql) => generic_database_url, + Some(DatabaseDriver::Postgres) | None => legacy_postgres_url.or(generic_database_url), + } +} + #[derive(ClapArgs, Debug, Clone)] struct GatewayUsageArgs { #[arg( @@ -1708,7 +1736,7 @@ fn validate_deployment_topology( { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, - "AETHER_GATEWAY_VIDEO_TASK_STORE_PATH must be unset when AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY=multi-node; use shared Postgres-backed state instead", + "AETHER_GATEWAY_VIDEO_TASK_STORE_PATH must be unset when AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY=multi-node; use shared SQL-backed state instead", )); } @@ -1764,7 +1792,6 @@ async fn run() -> Result<(), Box> { } init_service_runtime(args.runtime_config()?)?; let sql_database_config = args.data.effective_sql_database_config(); - let data_postgres_url = args.data.effective_postgres_url(); let data_redis_url = args.data.effective_redis_url(); let runtime_backend = args.effective_runtime_backend(sql_database_config.as_ref(), data_redis_url.as_deref()); @@ -1932,7 +1959,9 @@ async fn run() -> Result<(), Box> { .as_ref() .map(|database| database.pool.max_connections) .unwrap_or_default(), - data_postgres_configured = data_postgres_url.is_some(), + data_postgres_configured = sql_database_config + .as_ref() + .is_some_and(|database| database.driver == DatabaseDriver::Postgres), runtime_redis_configured = matches!(runtime_backend, RuntimeBackendArg::Redis), data_redis_url_supplied = data_redis_url.is_some(), data_has_encryption_key = data_config.encryption_key().is_some(), @@ -1989,7 +2018,7 @@ async fn run() -> Result<(), Box> { { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, - "usage persistence requires a configured Postgres data backend; set AETHER_GATEWAY_DATA_POSTGRES_URL before starting aether-gateway", + "usage persistence requires a configured shared SQL data backend; set AETHER_DATABASE_DRIVER and AETHER_DATABASE_URL before starting aether-gateway", ) .into()); } @@ -2751,6 +2780,42 @@ mod tests { assert_eq!(database.pool.max_connections, 1); } + #[test] + fn explicit_mysql_driver_accepts_generic_database_url() { + let url = super::resolve_database_url( + Some(DatabaseDriver::Mysql), + None, + Some("postgres://legacy/aether".to_string()), + Some("mysql://root:root@localhost/aether".to_string()), + ); + + assert_eq!(url.as_deref(), Some("mysql://root:root@localhost/aether")); + } + + #[test] + fn explicit_sqlite_driver_accepts_generic_database_url() { + let url = super::resolve_database_url( + Some(DatabaseDriver::Sqlite), + None, + Some("postgres://legacy/aether".to_string()), + Some("sqlite:///opt/aether/data/aether.db".to_string()), + ); + + assert_eq!(url.as_deref(), Some("sqlite:///opt/aether/data/aether.db")); + } + + #[test] + fn postgres_legacy_url_keeps_precedence_over_generic_database_url() { + let url = super::resolve_database_url( + Some(DatabaseDriver::Postgres), + None, + Some("postgres://legacy/aether".to_string()), + Some("postgres://generic/aether".to_string()), + ); + + assert_eq!(url.as_deref(), Some("postgres://legacy/aether")); + } + #[test] fn gateway_data_pool_auto_sizes_server_databases_from_runtime_cpu() { let mut args = test_args(); @@ -3278,6 +3343,23 @@ mod tests { .expect("single-node sqlite memory runtime should be accepted"); } + #[test] + fn multi_node_accepts_mysql_database_backend() { + let mut args = test_args(); + args.deployment_topology = DeploymentTopologyArg::MultiNode; + args.node_role = NodeRoleArg::Frontdoor; + args.video_task_store_path = None; + let database = test_database(DatabaseDriver::Mysql, 8); + + super::validate_deployment_topology( + &args, + Some(&database), + Some("redis://127.0.0.1/0"), + RuntimeBackendArg::Redis, + ) + .expect("multi-node mysql with shared redis should be accepted"); + } + #[test] fn multi_node_rejects_memory_runtime_backend() { let mut args = test_args(); diff --git a/apps/aether-gateway/src/state/bootstrap_admin.rs b/apps/aether-gateway/src/state/bootstrap_admin.rs index e5fe90e48..03db6e899 100644 --- a/apps/aether-gateway/src/state/bootstrap_admin.rs +++ b/apps/aether-gateway/src/state/bootstrap_admin.rs @@ -243,8 +243,7 @@ impl AppState { ) -> Result<(), GatewayError> { if !self.has_auth_user_write_capability() || !self.has_auth_wallet_write_capability() { return Err(GatewayError::Internal( - "bootstrap admin requires Postgres-backed user and wallet write capability" - .to_string(), + "bootstrap admin requires SQL-backed user and wallet write capability".to_string(), )); } diff --git a/apps/aether-gateway/src/tests/architecture/sql_and_data.rs b/apps/aether-gateway/src/tests/architecture/sql_and_data.rs index 80151e267..85bd7b132 100644 --- a/apps/aether-gateway/src/tests/architecture/sql_and_data.rs +++ b/apps/aether-gateway/src/tests/architecture/sql_and_data.rs @@ -1902,15 +1902,22 @@ fn lifecycle_backfills_are_partitioned_by_driver() { )); for required in [ format!("use crate::driver::{driver}::{pool}"), + format!("sqlx::migrate!(\"./backfills/{driver}\")"), + "static BACKFILL_MIGRATOR".to_string(), + "ENSURE_SCHEMA_BACKFILLS_TABLE_SQL".to_string(), + "LIST_APPLIED_BACKFILLS_SQL".to_string(), + "INSERT_APPLIED_BACKFILL_SQL".to_string(), "pub async fn run_backfills".to_string(), "pub async fn pending_backfills".to_string(), + "ensure_schema_backfills_table".to_string(), + "validate_applied_backfills".to_string(), ] { assert!( source.contains(&required), "backfill/{driver}.rs should own {required}" ); } - for forbidden in ["PgPool", "BACKFILL_MIGRATOR", "schema_backfills"] { + for forbidden in ["PgPool", "PgConnection", "crate::driver::postgres"] { assert!( !source.contains(forbidden), "backfill/{driver}.rs should not depend on PostgreSQL via {forbidden}" diff --git a/crates/aether-data/adapters/mysql/Cargo.toml b/crates/aether-data/adapters/mysql/Cargo.toml index 84436cf65..35a96b69b 100644 --- a/crates/aether-data/adapters/mysql/Cargo.toml +++ b/crates/aether-data/adapters/mysql/Cargo.toml @@ -13,6 +13,7 @@ aether-data-query.workspace = true async-trait.workspace = true chrono.workspace = true chrono-tz.workspace = true +flate2.workspace = true serde_json.workspace = true sha2.workspace = true sqlx = { workspace = true, features = ["mysql", "runtime-tokio-rustls", "chrono", "migrate", "macros"] } diff --git a/crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql b/crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql new file mode 100644 index 000000000..6ec4def0a --- /dev/null +++ b/crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql @@ -0,0 +1,230 @@ +-- Complete the portable schema contract that is already represented by the +-- logical/generated schema and PostgreSQL usage capture tables. + +ALTER TABLE provider_api_keys + ADD COLUMN `last_error_at` BIGINT, + ADD COLUMN `last_error_msg` LONGTEXT; + +CREATE TABLE IF NOT EXISTS api_key_provider_mappings ( + `id` VARCHAR(64) NOT NULL, + `api_key_id` VARCHAR(64) NOT NULL, + `provider_id` VARCHAR(64) NOT NULL, + `priority_adjustment` INT NOT NULL DEFAULT 0, + `weight_multiplier` DOUBLE NOT NULL DEFAULT 1, + `is_enabled` TINYINT(1) NOT NULL DEFAULT 1, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_apikey_provider (`api_key_id`, `provider_id`), + KEY api_key_provider_mappings_api_key_id_idx (`api_key_id`), + KEY api_key_provider_mappings_provider_id_idx (`provider_id`), + KEY idx_apikey_provider_enabled (`api_key_id`, `is_enabled`) +); + +CREATE TABLE IF NOT EXISTS provider_usage_tracking ( + `id` VARCHAR(64) NOT NULL, + `provider_id` VARCHAR(64) NOT NULL, + `window_start` BIGINT NOT NULL, + `window_end` BIGINT NOT NULL, + `total_requests` INT NOT NULL DEFAULT 0, + `successful_requests` INT NOT NULL DEFAULT 0, + `failed_requests` INT NOT NULL DEFAULT 0, + `avg_response_time_ms` DOUBLE NOT NULL DEFAULT 0, + `total_response_time_ms` DOUBLE NOT NULL DEFAULT 0, + `total_cost_usd` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + KEY provider_usage_tracking_provider_id_idx (`provider_id`), + KEY provider_usage_tracking_window_start_idx (`window_start`), + KEY idx_provider_window (`provider_id`, `window_start`), + KEY idx_window_time (`window_start`, `window_end`) +); + +ALTER TABLE video_tasks + ADD COLUMN `converted_request_body` JSON, + ADD COLUMN `max_retries` INT NOT NULL DEFAULT 3, + ADD COLUMN `video_urls` JSON, + ADD COLUMN `thumbnail_url` LONGTEXT, + ADD COLUMN `video_size_bytes` BIGINT, + ADD COLUMN `video_expires_at` BIGINT, + ADD COLUMN `stored_video_path` VARCHAR(500), + ADD COLUMN `storage_provider` VARCHAR(50), + ADD COLUMN `remixed_from_task_id` VARCHAR(64), + ADD COLUMN `webhook_url` VARCHAR(500), + ADD COLUMN `webhook_sent` TINYINT(1) NOT NULL DEFAULT 0, + ADD COLUMN `webhook_sent_at` BIGINT, + ADD COLUMN `video_duration_seconds` DOUBLE; + +-- Portable compatibility columns. New canonical HTTP payload writes use the +-- normalized usage_http_audits and usage_body_blobs tables below. +ALTER TABLE `usage` + ADD COLUMN `input_output_total_tokens` BIGINT NOT NULL DEFAULT 0, + ADD COLUMN `cache_creation_input_tokens_5m` BIGINT NOT NULL DEFAULT 0, + ADD COLUMN `cache_creation_input_tokens_1h` BIGINT NOT NULL DEFAULT 0, + ADD COLUMN `input_context_tokens` BIGINT NOT NULL DEFAULT 0, + ADD COLUMN `input_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `output_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `cache_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `cache_creation_cost_usd_5m` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `cache_creation_cost_usd_1h` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `request_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_input_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_output_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_cache_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_cache_creation_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_cache_creation_cost_usd_5m` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_cache_creation_cost_usd_1h` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_cache_read_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `actual_request_cost_usd` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `rate_multiplier` DOUBLE NOT NULL DEFAULT 1, + ADD COLUMN `input_price_per_1m` DOUBLE, + ADD COLUMN `cache_creation_price_per_1m` DOUBLE, + ADD COLUMN `cache_creation_price_per_1m_5m` DOUBLE, + ADD COLUMN `cache_creation_price_per_1m_1h` DOUBLE, + ADD COLUMN `cache_read_price_per_1m` DOUBLE, + ADD COLUMN `price_per_request` DOUBLE, + ADD COLUMN `request_headers` JSON, + ADD COLUMN `request_body` JSON, + ADD COLUMN `provider_request_headers` JSON, + ADD COLUMN `provider_request_body` JSON, + ADD COLUMN `response_headers` JSON, + ADD COLUMN `response_body` JSON, + ADD COLUMN `client_response_headers` JSON, + ADD COLUMN `client_response_body` JSON, + ADD COLUMN `request_body_compressed` LONGBLOB, + ADD COLUMN `provider_request_body_compressed` LONGBLOB, + ADD COLUMN `response_body_compressed` LONGBLOB, + ADD COLUMN `client_response_body_compressed` LONGBLOB, + ADD COLUMN `created_at` BIGINT, + ADD COLUMN `username` VARCHAR(255), + ADD COLUMN `api_key_name` VARCHAR(255); + +CREATE TABLE IF NOT EXISTS usage_body_blobs ( + `body_ref` VARCHAR(160) NOT NULL, + `request_id` VARCHAR(128) NOT NULL, + `body_field` VARCHAR(50) NOT NULL, + `payload_gzip` LONGBLOB NOT NULL, + `created_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + `updated_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + PRIMARY KEY (`body_ref`), + UNIQUE KEY usage_body_blobs_request_id_field_key (`request_id`, `body_field`), + KEY ix_usage_body_blobs_request_id (`request_id`), + CONSTRAINT usage_body_blobs_request_id_fkey + FOREIGN KEY (`request_id`) REFERENCES `usage` (`request_id`) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS usage_http_audits ( + `request_id` VARCHAR(128) NOT NULL, + `request_headers` JSON, + `provider_request_headers` JSON, + `response_headers` JSON, + `client_response_headers` JSON, + `request_body_ref` VARCHAR(160), + `provider_request_body_ref` VARCHAR(160), + `response_body_ref` VARCHAR(160), + `client_response_body_ref` VARCHAR(160), + `request_body_state` VARCHAR(32), + `provider_request_body_state` VARCHAR(32), + `response_body_state` VARCHAR(32), + `client_response_body_state` VARCHAR(32), + `body_capture_mode` VARCHAR(32) NOT NULL DEFAULT 'none', + `created_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + `updated_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + PRIMARY KEY (`request_id`), + KEY ix_usage_http_audits_updated_at (`updated_at`), + CONSTRAINT usage_http_audits_request_id_fkey + FOREIGN KEY (`request_id`) REFERENCES `usage` (`request_id`) ON DELETE CASCADE +); + +-- Billing V3 keeps the immutable pricing/token snapshot separate from the +-- mutable compatibility columns on usage. +ALTER TABLE usage_settlement_snapshots + ADD COLUMN `billing_snapshot_schema_version` VARCHAR(20), + ADD COLUMN `billing_snapshot_status` VARCHAR(20), + ADD COLUMN `rate_multiplier` DECIMAL(10,6), + ADD COLUMN `is_free_tier` TINYINT(1), + ADD COLUMN `input_price_per_1m` DECIMAL(20,8), + ADD COLUMN `output_price_per_1m` DECIMAL(20,8), + ADD COLUMN `cache_creation_price_per_1m` DECIMAL(20,8), + ADD COLUMN `cache_read_price_per_1m` DECIMAL(20,8), + ADD COLUMN `price_per_request` DECIMAL(20,8), + ADD COLUMN `settlement_snapshot_schema_version` VARCHAR(20), + ADD COLUMN `settlement_snapshot` JSON, + ADD COLUMN `billing_dimensions` JSON, + ADD COLUMN `billing_input_tokens` BIGINT, + ADD COLUMN `billing_effective_input_tokens` BIGINT, + ADD COLUMN `billing_output_tokens` BIGINT, + ADD COLUMN `billing_cache_creation_tokens` BIGINT, + ADD COLUMN `billing_cache_creation_5m_tokens` BIGINT, + ADD COLUMN `billing_cache_creation_1h_tokens` BIGINT, + ADD COLUMN `billing_cache_read_tokens` BIGINT, + ADD COLUMN `billing_total_input_context` BIGINT, + ADD COLUMN `billing_cache_creation_cost_usd` DECIMAL(20,8), + ADD COLUMN `billing_cache_read_cost_usd` DECIMAL(20,8), + ADD COLUMN `billing_total_cost_usd` DECIMAL(20,8), + ADD COLUMN `billing_actual_total_cost_usd` DECIMAL(20,8), + ADD COLUMN `billing_pricing_source` VARCHAR(50), + ADD COLUMN `billing_rule_id` VARCHAR(100), + ADD COLUMN `billing_rule_version` VARCHAR(50); + +CREATE INDEX ix_usage_settlement_snapshots_schema_version + ON usage_settlement_snapshots (`settlement_snapshot_schema_version`); +CREATE INDEX ix_usage_settlement_snapshots_pricing_source + ON usage_settlement_snapshots (`billing_pricing_source`); + +CREATE TABLE IF NOT EXISTS stats_summary ( + `id` VARCHAR(64) NOT NULL, + `cutoff_date` BIGINT NOT NULL, + `all_time_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_success_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_error_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_input_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_output_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cost` DOUBLE NOT NULL DEFAULT 0, + `all_time_actual_cost` DOUBLE NOT NULL DEFAULT 0, + `total_users` BIGINT NOT NULL DEFAULT 0, + `active_users` BIGINT NOT NULL DEFAULT 0, + `total_api_keys` BIGINT NOT NULL DEFAULT 0, + `active_api_keys` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`) +); + +CREATE TABLE IF NOT EXISTS user_model_usage_counts ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `model` VARCHAR(255) NOT NULL, + `usage_count` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_user_model_usage_count (`user_id`, `model`), + KEY idx_user_model_usage_user (`user_id`), + KEY idx_user_model_usage_model (`model`) +); + +ALTER TABLE stats_daily + ADD COLUMN `input_cost` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `output_cost` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN `p50_response_time_ms` BIGINT, + ADD COLUMN `p90_response_time_ms` BIGINT, + ADD COLUMN `p99_response_time_ms` BIGINT, + ADD COLUMN `p50_first_byte_time_ms` BIGINT, + ADD COLUMN `p90_first_byte_time_ms` BIGINT, + ADD COLUMN `p99_first_byte_time_ms` BIGINT; + +-- MySQL has no partial indexes, so keep the active status first and preserve +-- the cleanup query's ascending timestamp/request ordering. +CREATE INDEX idx_usage_stale_pending_created_request + ON `usage` (`status`, `created_at_unix_ms`, `request_id`); + +CREATE INDEX idx_provider_api_keys_provider_created_at_desc + ON provider_api_keys (`provider_id`, `created_at` DESC, `name`, `id`); +CREATE INDEX idx_provider_api_keys_provider_last_used_at_desc + ON provider_api_keys (`provider_id`, `last_used_at` DESC, `name`, `id`); diff --git a/crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql b/crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql new file mode 100644 index 000000000..a7b4d84a1 --- /dev/null +++ b/crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql @@ -0,0 +1,381 @@ +ALTER TABLE stats_user_daily + ADD COLUMN actual_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0, + ADD COLUMN effective_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN total_input_context BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN cache_read_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_output_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_first_finalized_at_unix_secs BIGINT, + ADD COLUMN settled_last_finalized_at_unix_secs BIGINT; + +ALTER TABLE stats_hourly_user + ADD COLUMN cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN actual_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_output_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_first_finalized_at_unix_secs BIGINT, + ADD COLUMN settled_last_finalized_at_unix_secs BIGINT; + +ALTER TABLE stats_daily + ADD COLUMN effective_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN total_input_context BIGINT NOT NULL DEFAULT 0, + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_hit_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_hit_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_hit_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_total_input_context BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_read_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_output_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_first_finalized_at_unix_secs BIGINT, + ADD COLUMN settled_last_finalized_at_unix_secs BIGINT; + +ALTER TABLE stats_hourly + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_hit_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_hit_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_hit_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_total_input_context BIGINT NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN completed_cache_read_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_cost DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN settled_total_requests BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_input_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_output_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN settled_first_finalized_at_unix_secs BIGINT, + ADD COLUMN settled_last_finalized_at_unix_secs BIGINT; + +ALTER TABLE stats_daily_model + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0; + +ALTER TABLE stats_hourly_model + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0; + +ALTER TABLE stats_hourly_user_model + ADD COLUMN response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + ADD COLUMN response_time_samples BIGINT NOT NULL DEFAULT 0; + +CREATE TABLE stats_user_summary ( + id VARCHAR(64) NOT NULL, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + cutoff_date BIGINT NOT NULL, + all_time_requests BIGINT NOT NULL DEFAULT 0, + all_time_success_requests BIGINT NOT NULL DEFAULT 0, + all_time_error_requests BIGINT NOT NULL DEFAULT 0, + all_time_input_tokens BIGINT NOT NULL DEFAULT 0, + all_time_output_tokens BIGINT NOT NULL DEFAULT 0, + all_time_cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + all_time_cache_read_tokens BIGINT NOT NULL DEFAULT 0, + all_time_cost DOUBLE NOT NULL DEFAULT 0, + all_time_actual_cost DOUBLE NOT NULL DEFAULT 0, + active_days BIGINT NOT NULL DEFAULT 0, + first_active_date BIGINT, + last_active_date BIGINT, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_user_summary_user_id (user_id), + KEY idx_stats_user_summary_cutoff_date (cutoff_date) +); + +CREATE TABLE stats_user_daily_model ( + id VARCHAR(64) NOT NULL, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + total_requests BIGINT NOT NULL DEFAULT 0, + success_requests BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + effective_input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_input_context BIGINT NOT NULL DEFAULT 0, + cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + total_cost DOUBLE NOT NULL DEFAULT 0, + actual_total_cost DOUBLE NOT NULL DEFAULT 0, + response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + response_time_samples BIGINT NOT NULL DEFAULT 0, + successful_response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + successful_response_time_samples BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_user_daily_model (`user_id`, `date`, `model`), + KEY idx_stats_user_daily_model_date (`date`), + KEY idx_stats_user_daily_model_user_id (`user_id`) +); + +CREATE TABLE stats_user_daily_provider ( + id VARCHAR(64) NOT NULL, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + provider_name VARCHAR(255) NOT NULL, + total_requests BIGINT NOT NULL DEFAULT 0, + success_requests BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + effective_input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_input_context BIGINT NOT NULL DEFAULT 0, + cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + total_cost DOUBLE NOT NULL DEFAULT 0, + actual_total_cost DOUBLE NOT NULL DEFAULT 0, + response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + response_time_samples BIGINT NOT NULL DEFAULT 0, + successful_response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + successful_response_time_samples BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_user_daily_provider (`user_id`, `date`, `provider_name`), + KEY idx_stats_user_daily_provider_date (`date`), + KEY idx_stats_user_daily_provider_user_id (`user_id`) +); + +CREATE TABLE stats_user_daily_api_format ( + id VARCHAR(64) NOT NULL, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + api_format VARCHAR(128) NOT NULL, + total_requests BIGINT NOT NULL DEFAULT 0, + success_requests BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + effective_input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_input_context BIGINT NOT NULL DEFAULT 0, + cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + total_cost DOUBLE NOT NULL DEFAULT 0, + actual_total_cost DOUBLE NOT NULL DEFAULT 0, + response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + response_time_samples BIGINT NOT NULL DEFAULT 0, + successful_response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + successful_response_time_samples BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_user_daily_api_format (`user_id`, `date`, `api_format`), + KEY idx_stats_user_daily_api_format_date (`date`), + KEY idx_stats_user_daily_api_format_user_id (`user_id`) +); + +CREATE TABLE stats_daily_model_provider ( + id VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + provider_name VARCHAR(255) NOT NULL, + total_requests BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_cost DOUBLE NOT NULL DEFAULT 0, + response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + response_time_samples BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_daily_model_provider (`date`, `model`, `provider_name`), + KEY idx_stats_daily_model_provider_date (`date`) +); + +CREATE TABLE stats_user_daily_model_provider ( + id VARCHAR(64) NOT NULL, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + provider_name VARCHAR(255) NOT NULL, + total_requests BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_cost DOUBLE NOT NULL DEFAULT 0, + response_time_sum_ms DOUBLE NOT NULL DEFAULT 0, + response_time_samples BIGINT NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (id), + UNIQUE KEY uq_stats_user_daily_model_provider (`user_id`, `date`, `model`, `provider_name`), + KEY idx_stats_user_daily_model_provider_date (`date`), + KEY idx_stats_user_daily_model_provider_user_date (`user_id`, `date`) +); + +CREATE TABLE stats_daily_cost_savings ( + id VARCHAR(64) PRIMARY KEY, + `date` BIGINT NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_daily_cost_savings_date (`date`) +); + +CREATE TABLE stats_daily_cost_savings_provider ( + id VARCHAR(64) PRIMARY KEY, + `date` BIGINT NOT NULL, + provider_name VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_daily_cost_savings_provider (`date`, `provider_name`) +); + +CREATE TABLE stats_daily_cost_savings_model ( + id VARCHAR(64) PRIMARY KEY, + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_daily_cost_savings_model (`date`, `model`) +); + +CREATE TABLE stats_daily_cost_savings_model_provider ( + id VARCHAR(64) PRIMARY KEY, + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + provider_name VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_daily_cost_savings_model_provider (`date`, `model`, `provider_name`) +); + +CREATE TABLE stats_user_daily_cost_savings ( + id VARCHAR(64) PRIMARY KEY, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_user_daily_cost_savings (`user_id`, `date`) +); + +CREATE TABLE stats_user_daily_cost_savings_provider ( + id VARCHAR(64) PRIMARY KEY, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + provider_name VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_user_daily_cost_savings_provider (`user_id`, `date`, `provider_name`) +); + +CREATE TABLE stats_user_daily_cost_savings_model ( + id VARCHAR(64) PRIMARY KEY, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_user_daily_cost_savings_model (`user_id`, `date`, `model`) +); + +CREATE TABLE stats_user_daily_cost_savings_model_provider ( + id VARCHAR(64) PRIMARY KEY, + user_id VARCHAR(64) NOT NULL, + username VARCHAR(255), + `date` BIGINT NOT NULL, + model VARCHAR(255) NOT NULL, + provider_name VARCHAR(255) NOT NULL, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_cost DOUBLE NOT NULL DEFAULT 0, + cache_creation_cost DOUBLE NOT NULL DEFAULT 0, + estimated_full_cost DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE KEY uq_stats_user_daily_cost_savings_model_provider (`user_id`, `date`, `model`, `provider_name`) +); + +CREATE INDEX idx_stats_daily_cost_savings_provider_date + ON stats_daily_cost_savings_provider (`date`); +CREATE INDEX idx_stats_daily_cost_savings_model_date + ON stats_daily_cost_savings_model (`date`); +CREATE INDEX idx_stats_daily_cost_savings_model_provider_date + ON stats_daily_cost_savings_model_provider (`date`); +CREATE INDEX idx_stats_user_daily_cost_savings_date + ON stats_user_daily_cost_savings (`date`); +CREATE INDEX idx_stats_user_daily_cost_savings_provider_date + ON stats_user_daily_cost_savings_provider (`date`); +CREATE INDEX idx_stats_user_daily_cost_savings_model_date + ON stats_user_daily_cost_savings_model (`date`); +CREATE INDEX idx_stats_user_daily_cost_savings_model_provider_date + ON stats_user_daily_cost_savings_model_provider (`date`); + +-- Existing completed buckets predate the enriched dimensions above. Preserve the rows for +-- reads, but make the bounded aggregation worker replay every historical bucket. +UPDATE stats_hourly SET is_complete = 0 WHERE is_complete <> 0; +UPDATE stats_daily SET is_complete = 0 WHERE is_complete <> 0; diff --git a/crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql b/crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql new file mode 100644 index 000000000..9dc553464 --- /dev/null +++ b/crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql @@ -0,0 +1,69 @@ +CREATE TABLE IF NOT EXISTS usage_routing_snapshots ( + `request_id` VARCHAR(128) NOT NULL, + `candidate_id` VARCHAR(160), + `candidate_index` BIGINT, + `key_name` VARCHAR(255), + `planner_kind` VARCHAR(120), + `route_family` VARCHAR(80), + `route_kind` VARCHAR(80), + `execution_path` VARCHAR(80), + `local_execution_runtime_miss_reason` VARCHAR(255), + `selected_provider_id` VARCHAR(100), + `selected_endpoint_id` VARCHAR(100), + `selected_provider_api_key_id` VARCHAR(100), + `has_format_conversion` TINYINT(1), + `created_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + `updated_at` BIGINT NOT NULL DEFAULT (UNIX_TIMESTAMP()), + PRIMARY KEY (`request_id`), + KEY ix_usage_routing_snapshots_route_family_kind (`route_family`, `route_kind`), + KEY ix_usage_routing_snapshots_candidate_id (`candidate_id`), + CONSTRAINT usage_routing_snapshots_request_id_fkey + FOREIGN KEY (`request_id`) REFERENCES `usage` (`request_id`) ON DELETE CASCADE +); + +INSERT INTO usage_routing_snapshots ( + request_id, + candidate_id, + candidate_index, + key_name, + planner_kind, + route_family, + route_kind, + execution_path, + local_execution_runtime_miss_reason, + selected_provider_id, + selected_endpoint_id, + selected_provider_api_key_id, + has_format_conversion, + created_at, + updated_at +) +SELECT + request_id, + candidate_id, + candidate_index, + key_name, + planner_kind, + route_family, + route_kind, + execution_path, + local_execution_runtime_miss_reason, + provider_id, + provider_endpoint_id, + provider_api_key_id, + has_format_conversion, + COALESCE(NULLIF(created_at_unix_ms, 0), NULLIF(updated_at_unix_secs, 0), UNIX_TIMESTAMP()), + COALESCE(NULLIF(updated_at_unix_secs, 0), NULLIF(created_at_unix_ms, 0), UNIX_TIMESTAMP()) +FROM `usage` +WHERE candidate_id IS NOT NULL + OR candidate_index IS NOT NULL + OR key_name IS NOT NULL + OR planner_kind IS NOT NULL + OR route_family IS NOT NULL + OR route_kind IS NOT NULL + OR execution_path IS NOT NULL + OR local_execution_runtime_miss_reason IS NOT NULL + OR provider_id IS NOT NULL + OR provider_endpoint_id IS NOT NULL + OR provider_api_key_id IS NOT NULL + OR has_format_conversion <> 0; diff --git a/crates/aether-data/adapters/mysql/src/billing.rs b/crates/aether-data/adapters/mysql/src/billing.rs index 1a9d2f3a6..0b38c435d 100644 --- a/crates/aether-data/adapters/mysql/src/billing.rs +++ b/crates/aether-data/adapters/mysql/src/billing.rs @@ -642,23 +642,8 @@ LIMIT 1 input: &PaymentGatewayConfigWriteInput, ) -> Result, DataLayerError> { let provider = input.provider.trim().to_ascii_lowercase(); - let existing_secret = if input.preserve_existing_secret { - sqlx::query_scalar::<_, String>( - "SELECT merchant_key_encrypted FROM payment_gateway_configs WHERE provider = ?", - ) - .bind(&provider) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - } else { - None - }; - let secret = if input.preserve_existing_secret { - existing_secret - } else { - input.merchant_key_encrypted.clone() - }; let now = current_unix_secs_i64(); + let mut tx = self.pool.begin().await.map_sql_err()?; sqlx::query( r#" INSERT INTO payment_gateway_configs ( @@ -672,7 +657,10 @@ ON DUPLICATE KEY UPDATE endpoint_url = VALUES(endpoint_url), callback_base_url = VALUES(callback_base_url), merchant_id = VALUES(merchant_id), - merchant_key_encrypted = VALUES(merchant_key_encrypted), + merchant_key_encrypted = CASE + WHEN ? THEN merchant_key_encrypted + ELSE VALUES(merchant_key_encrypted) + END, pay_currency = VALUES(pay_currency), usd_exchange_rate = VALUES(usd_exchange_rate), min_recharge_usd = VALUES(min_recharge_usd), @@ -685,22 +673,35 @@ ON DUPLICATE KEY UPDATE .bind(&input.endpoint_url) .bind(input.callback_base_url.as_deref()) .bind(&input.merchant_id) - .bind(secret.as_deref()) + .bind(input.merchant_key_encrypted.as_deref()) .bind(&input.pay_currency) .bind(input.usd_exchange_rate) .bind(input.min_recharge_usd) .bind(json_to_string(&input.channels_json)?) .bind(now) .bind(now) - .execute(&self.pool) + .bind(input.preserve_existing_secret) + .execute(&mut *tx) .await .map_sql_err()?; - match self.find_payment_gateway_config(&provider).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "upserted payment gateway config missing".to_string(), - )), - } + let row = sqlx::query( + r#" +SELECT + provider, enabled, endpoint_url, callback_base_url, merchant_id, + merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, + channels_json, created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs +FROM payment_gateway_configs +WHERE provider = ? +LIMIT 1 +"#, + ) + .bind(&provider) + .fetch_one(&mut *tx) + .await + .map_sql_err()?; + let record = map_payment_gateway_config_mysql(&row)?; + tx.commit().await.map_sql_err()?; + Ok(AdminBillingMutationOutcome::Applied(record)) } async fn list_billing_plans( @@ -1402,7 +1403,13 @@ fn parse_required_json(raw: String) -> Result #[cfg(test)] mod tests { + use serde_json::json; + use super::MysqlBillingReadRepository; + use crate::run_migrations; + use aether_data_contracts::repository::billing::{ + AdminBillingMutationOutcome, BillingReadRepository, PaymentGatewayConfigWriteInput, + }; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -1414,4 +1421,88 @@ mod tests { let _repository = MysqlBillingReadRepository::new(pool); } + + #[tokio::test] + async fn mysql_gateway_upsert_preserves_secret_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql billing parity test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + sqlx::query("DELETE FROM payment_gateway_configs WHERE provider = 'paritypay'") + .execute(&pool) + .await + .expect("gateway cleanup should succeed"); + let repository = MysqlBillingReadRepository::new(pool.clone()); + + let mut input = PaymentGatewayConfigWriteInput { + provider: " ParityPay ".to_string(), + enabled: true, + endpoint_url: "https://pay.example/first".to_string(), + callback_base_url: Some("https://api.example/callback".to_string()), + merchant_id: "merchant-1".to_string(), + merchant_key_encrypted: Some("secret-first".to_string()), + preserve_existing_secret: true, + pay_currency: "USD".to_string(), + usd_exchange_rate: 1.0, + min_recharge_usd: 2.0, + channels_json: json!([{"id": "card"}]), + }; + let inserted = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway insert should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway insert outcome: {other:?}"), + }; + assert_eq!( + inserted.merchant_key_encrypted.as_deref(), + Some("secret-first") + ); + + input.merchant_key_encrypted = Some("secret-ignored".to_string()); + let preserved = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway preserve update should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway preserve outcome: {other:?}"), + }; + assert_eq!( + preserved.merchant_key_encrypted.as_deref(), + Some("secret-first") + ); + + input.preserve_existing_secret = false; + input.merchant_key_encrypted = Some("secret-replaced".to_string()); + let replaced = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway replacement update should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway replace outcome: {other:?}"), + }; + assert_eq!( + replaced.merchant_key_encrypted.as_deref(), + Some("secret-replaced") + ); + + sqlx::query("DELETE FROM payment_gateway_configs WHERE provider = 'paritypay'") + .execute(&pool) + .await + .expect("gateway cleanup should succeed"); + } } diff --git a/crates/aether-data/adapters/mysql/src/candidates.rs b/crates/aether-data/adapters/mysql/src/candidates.rs index 33468fe97..b90cdcabc 100644 --- a/crates/aether-data/adapters/mysql/src/candidates.rs +++ b/crates/aether-data/adapters/mysql/src/candidates.rs @@ -1,7 +1,7 @@ use std::collections::{BTreeMap, BTreeSet}; use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; +use sqlx::{mysql::MySqlRow, MySql, MySqlConnection, QueryBuilder, Row}; use aether_data_contracts::repository::candidates::{ request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket, @@ -51,24 +51,6 @@ impl MysqlRequestCandidateRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } } - - async fn find_by_unique( - &self, - request_id: &str, - candidate_index: u32, - retry_index: u32, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1" - )) - .bind(request_id) - .bind(to_i32(candidate_index)?) - .bind(to_i32(retry_index)?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_candidate_row).transpose() - } } #[async_trait] @@ -237,16 +219,50 @@ impl RequestCandidateWriteRepository for MysqlRequestCandidateRepository { candidate: UpsertRequestCandidateRecord, ) -> Result { candidate.validate()?; - let existing = self - .find_by_unique( - &candidate.request_id, - candidate.candidate_index, - candidate.retry_index, - ) - .await?; - let merged = merge_candidate(candidate, existing)?; - upsert_merged_candidate(&self.pool, &merged).await?; - Ok(merged) + let mut tx = self.pool.begin().await.map_sql_err()?; + match upsert_candidate_in_transaction(&mut tx, candidate).await { + Ok(candidate) => { + tx.commit().await.map_sql_err()?; + Ok(candidate) + } + Err(err) => { + tx.rollback().await.map_sql_err()?; + Err(err) + } + } + } + + async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + if candidates.is_empty() { + return Ok(0); + } + for candidate in &candidates { + candidate.validate()?; + } + + let mut tx = self.pool.begin().await.map_sql_err()?; + let result: Result = async { + let mut persisted = 0usize; + for candidate in candidates { + upsert_candidate_in_transaction(&mut tx, candidate).await?; + persisted = persisted.saturating_add(1); + } + Ok(persisted) + } + .await; + match result { + Ok(persisted) => { + tx.commit().await.map_sql_err()?; + Ok(persisted) + } + Err(err) => { + tx.rollback().await.map_sql_err()?; + Err(err) + } + } } async fn delete_created_before( @@ -282,8 +298,90 @@ WHERE id IN ( } } +async fn upsert_candidate_in_transaction( + tx: &mut sqlx::Transaction<'_, MySql>, + candidate: UpsertRequestCandidateRecord, +) -> Result { + // Write first so both existing rows and previously empty unique keys are locked + // before the Rust merge reads their latest committed state. + let insert_candidate = merge_candidate(candidate.clone(), None)?; + insert_candidate_if_absent(tx, &insert_candidate).await?; + let existing = find_by_unique_for_update( + tx, + &candidate.request_id, + candidate.candidate_index, + candidate.retry_index, + ) + .await? + .ok_or_else(|| { + DataLayerError::UnexpectedValue( + "request candidate row was not locked after insert-if-absent".to_string(), + ) + })?; + let merged = merge_candidate(candidate, Some(existing))?; + upsert_merged_candidate(tx, &merged).await?; + find_by_unique_for_update( + tx, + &merged.request_id, + merged.candidate_index, + merged.retry_index, + ) + .await? + .ok_or_else(|| { + DataLayerError::UnexpectedValue( + "request candidate row disappeared after atomic upsert".to_string(), + ) + }) +} + +async fn insert_candidate_if_absent( + connection: &mut MySqlConnection, + candidate: &StoredRequestCandidate, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO request_candidates ( + id, request_id, candidate_index, retry_index, status, created_at +) +VALUES (?, ?, ?, ?, ?, ?) +ON DUPLICATE KEY UPDATE id = id +"#, + ) + .bind(&candidate.id) + .bind(&candidate.request_id) + .bind(to_i32(candidate.candidate_index)?) + .bind(to_i32(candidate.retry_index)?) + .bind(status_to_database(candidate.status)) + .bind(u64_to_i64( + candidate.created_at_unix_ms, + "request candidate created_at", + )?) + .execute(connection) + .await + .map_sql_err()?; + Ok(()) +} + +async fn find_by_unique_for_update( + connection: &mut MySqlConnection, + request_id: &str, + candidate_index: u32, + retry_index: u32, +) -> Result, DataLayerError> { + let row = sqlx::query(&format!( + "{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1 FOR UPDATE" + )) + .bind(request_id) + .bind(to_i32(candidate_index)?) + .bind(to_i32(retry_index)?) + .fetch_optional(connection) + .await + .map_sql_err()?; + row.as_ref().map(map_candidate_row).transpose() +} + async fn upsert_merged_candidate( - pool: &MysqlPool, + connection: &mut MySqlConnection, candidate: &StoredRequestCandidate, ) -> Result<(), DataLayerError> { sqlx::query( @@ -405,7 +503,7 @@ ON DUPLICATE KEY UPDATE candidate.finished_at_unix_ms, "request candidate finished_at", )?) - .execute(pool) + .execute(connection) .await .map_sql_err()?; Ok(()) @@ -776,9 +874,10 @@ mod tests { use super::MysqlRequestCandidateRepository; use crate::run_migrations; use aether_data_contracts::repository::candidates::{ - RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate, - UpsertRequestCandidateRecord, + RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, + StoredRequestCandidate, UpsertRequestCandidateRecord, }; + use serde_json::json; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -792,7 +891,7 @@ mod tests { } #[tokio::test] - async fn mysql_atomic_conflict_keeps_candidate_lifecycle_monotonic_when_configured() { + async fn mysql_concurrent_and_batch_upserts_are_atomic_when_configured() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") .ok() .filter(|value| !value.trim().is_empty()) @@ -803,7 +902,7 @@ mod tests { return; }; let pool = sqlx::mysql::MySqlPoolOptions::new() - .max_connections(2) + .max_connections(12) .connect(&database_url) .await .expect("mysql test pool should connect"); @@ -811,103 +910,172 @@ mod tests { .await .expect("mysql migrations should run"); let repository = MysqlRequestCandidateRepository::new(pool.clone()); - let request_id = format!("candidate-lifecycle-{}", uuid::Uuid::new_v4()); - - let terminal = stored_candidate( + let request_id = format!("candidate-concurrency-{}", uuid::Uuid::new_v4()); + let mut initial = sample_upsert( &request_id, - "terminal", - 0, - RequestCandidateStatus::Success, - Some(123), - Some(2_000_002), - ); - super::upsert_merged_candidate(&pool, &terminal) - .await - .expect("terminal candidate should insert"); - let stale_streaming = stored_candidate( - &request_id, - "stale-streaming", - 0, - RequestCandidateStatus::Streaming, - Some(9_999), - Some(9_999_999), - ); - super::upsert_merged_candidate(&pool, &stale_streaming) - .await - .expect("stale streaming conflict should execute"); - - let streaming = stored_candidate( - &request_id, - "streaming", - 1, - RequestCandidateStatus::Streaming, - None, - None, - ); - super::upsert_merged_candidate(&pool, &streaming) - .await - .expect("streaming candidate should insert"); - let stale_pending = stored_candidate( - &request_id, - "stale-pending", - 1, + "initial", RequestCandidateStatus::Pending, - None, - None, + Some(json!({"initial": true})), + 3_000_000, ); - super::upsert_merged_candidate(&pool, &stale_pending) + initial.is_cached = Some(false); + repository + .upsert(initial) .await - .expect("stale pending conflict should execute"); + .expect("initial candidate should insert"); - let pending = stored_candidate( - &request_id, - "pending", - 2, - RequestCandidateStatus::Pending, - Some(321), - None, - ); - super::upsert_merged_candidate(&pool, &pending) - .await - .expect("pending candidate should insert"); - let stale_available = stored_candidate( - &request_id, - "stale-available", - 2, - RequestCandidateStatus::Available, - Some(9_999), - Some(9_999_999), - ); - super::upsert_merged_candidate(&pool, &stale_available) - .await - .expect("stale available conflict should execute"); + const WRITERS: usize = 8; + let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(WRITERS)); + let mut tasks = Vec::new(); + for writer in 0..WRITERS { + let repository = repository.clone(); + let request_id = request_id.clone(); + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + let status = if writer == 0 { + RequestCandidateStatus::Success + } else { + RequestCandidateStatus::Streaming + }; + let mut extra_data = serde_json::Map::new(); + extra_data.insert(format!("writer_{writer}"), json!(writer)); + let mut candidate = sample_upsert( + &request_id, + format!("writer-{writer}").as_str(), + status, + Some(serde_json::Value::Object(extra_data)), + 3_100_000 + u64::try_from(writer).expect("writer index should fit") * 10, + ); + if writer != 0 { + candidate.latency_ms = Some(9_000 + writer as u64); + candidate.finished_at_unix_ms = Some(9_000_000 + writer as u64); + } + barrier.wait().await; + repository.upsert(candidate).await + })); + } + for task in tasks { + task.await + .expect("candidate writer should join") + .expect("candidate writer should persist"); + } let candidates = repository .list_by_request_id(&request_id) .await .expect("mysql request candidates should load"); - let terminal = candidates - .iter() - .find(|candidate| candidate.candidate_index == 0) - .expect("terminal candidate should remain"); - assert_eq!(terminal.status, RequestCandidateStatus::Success); - assert_eq!(terminal.latency_ms, Some(123)); - assert_eq!(terminal.finished_at_unix_ms, Some(2_000_002)); - let streaming = candidates - .iter() - .find(|candidate| candidate.candidate_index == 1) - .expect("streaming candidate should remain"); - assert_eq!(streaming.status, RequestCandidateStatus::Streaming); - let pending = candidates - .iter() - .find(|candidate| candidate.candidate_index == 2) - .expect("pending candidate should remain"); - assert_eq!(pending.status, RequestCandidateStatus::Pending); - assert_eq!(pending.latency_ms, Some(321)); - assert_eq!(pending.finished_at_unix_ms, None); + assert_eq!(candidates.len(), 1); + let candidate = &candidates[0]; + assert_eq!(candidate.id, "initial"); + assert_eq!(candidate.status, RequestCandidateStatus::Success); + assert_eq!(candidate.latency_ms, Some(123)); + assert_eq!(candidate.finished_at_unix_ms, Some(3_100_002)); + let extra_data = candidate + .extra_data + .as_ref() + .and_then(serde_json::Value::as_object) + .expect("merged extra data should be an object"); + assert_eq!(extra_data.get("initial"), Some(&json!(true))); + for writer in 0..WRITERS { + assert_eq!( + extra_data.get(format!("writer_{writer}").as_str()), + Some(&json!(writer)) + ); + } - sqlx::query("DELETE FROM request_candidates WHERE request_id = ?") + let batch_request_id = format!("candidate-batch-{}", uuid::Uuid::new_v4()); + let mut pending = sample_upsert( + &batch_request_id, + "batch-first", + RequestCandidateStatus::Pending, + Some(json!({"pending": true})), + 4_000_000, + ); + pending.is_cached = Some(false); + let mut streaming = sample_upsert( + &batch_request_id, + "batch-second", + RequestCandidateStatus::Streaming, + Some(json!({"streaming": true})), + 4_000_100, + ); + streaming.is_cached = None; + let mut success = sample_upsert( + &batch_request_id, + "batch-third", + RequestCandidateStatus::Success, + Some(json!({"success": true})), + 4_000_200, + ); + success.is_cached = Some(true); + let mut late_pending = sample_upsert( + &batch_request_id, + "batch-fourth", + RequestCandidateStatus::Pending, + Some(json!({"late": true})), + 4_000_300, + ); + late_pending.is_cached = None; + late_pending.latency_ms = Some(9_999); + late_pending.finished_at_unix_ms = Some(9_999_999); + assert_eq!( + repository + .upsert_many(vec![pending, streaming, success, late_pending]) + .await + .expect("ordered batch should persist"), + 4 + ); + let batch_candidates = repository + .list_by_request_id(&batch_request_id) + .await + .expect("batch candidate should load"); + assert_eq!(batch_candidates.len(), 1); + assert_eq!(batch_candidates[0].id, "batch-first"); + assert_eq!(batch_candidates[0].status, RequestCandidateStatus::Success); + assert!(batch_candidates[0].is_cached); + assert_eq!(batch_candidates[0].latency_ms, Some(123)); + assert_eq!(batch_candidates[0].finished_at_unix_ms, Some(4_000_202)); + assert_eq!( + batch_candidates[0].extra_data, + Some(json!({ + "pending": true, + "streaming": true, + "success": true, + "late": true + })) + ); + + let rollback_request_id = format!("candidate-rollback-{}", uuid::Uuid::new_v4()); + let valid = sample_upsert( + &rollback_request_id, + "rollback-valid", + RequestCandidateStatus::Pending, + None, + 5_000_000, + ); + let mut invalid = sample_upsert( + &rollback_request_id, + "rollback-invalid", + RequestCandidateStatus::Success, + None, + 5_000_100, + ); + invalid.candidate_index = 1; + invalid.latency_ms = Some(u64::MAX); + repository + .upsert_many(vec![valid, invalid]) + .await + .expect_err("invalid later row should roll back the batch"); + assert!(repository + .list_by_request_id(&rollback_request_id) + .await + .expect("rolled-back batch should be readable") + .is_empty()); + + sqlx::query("DELETE FROM request_candidates WHERE request_id IN (?, ?, ?)") .bind(&request_id) + .bind(&batch_request_id) + .bind(&rollback_request_id) .execute(&pool) .await .expect("mysql candidate test rows should clean up"); @@ -984,40 +1152,38 @@ mod tests { ); } - fn stored_candidate( + fn sample_upsert( request_id: &str, id: &str, - candidate_index: u32, status: RequestCandidateStatus, - latency_ms: Option, - finished_at_unix_ms: Option, - ) -> StoredRequestCandidate { - StoredRequestCandidate::new( - id.to_string(), - request_id.to_string(), - Some("user-1".to_string()), - Some("key-1".to_string()), - None, - None, - i32::try_from(candidate_index).expect("candidate index should fit"), - 0, - Some("provider-1".to_string()), - Some("endpoint-1".to_string()), - Some("provider-key-1".to_string()), + extra_data: Option, + created_at_unix_ms: u64, + ) -> UpsertRequestCandidateRecord { + UpsertRequestCandidateRecord { + id: id.to_string(), + request_id: request_id.to_string(), + user_id: Some("user-1".to_string()), + api_key_id: Some("key-1".to_string()), + username: Some("user".to_string()), + api_key_name: Some("Key".to_string()), + candidate_index: 0, + retry_index: 0, + provider_id: Some("provider-1".to_string()), + endpoint_id: Some("endpoint-1".to_string()), + key_id: Some("provider-key-1".to_string()), status, - None, - false, - Some(200), - None, - None, - latency_ms, - None, - None, - None, - 2_000_000, - Some(2_000_001), - finished_at_unix_ms, - ) - .expect("stored candidate should build") + skip_reason: None, + is_cached: Some(false), + status_code: Some(200), + error_type: None, + error_message: None, + latency_ms: Some(123), + concurrent_requests: Some(2), + extra_data, + required_capabilities: Some(json!({"streaming": true})), + created_at_unix_ms: Some(created_at_unix_ms), + started_at_unix_ms: Some(created_at_unix_ms + 1), + finished_at_unix_ms: Some(created_at_unix_ms + 2), + } } } diff --git a/crates/aether-data/adapters/mysql/src/global_models.rs b/crates/aether-data/adapters/mysql/src/global_models.rs index fd97cebf3..863418f06 100644 --- a/crates/aether-data/adapters/mysql/src/global_models.rs +++ b/crates/aether-data/adapters/mysql/src/global_models.rs @@ -1,20 +1,170 @@ use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; +use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use aether_data_contracts::repository::global_models::{ metadata_supports_embedding, AdminGlobalModelListQuery, AdminProviderModelListQuery, - CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelSnapshot, - GlobalModelWriteRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, - PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage, - StoredAdminProviderModel, StoredProviderActiveGlobalModel, StoredProviderModelStats, - StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage, - UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord, + CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelWriteRepository, + PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, PublicGlobalModelQuery, + StoredAdminGlobalModel, StoredAdminGlobalModelPage, StoredAdminProviderModel, + StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel, + StoredPublicGlobalModel, StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, + UpsertAdminProviderModelRecord, }; use aether_data_contracts::DataLayerError; use crate::error::SqlResultExt; use crate::MysqlPool; +const LIST_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#" +SELECT + id, + name, + display_name, + is_active, + CAST(default_price_per_request AS DOUBLE) AS default_price_per_request, + CAST(default_tiered_pricing AS CHAR) AS default_tiered_pricing, + CAST(supported_capabilities AS CHAR) AS supported_capabilities, + CAST(config AS CHAR) AS config, + CAST(0 AS SIGNED) AS usage_count +FROM global_models +"#; + +const COUNT_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#" +SELECT CAST(COUNT(id) AS SIGNED) AS total +FROM global_models +"#; + +const LIST_PUBLIC_CATALOG_MODELS_PREFIX: &str = r#" +SELECT + m.id, + m.provider_id, + p.name AS provider_name, + p.is_active AS provider_is_active, + m.provider_model_name, + COALESCE(gm.name, m.provider_model_name) AS name, + COALESCE(NULLIF(gm.display_name, ''), m.provider_model_name) AS display_name, + CAST(gm.config AS CHAR) AS global_model_config, + CAST(gm.supported_capabilities AS CHAR) AS global_model_supported_capabilities, + CAST(m.config AS CHAR) AS model_config, + CAST(m.tiered_pricing AS CHAR) AS tiered_pricing, + CAST(gm.default_tiered_pricing AS CHAR) AS default_tiered_pricing, + COALESCE( + m.supports_vision, + CASE + WHEN JSON_EXTRACT(gm.config, '$.vision') IS NULL THEN NULL + WHEN LOWER(JSON_UNQUOTE(JSON_EXTRACT(gm.config, '$.vision'))) IN ('true', '1') THEN 1 + ELSE 0 + END, + 0 + ) AS supports_vision, + COALESCE( + m.supports_function_calling, + CASE + WHEN JSON_EXTRACT(gm.config, '$.function_calling') IS NULL THEN NULL + WHEN LOWER(JSON_UNQUOTE(JSON_EXTRACT(gm.config, '$.function_calling'))) IN ('true', '1') THEN 1 + ELSE 0 + END, + 0 + ) AS supports_function_calling, + COALESCE( + m.supports_streaming, + CASE + WHEN JSON_EXTRACT(gm.config, '$.streaming') IS NULL THEN NULL + WHEN LOWER(JSON_UNQUOTE(JSON_EXTRACT(gm.config, '$.streaming'))) IN ('true', '1') THEN 1 + ELSE 0 + END, + 1 + ) AS supports_streaming, + m.is_active, + gm.is_active AS global_model_is_active +FROM models m +JOIN providers p ON p.id = m.provider_id +LEFT JOIN global_models gm ON gm.id = m.global_model_id +"#; + +const LIST_PROVIDER_MODEL_STATS_PREFIX: &str = r#" +SELECT + provider_id, + CAST(COUNT(id) AS SIGNED) AS total_models, + CAST(COALESCE(SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END), 0) AS SIGNED) AS active_models +FROM models +WHERE provider_id IN ( +"#; + +const LIST_ADMIN_PROVIDER_MODELS_PREFIX: &str = r#" +SELECT + m.id, + m.provider_id, + m.global_model_id, + m.provider_model_name, + CAST(m.provider_model_mappings AS CHAR) AS provider_model_mappings, + CAST(m.price_per_request AS DOUBLE) AS price_per_request, + CAST(m.tiered_pricing AS CHAR) AS tiered_pricing, + m.supports_vision, + m.supports_function_calling, + m.supports_streaming, + m.supports_extended_thinking, + m.supports_image_generation, + m.is_active, + COALESCE(m.is_available, 1) AS is_available, + CAST(m.config AS CHAR) AS config, + m.created_at AS created_at_unix_ms, + m.updated_at AS updated_at_unix_secs, + gm.name AS global_model_name, + gm.display_name AS global_model_display_name, + CAST(gm.default_price_per_request AS DOUBLE) AS global_model_default_price_per_request, + CAST(gm.default_tiered_pricing AS CHAR) AS global_model_default_tiered_pricing, + CAST(gm.supported_capabilities AS CHAR) AS global_model_supported_capabilities, + CAST(gm.config AS CHAR) AS global_model_config +FROM models m +LEFT JOIN global_models gm ON gm.id = m.global_model_id +"#; + +const LIST_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#" +SELECT + gm.id, + gm.name, + COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name, + gm.is_active, + CAST(gm.default_price_per_request AS DOUBLE) AS default_price_per_request, + CAST(gm.default_tiered_pricing AS CHAR) AS default_tiered_pricing, + CAST(gm.supported_capabilities AS CHAR) AS supported_capabilities, + CAST(gm.config AS CHAR) AS config, + COALESCE(gm_stats.provider_count, 0) AS provider_count, + COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count, + CAST(COALESCE(gm.usage_count, 0) AS SIGNED) AS usage_count, + gm.created_at AS created_at_unix_ms, + gm.updated_at AS updated_at_unix_secs +FROM global_models gm +LEFT JOIN ( + SELECT + m.global_model_id, + CAST(COUNT(DISTINCT m.provider_id) AS SIGNED) AS provider_count, + CAST(COUNT( + DISTINCT CASE + WHEN m.is_active = 1 AND COALESCE(m.is_available, 1) = 1 AND p.is_active = 1 THEN m.provider_id + ELSE NULL + END + ) AS SIGNED) AS active_provider_count + FROM models m + JOIN providers p ON p.id = m.provider_id + GROUP BY m.global_model_id +) gm_stats ON gm_stats.global_model_id = gm.id +"#; + +const COUNT_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#" +SELECT CAST(COUNT(id) AS SIGNED) AS total +FROM global_models gm +"#; + +const LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX: &str = r#" +SELECT DISTINCT + provider_id, + global_model_id +FROM models +WHERE provider_id IN ( +"#; + #[derive(Debug, Clone)] pub struct MysqlGlobalModelReadRepository { pool: MysqlPool, @@ -25,171 +175,6 @@ impl MysqlGlobalModelReadRepository { Self { pool } } - async fn load_snapshot(&self) -> Result { - Ok( - GlobalModelSnapshot::seed(self.load_public_global_models().await?) - .with_admin_global_models(self.load_admin_global_models().await?) - .with_admin_provider_models(self.load_admin_provider_models().await?) - .with_public_catalog_models(self.load_public_catalog_models().await?) - .with_provider_model_stats(self.load_provider_model_stats().await?) - .with_active_global_model_refs(self.load_active_global_model_refs().await?), - ) - } - - async fn load_public_global_models( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT id, name, display_name, is_active, default_price_per_request, - default_tiered_pricing, supported_capabilities, config, usage_count -FROM global_models -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_public_global_model_row).collect() - } - - async fn load_admin_global_models( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, - name, - COALESCE(NULLIF(display_name, ''), name) AS display_name, - is_active, - default_price_per_request, - default_tiered_pricing, - supported_capabilities, - config, - usage_count, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM global_models -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_admin_global_model_row).collect() - } - - async fn load_admin_provider_models( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - m.id, - m.provider_id, - m.global_model_id, - m.provider_model_name, - m.provider_model_mappings, - m.price_per_request, - m.tiered_pricing, - m.supports_vision, - m.supports_function_calling, - m.supports_streaming, - m.supports_extended_thinking, - m.supports_image_generation, - m.is_active, - m.is_available, - m.config, - m.created_at AS created_at_unix_ms, - m.updated_at AS updated_at_unix_secs, - gm.name AS global_model_name, - gm.display_name AS global_model_display_name, - gm.default_price_per_request AS global_model_default_price_per_request, - gm.default_tiered_pricing AS global_model_default_tiered_pricing, - gm.supported_capabilities AS global_model_supported_capabilities, - gm.config AS global_model_config -FROM models m -LEFT JOIN global_models gm ON gm.id = m.global_model_id -WHERE m.global_model_id IS NOT NULL -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_admin_provider_model_row).collect() - } - - async fn load_public_catalog_models( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - m.id, - m.provider_id, - p.name AS provider_name, - p.is_active AS provider_is_active, - m.provider_model_name, - COALESCE(gm.name, m.provider_model_name) AS name, - COALESCE(NULLIF(gm.display_name, ''), m.provider_model_name) AS display_name, - gm.config AS global_model_config, - gm.supported_capabilities AS global_model_supported_capabilities, - m.config AS model_config, - m.tiered_pricing, - gm.default_tiered_pricing, - m.supports_vision, - m.supports_function_calling, - m.supports_streaming, - m.is_active, - gm.is_active AS global_model_is_active -FROM models m -JOIN providers p ON p.id = m.provider_id -LEFT JOIN global_models gm ON gm.id = m.global_model_id -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_public_catalog_model_row).collect() - } - - async fn load_provider_model_stats( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - provider_id, - COUNT(id) AS total_models, - SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END) AS active_models -FROM models -GROUP BY provider_id -ORDER BY provider_id ASC -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_provider_model_stats_row).collect() - } - - async fn load_active_global_model_refs( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT DISTINCT provider_id, global_model_id -FROM models -WHERE is_active = 1 - AND global_model_id IS NOT NULL -ORDER BY provider_id ASC, global_model_id ASC -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_active_global_model_row).collect() - } - pub async fn create_admin_provider_model( &self, record: &UpsertAdminProviderModelRecord, @@ -474,64 +459,172 @@ impl GlobalModelReadRepository for MysqlGlobalModelReadRepository { &self, query: &PublicGlobalModelQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_public_models(query)) + let mut count_builder = QueryBuilder::::new(COUNT_PUBLIC_GLOBAL_MODELS_PREFIX); + apply_public_model_filters(&mut count_builder, query); + let count_row = count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?; + let total = count_row + .try_get::("total") + .map(|value| value.max(0) as usize) + .map_sql_err()?; + + let mut list_builder = QueryBuilder::::new(LIST_PUBLIC_GLOBAL_MODELS_PREFIX); + apply_public_model_filters(&mut list_builder, query); + list_builder + .push(" ORDER BY name ASC LIMIT ") + .push_bind(query.limit as i64) + .push(" OFFSET ") + .push_bind(query.offset as i64); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_public_global_model_row) + .collect::>()?; + + Ok(StoredPublicGlobalModelPage { items, total }) } async fn get_public_model_by_name( &self, model_name: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .get_public_model_by_name(model_name)) + let row = sqlx::query( + r#" +SELECT + id, + name, + display_name, + is_active, + CAST(default_price_per_request AS DOUBLE) AS default_price_per_request, + CAST(default_tiered_pricing AS CHAR) AS default_tiered_pricing, + CAST(supported_capabilities AS CHAR) AS supported_capabilities, + CAST(config AS CHAR) AS config, + CAST(0 AS SIGNED) AS usage_count +FROM global_models +WHERE name = ? AND is_active = 1 +LIMIT 1 + "#, + ) + .bind(model_name) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + + row.as_ref().map(map_public_global_model_row).transpose() } async fn list_public_catalog_models( &self, query: &PublicCatalogModelListQuery, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_public_catalog_models(query)) + let mut builder = QueryBuilder::::new(LIST_PUBLIC_CATALOG_MODELS_PREFIX); + apply_public_catalog_model_filters(&mut builder, query.provider_id.as_deref(), None); + builder + .push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC LIMIT ") + .push_bind(query.limit as i64) + .push(" OFFSET ") + .push_bind(query.offset as i64); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_public_catalog_model_row).collect() } async fn search_public_catalog_models( &self, query: &PublicCatalogModelSearchQuery, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .search_public_catalog_models(query)) + let mut builder = QueryBuilder::::new(LIST_PUBLIC_CATALOG_MODELS_PREFIX); + apply_public_catalog_model_filters( + &mut builder, + query.provider_id.as_deref(), + Some(query.search.as_str()), + ); + builder + .push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC LIMIT ") + .push_bind(query.limit as i64); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_public_catalog_model_row).collect() } async fn list_admin_global_models( &self, query: &AdminGlobalModelListQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_admin_global_models(query)) + let mut count_builder = QueryBuilder::::new(COUNT_ADMIN_GLOBAL_MODELS_PREFIX); + apply_admin_global_model_filters(&mut count_builder, query); + let count_row = count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?; + let total = count_row + .try_get::("total") + .map(|value| value.max(0) as usize) + .map_sql_err()?; + + let mut list_builder = QueryBuilder::::new(LIST_ADMIN_GLOBAL_MODELS_PREFIX); + apply_admin_global_model_filters(&mut list_builder, query); + list_builder + .push(" ORDER BY gm.name ASC LIMIT ") + .push_bind(query.limit as i64) + .push(" OFFSET ") + .push_bind(query.offset as i64); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_admin_global_model_row) + .collect::>()?; + Ok(StoredAdminGlobalModelPage { items, total }) } async fn list_admin_provider_models( &self, query: &AdminProviderModelListQuery, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_admin_provider_models(query)) + let mut builder = QueryBuilder::::new(LIST_ADMIN_PROVIDER_MODELS_PREFIX); + builder + .push(" WHERE m.provider_id = ") + .push_bind(query.provider_id.trim().to_string()); + if let Some(is_active) = query.is_active { + builder.push(" AND m.is_active = ").push_bind(is_active); + } + builder + .push(" ORDER BY m.created_at DESC, m.id ASC LIMIT ") + .push_bind(query.limit as i64) + .push(" OFFSET ") + .push_bind(query.offset as i64); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_admin_provider_model_row).collect() } async fn list_admin_provider_available_source_models( &self, provider_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_admin_provider_available_source_models(provider_id)) + let rows = sqlx::query(&format!( + r#" +{LIST_ADMIN_PROVIDER_MODELS_PREFIX} +WHERE m.provider_id = ? + AND m.is_active = 1 + AND gm.is_active = 1 +ORDER BY gm.name ASC, m.created_at DESC, m.id ASC + "# + )) + .bind(provider_id) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_admin_provider_model_row).collect() } async fn get_admin_provider_model( @@ -539,60 +632,111 @@ impl GlobalModelReadRepository for MysqlGlobalModelReadRepository { provider_id: &str, model_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .get_admin_provider_model(provider_id, model_id)) + let row = sqlx::query(&format!( + r#" +{LIST_ADMIN_PROVIDER_MODELS_PREFIX} +WHERE m.provider_id = ? + AND m.id = ? +LIMIT 1 + "# + )) + .bind(provider_id) + .bind(model_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + + row.as_ref().map(map_admin_provider_model_row).transpose() } async fn get_admin_global_model_by_id( &self, global_model_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .get_admin_global_model_by_id(global_model_id)) + let row = sqlx::query(&format!( + r#" +{LIST_ADMIN_GLOBAL_MODELS_PREFIX} +WHERE gm.id = ? +LIMIT 1 + "# + )) + .bind(global_model_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + + row.as_ref().map(map_admin_global_model_row).transpose() } async fn get_admin_global_model_by_name( &self, model_name: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .get_admin_global_model_by_name(model_name)) + let row = sqlx::query(&format!( + r#" +{LIST_ADMIN_GLOBAL_MODELS_PREFIX} +WHERE gm.name = ? +LIMIT 1 + "# + )) + .bind(model_name) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + + row.as_ref().map(map_admin_global_model_row).transpose() } async fn list_admin_provider_models_by_global_model_id( &self, global_model_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_admin_provider_models_by_global_model_id(global_model_id)) + let rows = sqlx::query(&format!( + r#" +{LIST_ADMIN_PROVIDER_MODELS_PREFIX} +WHERE m.global_model_id = ? +ORDER BY m.created_at DESC, m.id ASC + "# + )) + .bind(global_model_id) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_admin_provider_model_row).collect() } async fn list_provider_model_stats( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_provider_model_stats(provider_ids)) + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut builder = build_provider_id_list_query( + LIST_PROVIDER_MODEL_STATS_PREFIX, + provider_ids, + ")\nGROUP BY provider_id\nORDER BY provider_id ASC", + ); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_provider_model_stats_row).collect() } async fn list_active_global_model_ids_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_active_global_model_ids_by_provider_ids(provider_ids)) + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut builder = build_provider_id_list_query( + LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX, + provider_ids, + ")\nAND is_active = 1\nAND global_model_id IS NOT NULL\nORDER BY provider_id ASC, global_model_id ASC", + ); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_active_global_model_row).collect() } } @@ -696,6 +840,100 @@ fn first_tier_price(value: Option<&serde_json::Value>, key: &str) -> Option .and_then(serde_json::Value::as_f64) } +fn apply_public_model_filters( + builder: &mut QueryBuilder<'_, MySql>, + query: &PublicGlobalModelQuery, +) { + builder.push(" WHERE "); + match query.is_active { + Some(is_active) => { + builder.push("is_active = ").push_bind(is_active); + } + None => { + builder.push("is_active = 1"); + } + } + + if let Some(search) = query + .search + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let pattern = format!("%{}%", search.to_ascii_lowercase()); + builder + .push(" AND (LOWER(name) LIKE ") + .push_bind(pattern.clone()) + .push(" OR LOWER(display_name) LIKE ") + .push_bind(pattern) + .push(")"); + } +} + +fn apply_admin_global_model_filters( + builder: &mut QueryBuilder<'_, MySql>, + query: &AdminGlobalModelListQuery, +) { + builder.push(" WHERE 1=1"); + if let Some(is_active) = query.is_active { + builder.push(" AND gm.is_active = ").push_bind(is_active); + } + if let Some(search) = query + .search + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let pattern = format!("%{}%", search.to_ascii_lowercase()); + builder + .push(" AND (LOWER(gm.name) LIKE ") + .push_bind(pattern.clone()) + .push(" OR LOWER(gm.display_name) LIKE ") + .push_bind(pattern) + .push(")"); + } +} + +fn apply_public_catalog_model_filters( + builder: &mut QueryBuilder<'_, MySql>, + provider_id: Option<&str>, + search: Option<&str>, +) { + builder.push(" WHERE m.is_active = 1 AND COALESCE(m.is_available, 1) = 1 AND p.is_active = 1 AND COALESCE(gm.is_active, 1) = 1"); + + if let Some(provider_id) = provider_id.map(str::trim).filter(|value| !value.is_empty()) { + builder + .push(" AND m.provider_id = ") + .push_bind(provider_id.to_string()); + } + + if let Some(search) = search.map(str::trim).filter(|value| !value.is_empty()) { + let pattern = format!("%{}%", search.to_ascii_lowercase()); + builder + .push(" AND (LOWER(m.provider_model_name) LIKE ") + .push_bind(pattern.clone()) + .push(" OR LOWER(gm.name) LIKE ") + .push_bind(pattern.clone()) + .push(" OR LOWER(gm.display_name) LIKE ") + .push_bind(pattern) + .push(")"); + } +} + +fn build_provider_id_list_query<'a>( + prefix: &'static str, + provider_ids: &'a [String], + suffix: &'static str, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(prefix); + let mut separated = builder.separated(", "); + for provider_id in provider_ids { + separated.push_bind(provider_id); + } + separated.push_unseparated(suffix); + builder +} + fn map_public_global_model_row(row: &MySqlRow) -> Result { StoredPublicGlobalModel::new( row.try_get("id").map_sql_err()?, @@ -717,6 +955,16 @@ fn map_public_global_model_row(row: &MySqlRow) -> Result Result { + let provider_count = row + .try_get::("provider_count") + .map_sql_err()? + .max(0) as u64; + let active_provider_count = row + .try_get::("active_provider_count") + .map_sql_err()? + .max(0) as u64; + let usage_count = row.try_get::("usage_count").map_sql_err()?.max(0) as u64; + StoredAdminGlobalModel::new( row.try_get("id").map_sql_err()?, row.try_get("name").map_sql_err()?, @@ -732,9 +980,9 @@ fn map_admin_global_model_row(row: &MySqlRow) -> Result("usage_count").map_sql_err()?.max(0) as u64, + provider_count, + active_provider_count, + usage_count, optional_u64( row.try_get("created_at_unix_ms").map_sql_err()?, "global_models.created_at", @@ -847,8 +1095,8 @@ fn map_public_catalog_model_row( .and_then(|value| value.get("icon_url")) .and_then(serde_json::Value::as_str) .map(ToString::to_string), - first_tier_price(pricing, "input_price_per_1m"), - first_tier_price(pricing, "output_price_per_1m"), + Some(first_tier_price(pricing, "input_price_per_1m").unwrap_or(0.0)), + Some(first_tier_price(pricing, "output_price_per_1m").unwrap_or(0.0)), first_tier_price(pricing, "cache_creation_price_per_1m"), first_tier_price(pricing, "cache_read_price_per_1m"), row.try_get("supports_vision").map_sql_err()?, @@ -900,7 +1148,66 @@ fn optional_admin_global_model_usage_count_i64( #[cfg(test)] mod tests { - use super::MysqlGlobalModelReadRepository; + use sqlx::{MySql, QueryBuilder}; + + use aether_data_contracts::repository::global_models::{ + AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord, + GlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, + PublicGlobalModelQuery, UpsertAdminProviderModelRecord, + }; + use serde_json::json; + + use super::{ + apply_admin_global_model_filters, apply_public_model_filters, build_provider_id_list_query, + MysqlGlobalModelReadRepository, LIST_PROVIDER_MODEL_STATS_PREFIX, + }; + use crate::run_migrations; + + #[test] + fn mysql_global_model_builders_apply_filters_and_bound_lists() { + let mut public = QueryBuilder::::new("SELECT id FROM global_models"); + apply_public_model_filters( + &mut public, + &PublicGlobalModelQuery { + offset: 0, + limit: 25, + is_active: Some(false), + search: Some("GPT".to_string()), + }, + ); + assert_eq!( + public.sql(), + "SELECT id FROM global_models WHERE is_active = ? AND (LOWER(name) LIKE ? OR LOWER(display_name) LIKE ?)" + ); + + let mut admin = QueryBuilder::::new("SELECT gm.id FROM global_models gm"); + apply_admin_global_model_filters( + &mut admin, + &AdminGlobalModelListQuery { + offset: 0, + limit: 25, + is_active: None, + search: Some("embed".to_string()), + }, + ); + assert_eq!( + admin.sql(), + "SELECT gm.id FROM global_models gm WHERE 1=1 AND (LOWER(gm.name) LIKE ? OR LOWER(gm.display_name) LIKE ?)" + ); + + let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()]; + let provider_stats = build_provider_id_list_query( + LIST_PROVIDER_MODEL_STATS_PREFIX, + &provider_ids, + ") GROUP BY provider_id", + ); + let normalized_sql = provider_stats + .sql() + .split_whitespace() + .collect::>() + .join(" "); + assert!(normalized_sql.ends_with("WHERE provider_id IN ( ?, ?) GROUP BY provider_id")); + } #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -912,4 +1219,176 @@ mod tests { let _repository = MysqlGlobalModelReadRepository::new(pool); } + + #[tokio::test] + async fn mysql_global_model_queries_round_trip_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql global model test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + + let suffix = unique_suffix(); + let provider_id = format!("provider-{suffix}"); + let global_id = format!("global-{suffix}"); + let model_id = format!("model-{suffix}"); + let global_name = format!("global-name-{suffix}"); + sqlx::query( + r#" +INSERT INTO providers ( + id, name, provider_type, is_active, provider_priority, created_at, updated_at +) VALUES (?, ?, 'custom', 1, 10, 1, 1) +"#, + ) + .bind(&provider_id) + .bind(format!("Provider {suffix}")) + .execute(&pool) + .await + .expect("provider should seed"); + + let repository = MysqlGlobalModelReadRepository::new(pool.clone()); + let mut global_record = CreateAdminGlobalModelRecord::new( + global_id.clone(), + global_name.clone(), + format!("Global {suffix}"), + true, + None, + Some(json!({"tiers":[{"input_price_per_1m":2.0}]})), + Some(json!(["chat"])), + Some(json!({"description":"mysql parity"})), + ) + .expect("global model input should validate"); + global_record.usage_count = Some(7); + repository + .create_admin_global_model(&global_record) + .await + .expect("global model should create") + .expect("created global model should return"); + + let provider_model = UpsertAdminProviderModelRecord::new( + model_id.clone(), + provider_id.clone(), + global_id.clone(), + format!("upstream-{suffix}"), + Some(json!([global_name])), + None, + None, + Some(true), + Some(true), + Some(true), + Some(false), + Some(false), + true, + true, + Some(json!({"routing":"primary"})), + ) + .expect("provider model input should validate"); + repository + .create_admin_provider_model(&provider_model) + .await + .expect("provider model should create") + .expect("created provider model should return"); + + let public = repository + .list_public_models(&PublicGlobalModelQuery { + offset: 0, + limit: 10, + is_active: None, + search: Some(suffix.clone()), + }) + .await + .expect("public models should list"); + assert_eq!(public.total, 1); + assert_eq!(public.items[0].usage_count, 0); + + let catalog = repository + .search_public_catalog_models(&PublicCatalogModelSearchQuery { + search: suffix.clone(), + provider_id: Some(provider_id.clone()), + limit: 10, + }) + .await + .expect("catalog should search"); + assert_eq!(catalog.len(), 1); + assert_eq!(catalog[0].input_price_per_1m, Some(2.0)); + assert_eq!(catalog[0].output_price_per_1m, Some(0.0)); + + let catalog_page = repository + .list_public_catalog_models(&PublicCatalogModelListQuery { + provider_id: Some(provider_id.clone()), + offset: 0, + limit: 1, + }) + .await + .expect("catalog should list"); + assert_eq!(catalog_page.len(), 1); + + let admin = repository + .list_admin_global_models(&AdminGlobalModelListQuery { + offset: 0, + limit: 10, + is_active: None, + search: Some(suffix.clone()), + }) + .await + .expect("admin global models should list"); + assert_eq!(admin.total, 1); + assert_eq!(admin.items[0].provider_count, 1); + assert_eq!(admin.items[0].active_provider_count, 1); + assert_eq!(admin.items[0].usage_count, 7); + + let admin_models = repository + .list_admin_provider_models(&AdminProviderModelListQuery { + provider_id: provider_id.clone(), + is_active: Some(true), + offset: 0, + limit: 1, + }) + .await + .expect("admin provider models should list"); + assert_eq!(admin_models.len(), 1); + assert_eq!(admin_models[0].id, model_id); + + let stats = repository + .list_provider_model_stats(std::slice::from_ref(&provider_id)) + .await + .expect("provider stats should list"); + assert_eq!(stats.len(), 1); + assert_eq!(stats[0].total_models, 1); + let active_refs = repository + .list_active_global_model_ids_by_provider_ids(std::slice::from_ref(&provider_id)) + .await + .expect("active global refs should list"); + assert_eq!(active_refs.len(), 1); + assert_eq!(active_refs[0].global_model_id, global_id); + + assert!(repository + .delete_admin_global_model(&global_id) + .await + .expect("global model should delete")); + sqlx::query("DELETE FROM providers WHERE id = ?") + .bind(&provider_id) + .execute(&pool) + .await + .expect("provider should delete"); + } + + fn unique_suffix() -> String { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + format!("{}-{nanos}", std::process::id()) + } } diff --git a/crates/aether-data/adapters/mysql/src/lib.rs b/crates/aether-data/adapters/mysql/src/lib.rs index 247d430bb..5f5a21ac8 100644 --- a/crates/aether-data/adapters/mysql/src/lib.rs +++ b/crates/aether-data/adapters/mysql/src/lib.rs @@ -47,7 +47,7 @@ pub use proxy_nodes::MysqlProxyNodeReadRepository; pub use quota::MysqlProviderQuotaRepository; pub use routing_profiles::MysqlRoutingGroupRepository; pub use settlement::MysqlSettlementRepository; -pub use usage::{MysqlUsageStorage, MysqlUsageWriteRepository}; +pub use usage::{MysqlUsageReadFilter, MysqlUsageStorage, MysqlUsageWriteRepository}; pub use users::MysqlUserReadRepository; pub use video_tasks::MysqlVideoTaskRepository; pub use wallet::MysqlWalletReadRepository; diff --git a/crates/aether-data/adapters/mysql/src/migrations.rs b/crates/aether-data/adapters/mysql/src/migrations.rs index 197bc831b..68400ae74 100644 --- a/crates/aether-data/adapters/mysql/src/migrations.rs +++ b/crates/aether-data/adapters/mysql/src/migrations.rs @@ -1,5 +1,5 @@ use sqlx::{ - migrate::{Migrate, MigrateError, Migrator}, + migrate::{AppliedMigration, Migrate, MigrateError, Migrator}, MySqlPool, }; @@ -17,9 +17,15 @@ pub async fn pending_migrations( let mut conn = pool.acquire().await?; let applied_migrations = match conn.list_applied_migrations().await { Ok(applied_migrations) => applied_migrations, - Err(err) if is_missing_sqlx_migrations_table_error(&err) => Vec::new(), + Err(err) if is_missing_sqlx_migrations_table_error(&err) => { + return Ok(pending_migrations_from_applied(&[])); + } Err(err) => return Err(err), }; + if let Some(version) = conn.dirty_version().await? { + return Err(MigrateError::Dirty(version)); + } + validate_applied_migrations(&applied_migrations)?; Ok(pending_migrations_from_applied(&applied_migrations)) } @@ -56,9 +62,33 @@ fn pending_migrations_from_applied( .collect() } +fn validate_applied_migrations( + applied_migrations: &[AppliedMigration], +) -> Result<(), MigrateError> { + if MIGRATOR.ignore_missing { + return Ok(()); + } + let known_versions = MIGRATOR + .iter() + .map(|migration| migration.version) + .collect::>(); + if let Some(migration) = applied_migrations + .iter() + .find(|migration| !known_versions.contains(&migration.version)) + { + return Err(MigrateError::VersionMissing(migration.version)); + } + Ok(()) +} + #[cfg(test)] mod tests { - use super::MIGRATOR; + use std::borrow::Cow; + + use super::{ + pending_migrations, prepare_database_for_startup, validate_applied_migrations, MIGRATOR, + }; + use sqlx::migrate::{AppliedMigration, MigrateError}; #[test] fn embeds_mysql_migration_sources() { @@ -69,4 +99,152 @@ mod tests { assert!(!versions.is_empty()); assert!(versions.windows(2).all(|pair| pair[0] < pair[1])); } + + #[test] + fn embeds_cross_driver_schema_parity_migration() { + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == 20260725010000) + .expect("cross-driver schema parity migration should be embedded"); + let sql = migration.sql.as_ref(); + + for required_fragment in [ + "CREATE TABLE IF NOT EXISTS usage_body_blobs", + "CREATE TABLE IF NOT EXISTS usage_http_audits", + "CREATE TABLE IF NOT EXISTS stats_summary", + "CREATE TABLE IF NOT EXISTS user_model_usage_counts", + "CREATE TABLE IF NOT EXISTS api_key_provider_mappings", + "CREATE TABLE IF NOT EXISTS provider_usage_tracking", + "ADD COLUMN `settlement_snapshot_schema_version`", + "ADD COLUMN `billing_effective_input_tokens`", + "ADD COLUMN `converted_request_body`", + "ADD COLUMN `p99_first_byte_time_ms`", + "idx_usage_stale_pending_created_request", + ] { + assert!( + sql.contains(required_fragment), + "parity migration is missing {required_fragment}" + ); + } + } + + #[test] + fn embeds_advanced_stats_parity_migration() { + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == 20260725020000) + .expect("advanced stats parity migration should be embedded"); + let sql = migration.sql.as_ref(); + + for required_fragment in [ + "CREATE TABLE stats_user_summary", + "CREATE TABLE stats_user_daily_api_format", + "CREATE TABLE stats_user_daily_model_provider", + "CREATE TABLE stats_daily_model_provider", + "CREATE TABLE stats_daily_cost_savings", + "CREATE TABLE stats_user_daily_cost_savings_model_provider", + "ADD COLUMN completed_total_input_context", + "ADD COLUMN settled_total_cost", + "ADD COLUMN response_time_samples", + "UPDATE stats_hourly SET is_complete = 0", + "UPDATE stats_daily SET is_complete = 0", + ] { + assert!( + sql.contains(required_fragment), + "advanced stats migration is missing {required_fragment}" + ); + } + } + + #[test] + fn rejects_applied_migration_versions_unknown_to_this_binary() { + let version = MIGRATOR + .iter() + .map(|migration| migration.version) + .max() + .expect("mysql migrations should not be empty") + + 1; + let error = validate_applied_migrations(&[AppliedMigration { + version, + checksum: Cow::Borrowed(&[]), + }]) + .expect_err("unknown applied migration should block startup"); + + assert!(matches!(error, MigrateError::VersionMissing(found) if found == version)); + } + + #[tokio::test] + async fn pending_and_startup_preparation_reject_dirty_mysql_migration_state_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql dirty migration test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + let dirty_version = MIGRATOR + .iter() + .next() + .expect("mysql migrations should not be empty") + .version; + + let mut conn = pool + .acquire() + .await + .expect("mysql connection should acquire"); + sqlx::query( + r#" +CREATE TEMPORARY TABLE _sqlx_migrations ( + version BIGINT PRIMARY KEY, + description TEXT NOT NULL, + installed_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + success BOOLEAN NOT NULL, + checksum BLOB NOT NULL, + execution_time BIGINT NOT NULL +) +"#, + ) + .execute(&mut *conn) + .await + .expect("temporary mysql migrations table should create"); + sqlx::query( + r#" +INSERT INTO _sqlx_migrations ( + version, + description, + success, + checksum, + execution_time +) VALUES (?, 'dirty test migration', FALSE, ?, 0) +"#, + ) + .bind(dirty_version) + .bind(Vec::::new()) + .execute(&mut *conn) + .await + .expect("dirty mysql migration should insert"); + drop(conn); + + let pending_error = pending_migrations(&pool) + .await + .expect_err("dirty mysql migration should fail pending inspection"); + assert!( + matches!(&pending_error, MigrateError::Dirty(version) if *version == dirty_version), + "unexpected pending migration error: {pending_error}" + ); + + let preparation_error = prepare_database_for_startup(&pool) + .await + .expect_err("dirty mysql migration should fail startup preparation"); + assert!( + matches!(&preparation_error, MigrateError::Dirty(version) if *version == dirty_version), + "unexpected startup preparation error: {preparation_error}" + ); + } } diff --git a/crates/aether-data/adapters/mysql/src/pool.rs b/crates/aether-data/adapters/mysql/src/pool.rs index 26a2e1539..a6826eb22 100644 --- a/crates/aether-data/adapters/mysql/src/pool.rs +++ b/crates/aether-data/adapters/mysql/src/pool.rs @@ -62,6 +62,14 @@ impl MysqlPoolFactory { .acquire_timeout(Duration::from_millis(acquire_timeout_ms)) .idle_timeout(Duration::from_millis(idle_timeout_ms)) .max_lifetime(Duration::from_millis(max_lifetime_ms)) + .after_connect(|connection, _metadata| { + Box::pin(async move { + sqlx::query("SET time_zone = '+00:00'") + .execute(connection) + .await?; + Ok(()) + }) + }) .connect_lazy_with(self.connect_options()?)) } } @@ -90,4 +98,34 @@ mod tests { let factory = MysqlPoolFactory::new(config).expect("factory should build"); let _pool = factory.connect_lazy().expect("lazy pool should build"); } + + #[tokio::test] + async fn factory_configures_utc_session_timezone_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql timezone test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + let config = SqlDatabaseConfig { + driver: DatabaseDriver::Mysql, + url: database_url, + pool: SqlPoolConfig { + max_connections: 1, + ..SqlPoolConfig::default() + }, + }; + + let pool = MysqlPoolFactory::new(config) + .expect("factory should build") + .connect_lazy() + .expect("lazy pool should build"); + let timezone: String = sqlx::query_scalar("SELECT @@session.time_zone") + .fetch_one(&pool) + .await + .expect("mysql session timezone should load"); + + assert_eq!(timezone, "+00:00"); + } } diff --git a/crates/aether-data/adapters/mysql/src/provider_catalog.rs b/crates/aether-data/adapters/mysql/src/provider_catalog.rs index 46d69b4f2..1efa21335 100644 --- a/crates/aether-data/adapters/mysql/src/provider_catalog.rs +++ b/crates/aether-data/adapters/mysql/src/provider_catalog.rs @@ -9,15 +9,19 @@ use sqlx::{ use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - ProviderCatalogReadRepository, ProviderCatalogSnapshot, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; use aether_data_contracts::DataLayerError; +use aether_data_query::{ + push_ci_contains_any, push_eq, push_in, push_limit_offset, push_optional_eq, SqlDialect, + WhereClause, +}; use crate::error::SqlResultExt; use crate::MysqlPool; @@ -27,6 +31,30 @@ pub struct MysqlProviderCatalogReadRepository { pool: MysqlPool, } +const PROVIDER_SELECT_SQL: &str = r#" +SELECT + id, name, description, website, provider_type, billing_type, + monthly_quota_usd, monthly_used_usd, quota_reset_day, + quota_last_reset_at AS quota_last_reset_at_unix_secs, + quota_expires_at AS quota_expires_at_unix_secs, + provider_priority, is_active, keep_priority_on_conversion, + enable_format_conversion, concurrent_limit, max_retries, proxy, + request_timeout, stream_first_byte_timeout, config, + created_at AS created_at_unix_ms, + updated_at AS updated_at_unix_secs +FROM providers +"#; + +const ENDPOINT_SELECT_SQL: &str = r#" +SELECT + id, provider_id, api_format, api_family, endpoint_kind, is_active, + health_score, base_url, header_rules, body_rules, max_retries, + custom_path, config, format_acceptance_config, proxy, + created_at AS created_at_unix_ms, + updated_at AS updated_at_unix_secs +FROM provider_endpoints +"#; + const KEY_SELECT_SQL: &str = r#" SELECT id, provider_id, name, auth_type, capabilities, is_active, api_formats, @@ -54,69 +82,176 @@ SELECT FROM provider_api_keys "#; +const KEY_SUMMARY_SELECT_SQL: &str = r#" +SELECT + id, + provider_id, + COALESCE(NULLIF(name, ''), id) AS name, + COALESCE(NULLIF(auth_type, ''), 'summary') AS auth_type, + NULL AS capabilities, + is_active, + api_formats, + NULL AS auth_type_by_format, + NULL AS allow_auth_channel_mismatch_formats, + 'summary' AS api_key, + CASE WHEN auth_config IS NULL THEN NULL ELSE '{}' END AS auth_config, + NULL AS note, + NULL AS internal_priority, + NULL AS rate_multipliers, + NULL AS global_priority_by_format, + NULL AS allowed_models, + NULL AS expires_at_unix_secs, + NULL AS cache_ttl_minutes, + NULL AS max_probe_interval_minutes, + NULL AS proxy, + NULL AS fingerprint, + NULL AS rpm_limit, + NULL AS concurrent_limit, + NULL AS learned_rpm_limit, + NULL AS concurrent_429_count, + NULL AS rpm_429_count, + NULL AS last_429_at_unix_secs, + NULL AS last_429_type, + NULL AS adjustment_history, + NULL AS utilization_samples, + NULL AS last_probe_increase_at_unix_secs, + NULL AS last_rpm_peak, + NULL AS request_count, + CAST(0 AS SIGNED) AS total_tokens, + CAST(0 AS DOUBLE) AS total_cost_usd, + NULL AS success_count, + NULL AS error_count, + NULL AS total_response_time_ms, + NULL AS last_used_at_unix_secs, + FALSE AS auto_fetch_models, + NULL AS last_models_fetch_at_unix_secs, + NULL AS last_models_fetch_error, + NULL AS locked_models, + NULL AS model_include_patterns, + NULL AS model_exclude_patterns, + NULL AS upstream_metadata, + NULL AS oauth_invalid_at_unix_secs, + NULL AS oauth_invalid_reason, + NULL AS status_snapshot, + NULL AS created_at_unix_ms, + NULL AS updated_at_unix_secs, + health_by_format, + NULL AS circuit_breaker_by_format +FROM provider_api_keys +"#; + +const KEY_MAINTENANCE_SUMMARY_SELECT_SQL: &str = r#" +SELECT id, provider_id, is_active, upstream_metadata +FROM provider_api_keys +"#; + +const KEY_STATS_SELECT_SQL: &str = r#" +SELECT + provider_id, + CAST(COUNT(*) AS SIGNED) AS total_keys, + CAST(SUM(CASE WHEN is_active THEN 1 ELSE 0 END) AS SIGNED) AS active_keys +FROM provider_api_keys +"#; + impl MysqlProviderCatalogReadRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } } - async fn load_snapshot(&self) -> Result { - Ok(ProviderCatalogSnapshot::new( - self.load_providers().await?, - self.load_endpoints().await?, - self.load_keys().await?, - )) - } + pub async fn list_providers_by_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } - async fn load_providers(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, name, description, website, provider_type, billing_type, - monthly_quota_usd, monthly_used_usd, quota_reset_day, - quota_last_reset_at AS quota_last_reset_at_unix_secs, - quota_expires_at AS quota_expires_at_unix_secs, - provider_priority, is_active, keep_priority_on_conversion, - enable_format_conversion, concurrent_limit, max_retries, proxy, - request_timeout, stream_first_byte_timeout, config, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM providers -"#, + let rows = build_in_query( + PROVIDER_SELECT_SQL, + "id", + provider_ids, + " ORDER BY name ASC", ) + .build() .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_provider_row).collect() } - async fn load_endpoints(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, provider_id, api_format, api_family, endpoint_kind, is_active, - health_score, base_url, header_rules, body_rules, max_retries, - custom_path, config, format_acceptance_config, proxy, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM provider_endpoints -WHERE api_format IS NOT NULL -"#, + pub async fn list_providers( + &self, + active_only: bool, + ) -> Result, DataLayerError> { + let mut builder = QueryBuilder::::new(PROVIDER_SELECT_SQL); + let mut where_clause = WhereClause::new(); + if active_only { + push_eq(&mut builder, &mut where_clause, "is_active", true); + } + builder.push(" ORDER BY provider_priority ASC, name ASC"); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_provider_row).collect() + } + + pub async fn list_endpoints_by_ids( + &self, + endpoint_ids: &[String], + ) -> Result, DataLayerError> { + if endpoint_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query( + ENDPOINT_SELECT_SQL, + "id", + endpoint_ids, + " ORDER BY api_format ASC, id ASC", ) + .build() .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_endpoint_row).collect() } - async fn load_keys(&self) -> Result, DataLayerError> { - let rows = sqlx::query(KEY_SELECT_SQL) + pub async fn list_endpoints_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query( + ENDPOINT_SELECT_SQL, + "provider_id", + provider_ids, + " ORDER BY provider_id ASC, api_format ASC, id ASC", + ) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_endpoint_row).collect() + } + + pub async fn list_keys_by_ids( + &self, + key_ids: &[String], + ) -> Result, DataLayerError> { + if key_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query(KEY_SELECT_SQL, "id", key_ids, " ORDER BY name ASC, id ASC") + .build() .fetch_all(&self.pool) .await .map_sql_err()?; rows.iter().map(map_key_row).collect() } - async fn list_keys_by_provider_ids_direct( + pub async fn list_keys_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { @@ -124,12 +259,120 @@ WHERE api_format IS NOT NULL return Ok(Vec::new()); } - let rows = build_list_keys_by_provider_ids_query(provider_ids) + let rows = build_in_query( + KEY_SELECT_SQL, + "provider_id", + provider_ids, + " ORDER BY provider_id ASC, name ASC, id ASC", + ) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_key_row).collect() + } + + pub async fn list_key_summaries_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query( + KEY_SUMMARY_SELECT_SQL, + "provider_id", + provider_ids, + " ORDER BY provider_id ASC, id ASC", + ) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_key_row).collect() + } + + pub async fn list_key_maintenance_summaries_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query( + KEY_MAINTENANCE_SUMMARY_SELECT_SQL, + "provider_id", + provider_ids, + " ORDER BY provider_id ASC, id ASC", + ) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_key_maintenance_summary_row).collect() + } + + pub async fn list_keys_page( + &self, + query: &ProviderCatalogKeyListQuery, + ) -> Result { + validate_non_empty(&query.provider_id, "provider catalog provider_id")?; + let offset = i64::try_from(query.offset).map_err(|_| { + DataLayerError::InvalidInput(format!( + "invalid provider catalog key offset: {}", + query.offset + )) + })?; + let limit = i64::try_from(query.limit).map_err(|_| { + DataLayerError::InvalidInput(format!( + "invalid provider catalog key limit: {}", + query.limit + )) + })?; + + let mut count_builder = build_key_page_count_query(query); + let total = count_builder + .build_query_scalar::() + .fetch_one(&self.pool) + .await + .map_sql_err()? + .max(0) as usize; + + let mut list_builder = build_key_page_list_query(query, limit, offset); + let rows = list_builder .build() .fetch_all(&self.pool) .await .map_sql_err()?; - rows.iter().map(map_key_row).collect() + let items = rows + .iter() + .map(map_key_row) + .collect::, _>>()?; + + Ok(StoredProviderCatalogKeyPage { items, total }) + } + + pub async fn list_key_stats_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_in_query( + KEY_STATS_SELECT_SQL, + "provider_id", + provider_ids, + " GROUP BY provider_id ORDER BY provider_id ASC", + ) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_key_stats_row).collect() } pub async fn create_provider( @@ -1443,84 +1686,70 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository { &self, active_only: bool, ) -> Result, DataLayerError> { - Ok(self.load_snapshot().await?.list_providers(active_only)) + Self::list_providers(self, active_only).await } async fn list_providers_by_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_providers_by_ids(provider_ids)) + Self::list_providers_by_ids(self, provider_ids).await } async fn list_endpoints_by_ids( &self, endpoint_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_endpoints_by_ids(endpoint_ids)) + Self::list_endpoints_by_ids(self, endpoint_ids).await } async fn list_endpoints_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_endpoints_by_provider_ids(provider_ids)) + Self::list_endpoints_by_provider_ids(self, provider_ids).await } async fn list_keys_by_ids( &self, key_ids: &[String], ) -> Result, DataLayerError> { - Ok(self.load_snapshot().await?.list_keys_by_ids(key_ids)) + Self::list_keys_by_ids(self, key_ids).await } async fn list_keys_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - self.list_keys_by_provider_ids_direct(provider_ids).await + Self::list_keys_by_provider_ids(self, provider_ids).await } async fn list_key_summaries_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - self.list_keys_by_provider_ids_direct(provider_ids).await + Self::list_key_summaries_by_provider_ids(self, provider_ids).await } async fn list_key_maintenance_summaries_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_key_maintenance_summaries_by_provider_ids(provider_ids)) + Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await } async fn list_keys_page( &self, query: &ProviderCatalogKeyListQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_keys_page(query)) + Self::list_keys_page(self, query).await } async fn list_key_stats_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - self.load_snapshot() - .await? - .list_key_stats_by_provider_ids(provider_ids) + Self::list_key_stats_by_provider_ids(self, provider_ids).await } } @@ -1975,17 +2204,87 @@ fn optional_json_to_string( optional_json_ref_to_string(value.as_ref(), field_name) } -fn build_list_keys_by_provider_ids_query(provider_ids: &[String]) -> QueryBuilder<'_, MySql> { +fn build_in_query<'a>( + select_sql: &'static str, + column: &'static str, + values: &'a [String], + suffix: &'static str, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(select_sql); + let mut where_clause = WhereClause::new(); + push_in(&mut builder, &mut where_clause, column, values); + builder.push(suffix); + builder +} + +fn build_key_page_count_query<'a>( + query: &'a ProviderCatalogKeyListQuery, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new( + "SELECT CAST(COUNT(*) AS SIGNED) AS total FROM provider_api_keys", + ); + let mut where_clause = WhereClause::new(); + apply_key_page_filters(&mut builder, &mut where_clause, query); + builder +} + +fn build_key_page_list_query<'a>( + query: &'a ProviderCatalogKeyListQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { let mut builder = QueryBuilder::::new(KEY_SELECT_SQL); - builder.push("WHERE provider_id IN ("); - { - let mut separated = builder.separated(", "); - for provider_id in provider_ids { - separated.push_bind(provider_id.clone()); + let mut where_clause = WhereClause::new(); + apply_key_page_filters(&mut builder, &mut where_clause, query); + builder + .push(" ORDER BY ") + .push(key_page_order_by(&query.order)); + push_limit_offset(&mut builder, limit, offset); + builder +} + +fn key_page_order_by(order: &ProviderCatalogKeyListOrder) -> &'static str { + match order { + ProviderCatalogKeyListOrder::Name => "internal_priority ASC, name ASC, id ASC", + ProviderCatalogKeyListOrder::CreatedAt => { + "internal_priority ASC, COALESCE(created_at, 0) ASC, id ASC" + } + ProviderCatalogKeyListOrder::CreatedAtAsc => { + "created_at IS NULL ASC, created_at ASC, name ASC, id ASC" + } + ProviderCatalogKeyListOrder::CreatedAtDesc => { + "created_at IS NULL ASC, created_at DESC, name ASC, id ASC" + } + ProviderCatalogKeyListOrder::LastUsedAtAsc => { + "last_used_at IS NULL ASC, last_used_at ASC, name ASC, id ASC" + } + ProviderCatalogKeyListOrder::LastUsedAtDesc => { + "last_used_at IS NULL ASC, last_used_at DESC, name ASC, id ASC" } } - builder.push(") ORDER BY provider_id ASC, name ASC, id ASC"); - builder +} + +fn apply_key_page_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + where_clause: &mut WhereClause, + query: &'a ProviderCatalogKeyListQuery, +) { + push_eq( + builder, + where_clause, + "provider_id", + query.provider_id.clone(), + ); + if let Some(search) = query.search.as_deref() { + push_ci_contains_any( + builder, + where_clause, + SqlDialect::MySql, + &["name", "id"], + search, + ); + } + push_optional_eq(builder, where_clause, "is_active", query.is_active); } fn key_insert_sql() -> &'static str { @@ -2258,6 +2557,28 @@ fn map_endpoint_row(row: &MySqlRow) -> Result Result { + StoredProviderCatalogKeyStats::new( + row.try_get("provider_id").map_sql_err()?, + row.try_get("total_keys").map_sql_err()?, + row.try_get("active_keys").map_sql_err()?, + ) +} + +fn map_key_maintenance_summary_row( + row: &MySqlRow, +) -> Result { + Ok(StoredProviderCatalogKeyMaintenanceSummary { + id: row.try_get("id").map_sql_err()?, + provider_id: row.try_get("provider_id").map_sql_err()?, + is_active: row.try_get("is_active").map_sql_err()?, + upstream_metadata: optional_json_from_string( + row.try_get("upstream_metadata").map_sql_err()?, + "provider_api_keys.upstream_metadata", + )?, + }) +} + fn map_key_row(row: &MySqlRow) -> Result { let total_cost_usd = row .try_get::, _>("total_cost_usd") @@ -2387,9 +2708,9 @@ fn map_key_row(row: &MySqlRow) -> Result Result { .map_err(|_| DataLayerError::InvalidInput("timestamp overflow".to_string())) } +async fn enqueue_provider_monthly_usage_delta_mysql( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, + provider_id: &str, + total_cost_usd_delta: f64, + created_at: i64, +) -> Result<(), DataLayerError> { + let request_id = request_id.trim(); + let provider_id = provider_id.trim(); + if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 { + return Ok(()); + } + if !total_cost_usd_delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "provider monthly usage delta is not finite for {provider_id}" + ))); + } + + sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL) + .bind(uuid::Uuid::new_v4().to_string()) + .bind(request_id) + .bind(provider_id) + .bind(total_cost_usd_delta) + .bind(created_at) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + #[derive(Debug, Default)] struct DailyQuotaDebitResult { debited_usd: f64, @@ -622,30 +659,14 @@ WHERE id = ? .as_deref() .filter(|value| !value.is_empty()) { - sqlx::query( - r#" -UPDATE providers -SET - monthly_used_usd = COALESCE(monthly_used_usd, 0) + ?, - updated_at = ? -WHERE id = ? -"#, + enqueue_provider_monthly_usage_delta_mysql( + &mut tx, + &input.request_id, + provider_id, + input.actual_total_cost_usd, + updated_at, ) - .bind(input.actual_total_cost_usd) - .bind(updated_at) - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - settlement.provider_monthly_used_usd = sqlx::query_scalar::<_, Option>( - "SELECT monthly_used_usd FROM providers WHERE id = ? LIMIT 1", - ) - .bind(provider_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .flatten(); + .await?; } } @@ -683,6 +704,10 @@ WHERE id = ? #[cfg(test)] mod tests { use super::MysqlSettlementRepository; + use crate::run_migrations; + use aether_data_contracts::repository::settlement::{ + SettlementWriteRepository, UsageSettlementInput, + }; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -694,4 +719,132 @@ mod tests { let _repository = MysqlSettlementRepository::new(pool); } + + #[tokio::test] + async fn mysql_repository_settles_once_and_enqueues_provider_delta_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!( + "skipping mysql settlement parity test because AETHER_TEST_MYSQL_URL is unset" + ); + return; + }; + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + cleanup_settlement_rows(&pool).await; + + sqlx::query( + r#" +INSERT INTO providers (id, name, provider_type, monthly_used_usd, created_at, updated_at) +VALUES ('settlement-provider-1', 'Settlement Provider', 'openai', 5.0, 1, 1) +"#, + ) + .execute(&pool) + .await + .expect("provider should seed"); + sqlx::query( + r#" +INSERT INTO wallets (id, user_id, balance, gift_balance, limit_mode, created_at, updated_at) +VALUES ('settlement-wallet-1', 'settlement-user-1', 10.0, 2.0, 'finite', 1, 1) +"#, + ) + .execute(&pool) + .await + .expect("wallet should seed"); + sqlx::query( + r#" +INSERT INTO `usage` ( + request_id, user_id, provider_id, status, billing_status, + total_cost_usd, actual_total_cost_usd +) +VALUES ( + 'settlement-request-1', 'settlement-user-1', 'settlement-provider-1', + 'completed', 'pending', 3.0, 6.0 +) +"#, + ) + .execute(&pool) + .await + .expect("usage should seed"); + + let repository = MysqlSettlementRepository::new(pool.clone()); + let input = UsageSettlementInput { + request_id: "settlement-request-1".to_string(), + user_id: Some("settlement-user-1".to_string()), + api_key_id: None, + api_key_is_standalone: false, + provider_id: Some("settlement-provider-1".to_string()), + status: "completed".to_string(), + billing_status: "pending".to_string(), + total_cost_usd: 3.0, + actual_total_cost_usd: 6.0, + finalized_at_unix_secs: Some(1_234), + }; + let first = repository + .settle_usage(input.clone()) + .await + .expect("settlement should run") + .expect("usage should exist"); + let second = repository + .settle_usage(input) + .await + .expect("second settlement should run") + .expect("usage should exist"); + assert_eq!(first.billing_status, "settled"); + assert_eq!(first.provider_monthly_used_usd, None); + assert_eq!(second.finalized_at_unix_secs, Some(1_234)); + + let wallet: (f64, f64, f64) = sqlx::query_as( + "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'settlement-wallet-1'", + ) + .fetch_one(&pool) + .await + .expect("wallet should load"); + assert_eq!(wallet, (4.0, 2.0, 6.0)); + let provider_used: f64 = sqlx::query_scalar( + "SELECT monthly_used_usd FROM providers WHERE id = 'settlement-provider-1'", + ) + .fetch_one(&pool) + .await + .expect("provider should load"); + assert_eq!(provider_used, 5.0); + let provider_delta: (i64, f64) = sqlx::query_as( + r#" +SELECT CAST(COUNT(*) AS SIGNED), COALESCE(SUM(total_cost_usd_delta), 0) +FROM usage_counter_deltas +WHERE request_id = 'settlement-request-1' + AND kind = 'provider_monthly' + AND target_id = 'settlement-provider-1' +"#, + ) + .fetch_one(&pool) + .await + .expect("provider delta should load"); + assert_eq!(provider_delta, (1, 6.0)); + + cleanup_settlement_rows(&pool).await; + } + + async fn cleanup_settlement_rows(pool: &sqlx::MySqlPool) { + for sql in [ + "DELETE FROM usage_counter_deltas WHERE request_id = 'settlement-request-1'", + "DELETE FROM usage_settlement_snapshots WHERE request_id = 'settlement-request-1'", + "DELETE FROM `usage` WHERE request_id = 'settlement-request-1'", + "DELETE FROM wallets WHERE id = 'settlement-wallet-1'", + "DELETE FROM providers WHERE id = 'settlement-provider-1'", + ] { + sqlx::query(sql) + .execute(pool) + .await + .expect("settlement cleanup should succeed"); + } + } } diff --git a/crates/aether-data/adapters/mysql/src/usage.rs b/crates/aether-data/adapters/mysql/src/usage.rs index b8ca918f6..20869e968 100644 --- a/crates/aether-data/adapters/mysql/src/usage.rs +++ b/crates/aether-data/adapters/mysql/src/usage.rs @@ -6,11 +6,11 @@ use async_trait::async_trait; use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use aether_data_contracts::repository::usage::{ - provider_api_key_usage_is_error, provider_api_key_usage_is_success, strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure, usage_request_metadata_client_family, PendingUsageCleanupSummary, StoredRequestUsageAudit, StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardSummary, - StoredUsageUserTotals, UpsertUsageRecord, UsageDailyHeatmapQuery, + StoredUsageUserTotals, UpsertUsageRecord, UsageCleanupExecutionMode, UsageCleanupPreviewCounts, + UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageWriteRepository, }; use aether_data_contracts::DataLayerError; @@ -18,18 +18,40 @@ use aether_data_contracts::DataLayerError; use crate::error::SqlResultExt; use crate::MysqlPool; +mod cleanup; +mod counters; +mod http_capture; +mod read; +mod snapshots; + +pub use read::MysqlUsageReadFilter; + const USAGE_COLUMNS: &str = r#" SELECT id, - request_id, + `usage`.request_id, user_id, api_key_id, + `usage`.username, + `usage`.api_key_name, provider_name, model, target_model, - provider_id, - provider_endpoint_id, - provider_api_key_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_id + ELSE `usage`.provider_id + END AS provider_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_endpoint_id + ELSE `usage`.provider_endpoint_id + END AS provider_endpoint_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_api_key_id + ELSE `usage`.provider_api_key_id + END AS provider_api_key_id, request_type, api_format, api_family, @@ -37,41 +59,198 @@ SELECT endpoint_api_format, provider_api_family, provider_endpoint_kind, - has_format_conversion, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN COALESCE(usage_routing_snapshots.has_format_conversion, FALSE) + ELSE COALESCE(`usage`.has_format_conversion, FALSE) + END AS has_format_conversion, is_stream, upstream_is_stream, input_tokens, - output_tokens, + COALESCE(usage_settlement_snapshots.billing_output_tokens, `usage`.output_tokens, 0) + AS output_tokens, total_tokens, - cache_creation_input_tokens, - cache_creation_ephemeral_5m_input_tokens, - cache_creation_ephemeral_1h_input_tokens, - cache_read_input_tokens, - cache_creation_cost_usd, - cache_read_cost_usd, - output_price_per_1m, - total_cost_usd, - actual_total_cost_usd, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_tokens, + CASE + WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL + OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0) + + COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0) + END, + `usage`.cache_creation_input_tokens, + 0 + ) AS cache_creation_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_5m_tokens, + `usage`.cache_creation_ephemeral_5m_input_tokens, + 0 + ) AS cache_creation_ephemeral_5m_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_1h_tokens, + `usage`.cache_creation_ephemeral_1h_input_tokens, + 0 + ) AS cache_creation_ephemeral_1h_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_read_tokens, + `usage`.cache_read_input_tokens, + 0 + ) AS cache_read_input_tokens, + CAST(COALESCE( + usage_settlement_snapshots.billing_cache_creation_cost_usd, + `usage`.cache_creation_cost_usd, + 0 + ) AS DOUBLE) AS cache_creation_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.billing_cache_read_cost_usd, + `usage`.cache_read_cost_usd, + 0 + ) AS DOUBLE) AS cache_read_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.output_price_per_1m, + `usage`.output_price_per_1m + ) AS DOUBLE) AS output_price_per_1m, + CAST(COALESCE( + usage_settlement_snapshots.billing_total_cost_usd, + `usage`.total_cost_usd, + 0 + ) AS DOUBLE) AS total_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.billing_actual_total_cost_usd, + `usage`.actual_total_cost_usd, + 0 + ) AS DOUBLE) AS actual_total_cost_usd, status_code, error_message, error_category, response_time_ms, first_byte_time_ms, status, - billing_status, + COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) + AS billing_status, + CAST(COALESCE(usage_http_audits.request_headers, `usage`.request_headers) AS CHAR) AS request_headers, + CAST(`usage`.request_body AS CHAR) AS request_body, + `usage`.request_body_compressed, + CAST(COALESCE( + usage_http_audits.provider_request_headers, + `usage`.provider_request_headers + ) AS CHAR) AS provider_request_headers, + CAST(`usage`.provider_request_body AS CHAR) AS provider_request_body, + `usage`.provider_request_body_compressed, + CAST(COALESCE(usage_http_audits.response_headers, `usage`.response_headers) AS CHAR) AS response_headers, + CAST(`usage`.response_body AS CHAR) AS response_body, + `usage`.response_body_compressed, + CAST(COALESCE( + usage_http_audits.client_response_headers, + `usage`.client_response_headers + ) AS CHAR) AS client_response_headers, + CAST(`usage`.client_response_body AS CHAR) AS client_response_body, + `usage`.client_response_body_compressed, + usage_http_audits.request_body_ref AS http_request_body_ref, + usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref, + usage_http_audits.response_body_ref AS http_response_body_ref, + usage_http_audits.client_response_body_ref AS http_client_response_body_ref, + usage_http_audits.request_body_state AS http_request_body_state, + usage_http_audits.provider_request_body_state AS http_provider_request_body_state, + usage_http_audits.response_body_state AS http_response_body_state, + usage_http_audits.client_response_body_state AS http_client_response_body_state, request_metadata, - candidate_id, - candidate_index, - key_name, - planner_kind, - route_family, - route_kind, - execution_path, - local_execution_runtime_miss_reason, - finalized_at AS finalized_at_unix_secs, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.candidate_id + ELSE `usage`.candidate_id + END AS routing_candidate_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.candidate_index + ELSE `usage`.candidate_index + END AS routing_candidate_index, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.key_name + ELSE `usage`.key_name + END AS routing_key_name, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.planner_kind + ELSE `usage`.planner_kind + END AS routing_planner_kind, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.route_family + ELSE `usage`.route_family + END AS routing_route_family, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.route_kind + ELSE `usage`.route_kind + END AS routing_route_kind, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.execution_path + ELSE `usage`.execution_path + END AS routing_execution_path, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.local_execution_runtime_miss_reason + ELSE `usage`.local_execution_runtime_miss_reason + END AS routing_local_execution_runtime_miss_reason, + usage_settlement_snapshots.billing_snapshot_schema_version + AS settlement_billing_snapshot_schema_version, + usage_settlement_snapshots.billing_snapshot_status AS settlement_billing_snapshot_status, + CAST(usage_settlement_snapshots.rate_multiplier AS DOUBLE) AS settlement_rate_multiplier, + usage_settlement_snapshots.is_free_tier AS settlement_is_free_tier, + CAST(usage_settlement_snapshots.input_price_per_1m AS DOUBLE) + AS settlement_input_price_per_1m, + CAST(usage_settlement_snapshots.output_price_per_1m AS DOUBLE) + AS settlement_output_price_per_1m, + CAST(usage_settlement_snapshots.cache_creation_price_per_1m AS DOUBLE) + AS settlement_cache_creation_price_per_1m, + CAST(usage_settlement_snapshots.cache_read_price_per_1m AS DOUBLE) + AS settlement_cache_read_price_per_1m, + CAST(usage_settlement_snapshots.price_per_request AS DOUBLE) + AS settlement_price_per_request, + usage_settlement_snapshots.settlement_snapshot_schema_version + AS settlement_snapshot_schema_version, + CAST(usage_settlement_snapshots.settlement_snapshot AS CHAR) AS settlement_snapshot, + CAST(usage_settlement_snapshots.billing_dimensions AS CHAR) + AS settlement_billing_dimensions, + usage_settlement_snapshots.billing_input_tokens AS settlement_billing_input_tokens, + usage_settlement_snapshots.billing_effective_input_tokens + AS settlement_billing_effective_input_tokens, + usage_settlement_snapshots.billing_output_tokens AS settlement_billing_output_tokens, + usage_settlement_snapshots.billing_cache_creation_tokens + AS settlement_billing_cache_creation_tokens, + usage_settlement_snapshots.billing_cache_creation_5m_tokens + AS settlement_billing_cache_creation_5m_tokens, + usage_settlement_snapshots.billing_cache_creation_1h_tokens + AS settlement_billing_cache_creation_1h_tokens, + usage_settlement_snapshots.billing_cache_read_tokens + AS settlement_billing_cache_read_tokens, + usage_settlement_snapshots.billing_total_input_context + AS settlement_billing_total_input_context, + CAST(usage_settlement_snapshots.billing_cache_creation_cost_usd AS DOUBLE) + AS settlement_billing_cache_creation_cost_usd, + CAST(usage_settlement_snapshots.billing_cache_read_cost_usd AS DOUBLE) + AS settlement_billing_cache_read_cost_usd, + CAST(usage_settlement_snapshots.billing_total_cost_usd AS DOUBLE) + AS settlement_billing_total_cost_usd, + CAST(usage_settlement_snapshots.billing_actual_total_cost_usd AS DOUBLE) + AS settlement_billing_actual_total_cost_usd, + usage_settlement_snapshots.billing_pricing_source AS settlement_billing_pricing_source, + usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id, + usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version, + COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at) + AS finalized_at_unix_secs, created_at_unix_ms, updated_at_unix_secs FROM `usage` +LEFT JOIN usage_http_audits + ON usage_http_audits.request_id = `usage`.request_id +LEFT JOIN usage_routing_snapshots + ON usage_routing_snapshots.request_id = `usage`.request_id +LEFT JOIN usage_settlement_snapshots + ON usage_settlement_snapshots.request_id = `usage`.request_id "#; const UPSERT_USAGE_SQL: &str = r#" @@ -276,6 +455,111 @@ WHERE request_id = ? AND status IN ('streaming', 'success') "#; +const MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR: &str = r#" +GREATEST( + COALESCE( + CASE + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + GREATEST(settlement.billing_effective_input_tokens, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + + GREATEST( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) > 0 + THEN COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ELSE COALESCE(`usage`.cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + GREATEST( + COALESCE( + settlement.billing_cache_read_tokens, + `usage`.cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + GREATEST(settlement.billing_total_input_context, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + END, + NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0), + ( + CASE + WHEN ( + LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'openai' + OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'openai:%' + OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'gemini' + OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'gemini:%' + OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'google' + OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'google:%' + ) + AND COALESCE(`usage`.input_tokens, 0) > 0 + AND COALESCE(`usage`.cache_read_input_tokens, 0) > 0 + THEN GREATEST( + COALESCE(`usage`.input_tokens, 0) - COALESCE(`usage`.cache_read_input_tokens, 0), + 0 + ) + ELSE GREATEST(COALESCE(`usage`.input_tokens, 0), 0) + END + + GREATEST(COALESCE(`usage`.output_tokens, 0), 0) + + ( + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) > 0 + THEN COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ELSE GREATEST(COALESCE(`usage`.cache_creation_input_tokens, 0), 0) + END + ) + + GREATEST(COALESCE(`usage`.cache_read_input_tokens, 0), 0) + ), + 0 + ), + 0 +) +"#; + +const MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR: &str = r#" +CASE + WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled') + AND (status_code IS NULL OR status_code < 400) + AND (error_message IS NULL OR TRIM(error_message) = '') + THEN 1 + ELSE 0 +END +"#; + +const MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR: &str = r#" +CASE + WHEN status NOT IN ('pending', 'streaming') + AND NOT ( + status IN ('completed', 'success', 'ok', 'billed', 'settled') + AND (status_code IS NULL OR status_code < 400) + AND (error_message IS NULL OR TRIM(error_message) = '') + ) + THEN 1 + ELSE 0 +END +"#; + #[derive(Debug, Clone)] pub struct MysqlUsageWriteRepository { pool: MysqlPool, @@ -291,19 +575,20 @@ impl MysqlUsageStorage { Self { pool } } - pub async fn load_usage_records(&self) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{USAGE_COLUMNS} ORDER BY created_at_unix_ms ASC, request_id ASC" - )) - .fetch_all(&self.pool) - .await - .map_sql_err()?; + pub async fn read_usage_counter_health( + &self, + ) -> Result + { + counters::read_health(&self.pool).await + } - let items = rows - .iter() - .map(map_usage_row) - .collect::, _>>()?; - Ok(items) + pub async fn read_usage_counter_pending_health( + &self, + ) -> Result< + aether_data_contracts::repository::usage::UsageCounterPendingHealthSnapshot, + DataLayerError, + > { + counters::read_pending_health(&self.pool).await } async fn summarize_usage_daily_heatmap_raw_from_range( @@ -312,30 +597,23 @@ impl MysqlUsageStorage { created_until_unix_secs: u64, user_id: Option<&str>, ) -> Result, DataLayerError> { - let mut sql = String::from( + let mut sql = format!( r#" SELECT DATE_FORMAT(FROM_UNIXTIME(created_at_unix_ms), '%Y-%m-%d') AS date, CAST(COUNT(*) AS SIGNED) AS requests, - CAST(COALESCE(SUM( - GREATEST(COALESCE(input_tokens, 0), 0) - + GREATEST(COALESCE(output_tokens, 0), 0) - + CASE - WHEN COALESCE(cache_creation_input_tokens, 0) = 0 - AND (COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)) > 0 - THEN COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0) - ELSE GREATEST(COALESCE(cache_creation_input_tokens, 0), 0) - END - + GREATEST(COALESCE(cache_read_input_tokens, 0), 0) - ), 0) AS SIGNED) AS total_tokens, + CAST(COALESCE(SUM({canonical_total_tokens_expr}), 0) AS SIGNED) AS total_tokens, CAST(COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS DOUBLE) AS total_cost_usd, CAST(COALESCE(SUM(COALESCE(actual_total_cost_usd, 0)), 0) AS DOUBLE) AS actual_total_cost_usd FROM `usage` +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? AND status NOT IN ('pending', 'streaming') AND provider_name NOT IN ('unknown', 'pending') "#, + canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, ); if user_id.is_some() { sql.push_str(" AND user_id = ?\n"); @@ -671,16 +949,17 @@ WHERE user_id IN ( ); } - let mut raw_builder = QueryBuilder::::new( + let mut raw_builder = QueryBuilder::::new(format!( r#" SELECT `usage`.user_id, CAST(COUNT(*) AS SIGNED) AS request_count, - CAST(COALESCE(SUM(GREATEST(COALESCE(`usage`.total_tokens, 0), 0)), 0) AS SIGNED) AS total_tokens + CAST(COALESCE(SUM({canonical_total_tokens_expr}), 0) AS SIGNED) AS total_tokens FROM `usage` JOIN ( "#, - ); + canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, + )); for (index, user_id) in unique_user_ids.iter().enumerate() { if index > 0 { raw_builder.push(" UNION ALL "); @@ -696,6 +975,8 @@ JOIN ( raw_builder.push( r#" ) AS requested ON requested.user_id = `usage`.user_id +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id WHERE `usage`.created_at_unix_ms >= requested.cutoff_unix_secs AND `usage`.status NOT IN ('pending', 'streaming') AND `usage`.provider_name NOT IN ('unknown', 'pending') @@ -739,12 +1020,23 @@ impl MysqlUsageWriteRepository { &self, request_id: &str, ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE request_id = ? LIMIT 1")) - .bind(request_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_usage_row).transpose() + let row = sqlx::query(&format!( + "{USAGE_COLUMNS} WHERE `usage`.request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + let usage = row + .as_ref() + .map(|row| map_usage_row(row, true)) + .transpose()?; + match usage { + Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage) + .await + .map(Some), + None => Ok(None), + } } } @@ -754,28 +1046,63 @@ impl UsageWriteRepository for MysqlUsageWriteRepository { &self, usage: UpsertUsageRecord, ) -> Result { - let usage = strip_deprecated_usage_display_fields(usage); + let mut usage = strip_deprecated_usage_display_fields(usage); usage.validate()?; - - if let Some(existing) = self.find_by_request_id(&usage.request_id).await? { + let prepared_capture = http_capture::prepare_usage_http_capture(&mut usage)?; + let mut tx = self.pool.begin().await.map_sql_err()?; + let existing = counters::lock_and_load_usage(&mut tx, &usage.request_id).await?; + let recovers_terminal_failure = existing.as_ref().is_some_and(|existing| { + usage_can_recover_terminal_failure( + &existing.status, + &existing.billing_status, + &usage.status, + &usage.billing_status, + ) + }); + if let Some(existing) = existing.as_ref() { if (existing.billing_status == "settled" || existing.billing_status == "void") - && !usage_can_recover_terminal_failure( - &existing.status, - &existing.billing_status, - &usage.status, - &usage.billing_status, - ) + && !recovers_terminal_failure { - return Ok(existing); + let existing = existing.clone(); + tx.rollback().await.map_sql_err()?; + return http_capture::hydrate_usage_body_refs(&self.pool, existing).await; } } + let capture_update_allowed = recovers_terminal_failure + || http_capture::capture_update_allowed(existing.as_ref(), &usage.status); + if capture_update_allowed { + http_capture::apply_previous_metadata_tombstones(&mut usage, existing.as_ref()); + } + let prepared_snapshots = capture_update_allowed + .then(|| snapshots::from_usage(&usage)) + .transpose()?; bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? - .execute(&self.pool) + .execute(&mut *tx) .await .map_sql_err()?; - self.rebuild_api_key_usage_stats().await?; - self.rebuild_provider_api_key_usage_stats().await?; + if capture_update_allowed { + http_capture::sync_usage_http_capture(&mut tx, &usage.request_id, &prepared_capture) + .await?; + let (routing_snapshot, settlement_snapshot) = prepared_snapshots + .as_ref() + .expect("capture-allowed usage has prepared snapshots"); + snapshots::sync( + &mut tx, + &usage.request_id, + routing_snapshot, + settlement_snapshot, + matches!(usage.status.as_str(), "completed" | "failed" | "cancelled"), + ) + .await?; + } + counters::enqueue_usage_transition_for_request( + &mut tx, + &usage.request_id, + existing.as_ref(), + ) + .await?; + tx.commit().await.map_sql_err()?; self.find_by_request_id(&usage.request_id) .await? .ok_or_else(|| { @@ -784,6 +1111,7 @@ impl UsageWriteRepository for MysqlUsageWriteRepository { } async fn rebuild_api_key_usage_stats(&self) -> Result { + let mut tx = self.pool.begin().await.map_sql_err()?; sqlx::query( r#" UPDATE api_keys @@ -793,55 +1121,45 @@ SET total_requests = 0, last_used_at = NULL "#, ) - .execute(&self.pool) + .execute(&mut *tx) .await .map_sql_err()?; - let rows = sqlx::query( + let rows_affected = sqlx::query(&format!( r#" -SELECT - api_key_id, - COUNT(*) AS total_requests, - CAST(COALESCE(SUM(total_tokens), 0) AS SIGNED) AS total_tokens, - CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE) AS total_cost_usd, - MAX(updated_at_unix_secs) AS last_used_at -FROM `usage` -WHERE api_key_id IS NOT NULL AND api_key_id <> '' -GROUP BY api_key_id -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - for row in &rows { - sqlx::query( - r#" UPDATE api_keys -SET total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - last_used_at = ? -WHERE id = ? +JOIN ( + SELECT + api_key_id, + COUNT(*) AS total_requests, + COALESCE(SUM({canonical_total_tokens_expr}), 0) AS total_tokens, + COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS total_cost_usd, + MAX(created_at_unix_ms) AS last_used_at + FROM `usage` + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id + WHERE api_key_id IS NOT NULL + AND TRIM(api_key_id) <> '' + AND status NOT IN ('pending', 'streaming') + GROUP BY api_key_id +) AS aggregated ON aggregated.api_key_id = api_keys.id +SET api_keys.total_requests = aggregated.total_requests, + api_keys.total_tokens = aggregated.total_tokens, + api_keys.total_cost_usd = aggregated.total_cost_usd, + api_keys.last_used_at = aggregated.last_used_at "#, - ) - .bind(row.try_get::("total_requests").map_sql_err()?) - .bind(row.try_get::("total_tokens").map_sql_err()?) - .bind(row.try_get::("total_cost_usd").map_sql_err()?) - .bind( - row.try_get::, _>("last_used_at") - .map_sql_err()?, - ) - .bind(row.try_get::("api_key_id").map_sql_err()?) - .execute(&self.pool) - .await - .map_sql_err()?; - } - - Ok(rows.len() as u64) + canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, + )) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + tx.commit().await.map_sql_err()?; + Ok(rows_affected) } async fn rebuild_provider_api_key_usage_stats(&self) -> Result { + let mut tx = self.pool.begin().await.map_sql_err()?; sqlx::query( r#" UPDATE provider_api_keys @@ -854,94 +1172,58 @@ SET request_count = 0, last_used_at = NULL "#, ) - .execute(&self.pool) + .execute(&mut *tx) .await .map_sql_err()?; - let rows = sqlx::query( + let rows_affected = sqlx::query(&format!( r#" -SELECT - provider_api_key_id, - status, - status_code, - error_message, - total_tokens, - total_cost_usd, - response_time_ms, - updated_at_unix_secs -FROM `usage` -WHERE provider_api_key_id IS NOT NULL AND provider_api_key_id <> '' -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let mut stats = BTreeMap::::new(); - for row in rows { - let key_id: String = row.try_get("provider_api_key_id").map_sql_err()?; - let status: String = row.try_get("status").map_sql_err()?; - let status_code = row.try_get::, _>("status_code").map_sql_err()?; - let status_code_u16 = status_code.and_then(|value| u16::try_from(value).ok()); - let error_message: Option = row.try_get("error_message").map_sql_err()?; - let entry = stats.entry(key_id).or_default(); - entry.request_count += 1; - let is_success = provider_api_key_usage_is_success( - &status, - status_code_u16, - error_message.as_deref(), - ); - let is_in_flight = matches!(status.as_str(), "pending" | "streaming"); - if is_success { - entry.success_count += 1; - } - if provider_api_key_usage_is_error(&status, status_code_u16, error_message.as_deref()) { - entry.error_count += 1; - } - if !is_in_flight { - entry.total_tokens += row.try_get::("total_tokens").map_sql_err()?; - entry.total_cost_usd += row.try_get::("total_cost_usd").map_sql_err()?; - } - if is_success { - entry.total_response_time_ms += row - .try_get::, _>("response_time_ms") - .map_sql_err()? - .unwrap_or_default(); - } - entry.last_used_at = entry.last_used_at.max( - row.try_get::, _>("updated_at_unix_secs") - .map_sql_err()?, - ); - } - - for (key_id, stat) in &stats { - sqlx::query( - r#" UPDATE provider_api_keys -SET request_count = ?, - success_count = ?, - error_count = ?, - total_tokens = ?, - total_cost_usd = ?, - total_response_time_ms = ?, - last_used_at = ? -WHERE id = ? +JOIN ( + SELECT + provider_api_key_id, + COUNT(*) AS request_count, + COALESCE(SUM({success_flag_expr}), 0) AS success_count, + COALESCE(SUM({error_flag_expr}), 0) AS error_count, + COALESCE(SUM(CASE + WHEN status IN ('pending', 'streaming') THEN 0 + ELSE {canonical_total_tokens_expr} + END), 0) AS total_tokens, + COALESCE(SUM(CASE + WHEN status IN ('pending', 'streaming') THEN 0 + ELSE COALESCE(total_cost_usd, 0) + END), 0) AS total_cost_usd, + COALESCE(SUM(CASE + WHEN {success_flag_expr} = 1 AND response_time_ms IS NOT NULL + THEN GREATEST(COALESCE(response_time_ms, 0), 0) + ELSE 0 + END), 0) AS total_response_time_ms, + MAX(created_at_unix_ms) AS last_used_at + FROM `usage` + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id + WHERE provider_api_key_id IS NOT NULL + AND TRIM(provider_api_key_id) <> '' + GROUP BY provider_api_key_id +) AS aggregated ON aggregated.provider_api_key_id = provider_api_keys.id +SET provider_api_keys.request_count = aggregated.request_count, + provider_api_keys.success_count = aggregated.success_count, + provider_api_keys.error_count = aggregated.error_count, + provider_api_keys.total_tokens = aggregated.total_tokens, + provider_api_keys.total_cost_usd = aggregated.total_cost_usd, + provider_api_keys.total_response_time_ms = aggregated.total_response_time_ms, + provider_api_keys.last_used_at = aggregated.last_used_at "#, - ) - .bind(stat.request_count) - .bind(stat.success_count) - .bind(stat.error_count) - .bind(stat.total_tokens) - .bind(stat.total_cost_usd) - .bind(stat.total_response_time_ms) - .bind(stat.last_used_at) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - } - - Ok(stats.len() as u64) + success_flag_expr = MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR, + error_flag_expr = MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR, + canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, + )) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + tx.commit().await.map_sql_err()?; + Ok(rows_affected) } async fn cleanup_stale_pending_requests( @@ -955,7 +1237,6 @@ WHERE id = ? return Ok(PendingUsageCleanupSummary::default()); } - let cutoff_unix_ms = cutoff_unix_secs.saturating_mul(1000); let now_unix_ms = now_unix_secs.saturating_mul(1000); let mut summary = PendingUsageCleanupSummary::default(); let batch_size_u64 = u64::try_from(batch_size).map_err(|_| { @@ -967,7 +1248,7 @@ WHERE id = ? loop { let mut tx = self.pool.begin().await.map_sql_err()?; let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL) - .bind(to_i64(cutoff_unix_ms, "stale pending usage cutoff")?) + .bind(to_i64(cutoff_unix_secs, "stale pending usage cutoff")?) .bind(to_i64(batch_size_u64, "stale pending usage batch size")?) .fetch_all(&mut *tx) .await @@ -1101,6 +1382,71 @@ WHERE request_id = ? Ok(summary) } + + async fn flush_usage_counter_deltas( + &self, + batch_size: usize, + ) -> Result + { + counters::flush(&self.pool, batch_size).await + } + + async fn enqueue_proxy_node_counter_delta( + &self, + delta: aether_data_contracts::repository::usage::ProxyNodeCounterDelta, + ) -> Result { + counters::enqueue_proxy_node(&self.pool, delta).await + } + + async fn enqueue_management_token_counter_delta( + &self, + delta: aether_data_contracts::repository::usage::ManagementTokenCounterDelta, + ) -> Result { + counters::enqueue_management_token(&self.pool, delta).await + } + + async fn enqueue_api_key_last_used_delta( + &self, + delta: aether_data_contracts::repository::usage::ApiKeyLastUsedDelta, + ) -> Result { + counters::enqueue_api_key_last_used(&self.pool, delta).await + } + + async fn cleanup_processed_usage_counter_deltas( + &self, + cutoff_unix_secs: u64, + batch_size: usize, + ) -> Result { + counters::cleanup_processed(&self.pool, cutoff_unix_secs, batch_size).await + } + + async fn cleanup_usage( + &self, + window: &UsageCleanupWindow, + batch_size: usize, + auto_delete_expired_keys: bool, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, + ) -> Result { + cleanup::cleanup_usage( + &self.pool, + window, + batch_size, + auto_delete_expired_keys, + targets, + mode, + ) + .await + } + + async fn preview_usage_cleanup( + &self, + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, + ) -> Result { + cleanup::preview_usage_cleanup(&self.pool, window, targets, mode).await + } } struct StalePendingUsageRow { @@ -1109,17 +1455,6 @@ struct StalePendingUsageRow { billing_status: String, } -#[derive(Default)] -struct ProviderKeyStats { - request_count: i64, - success_count: i64, - error_count: i64, - total_tokens: i64, - total_cost_usd: f64, - total_response_time_ms: i64, - last_used_at: Option, -} - async fn completed_request_ids_mysql<'a>( tx: &mut sqlx::Transaction<'_, sqlx::MySql>, request_ids: impl Iterator, @@ -1185,6 +1520,33 @@ INSERT INTO usage_settlement_snapshots ( ) VALUES (?, 'void', ?, ?, ?) ON DUPLICATE KEY UPDATE billing_status = VALUES(billing_status), + billing_snapshot_schema_version = NULL, + billing_snapshot_status = NULL, + settlement_snapshot_schema_version = NULL, + settlement_snapshot = NULL, + billing_dimensions = NULL, + billing_input_tokens = NULL, + billing_effective_input_tokens = NULL, + billing_output_tokens = NULL, + billing_cache_creation_tokens = NULL, + billing_cache_creation_5m_tokens = NULL, + billing_cache_creation_1h_tokens = NULL, + billing_cache_read_tokens = NULL, + billing_total_input_context = NULL, + billing_cache_creation_cost_usd = NULL, + billing_cache_read_cost_usd = NULL, + billing_total_cost_usd = NULL, + billing_actual_total_cost_usd = NULL, + billing_pricing_source = NULL, + billing_rule_id = NULL, + billing_rule_version = NULL, + rate_multiplier = NULL, + is_free_tier = NULL, + input_price_per_1m = NULL, + output_price_per_1m = NULL, + cache_creation_price_per_1m = NULL, + cache_read_price_per_1m = NULL, + price_per_request = NULL, finalized_at = COALESCE(usage_settlement_snapshots.finalized_at, VALUES(finalized_at)), updated_at = VALUES(updated_at) "#, @@ -1289,7 +1651,7 @@ fn bind_upsert<'q>( .unwrap_or(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens); let created_at = usage .created_at_unix_ms - .unwrap_or(usage.updated_at_unix_secs.saturating_mul(1000)); + .unwrap_or(usage.updated_at_unix_secs); let request_metadata = usage .request_metadata .as_ref() @@ -1365,7 +1727,10 @@ fn bind_upsert<'q>( Ok(query) } -fn map_usage_row(row: &MySqlRow) -> Result { +fn map_usage_row( + row: &MySqlRow, + resolve_legacy_compressed: bool, +) -> Result { let id = row .try_get::, _>("id") .map_sql_err()? @@ -1378,8 +1743,8 @@ fn map_usage_row(row: &MySqlRow) -> Result Result, _>("upstream_is_stream") .map_sql_err()?; merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream); - audit.candidate_id = row.try_get("candidate_id").map_sql_err()?; - audit.candidate_index = row - .try_get::, _>("candidate_index") - .map_sql_err()? - .map(|value| value as u64); - audit.key_name = row.try_get("key_name").map_sql_err()?; - audit.planner_kind = row.try_get("planner_kind").map_sql_err()?; - audit.route_family = row.try_get("route_family").map_sql_err()?; - audit.route_kind = row.try_get("route_kind").map_sql_err()?; - audit.execution_path = row.try_get("execution_path").map_sql_err()?; - audit.local_execution_runtime_miss_reason = row - .try_get("local_execution_runtime_miss_reason") - .map_sql_err()?; + snapshots::hydrate_row(row, &mut audit)?; Ok(audit) } diff --git a/crates/aether-data/adapters/mysql/src/usage/cleanup.rs b/crates/aether-data/adapters/mysql/src/usage/cleanup.rs new file mode 100644 index 000000000..fc98ee971 --- /dev/null +++ b/crates/aether-data/adapters/mysql/src/usage/cleanup.rs @@ -0,0 +1,1177 @@ +use std::io::Write; + +use aether_data_contracts::repository::usage::{ + parse_usage_body_ref, usage_body_ref, UsageBodyField, UsageCleanupExecutionMode, + UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, +}; +use chrono::{DateTime, Utc}; +use flate2::{write::GzEncoder, Compression}; +use serde_json::Value; +use sqlx::Row; +use tracing::warn; + +use crate::error::SqlResultExt; +use crate::{DataLayerError, MysqlPool}; + +const RAW_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +"#; + +const COMPRESSED_BODY_PREDICATE: &str = r#" +request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_body_blobs + WHERE usage_body_blobs.request_id = `usage`.request_id +) +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = `usage`.request_id + AND ( + usage_http_audits.request_body_ref IS NOT NULL + OR usage_http_audits.provider_request_body_ref IS NOT NULL + OR usage_http_audits.response_body_ref IS NOT NULL + OR usage_http_audits.client_response_body_ref IS NOT NULL + ) +) +"#; + +const ALL_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_body_blobs + WHERE usage_body_blobs.request_id = `usage`.request_id +) +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = `usage`.request_id + AND ( + usage_http_audits.request_body_ref IS NOT NULL + OR usage_http_audits.provider_request_body_ref IS NOT NULL + OR usage_http_audits.response_body_ref IS NOT NULL + OR usage_http_audits.client_response_body_ref IS NOT NULL + ) +) +"#; + +const INLINE_OR_COMPRESSED_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +"#; + +const HEADER_PREDICATE: &str = r#" +request_headers IS NOT NULL +OR response_headers IS NOT NULL +OR provider_request_headers IS NOT NULL +OR client_response_headers IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = `usage`.request_id + AND ( + usage_http_audits.request_headers IS NOT NULL + OR usage_http_audits.response_headers IS NOT NULL + OR usage_http_audits.provider_request_headers IS NOT NULL + OR usage_http_audits.client_response_headers IS NOT NULL + ) +) +"#; + +const LEGACY_BODY_REF_PREDICATE: &str = r#" +request_metadata IS NOT NULL +AND JSON_VALID(request_metadata) +AND JSON_CONTAINS_PATH( + CASE WHEN JSON_VALID(request_metadata) THEN request_metadata ELSE JSON_OBJECT() END, + 'one', + '$.request_body_ref', + '$.provider_request_body_ref', + '$.response_body_ref', + '$.client_response_body_ref' +) = 1 +"#; + +const DETAIL_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR ( + request_metadata IS NOT NULL + AND JSON_VALID(request_metadata) + AND JSON_CONTAINS_PATH( + CASE WHEN JSON_VALID(request_metadata) THEN request_metadata ELSE JSON_OBJECT() END, + 'one', + '$.request_body_ref', + '$.provider_request_body_ref', + '$.response_body_ref', + '$.client_response_body_ref' + ) = 1 +) +"#; + +#[derive(Debug)] +struct CleanupRow { + id: String, + request_id: String, +} + +#[derive(Debug)] +struct BodyRow { + id: String, + request_id: String, + request_body: Option, + request_body_compressed: Option>, + provider_request_body: Option, + provider_request_body_compressed: Option>, + response_body: Option, + response_body_compressed: Option>, + client_response_body: Option, + client_response_body_compressed: Option>, +} + +#[derive(Debug, Default)] +struct DetachedRefs { + request_body_ref: Option, + provider_request_body_ref: Option, + response_body_ref: Option, + client_response_body_ref: Option, +} + +impl DetachedRefs { + fn any_present(&self) -> bool { + self.request_body_ref.is_some() + || self.provider_request_body_ref.is_some() + || self.response_body_ref.is_some() + || self.client_response_body_ref.is_some() + } +} + +struct DetachedBlob { + body_ref: String, + body_field: &'static str, + payload_gzip: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BodyCleanupKind { + Raw, + Compressed, + All, +} + +impl BodyCleanupKind { + fn predicate(self) -> &'static str { + match self { + Self::Raw => RAW_BODY_PREDICATE, + Self::Compressed => COMPRESSED_BODY_PREDICATE, + Self::All => ALL_BODY_PREDICATE, + } + } + + fn clears_detached(self) -> bool { + self != Self::Raw + } +} + +pub(crate) async fn cleanup_usage( + pool: &MysqlPool, + window: &UsageCleanupWindow, + batch_size: usize, + auto_delete_expired_keys: bool, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, +) -> Result { + if batch_size == 0 || !targets.any_selected() { + return Ok(UsageCleanupSummary::default()); + } + + if mode == UsageCleanupExecutionMode::BeforeNowBodyFields { + let body_externalized = if targets.detail_body { + cleanup_body_fields( + pool, + window.detail_cutoff, + None, + batch_size, + BodyCleanupKind::Raw, + ) + .await? + } else { + 0 + }; + let body_cleaned = if targets.compressed_body { + cleanup_body_fields( + pool, + window.compressed_cutoff, + None, + batch_size, + BodyCleanupKind::Compressed, + ) + .await? + } else { + 0 + }; + return Ok(UsageCleanupSummary { + body_externalized, + body_cleaned, + ..UsageCleanupSummary::default() + }); + } + + let records_deleted = if targets.records { + delete_old_usage_records(pool, window.log_cutoff, batch_size).await? + } else { + 0 + }; + let record_cutoff = targets.records.then_some(window.log_cutoff); + let header_cleaned = if targets.headers { + cleanup_headers(pool, window.header_cutoff, record_cutoff, batch_size).await? + } else { + 0 + }; + let body_cleaned = if targets.compressed_body { + cleanup_body_fields( + pool, + window.compressed_cutoff, + record_cutoff, + batch_size, + BodyCleanupKind::All, + ) + .await? + } else { + 0 + }; + let detail_newer_than = detail_body_newer_than(window, targets); + let legacy_body_refs_migrated = if targets.detail_body { + migrate_legacy_body_refs(pool, window.detail_cutoff, detail_newer_than, batch_size).await? + } else { + 0 + }; + let body_externalized = if targets.detail_body { + externalize_detail_bodies(pool, window.detail_cutoff, detail_newer_than, batch_size).await? + } else { + 0 + }; + let keys_cleaned = if targets.expired_keys { + match cleanup_expired_api_keys(pool, auto_delete_expired_keys).await { + Ok(count) => count, + Err(err) => { + warn!(error = %err, "MySQL usage cleanup expired api key sweep failed"); + 0 + } + } + } else { + 0 + }; + + Ok(UsageCleanupSummary { + body_externalized, + legacy_body_refs_migrated, + body_cleaned, + header_cleaned, + keys_cleaned, + records_deleted, + }) +} + +pub(crate) async fn preview_usage_cleanup( + pool: &MysqlPool, + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, +) -> Result { + if mode == UsageCleanupExecutionMode::BeforeNowBodyFields { + let detail = if targets.detail_body { + count_candidates(pool, RAW_BODY_PREDICATE, window.detail_cutoff, None).await? + } else { + 0 + }; + let compressed = if targets.compressed_body { + count_candidates( + pool, + COMPRESSED_BODY_PREDICATE, + window.compressed_cutoff, + None, + ) + .await? + } else { + 0 + }; + return Ok(UsageCleanupPreviewCounts { + detail, + compressed, + header: 0, + log: 0, + }); + } + + let record_cutoff = targets.records.then_some(window.log_cutoff); + let detail = if targets.detail_body { + count_candidates( + pool, + DETAIL_BODY_PREDICATE, + window.detail_cutoff, + detail_body_newer_than(window, targets), + ) + .await? + } else { + 0 + }; + let compressed = if targets.compressed_body { + count_candidates( + pool, + ALL_BODY_PREDICATE, + window.compressed_cutoff, + record_cutoff, + ) + .await? + } else { + 0 + }; + let header = if targets.headers { + count_candidates(pool, HEADER_PREDICATE, window.header_cutoff, record_cutoff).await? + } else { + 0 + }; + let log = if targets.records { + let count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM `usage` WHERE created_at_unix_ms < ?") + .bind(window.log_cutoff.timestamp()) + .fetch_one(pool) + .await + .map_sql_err()?; + u64::try_from(count).unwrap_or(0) + } else { + 0 + }; + + Ok(UsageCleanupPreviewCounts { + detail, + compressed, + header, + log, + }) +} + +fn detail_body_newer_than( + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, +) -> Option> { + [ + targets.compressed_body.then_some(window.compressed_cutoff), + targets.records.then_some(window.log_cutoff), + ] + .into_iter() + .flatten() + .max() +} + +fn invalid_window(cutoff: DateTime, newer_than: Option>) -> bool { + matches!(newer_than, Some(value) if value >= cutoff) +} + +async fn count_candidates( + pool: &MysqlPool, + predicate: &str, + cutoff: DateTime, + newer_than: Option>, +) -> Result { + if invalid_window(cutoff, newer_than) { + return Ok(0); + } + let sql = format!( + r#" +SELECT COUNT(*) +FROM `usage` +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({predicate}) +"# + ); + let newer_than = newer_than.map(|value| value.timestamp()); + let count: i64 = sqlx::query_scalar(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .fetch_one(pool) + .await + .map_sql_err()?; + Ok(u64::try_from(count).unwrap_or(0)) +} + +async fn fetch_cleanup_rows( + pool: &MysqlPool, + predicate: &str, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result, DataLayerError> { + if invalid_window(cutoff, newer_than) { + return Ok(Vec::new()); + } + let sql = format!( + r#" +SELECT id, request_id +FROM `usage` +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({predicate}) +ORDER BY created_at_unix_ms ASC, id ASC +LIMIT ? +"# + ); + let newer_than = newer_than.map(|value| value.timestamp()); + sqlx::query(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .bind(i64::try_from(batch_size).unwrap_or(i64::MAX)) + .fetch_all(pool) + .await + .map_sql_err()? + .into_iter() + .map(|row| { + Ok(CleanupRow { + id: row.try_get("id").map_sql_err()?, + request_id: row.try_get("request_id").map_sql_err()?, + }) + }) + .collect() +} + +async fn delete_old_usage_records( + pool: &MysqlPool, + cutoff: DateTime, + batch_size: usize, +) -> Result { + let mut total = 0usize; + loop { + let rows = fetch_cleanup_rows(pool, "1 = 1", cutoff, None, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + let mut deleted = 0usize; + for row in rows { + deleted += usize::try_from( + sqlx::query("DELETE FROM `usage` WHERE id = ?") + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(), + ) + .unwrap_or(usize::MAX); + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(deleted); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn cleanup_headers( + pool: &MysqlPool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "MySQL usage header cleanup skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = + fetch_cleanup_rows(pool, HEADER_PREDICATE, cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + for row in rows { + sqlx::query( + r#" +UPDATE `usage` +SET request_headers = NULL, + response_headers = NULL, + provider_request_headers = NULL, + client_response_headers = NULL +WHERE id = ? +"#, + ) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query( + r#" +UPDATE usage_http_audits +SET request_headers = NULL, + response_headers = NULL, + provider_request_headers = NULL, + client_response_headers = NULL, + updated_at = UNIX_TIMESTAMP() +WHERE request_id = ? +"#, + ) + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + delete_empty_http_audit(&mut tx, &row.request_id).await?; + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(row_count); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn cleanup_body_fields( + pool: &MysqlPool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, + kind: BodyCleanupKind, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "MySQL usage body cleanup skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = + fetch_cleanup_rows(pool, kind.predicate(), cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + for row in rows { + let update = if kind == BodyCleanupKind::All { + r#" +UPDATE `usage` +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL, + request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"# + } else if kind == BodyCleanupKind::Compressed { + r#" +UPDATE `usage` +SET request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"# + } else { + r#" +UPDATE `usage` +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL +WHERE id = ? +"# + }; + sqlx::query(update) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + if kind.clears_detached() { + sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query( + r#" +UPDATE usage_http_audits +SET request_body_ref = NULL, + provider_request_body_ref = NULL, + response_body_ref = NULL, + client_response_body_ref = NULL, + body_capture_mode = 'none', + updated_at = UNIX_TIMESTAMP() +WHERE request_id = ? +"#, + ) + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + delete_empty_http_audit(&mut tx, &row.request_id).await?; + } + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(row_count); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn delete_empty_http_audit( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +DELETE FROM usage_http_audits +WHERE request_id = ? + AND request_headers IS NULL + AND response_headers IS NULL + AND provider_request_headers IS NULL + AND client_response_headers IS NULL + AND request_body_ref IS NULL + AND provider_request_body_ref IS NULL + AND response_body_ref IS NULL + AND client_response_body_ref IS NULL +"#, + ) + .bind(request_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn migrate_legacy_body_refs( + pool: &MysqlPool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "MySQL usage legacy body-ref migration skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = fetch_cleanup_rows( + pool, + LEGACY_BODY_REF_PREDICATE, + cutoff, + newer_than, + batch_size, + ) + .await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + let mut migrated = 0usize; + for row in rows { + let metadata: Option = + sqlx::query_scalar("SELECT request_metadata FROM `usage` WHERE id = ? LIMIT 1") + .bind(&row.id) + .fetch_optional(&mut *tx) + .await + .map_sql_err()? + .flatten(); + let Some((refs, metadata)) = + legacy_body_ref_plan(&row.request_id, metadata.as_deref())? + else { + continue; + }; + if refs.any_present() { + upsert_http_audit_refs(&mut tx, &row.request_id, &refs).await?; + } + let updated = sqlx::query( + r#" +UPDATE `usage` +SET request_metadata = ?, + updated_at_unix_secs = UNIX_TIMESTAMP() +WHERE id = ? +"#, + ) + .bind(metadata) + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + if updated > 0 { + migrated += 1; + } + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(migrated); + if row_count < batch_size || migrated == 0 { + break; + } + } + Ok(total) +} + +fn legacy_body_ref_plan( + request_id: &str, + metadata: Option<&str>, +) -> Result)>, DataLayerError> { + let Some(metadata) = metadata else { + return Ok(None); + }; + let value: Value = serde_json::from_str(metadata).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid usage request_metadata JSON: {err}")) + })?; + let Value::Object(mut object) = value else { + return Ok(None); + }; + let mut refs = DetachedRefs::default(); + let mut removed = false; + for field in [ + UsageBodyField::RequestBody, + UsageBodyField::ProviderRequestBody, + UsageBodyField::ResponseBody, + UsageBodyField::ClientResponseBody, + ] { + let Some(value) = object.remove(field.as_ref_key()) else { + continue; + }; + removed = true; + let parsed = value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(parse_usage_body_ref) + .filter(|(parsed_request_id, parsed_field)| { + parsed_request_id == request_id && *parsed_field == field + }) + .map(|(parsed_request_id, parsed_field)| { + usage_body_ref(&parsed_request_id, parsed_field) + }); + set_ref(&mut refs, field, parsed); + } + if !removed { + return Ok(None); + } + let metadata = if object.is_empty() { + None + } else { + Some( + serde_json::to_string(&Value::Object(object)).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "failed to serialize request_metadata: {err}" + )) + })?, + ) + }; + Ok(Some((refs, metadata))) +} + +async fn externalize_detail_bodies( + pool: &MysqlPool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "MySQL usage body externalization skipped due to invalid window"); + return Ok(0); + } + let batch_size = batch_size.clamp(1, 25); + let mut total = 0usize; + loop { + let rows = fetch_body_rows(pool, cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut externalized = 0usize; + for row in rows { + let (blobs, refs) = build_detached_bodies(&row)?; + let mut tx = pool.begin().await.map_sql_err()?; + for blob in blobs { + sqlx::query( + r#" +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES (?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + request_id = VALUES(request_id), + body_field = VALUES(body_field), + payload_gzip = VALUES(payload_gzip), + updated_at = UNIX_TIMESTAMP() +"#, + ) + .bind(blob.body_ref) + .bind(&row.request_id) + .bind(blob.body_field) + .bind(blob.payload_gzip) + .execute(&mut *tx) + .await + .map_sql_err()?; + } + if refs.any_present() { + upsert_http_audit_refs(&mut tx, &row.request_id, &refs).await?; + } + let updated = sqlx::query( + r#" +UPDATE `usage` +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL, + request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"#, + ) + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + tx.commit().await.map_sql_err()?; + if updated > 0 { + externalized += 1; + } + } + total = total.saturating_add(externalized); + if row_count < batch_size || externalized == 0 { + break; + } + } + Ok(total) +} + +async fn fetch_body_rows( + pool: &MysqlPool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result, DataLayerError> { + let newer_than = newer_than.map(|value| value.timestamp()); + let sql = format!( + r#" +SELECT id, + request_id, + CAST(request_body AS CHAR) AS request_body, + request_body_compressed, + CAST(provider_request_body AS CHAR) AS provider_request_body, + provider_request_body_compressed, + CAST(response_body AS CHAR) AS response_body, + response_body_compressed, + CAST(client_response_body AS CHAR) AS client_response_body, + client_response_body_compressed +FROM `usage` +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({INLINE_OR_COMPRESSED_BODY_PREDICATE}) +ORDER BY created_at_unix_ms ASC, id ASC +LIMIT ? +"# + ); + sqlx::query(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .bind(i64::try_from(batch_size).unwrap_or(i64::MAX)) + .fetch_all(pool) + .await + .map_sql_err()? + .into_iter() + .map(|row| { + Ok(BodyRow { + id: row.try_get("id").map_sql_err()?, + request_id: row.try_get("request_id").map_sql_err()?, + request_body: parse_optional_json(row.try_get("request_body").map_sql_err()?)?, + request_body_compressed: row.try_get("request_body_compressed").map_sql_err()?, + provider_request_body: parse_optional_json( + row.try_get("provider_request_body").map_sql_err()?, + )?, + provider_request_body_compressed: row + .try_get("provider_request_body_compressed") + .map_sql_err()?, + response_body: parse_optional_json(row.try_get("response_body").map_sql_err()?)?, + response_body_compressed: row.try_get("response_body_compressed").map_sql_err()?, + client_response_body: parse_optional_json( + row.try_get("client_response_body").map_sql_err()?, + )?, + client_response_body_compressed: row + .try_get("client_response_body_compressed") + .map_sql_err()?, + }) + }) + .collect() +} + +fn parse_optional_json(raw: Option) -> Result, DataLayerError> { + raw.map(|raw| { + serde_json::from_str(&raw).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid inline usage body JSON: {err}")) + }) + }) + .transpose() +} + +fn build_detached_bodies( + row: &BodyRow, +) -> Result<(Vec, DetachedRefs), DataLayerError> { + let mut blobs = Vec::new(); + let mut refs = DetachedRefs::default(); + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::RequestBody, + row.request_body.as_ref(), + row.request_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ProviderRequestBody, + row.provider_request_body.as_ref(), + row.provider_request_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ResponseBody, + row.response_body.as_ref(), + row.response_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ClientResponseBody, + row.client_response_body.as_ref(), + row.client_response_body_compressed.as_deref(), + )?; + Ok((blobs, refs)) +} + +fn add_detached_body( + blobs: &mut Vec, + refs: &mut DetachedRefs, + request_id: &str, + field: UsageBodyField, + raw: Option<&Value>, + compressed: Option<&[u8]>, +) -> Result<(), DataLayerError> { + let payload_gzip = match raw { + Some(value) => Some(compress_json(value)?), + None => compressed.map(ToOwned::to_owned), + }; + let Some(payload_gzip) = payload_gzip else { + return Ok(()); + }; + let body_ref = usage_body_ref(request_id, field); + blobs.push(DetachedBlob { + body_ref: body_ref.clone(), + body_field: field.as_storage_field(), + payload_gzip, + }); + set_ref(refs, field, Some(body_ref)); + Ok(()) +} + +fn compress_json(value: &Value) -> Result, DataLayerError> { + let bytes = serde_json::to_vec(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage body: {err}")) + })?; + let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6)); + encoder.write_all(&bytes).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to gzip usage body: {err}")) + })?; + encoder.finish().map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to finish usage body gzip: {err}")) + }) +} + +fn set_ref(refs: &mut DetachedRefs, field: UsageBodyField, value: Option) { + match field { + UsageBodyField::RequestBody => refs.request_body_ref = value, + UsageBodyField::ProviderRequestBody => refs.provider_request_body_ref = value, + UsageBodyField::ResponseBody => refs.response_body_ref = value, + UsageBodyField::ClientResponseBody => refs.client_response_body_ref = value, + } +} + +async fn upsert_http_audit_refs( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, + refs: &DetachedRefs, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO usage_http_audits ( + request_id, + request_body_ref, + provider_request_body_ref, + response_body_ref, + client_response_body_ref, + body_capture_mode +) +VALUES (?, ?, ?, ?, ?, 'ref_backed') +ON DUPLICATE KEY UPDATE + request_body_ref = COALESCE(VALUES(request_body_ref), request_body_ref), + provider_request_body_ref = COALESCE(VALUES(provider_request_body_ref), provider_request_body_ref), + response_body_ref = COALESCE(VALUES(response_body_ref), response_body_ref), + client_response_body_ref = COALESCE(VALUES(client_response_body_ref), client_response_body_ref), + body_capture_mode = 'ref_backed', + updated_at = UNIX_TIMESTAMP() +"#, + ) + .bind(request_id) + .bind(refs.request_body_ref.as_deref()) + .bind(refs.provider_request_body_ref.as_deref()) + .bind(refs.response_body_ref.as_deref()) + .bind(refs.client_response_body_ref.as_deref()) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn cleanup_expired_api_keys( + pool: &MysqlPool, + auto_delete_expired_keys: bool, +) -> Result { + let now = Utc::now().timestamp(); + let rows = sqlx::query( + r#" +SELECT id, auto_delete_on_expiry +FROM api_keys +WHERE expires_at <= ? + AND is_active = TRUE +ORDER BY expires_at ASC, id ASC +"#, + ) + .bind(now) + .fetch_all(pool) + .await + .map_sql_err()?; + let mut cleaned = 0usize; + for row in rows { + let id: String = row.try_get("id").map_sql_err()?; + let auto_delete = row + .try_get::, _>("auto_delete_on_expiry") + .map_sql_err()? + .map(|value| value != 0) + .unwrap_or(auto_delete_expired_keys); + let mut tx = pool.begin().await.map_sql_err()?; + let affected = if auto_delete { + sqlx::query( + "UPDATE wallets SET status = 'disabled', updated_at = ? WHERE api_key_id = ? AND status <> 'disabled'", + ) + .bind(now) + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query("DELETE FROM api_keys WHERE id = ?") + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected() + } else { + sqlx::query( + "UPDATE api_keys SET is_active = FALSE, updated_at = ? WHERE id = ? AND is_active = TRUE", + ) + .bind(now) + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected() + }; + tx.commit().await.map_sql_err()?; + if affected > 0 { + cleaned += 1; + } + } + Ok(cleaned) +} + +#[cfg(test)] +mod tests { + use std::io::Read; + + use flate2::read::GzDecoder; + use serde_json::json; + + use super::{compress_json, legacy_body_ref_plan}; + + #[test] + fn mysql_cleanup_legacy_body_ref_plan_preserves_unrelated_metadata() { + let metadata = json!({ + "trace": "kept", + "request_body_ref": "usage://request/request-1/request_body", + "response_body_ref": "usage://request/other/response_body" + }) + .to_string(); + let (refs, metadata) = legacy_body_ref_plan("request-1", Some(&metadata)) + .expect("legacy plan should build") + .expect("legacy refs should be present"); + assert_eq!( + refs.request_body_ref.as_deref(), + Some("usage://request/request-1/request_body") + ); + assert!(refs.response_body_ref.is_none()); + assert_eq!( + serde_json::from_str::( + metadata.as_deref().expect("trace metadata should remain") + ) + .expect("metadata should remain valid JSON"), + json!({"trace": "kept"}) + ); + } + + #[test] + fn mysql_cleanup_gzip_payload_round_trips() { + let value = json!({"hello": "world"}); + let payload = compress_json(&value).expect("body should compress"); + let mut decoder = GzDecoder::new(payload.as_slice()); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .expect("body should decompress"); + assert_eq!( + serde_json::from_slice::(&decoded) + .expect("decoded body should be JSON"), + value + ); + } +} diff --git a/crates/aether-data/adapters/mysql/src/usage/counters.rs b/crates/aether-data/adapters/mysql/src/usage/counters.rs new file mode 100644 index 000000000..496fadc77 --- /dev/null +++ b/crates/aether-data/adapters/mysql/src/usage/counters.rs @@ -0,0 +1,1153 @@ +use std::collections::BTreeMap; + +use aether_data_contracts::repository::usage::{ + api_key_usage_contribution, model_usage_contribution, provider_api_key_usage_contribution, + ApiKeyLastUsedDelta, ApiKeyUsageDelta, ManagementTokenCounterDelta, ModelUsageDelta, + ProviderApiKeyUsageDelta, ProxyNodeCounterDelta, StoredRequestUsageAudit, + UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, +}; +use aether_data_contracts::DataLayerError; +use sqlx::{MySql, MySqlPool, QueryBuilder, Row}; + +use crate::error::SqlResultExt; + +const KIND_API_KEY: &str = "api_key"; +const KIND_PROVIDER_API_KEY: &str = "provider_api_key"; +const KIND_MODEL: &str = "model"; +const KIND_PROVIDER_MONTHLY: &str = "provider_monthly"; +const KIND_PROXY_NODE: &str = "proxy_node"; +const KIND_MANAGEMENT_TOKEN: &str = "management_token"; +const KIND_API_KEY_LAST_USED: &str = "api_key_last_used"; + +const CLAIM_SQL: &str = r#" +SELECT + id, + kind, + target_id, + request_count_delta, + total_requests_delta, + success_count_delta, + error_count_delta, + dns_failures_delta, + stream_errors_delta, + total_tokens_delta, + total_cost_usd_delta, + total_response_time_ms_delta, + last_used_at_unix_secs, + last_used_ip, + candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs, + usage_created_at_unix_secs +FROM usage_counter_deltas +WHERE processed_at IS NULL +ORDER BY created_at ASC, id ASC +LIMIT ? +FOR UPDATE SKIP LOCKED +"#; + +struct DeltaRow { + id: String, + kind: String, + target_id: String, + request_count_delta: i64, + total_requests_delta: i64, + success_count_delta: i64, + error_count_delta: i64, + dns_failures_delta: i64, + stream_errors_delta: i64, + total_tokens_delta: i64, + total_cost_usd_delta: f64, + total_response_time_ms_delta: i64, + last_used_at_unix_secs: Option, + last_used_ip: Option, + candidate_last_used_at_unix_secs: Option, + removed_last_used_at_unix_secs: Option, + usage_created_at_unix_secs: Option, +} + +#[derive(Default)] +struct Aggregates { + api_keys: BTreeMap, + provider_api_keys: BTreeMap, + models: BTreeMap, + provider_monthly: BTreeMap, + proxy_nodes: BTreeMap, + management_tokens: BTreeMap, + api_key_last_used: BTreeMap, +} + +impl Aggregates { + fn from_rows(rows: &[DeltaRow]) -> Result { + let mut aggregates = Self::default(); + for row in rows { + if !row.total_cost_usd_delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "usage_counter_deltas.total_cost_usd_delta is not finite for {}", + row.id + ))); + } + match row.kind.as_str() { + KIND_API_KEY => { + let entry = aggregates + .api_keys + .entry(row.target_id.clone()) + .or_default(); + entry.total_requests += row.total_requests_delta; + entry.total_tokens += row.total_tokens_delta; + entry.total_cost_usd += row.total_cost_usd_delta; + merge_optional_max( + &mut entry.candidate_last_used_at_unix_secs, + row.candidate_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.removed_last_used_at_unix_secs, + row.removed_last_used_at_unix_secs, + ); + } + KIND_PROVIDER_API_KEY => { + let entry = aggregates + .provider_api_keys + .entry(row.target_id.clone()) + .or_default(); + entry.request_count += row.request_count_delta; + entry.success_count += row.success_count_delta; + entry.error_count += row.error_count_delta; + entry.total_tokens += row.total_tokens_delta; + entry.total_cost_usd += row.total_cost_usd_delta; + entry.total_response_time_ms += row.total_response_time_ms_delta; + merge_optional_max( + &mut entry.candidate_last_used_at_unix_secs, + row.candidate_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.removed_last_used_at_unix_secs, + row.removed_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.usage_created_at_unix_secs, + row.usage_created_at_unix_secs, + ); + } + KIND_MODEL => { + aggregates + .models + .entry(row.target_id.clone()) + .or_default() + .request_count += row.request_count_delta; + } + KIND_PROVIDER_MONTHLY => { + *aggregates + .provider_monthly + .entry(row.target_id.clone()) + .or_default() += row.total_cost_usd_delta; + } + KIND_PROXY_NODE => { + let entry = aggregates + .proxy_nodes + .entry(row.target_id.clone()) + .or_insert(ProxyNodeCounterDelta { + node_id: row.target_id.clone(), + total_requests_delta: 0, + failed_requests_delta: 0, + dns_failures_delta: 0, + stream_errors_delta: 0, + }); + entry.total_requests_delta += row.total_requests_delta; + entry.failed_requests_delta += row.error_count_delta; + entry.dns_failures_delta += row.dns_failures_delta; + entry.stream_errors_delta += row.stream_errors_delta; + } + KIND_MANAGEMENT_TOKEN => { + let entry = aggregates + .management_tokens + .entry(row.target_id.clone()) + .or_insert(ManagementTokenCounterDelta { + token_id: row.target_id.clone(), + usage_count_delta: 0, + last_used_at_unix_secs: None, + last_used_ip: None, + }); + entry.usage_count_delta += row.request_count_delta; + merge_latest_timestamp_with_value( + &mut entry.last_used_at_unix_secs, + &mut entry.last_used_ip, + row.last_used_at_unix_secs, + row.last_used_ip.clone(), + ); + } + KIND_API_KEY_LAST_USED => { + let Some(last_used_at_unix_secs) = row.last_used_at_unix_secs else { + continue; + }; + let entry = aggregates + .api_key_last_used + .entry(row.target_id.clone()) + .or_insert(ApiKeyLastUsedDelta { + api_key_id: row.target_id.clone(), + last_used_at_unix_secs, + }); + if last_used_at_unix_secs > entry.last_used_at_unix_secs { + entry.last_used_at_unix_secs = last_used_at_unix_secs; + } + } + other => { + return Err(DataLayerError::UnexpectedValue(format!( + "unknown usage counter delta kind: {other}" + ))); + } + } + } + Ok(aggregates) + } +} + +pub(super) async fn flush( + pool: &MySqlPool, + batch_size: usize, +) -> Result { + if batch_size == 0 { + return Ok(UsageCounterFlushSummary::default()); + } + let limit = i64::try_from(batch_size).map_err(|_| { + DataLayerError::InvalidInput(format!( + "usage counter flush batch size is out of range: {batch_size}" + )) + })?; + + let mut tx = pool.begin().await.map_sql_err()?; + let rows = sqlx::query(CLAIM_SQL) + .bind(limit) + .fetch_all(&mut *tx) + .await + .map_sql_err()? + .iter() + .map(map_row) + .collect::, _>>()?; + if rows.is_empty() { + tx.rollback().await.map_sql_err()?; + return Ok(UsageCounterFlushSummary::default()); + } + + let aggregates = Aggregates::from_rows(&rows)?; + for (target_id, delta) in &aggregates.api_keys { + apply_api_key(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.models { + apply_model(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.provider_api_keys { + apply_provider_api_key(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.provider_monthly { + apply_provider_monthly(&mut tx, target_id, *delta).await?; + } + for (target_id, delta) in &aggregates.proxy_nodes { + apply_proxy_node(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.management_tokens { + apply_management_token(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.api_key_last_used { + apply_api_key_last_used(&mut tx, target_id, delta).await?; + } + + let now = current_unix_secs(); + let mut mark = QueryBuilder::::new("UPDATE usage_counter_deltas SET processed_at = "); + mark.push_bind(now).push(" WHERE id IN ("); + { + let mut ids = mark.separated(", "); + for row in &rows { + ids.push_bind(&row.id); + } + } + mark.push(")"); + mark.build().execute(&mut *tx).await.map_sql_err()?; + tx.commit().await.map_sql_err()?; + + Ok(UsageCounterFlushSummary { + rows_claimed: rows.len(), + api_key_targets: aggregates.api_keys.len(), + provider_api_key_targets: aggregates.provider_api_keys.len(), + model_targets: aggregates.models.len(), + provider_monthly_targets: aggregates.provider_monthly.len(), + proxy_node_targets: aggregates.proxy_nodes.len(), + management_token_targets: aggregates.management_tokens.len(), + api_key_last_used_targets: aggregates.api_key_last_used.len(), + }) +} + +pub(super) async fn enqueue_proxy_node( + pool: &MySqlPool, + delta: ProxyNodeCounterDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let node_id = delta.node_id.trim().to_string(); + let request_id = format!("proxy_node:{node_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_PROXY_NODE, + target_id: &node_id, + total_requests_delta: delta.total_requests_delta, + error_count_delta: delta.failed_requests_delta, + dns_failures_delta: delta.dns_failures_delta, + stream_errors_delta: delta.stream_errors_delta, + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn enqueue_management_token( + pool: &MySqlPool, + delta: ManagementTokenCounterDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let token_id = delta.token_id.trim().to_string(); + let last_used_ip = delta + .last_used_ip + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let last_used_at = delta + .last_used_at_unix_secs + .unwrap_or_else(|| current_unix_secs().max(0) as u64); + let request_id = format!("management_token:{token_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_MANAGEMENT_TOKEN, + target_id: &token_id, + request_count_delta: delta.usage_count_delta, + last_used_at_unix_secs: Some(last_used_at), + last_used_ip: last_used_ip.as_deref(), + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn enqueue_api_key_last_used( + pool: &MySqlPool, + delta: ApiKeyLastUsedDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let api_key_id = delta.api_key_id.trim().to_string(); + let request_id = format!("api_key_last_used:{api_key_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_API_KEY_LAST_USED, + target_id: &api_key_id, + last_used_at_unix_secs: Some(delta.last_used_at_unix_secs), + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn cleanup_processed( + pool: &MySqlPool, + cutoff_unix_secs: u64, + batch_size: usize, +) -> Result { + if batch_size == 0 { + return Ok(0); + } + let cutoff = to_i64(cutoff_unix_secs, "usage counter cleanup cutoff")?; + let limit = i64::try_from(batch_size).map_err(|_| { + DataLayerError::InvalidInput(format!( + "usage counter cleanup batch size is out of range: {batch_size}" + )) + })?; + let deleted = sqlx::query( + r#" +DELETE FROM usage_counter_deltas +WHERE id IN ( + SELECT id FROM ( + SELECT id + FROM usage_counter_deltas + WHERE processed_at IS NOT NULL AND processed_at < ? + ORDER BY processed_at ASC, created_at ASC, id ASC + LIMIT ? + ) AS doomed +) +"#, + ) + .bind(cutoff) + .bind(limit) + .execute(pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(usize::try_from(deleted).unwrap_or(usize::MAX)) +} + +pub(super) async fn read_health( + pool: &MySqlPool, +) -> Result { + let row = sqlx::query( + r#" +SELECT + (SELECT CAST(COUNT(*) AS SIGNED) FROM usage_counter_deltas WHERE processed_at IS NULL) + AS pending_rows, + (SELECT CAST(COUNT(*) AS SIGNED) FROM usage_counter_deltas WHERE processed_at IS NOT NULL) + AS processed_rows, + (SELECT MIN(created_at) FROM usage_counter_deltas WHERE processed_at IS NULL) + AS oldest_pending_created_at_unix_secs, + (SELECT MAX(processed_at) FROM usage_counter_deltas WHERE processed_at IS NOT NULL) + AS latest_processed_at_unix_secs +"#, + ) + .fetch_one(pool) + .await + .map_sql_err()?; + let mut snapshot = UsageCounterHealthSnapshot { + pending_rows: nonnegative_u64(row.try_get("pending_rows").map_sql_err()?), + processed_rows: nonnegative_u64(row.try_get("processed_rows").map_sql_err()?), + oldest_pending_created_at_unix_secs: optional_nonnegative_u64( + row.try_get("oldest_pending_created_at_unix_secs") + .map_sql_err()?, + ), + latest_processed_at_unix_secs: optional_nonnegative_u64( + row.try_get("latest_processed_at_unix_secs").map_sql_err()?, + ), + pending_by_kind: BTreeMap::new(), + }; + for row in pending_health_rows(pool).await? { + snapshot.pending_by_kind.insert(row.0, row.1); + } + Ok(snapshot) +} + +pub(super) async fn read_pending_health( + pool: &MySqlPool, +) -> Result { + let mut snapshot = UsageCounterPendingHealthSnapshot::default(); + for (kind, pending_rows, oldest) in pending_health_rows(pool).await? { + snapshot.pending_rows = snapshot.pending_rows.saturating_add(pending_rows); + if let Some(oldest) = oldest { + snapshot.oldest_pending_created_at_unix_secs = Some( + snapshot + .oldest_pending_created_at_unix_secs + .map_or(oldest, |current| current.min(oldest)), + ); + } + snapshot.pending_by_kind.insert(kind, pending_rows); + } + Ok(snapshot) +} + +async fn pending_health_rows( + pool: &MySqlPool, +) -> Result)>, DataLayerError> { + let rows = sqlx::query( + r#" +SELECT + kind, + CAST(COUNT(*) AS SIGNED) AS pending_rows, + MIN(created_at) AS oldest_pending_created_at_unix_secs +FROM usage_counter_deltas +WHERE processed_at IS NULL +GROUP BY kind +ORDER BY kind ASC +"#, + ) + .fetch_all(pool) + .await + .map_sql_err()?; + rows.iter() + .map(|row| { + Ok(( + row.try_get("kind").map_sql_err()?, + nonnegative_u64(row.try_get("pending_rows").map_sql_err()?), + optional_nonnegative_u64( + row.try_get("oldest_pending_created_at_unix_secs") + .map_sql_err()?, + ), + )) + }) + .collect() +} + +pub(super) async fn enqueue_usage_transition( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, + before: Option<&StoredRequestUsageAudit>, + after: &StoredRequestUsageAudit, +) -> Result<(), DataLayerError> { + let before_api_key = before.and_then(api_key_usage_contribution); + let after_api_key = api_key_usage_contribution(after); + match (before_api_key.as_ref(), after_api_key.as_ref()) { + (Some(before), Some(after)) if before.api_key_id == after.api_key_id => { + enqueue_api_key_delta( + tx, + request_id, + &before.api_key_id, + &ApiKeyUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_api_key.as_ref() { + enqueue_api_key_delta( + tx, + request_id, + &before.api_key_id, + &ApiKeyUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_api_key.as_ref() { + enqueue_api_key_delta( + tx, + request_id, + &after.api_key_id, + &ApiKeyUsageDelta::addition(after), + ) + .await?; + } + } + } + + let before_model = before.and_then(model_usage_contribution); + let after_model = model_usage_contribution(after); + match (before_model.as_ref(), after_model.as_ref()) { + (Some(before), Some(after)) if before.model == after.model => { + enqueue_model_delta( + tx, + request_id, + &before.model, + &ModelUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_model.as_ref() { + enqueue_model_delta( + tx, + request_id, + &before.model, + &ModelUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_model.as_ref() { + enqueue_model_delta( + tx, + request_id, + &after.model, + &ModelUsageDelta::addition(after), + ) + .await?; + } + } + } + + let before_provider = before.and_then(provider_api_key_usage_contribution); + let after_provider = provider_api_key_usage_contribution(after); + match (before_provider.as_ref(), after_provider.as_ref()) { + (Some(before), Some(after)) if before.key_id == after.key_id => { + enqueue_provider_api_key_delta( + tx, + request_id, + &before.key_id, + &ProviderApiKeyUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_provider.as_ref() { + enqueue_provider_api_key_delta( + tx, + request_id, + &before.key_id, + &ProviderApiKeyUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_provider.as_ref() { + enqueue_provider_api_key_delta( + tx, + request_id, + &after.key_id, + &ProviderApiKeyUsageDelta::addition(after), + ) + .await?; + } + } + } + Ok(()) +} + +pub(super) async fn enqueue_usage_transition_for_request( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, + before: Option<&StoredRequestUsageAudit>, +) -> Result<(), DataLayerError> { + let row = sqlx::query(&format!( + "{} WHERE `usage`.request_id = ? LIMIT 1", + super::USAGE_COLUMNS + )) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()? + .ok_or_else(|| { + DataLayerError::UnexpectedValue(format!( + "usage row missing while preparing counter delta: {request_id}" + )) + })?; + let after = super::map_usage_row(&row, false)?; + enqueue_usage_transition(tx, request_id, before, &after).await +} + +pub(super) async fn lock_and_load_usage( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, +) -> Result, DataLayerError> { + let exists = sqlx::query_scalar::<_, String>( + "SELECT request_id FROM `usage` WHERE request_id = ? FOR UPDATE", + ) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + if exists.is_none() { + return Ok(None); + } + let row = sqlx::query(&format!( + "{} WHERE `usage`.request_id = ? LIMIT 1", + super::USAGE_COLUMNS + )) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + row.as_ref() + .map(|row| super::map_usage_row(row, false)) + .transpose() +} + +async fn enqueue_api_key_delta( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, + target_id: &str, + delta: &ApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_API_KEY, + target_id, + total_requests_delta: delta.total_requests, + total_tokens_delta: delta.total_tokens, + total_cost_usd_delta: finite_or_zero(delta.total_cost_usd), + candidate_last_used_at_unix_secs: delta.candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs: delta.removed_last_used_at_unix_secs, + ..DeltaInsert::default() + }, + ) + .await +} + +async fn enqueue_model_delta( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, + target_id: &str, + delta: &ModelUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_MODEL, + target_id, + request_count_delta: delta.request_count, + ..DeltaInsert::default() + }, + ) + .await +} + +async fn enqueue_provider_api_key_delta( + tx: &mut sqlx::Transaction<'_, MySql>, + request_id: &str, + target_id: &str, + delta: &ProviderApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_PROVIDER_API_KEY, + target_id, + request_count_delta: delta.request_count, + success_count_delta: delta.success_count, + error_count_delta: delta.error_count, + total_tokens_delta: delta.total_tokens, + total_cost_usd_delta: finite_or_zero(delta.total_cost_usd), + total_response_time_ms_delta: delta.total_response_time_ms, + candidate_last_used_at_unix_secs: delta.candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs: delta.removed_last_used_at_unix_secs, + usage_created_at_unix_secs: delta.usage_created_at_unix_secs, + ..DeltaInsert::default() + }, + ) + .await +} + +#[derive(Default)] +struct DeltaInsert<'a> { + request_id: &'a str, + kind: &'a str, + target_id: &'a str, + request_count_delta: i64, + total_requests_delta: i64, + success_count_delta: i64, + error_count_delta: i64, + dns_failures_delta: i64, + stream_errors_delta: i64, + total_tokens_delta: i64, + total_cost_usd_delta: f64, + total_response_time_ms_delta: i64, + last_used_at_unix_secs: Option, + last_used_ip: Option<&'a str>, + candidate_last_used_at_unix_secs: Option, + removed_last_used_at_unix_secs: Option, + usage_created_at_unix_secs: Option, +} + +async fn insert_delta( + tx: &mut sqlx::Transaction<'_, MySql>, + input: DeltaInsert<'_>, +) -> Result<(), DataLayerError> { + let request_id = input.request_id.trim(); + let target_id = input.target_id.trim(); + if request_id.is_empty() || target_id.is_empty() { + return Ok(()); + } + sqlx::query( + r#" +INSERT INTO usage_counter_deltas ( + id, request_id, kind, target_id, request_count_delta, total_requests_delta, + success_count_delta, error_count_delta, dns_failures_delta, stream_errors_delta, + total_tokens_delta, total_cost_usd_delta, total_response_time_ms_delta, + last_used_at_unix_secs, last_used_ip, candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs, usage_created_at_unix_secs, created_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +"#, + ) + .bind(uuid::Uuid::new_v4().to_string()) + .bind(request_id) + .bind(input.kind) + .bind(target_id) + .bind(input.request_count_delta) + .bind(input.total_requests_delta) + .bind(input.success_count_delta) + .bind(input.error_count_delta) + .bind(input.dns_failures_delta) + .bind(input.stream_errors_delta) + .bind(input.total_tokens_delta) + .bind(finite_or_zero(input.total_cost_usd_delta)) + .bind(input.total_response_time_ms_delta) + .bind(optional_to_i64( + input.last_used_at_unix_secs, + "usage counter last_used_at_unix_secs", + )?) + .bind( + input + .last_used_ip + .map(str::trim) + .filter(|value| !value.is_empty()), + ) + .bind(optional_to_i64( + input.candidate_last_used_at_unix_secs, + "usage counter candidate_last_used_at_unix_secs", + )?) + .bind(optional_to_i64( + input.removed_last_used_at_unix_secs, + "usage counter removed_last_used_at_unix_secs", + )?) + .bind(optional_to_i64( + input.usage_created_at_unix_secs, + "usage counter usage_created_at_unix_secs", + )?) + .bind(current_unix_secs()) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +fn map_row(row: &sqlx::mysql::MySqlRow) -> Result { + Ok(DeltaRow { + id: row.try_get("id").map_sql_err()?, + kind: row.try_get("kind").map_sql_err()?, + target_id: row.try_get("target_id").map_sql_err()?, + request_count_delta: row.try_get("request_count_delta").map_sql_err()?, + total_requests_delta: row.try_get("total_requests_delta").map_sql_err()?, + success_count_delta: row.try_get("success_count_delta").map_sql_err()?, + error_count_delta: row.try_get("error_count_delta").map_sql_err()?, + dns_failures_delta: row.try_get("dns_failures_delta").map_sql_err()?, + stream_errors_delta: row.try_get("stream_errors_delta").map_sql_err()?, + total_tokens_delta: row.try_get("total_tokens_delta").map_sql_err()?, + total_cost_usd_delta: row.try_get("total_cost_usd_delta").map_sql_err()?, + total_response_time_ms_delta: row.try_get("total_response_time_ms_delta").map_sql_err()?, + last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.last_used_at_unix_secs", + row.try_get("last_used_at_unix_secs").map_sql_err()?, + )?, + last_used_ip: row.try_get("last_used_ip").map_sql_err()?, + candidate_last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.candidate_last_used_at_unix_secs", + row.try_get("candidate_last_used_at_unix_secs") + .map_sql_err()?, + )?, + removed_last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.removed_last_used_at_unix_secs", + row.try_get("removed_last_used_at_unix_secs") + .map_sql_err()?, + )?, + usage_created_at_unix_secs: optional_u64( + "usage_counter_deltas.usage_created_at_unix_secs", + row.try_get("usage_created_at_unix_secs").map_sql_err()?, + )?, + }) +} + +async fn apply_api_key( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let candidate = optional_to_i64( + delta.candidate_last_used_at_unix_secs, + "api key candidate last used at", + )?; + let removed = optional_to_i64( + delta.removed_last_used_at_unix_secs, + "api key removed last used at", + )?; + sqlx::query( + r#" +UPDATE api_keys +SET total_requests = GREATEST(COALESCE(total_requests, 0) + ?, 0), + total_tokens = GREATEST(COALESCE(total_tokens, 0) + ?, 0), + total_cost_usd = GREATEST(COALESCE(total_cost_usd, 0) + ?, 0), + last_used_at = CASE + WHEN ? IS NOT NULL THEN GREATEST(COALESCE(last_used_at, 0), ?) + WHEN ? IS NOT NULL AND last_used_at = ? THEN ( + SELECT MAX(created_at_unix_ms) + FROM `usage` + WHERE api_key_id = ? AND status NOT IN ('pending', 'streaming') + ) + ELSE last_used_at + END +WHERE id = ? +"#, + ) + .bind(delta.total_requests) + .bind(delta.total_tokens) + .bind(finite_or_zero(delta.total_cost_usd)) + .bind(candidate) + .bind(candidate) + .bind(removed) + .bind(removed) + .bind(target_id) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_model( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ModelUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + "UPDATE global_models SET usage_count = GREATEST(COALESCE(usage_count, 0) + ?, 0), updated_at = ? WHERE name = ?", + ) + .bind(delta.request_count) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_provider_api_key( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ProviderApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let candidate = optional_to_i64( + delta.candidate_last_used_at_unix_secs, + "provider api key candidate last used at", + )?; + let removed = optional_to_i64( + delta.removed_last_used_at_unix_secs, + "provider api key removed last used at", + )?; + sqlx::query( + r#" +UPDATE provider_api_keys +SET request_count = GREATEST(COALESCE(request_count, 0) + ?, 0), + success_count = GREATEST(COALESCE(success_count, 0) + ?, 0), + error_count = GREATEST(COALESCE(error_count, 0) + ?, 0), + total_tokens = GREATEST(COALESCE(total_tokens, 0) + ?, 0), + total_cost_usd = GREATEST(COALESCE(total_cost_usd, 0) + ?, 0), + total_response_time_ms = GREATEST(COALESCE(total_response_time_ms, 0) + ?, 0), + last_used_at = CASE + WHEN ? IS NOT NULL THEN GREATEST(COALESCE(last_used_at, 0), ?) + WHEN ? IS NOT NULL AND last_used_at = ? THEN ( + SELECT MAX(created_at_unix_ms) + FROM `usage` + WHERE provider_api_key_id = ? AND status NOT IN ('pending', 'streaming') + ) + ELSE last_used_at + END +WHERE id = ? +"#, + ) + .bind(delta.request_count) + .bind(delta.success_count) + .bind(delta.error_count) + .bind(delta.total_tokens) + .bind(finite_or_zero(delta.total_cost_usd)) + .bind(delta.total_response_time_ms) + .bind(candidate) + .bind(candidate) + .bind(removed) + .bind(removed) + .bind(target_id) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_provider_monthly( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: f64, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta == 0.0 { + return Ok(()); + } + if !delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "providers.monthly_used_usd delta is not finite for {target_id}" + ))); + } + sqlx::query( + "UPDATE providers SET monthly_used_usd = COALESCE(monthly_used_usd, 0) + ?, updated_at = ? WHERE id = ?", + ) + .bind(delta) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_proxy_node( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ProxyNodeCounterDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + r#" +UPDATE proxy_nodes +SET total_requests = total_requests + GREATEST(?, 0), + failed_requests = failed_requests + GREATEST(?, 0), + dns_failures = dns_failures + GREATEST(?, 0), + stream_errors = stream_errors + GREATEST(?, 0), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(delta.total_requests_delta) + .bind(delta.failed_requests_delta) + .bind(delta.dns_failures_delta) + .bind(delta.stream_errors_delta) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_management_token( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ManagementTokenCounterDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let last_used_at = optional_to_i64( + delta.last_used_at_unix_secs, + "management token last used at", + )?; + sqlx::query( + r#" +UPDATE management_tokens +SET usage_count = COALESCE(usage_count, 0) + GREATEST(?, 0), + last_used_at = CASE + WHEN ? IS NULL THEN last_used_at + ELSE GREATEST(COALESCE(last_used_at, 0), ?) + END, + last_used_ip = COALESCE(?, last_used_ip), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(delta.usage_count_delta) + .bind(last_used_at) + .bind(last_used_at) + .bind(delta.last_used_ip.as_deref()) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_api_key_last_used( + tx: &mut sqlx::Transaction<'_, MySql>, + target_id: &str, + delta: &ApiKeyLastUsedDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + "UPDATE api_keys SET last_used_at = GREATEST(COALESCE(last_used_at, 0), ?) WHERE id = ?", + ) + .bind(to_i64( + delta.last_used_at_unix_secs, + "api key last used at", + )?) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +fn merge_optional_max(target: &mut Option, value: Option) { + if let Some(value) = value { + if target.is_none_or(|current| value > current) { + *target = Some(value); + } + } +} + +fn merge_latest_timestamp_with_value( + target_timestamp: &mut Option, + target_value: &mut Option, + timestamp: Option, + value: Option, +) { + let Some(timestamp) = timestamp else { + return; + }; + if target_timestamp.is_none_or(|current| timestamp >= current) { + *target_timestamp = Some(timestamp); + if value + .as_deref() + .map(str::trim) + .is_some_and(|v| !v.is_empty()) + { + *target_value = value; + } + } +} + +fn finite_or_zero(value: f64) -> f64 { + if value.is_finite() { + value + } else { + 0.0 + } +} + +fn current_unix_secs() -> i64 { + chrono::Utc::now().timestamp().max(0) +} + +fn to_i64(value: u64, field: &str) -> Result { + i64::try_from(value) + .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds i64: {value}"))) +} + +fn optional_to_i64(value: Option, field: &str) -> Result, DataLayerError> { + value.map(|value| to_i64(value, field)).transpose() +} + +fn optional_u64(field: &str, value: Option) -> Result, DataLayerError> { + value + .map(|value| { + u64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!("{field} is negative: {value}")) + }) + }) + .transpose() +} + +fn nonnegative_u64(value: i64) -> u64 { + value.max(0) as u64 +} + +fn optional_nonnegative_u64(value: Option) -> Option { + value.map(nonnegative_u64) +} diff --git a/crates/aether-data/adapters/mysql/src/usage/http_capture.rs b/crates/aether-data/adapters/mysql/src/usage/http_capture.rs new file mode 100644 index 000000000..21a91bda5 --- /dev/null +++ b/crates/aether-data/adapters/mysql/src/usage/http_capture.rs @@ -0,0 +1,830 @@ +use std::io::{Read, Write}; + +use aether_data_contracts::repository::usage::{ + parse_usage_body_ref, usage_body_ref, StoredRequestUsageAudit, UpsertUsageRecord, + UsageBodyCaptureState, UsageBodyField, +}; +use flate2::{read::GzDecoder, write::GzEncoder, Compression}; +use serde_json::{Map, Value}; +use sqlx::{mysql::MySqlRow, Row}; + +use crate::error::SqlResultExt; +use crate::{DataLayerError, MysqlPool}; + +#[derive(Debug)] +pub(crate) struct PreparedUsageHttpCapture { + request_headers: Option, + provider_request_headers: Option, + response_headers: Option, + client_response_headers: Option, + request_body: PreparedBody, + provider_request_body: PreparedBody, + response_body: PreparedBody, + client_response_body: PreparedBody, + refs: HttpAuditRefs, + states: HttpAuditStates, + capture_mode: &'static str, +} + +#[derive(Debug)] +struct PreparedBody { + field: UsageBodyField, + payload_gzip: Option>, + clear_existing: bool, +} + +#[derive(Debug, Default)] +struct HttpAuditRefs { + request_body_ref: Option, + provider_request_body_ref: Option, + response_body_ref: Option, + client_response_body_ref: Option, +} + +impl HttpAuditRefs { + fn any_present(&self) -> bool { + self.request_body_ref.is_some() + || self.provider_request_body_ref.is_some() + || self.response_body_ref.is_some() + || self.client_response_body_ref.is_some() + } +} + +#[derive(Debug, Default)] +struct HttpAuditStates { + request_body_state: Option, + provider_request_body_state: Option, + response_body_state: Option, + client_response_body_state: Option, +} + +impl HttpAuditStates { + fn any_present(&self) -> bool { + self.request_body_state.is_some() + || self.provider_request_body_state.is_some() + || self.response_body_state.is_some() + || self.client_response_body_state.is_some() + } +} + +pub(crate) fn capture_update_allowed( + previous: Option<&StoredRequestUsageAudit>, + incoming_status: &str, +) -> bool { + let Some(previous) = previous else { + return true; + }; + if previous.billing_status != "pending" { + return false; + } + if matches!( + previous.status.as_str(), + "completed" | "failed" | "cancelled" + ) && matches!(incoming_status, "pending" | "streaming") + { + return false; + } + !(previous.status == "streaming" && incoming_status == "pending") +} + +pub(crate) fn apply_previous_metadata_tombstones( + usage: &mut UpsertUsageRecord, + previous: Option<&StoredRequestUsageAudit>, +) { + if usage.request_metadata.is_some() { + return; + } + let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None); + let clear_provider_request = + usage.provider_request_body_state == Some(UsageBodyCaptureState::None); + if !clear_request && !clear_provider_request { + return; + } + let mut metadata = previous + .and_then(|previous| previous.request_metadata.as_ref()) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + if clear_request { + metadata.remove("requested_reasoning_effort"); + metadata.remove("request_body_ref"); + } + if clear_provider_request { + metadata.remove("provider_reasoning_effort"); + metadata.remove("provider_service_tier"); + metadata.remove("provider_cache_ttl_minutes"); + metadata.remove("provider_request_body_ref"); + } + usage.request_metadata = Some(Value::Object(metadata)); +} + +pub(crate) fn prepare_usage_http_capture( + usage: &mut UpsertUsageRecord, +) -> Result { + let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None); + let clear_provider_request = + usage.provider_request_body_state == Some(UsageBodyCaptureState::None); + let clear_response = usage.response_body_state == Some(UsageBodyCaptureState::None); + let clear_client_response = + usage.client_response_body_state == Some(UsageBodyCaptureState::None); + + let request_body_value = (!clear_request) + .then_some(usage.request_body.as_ref()) + .flatten(); + let provider_request_body_value = (!clear_provider_request) + .then_some(usage.provider_request_body.as_ref()) + .flatten(); + let response_body_value = (!clear_response) + .then_some(usage.response_body.as_ref()) + .flatten(); + let client_response_body_value = (!clear_client_response) + .then_some(usage.client_response_body.as_ref()) + .flatten(); + + let request_body = prepare_body( + UsageBodyField::RequestBody, + request_body_value, + clear_request, + )?; + let provider_request_body = prepare_body( + UsageBodyField::ProviderRequestBody, + provider_request_body_value, + clear_provider_request, + )?; + let response_body = prepare_body( + UsageBodyField::ResponseBody, + response_body_value, + clear_response, + )?; + let client_response_body = prepare_body( + UsageBodyField::ClientResponseBody, + client_response_body_value, + clear_client_response, + )?; + + let refs = HttpAuditRefs { + request_body_ref: resolved_write_ref( + (!clear_request) + .then_some(usage.request_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::RequestBody, + request_body.payload_gzip.is_some(), + ), + provider_request_body_ref: resolved_write_ref( + (!clear_provider_request) + .then_some(usage.provider_request_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ProviderRequestBody, + provider_request_body.payload_gzip.is_some(), + ), + response_body_ref: resolved_write_ref( + (!clear_response) + .then_some(usage.response_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ResponseBody, + response_body.payload_gzip.is_some(), + ), + client_response_body_ref: resolved_write_ref( + (!clear_client_response) + .then_some(usage.client_response_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ClientResponseBody, + client_response_body.payload_gzip.is_some(), + ), + }; + let states = HttpAuditStates { + request_body_state: state_for_storage( + usage.request_body_state, + &request_body, + refs.request_body_ref.as_deref(), + ), + provider_request_body_state: state_for_storage( + usage.provider_request_body_state, + &provider_request_body, + refs.provider_request_body_ref.as_deref(), + ), + response_body_state: state_for_storage( + usage.response_body_state, + &response_body, + refs.response_body_ref.as_deref(), + ), + client_response_body_state: state_for_storage( + usage.client_response_body_state, + &client_response_body, + refs.client_response_body_ref.as_deref(), + ), + }; + + usage.request_metadata = prepare_metadata_for_body_storage( + usage.request_metadata.take(), + [ + ( + UsageBodyField::RequestBody, + request_body_value.is_some(), + usage.request_body_ref.as_deref(), + ), + ( + UsageBodyField::ProviderRequestBody, + provider_request_body_value.is_some(), + usage.provider_request_body_ref.as_deref(), + ), + ( + UsageBodyField::ResponseBody, + response_body_value.is_some(), + usage.response_body_ref.as_deref(), + ), + ( + UsageBodyField::ClientResponseBody, + client_response_body_value.is_some(), + usage.client_response_body_ref.as_deref(), + ), + ], + ); + + let capture_mode = if refs.any_present() { + "ref_backed" + } else if [ + request_body_value, + provider_request_body_value, + response_body_value, + client_response_body_value, + ] + .iter() + .any(Option::is_some) + { + "inline_legacy" + } else { + "none" + }; + + Ok(PreparedUsageHttpCapture { + request_headers: json_text(usage.request_headers.as_ref())?, + provider_request_headers: json_text(usage.provider_request_headers.as_ref())?, + response_headers: json_text(usage.response_headers.as_ref())?, + client_response_headers: json_text(usage.client_response_headers.as_ref())?, + request_body, + provider_request_body, + response_body, + client_response_body, + refs, + states, + capture_mode, + }) +} + +fn prepare_body( + field: UsageBodyField, + value: Option<&Value>, + clear_existing: bool, +) -> Result { + Ok(PreparedBody { + field, + payload_gzip: value.map(compress_json).transpose()?, + clear_existing, + }) +} + +fn compress_json(value: &Value) -> Result, DataLayerError> { + let bytes = serde_json::to_vec(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage body: {err}")) + })?; + let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6)); + encoder.write_all(&bytes).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to gzip usage body: {err}")) + })?; + encoder.finish().map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to finish usage body gzip: {err}")) + }) +} + +fn resolved_write_ref( + explicit_ref: Option<&str>, + request_id: &str, + field: UsageBodyField, + has_blob: bool, +) -> Option { + explicit_ref + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .or_else(|| has_blob.then(|| usage_body_ref(request_id, field))) +} + +fn state_for_storage( + incoming: Option, + body: &PreparedBody, + body_ref: Option<&str>, +) -> Option { + if matches!( + incoming, + Some( + UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable + | UsageBodyCaptureState::None + ) + ) { + return incoming; + } + if body.payload_gzip.is_some() || body_ref.is_some() { + return Some(UsageBodyCaptureState::Reference); + } + incoming +} + +fn json_text(value: Option<&Value>) -> Result, DataLayerError> { + value + .map(|value| { + serde_json::to_string(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage JSON: {err}")) + }) + }) + .transpose() +} + +fn prepare_metadata_for_body_storage( + metadata: Option, + body_fields: [(UsageBodyField, bool, Option<&str>); 4], +) -> Option { + let mut object = match metadata { + Some(Value::Object(object)) => object, + Some(value) => { + let mut object = Map::new(); + object.insert("request_metadata".to_string(), value); + object + } + None => Map::new(), + }; + let should_replace = !object.is_empty() + || body_fields + .iter() + .any(|(_, has_value, explicit_ref)| *has_value || explicit_ref.is_some()); + if !should_replace { + return None; + } + for (field, has_value, explicit_ref) in body_fields { + if has_value || explicit_ref.is_some() { + object.remove(field.as_ref_key()); + } + } + Some(Value::Object(object)) +} + +pub(crate) async fn sync_usage_http_capture( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, + prepared: &PreparedUsageHttpCapture, +) -> Result<(), DataLayerError> { + for body in [ + &prepared.request_body, + &prepared.provider_request_body, + &prepared.response_body, + &prepared.client_response_body, + ] { + sync_body(tx, request_id, body).await?; + } + let headers_present = prepared.request_headers.is_some() + || prepared.provider_request_headers.is_some() + || prepared.response_headers.is_some() + || prepared.client_response_headers.is_some(); + if !headers_present + && !prepared.refs.any_present() + && !prepared.states.any_present() + && prepared.capture_mode == "none" + { + return Ok(()); + } + + sqlx::query( + r#" +INSERT INTO usage_http_audits ( + request_id, + request_headers, + provider_request_headers, + response_headers, + client_response_headers, + request_body_ref, + provider_request_body_ref, + response_body_ref, + client_response_body_ref, + request_body_state, + provider_request_body_state, + response_body_state, + client_response_body_state, + body_capture_mode +) +VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + request_headers = COALESCE(VALUES(request_headers), request_headers), + provider_request_headers = COALESCE( + VALUES(provider_request_headers), + provider_request_headers + ), + response_headers = COALESCE(VALUES(response_headers), response_headers), + client_response_headers = COALESCE( + VALUES(client_response_headers), + client_response_headers + ), + request_body_ref = CASE + WHEN VALUES(request_body_state) = 'none' THEN NULL + ELSE COALESCE(VALUES(request_body_ref), request_body_ref) + END, + provider_request_body_ref = CASE + WHEN VALUES(provider_request_body_state) = 'none' THEN NULL + ELSE COALESCE(VALUES(provider_request_body_ref), provider_request_body_ref) + END, + response_body_ref = CASE + WHEN VALUES(response_body_state) = 'none' THEN NULL + ELSE COALESCE(VALUES(response_body_ref), response_body_ref) + END, + client_response_body_ref = CASE + WHEN VALUES(client_response_body_state) = 'none' THEN NULL + ELSE COALESCE(VALUES(client_response_body_ref), client_response_body_ref) + END, + request_body_state = COALESCE(VALUES(request_body_state), request_body_state), + provider_request_body_state = COALESCE( + VALUES(provider_request_body_state), + provider_request_body_state + ), + response_body_state = COALESCE(VALUES(response_body_state), response_body_state), + client_response_body_state = COALESCE( + VALUES(client_response_body_state), + client_response_body_state + ), + body_capture_mode = CASE + WHEN VALUES(body_capture_mode) = 'none' + AND ( + VALUES(request_body_state) = 'none' + OR VALUES(provider_request_body_state) = 'none' + OR VALUES(response_body_state) = 'none' + OR VALUES(client_response_body_state) = 'none' + ) + THEN 'none' + ELSE COALESCE( + NULLIF(VALUES(body_capture_mode), 'none'), + body_capture_mode, + 'none' + ) + END, + updated_at = UNIX_TIMESTAMP() +"#, + ) + .bind(request_id) + .bind(&prepared.request_headers) + .bind(&prepared.provider_request_headers) + .bind(&prepared.response_headers) + .bind(&prepared.client_response_headers) + .bind(prepared.refs.request_body_ref.as_deref()) + .bind(prepared.refs.provider_request_body_ref.as_deref()) + .bind(prepared.refs.response_body_ref.as_deref()) + .bind(prepared.refs.client_response_body_ref.as_deref()) + .bind( + prepared + .states + .request_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .provider_request_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .response_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .client_response_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind(prepared.capture_mode) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn sync_body( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, + body: &PreparedBody, +) -> Result<(), DataLayerError> { + let body_ref = usage_body_ref(request_id, body.field); + if body.clear_existing || body.payload_gzip.is_some() { + sqlx::query(clear_legacy_body_sql(body.field)) + .bind(request_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + if body.clear_existing { + sqlx::query("DELETE FROM usage_body_blobs WHERE body_ref = ?") + .bind(body_ref) + .execute(&mut **tx) + .await + .map_sql_err()?; + return Ok(()); + } + if let Some(payload_gzip) = body.payload_gzip.as_deref() { + sqlx::query( + r#" +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES (?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + request_id = VALUES(request_id), + body_field = VALUES(body_field), + payload_gzip = VALUES(payload_gzip), + updated_at = UNIX_TIMESTAMP() +"#, + ) + .bind(body_ref) + .bind(request_id) + .bind(body.field.as_storage_field()) + .bind(payload_gzip) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + Ok(()) +} + +fn clear_legacy_body_sql(field: UsageBodyField) -> &'static str { + match field { + UsageBodyField::RequestBody => { + "UPDATE `usage` SET request_body = NULL, request_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ProviderRequestBody => { + "UPDATE `usage` SET provider_request_body = NULL, provider_request_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ResponseBody => { + "UPDATE `usage` SET response_body = NULL, response_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ClientResponseBody => { + "UPDATE `usage` SET client_response_body = NULL, client_response_body_compressed = NULL WHERE request_id = ?" + } + } +} + +pub(crate) fn hydrate_usage_row( + row: &MySqlRow, + usage: &mut StoredRequestUsageAudit, + resolve_legacy_compressed: bool, +) -> Result<(), DataLayerError> { + usage.request_headers = optional_json(row, "request_headers")?; + usage.provider_request_headers = optional_json(row, "provider_request_headers")?; + usage.response_headers = optional_json(row, "response_headers")?; + usage.client_response_headers = optional_json(row, "client_response_headers")?; + + let request_body = legacy_body_column( + row, + "request_body", + "request_body_compressed", + resolve_legacy_compressed, + )?; + let provider_request_body = legacy_body_column( + row, + "provider_request_body", + "provider_request_body_compressed", + resolve_legacy_compressed, + )?; + let response_body = legacy_body_column( + row, + "response_body", + "response_body_compressed", + resolve_legacy_compressed, + )?; + let client_response_body = legacy_body_column( + row, + "client_response_body", + "client_response_body_compressed", + resolve_legacy_compressed, + )?; + usage.request_body = request_body.0; + usage.provider_request_body = provider_request_body.0; + usage.response_body = response_body.0; + usage.client_response_body = client_response_body.0; + + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + usage.request_body_ref = resolved_read_ref( + row.try_get("http_request_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::RequestBody, + request_body.1, + ); + usage.provider_request_body_ref = resolved_read_ref( + row.try_get("http_provider_request_body_ref") + .map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ProviderRequestBody, + provider_request_body.1, + ); + usage.response_body_ref = resolved_read_ref( + row.try_get("http_response_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ResponseBody, + response_body.1, + ); + usage.client_response_body_ref = resolved_read_ref( + row.try_get("http_client_response_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ClientResponseBody, + client_response_body.1, + ); + usage.request_body_state = optional_state(row, "http_request_body_state")?; + usage.provider_request_body_state = optional_state(row, "http_provider_request_body_state")?; + usage.response_body_state = optional_state(row, "http_response_body_state")?; + usage.client_response_body_state = optional_state(row, "http_client_response_body_state")?; + Ok(()) +} + +fn optional_json(row: &MySqlRow, column: &str) -> Result, DataLayerError> { + row.try_get::, _>(column) + .map_sql_err()? + .map(|raw| { + serde_json::from_str(&raw).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid usage JSON in {column}: {err}")) + }) + }) + .transpose() +} + +fn legacy_body_column( + row: &MySqlRow, + inline_column: &str, + compressed_column: &str, + resolve_compressed: bool, +) -> Result<(Option, bool), DataLayerError> { + let inline = optional_json(row, inline_column)?; + if inline.is_some() { + return Ok((inline, false)); + } + let compressed = row + .try_get::>, _>(compressed_column) + .map_sql_err()?; + let has_compressed = compressed.is_some(); + let value = if resolve_compressed { + compressed.map(|bytes| inflate_json(&bytes)).transpose()? + } else { + None + }; + Ok((value, has_compressed)) +} + +fn resolved_read_ref( + audit_ref: Option, + metadata: Option<&Map>, + request_id: &str, + field: UsageBodyField, + has_compressed: bool, +) -> Option { + audit_ref + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .or_else(|| has_compressed.then(|| usage_body_ref(request_id, field))) + .or_else(|| metadata_body_ref(metadata, request_id, field)) +} + +fn metadata_body_ref( + metadata: Option<&Map>, + request_id: &str, + field: UsageBodyField, +) -> Option { + metadata + .and_then(|metadata| metadata.get(field.as_ref_key())) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(parse_usage_body_ref) + .filter(|(parsed_request_id, parsed_field)| { + parsed_request_id == request_id && *parsed_field == field + }) + .map(|(parsed_request_id, parsed_field)| usage_body_ref(&parsed_request_id, parsed_field)) +} + +fn optional_state( + row: &MySqlRow, + column: &str, +) -> Result, DataLayerError> { + Ok(row + .try_get::, _>(column) + .map_sql_err()? + .as_deref() + .and_then(parse_state)) +} + +fn parse_state(value: &str) -> Option { + match value.trim() { + "none" => Some(UsageBodyCaptureState::None), + "inline" => Some(UsageBodyCaptureState::Inline), + "reference" => Some(UsageBodyCaptureState::Reference), + "truncated" => Some(UsageBodyCaptureState::Truncated), + "disabled" => Some(UsageBodyCaptureState::Disabled), + "unavailable" => Some(UsageBodyCaptureState::Unavailable), + _ => None, + } +} + +pub(crate) async fn hydrate_usage_body_refs( + pool: &MysqlPool, + mut usage: StoredRequestUsageAudit, +) -> Result { + for field in [ + UsageBodyField::RequestBody, + UsageBodyField::ProviderRequestBody, + UsageBodyField::ResponseBody, + UsageBodyField::ClientResponseBody, + ] { + if usage.body_value(field).is_some() { + continue; + } + let Some(body_ref) = usage.body_ref(field) else { + continue; + }; + let value = resolve_body_ref(pool, body_ref).await?; + match field { + UsageBodyField::RequestBody => usage.request_body = value, + UsageBodyField::ProviderRequestBody => usage.provider_request_body = value, + UsageBodyField::ResponseBody => usage.response_body = value, + UsageBodyField::ClientResponseBody => usage.client_response_body = value, + } + } + Ok(usage) +} + +pub(crate) async fn resolve_body_ref( + pool: &MysqlPool, + body_ref: &str, +) -> Result, DataLayerError> { + if let Some(payload_gzip) = sqlx::query_scalar::<_, Vec>( + "SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? LIMIT 1", + ) + .bind(body_ref) + .fetch_optional(pool) + .await + .map_sql_err()? + { + return inflate_json(&payload_gzip).map(Some); + } + let Some((request_id, field)) = parse_usage_body_ref(body_ref) else { + return Ok(None); + }; + let (inline_column, compressed_column) = usage_body_sql_columns(field); + let row = sqlx::query(&format!( + "SELECT CAST({inline_column} AS CHAR) AS inline_body, {compressed_column} AS compressed_body FROM `usage` WHERE request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(pool) + .await + .map_sql_err()?; + let Some(row) = row.as_ref() else { + return Ok(None); + }; + if let Some(raw) = row + .try_get::, _>("inline_body") + .map_sql_err()? + { + return serde_json::from_str(&raw).map(Some).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid inline usage body JSON: {err}")) + }); + } + row.try_get::>, _>("compressed_body") + .map_sql_err()? + .map(|bytes| inflate_json(&bytes)) + .transpose() +} + +fn usage_body_sql_columns(field: UsageBodyField) -> (&'static str, &'static str) { + match field { + UsageBodyField::RequestBody => ("request_body", "request_body_compressed"), + UsageBodyField::ProviderRequestBody => { + ("provider_request_body", "provider_request_body_compressed") + } + UsageBodyField::ResponseBody => ("response_body", "response_body_compressed"), + UsageBodyField::ClientResponseBody => { + ("client_response_body", "client_response_body_compressed") + } + } +} + +fn inflate_json(bytes: &[u8]) -> Result { + let mut decoder = GzDecoder::new(bytes); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to decompress usage body: {err}")) + })?; + serde_json::from_slice(&decoded).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to decode usage body JSON: {err}")) + }) +} diff --git a/crates/aether-data/adapters/mysql/src/usage/read.rs b/crates/aether-data/adapters/mysql/src/usage/read.rs new file mode 100644 index 000000000..bca791beb --- /dev/null +++ b/crates/aether-data/adapters/mysql/src/usage/read.rs @@ -0,0 +1,811 @@ +use aether_data_contracts::repository::usage::{ + ProviderApiKeyWindowUsageRequest, StoredRequestUsageAudit, UsageAuditKeywordSearchQuery, + UsageAuditListQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, +}; +use aether_data_contracts::DataLayerError; +use sqlx::{MySql, QueryBuilder}; + +use crate::error::SqlResultExt; + +use super::{http_capture, map_usage_row, row_u64, to_i64, MysqlUsageStorage, USAGE_COLUMNS}; + +const EFFECTIVE_PROVIDER_ID_EXPR: &str = r#"CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_id + ELSE `usage`.provider_id +END"#; + +const EFFECTIVE_PROVIDER_API_KEY_ID_EXPR: &str = r#"CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_api_key_id + ELSE `usage`.provider_api_key_id +END"#; + +const MONITORING_ERROR_PREDICATE: &str = r#"( + LOWER(TRIM(COALESCE(`usage`.status, ''))) IN ('failed', 'error') + OR (`usage`.error_category IS NOT NULL AND TRIM(`usage`.error_category) <> '') + OR ( + TRIM(COALESCE(`usage`.status, '')) = '' + AND ( + COALESCE(`usage`.status_code, 0) >= 400 + OR (`usage`.error_message IS NOT NULL AND TRIM(`usage`.error_message) <> '') + ) + ) +)"#; + +/// A SQL-side superset filter used before the runtime applies complex usage analytics. +/// +/// Every scan has explicit time bounds; an empty range returns no rows. Optional dimensions +/// further reduce the rows sent to the in-memory analytics implementation without changing its +/// result. +#[derive(Debug, Clone)] +pub struct MysqlUsageReadFilter { + created_from_unix_secs: u64, + created_until_unix_secs: u64, + user_id: Option, + api_key_id: Option, + provider_name: Option, + provider_id: Option, + model: Option, + api_format: Option, + endpoint_kind: Option, + is_stream: Option, + has_format_conversion: Option, + finalized_only: bool, + completed_only: bool, +} + +impl MysqlUsageReadFilter { + pub fn new(created_from_unix_secs: u64, created_until_unix_secs: u64) -> Self { + Self { + created_from_unix_secs, + created_until_unix_secs, + user_id: None, + api_key_id: None, + provider_name: None, + provider_id: None, + model: None, + api_format: None, + endpoint_kind: None, + is_stream: None, + has_format_conversion: None, + finalized_only: false, + completed_only: false, + } + } + + pub fn with_user_id(mut self, value: Option<&str>) -> Self { + self.user_id = value.map(ToOwned::to_owned); + self + } + + pub fn with_api_key_id(mut self, value: Option<&str>) -> Self { + self.api_key_id = value.map(ToOwned::to_owned); + self + } + + pub fn with_provider_name(mut self, value: Option<&str>) -> Self { + self.provider_name = value.map(ToOwned::to_owned); + self + } + + pub fn with_provider_id(mut self, value: Option<&str>) -> Self { + self.provider_id = value.map(ToOwned::to_owned); + self + } + + pub fn with_model(mut self, value: Option<&str>) -> Self { + self.model = value.map(ToOwned::to_owned); + self + } + + pub fn with_api_format(mut self, value: Option<&str>) -> Self { + self.api_format = value.map(ToOwned::to_owned); + self + } + + pub fn with_endpoint_kind(mut self, value: Option<&str>) -> Self { + self.endpoint_kind = value.map(ToOwned::to_owned); + self + } + + pub fn with_is_stream(mut self, value: Option) -> Self { + self.is_stream = value; + self + } + + pub fn with_has_format_conversion(mut self, value: Option) -> Self { + self.has_format_conversion = value; + self + } + + pub fn finalized_only(mut self) -> Self { + self.finalized_only = true; + self + } + + pub fn completed_only(mut self) -> Self { + self.completed_only = true; + self + } + + fn is_empty(&self) -> bool { + self.created_from_unix_secs >= self.created_until_unix_secs + } +} + +impl MysqlUsageStorage { + pub async fn find_by_id( + &self, + id: &str, + ) -> Result, DataLayerError> { + let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE `usage`.id = ? LIMIT 1")) + .bind(id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref() + .map(|row| map_usage_row(row, false)) + .transpose() + } + + pub async fn list_by_ids( + &self, + ids: &[String], + ) -> Result, DataLayerError> { + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder.push(" WHERE `usage`.id IN ("); + { + let mut separated = builder.separated(", "); + for id in ids { + separated.push_bind(id.clone()); + } + } + builder.push(") ORDER BY created_at_unix_ms DESC, `usage`.id ASC"); + self.fetch_usage_items(builder).await + } + + pub async fn find_by_request_id( + &self, + request_id: &str, + ) -> Result, DataLayerError> { + let row = sqlx::query(&format!( + "{USAGE_COLUMNS} WHERE `usage`.request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + let usage = row + .as_ref() + .map(|row| map_usage_row(row, true)) + .transpose()?; + match usage { + Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage) + .await + .map(Some), + None => Ok(None), + } + } + + pub async fn resolve_body_ref( + &self, + body_ref: &str, + ) -> Result, DataLayerError> { + http_capture::resolve_body_ref(&self.pool, body_ref).await + } + + pub async fn list_usage_audits( + &self, + query: &UsageAuditListQuery, + ) -> Result, DataLayerError> { + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + let mut has_where = false; + push_list_filters(&mut builder, query, &mut has_where)?; + push_order_limit_offset(&mut builder, query.newest_first, query.limit, query.offset)?; + self.fetch_usage_items(builder).await + } + + pub async fn count_usage_audits( + &self, + query: &UsageAuditListQuery, + ) -> Result { + let mut builder = + QueryBuilder::::new("SELECT CAST(COUNT(*) AS SIGNED) AS total FROM `usage`"); + let mut has_where = false; + push_list_filters(&mut builder, query, &mut has_where)?; + let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; + row_u64(&row, "total") + } + + pub async fn list_usage_audits_by_keyword_search( + &self, + query: &UsageAuditKeywordSearchQuery, + ) -> Result, DataLayerError> { + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + let mut has_where = false; + push_keyword_filters(&mut builder, query, &mut has_where)?; + push_order_limit_offset(&mut builder, query.newest_first, query.limit, query.offset)?; + self.fetch_usage_items(builder).await + } + + pub async fn count_usage_audits_by_keyword_search( + &self, + query: &UsageAuditKeywordSearchQuery, + ) -> Result { + let mut builder = + QueryBuilder::::new("SELECT CAST(COUNT(*) AS SIGNED) AS total FROM `usage`"); + let mut has_where = false; + push_keyword_filters(&mut builder, query, &mut has_where)?; + let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; + row_u64(&row, "total") + } + + pub async fn load_usage_records_in_range( + &self, + filter: &MysqlUsageReadFilter, + ) -> Result, DataLayerError> { + if filter.is_empty() { + return Ok(Vec::new()); + } + let mut builder = build_range_query(filter)?; + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(|row| map_usage_row(row, false)).collect() + } + + pub async fn count_monitoring_usage_errors( + &self, + query: &UsageMonitoringErrorCountQuery, + ) -> Result { + if query.created_from_unix_secs >= query.created_until_unix_secs { + return Ok(0); + } + let row = sqlx::query(&format!( + r#" +SELECT CAST(COUNT(*) AS SIGNED) AS total +FROM `usage` +WHERE created_at_unix_ms >= ? + AND created_at_unix_ms < ? + AND {MONITORING_ERROR_PREDICATE} +"# + )) + .bind(to_i64( + query.created_from_unix_secs, + "usage.created_at_unix_ms", + )?) + .bind(to_i64( + query.created_until_unix_secs, + "usage.created_at_unix_ms", + )?) + .fetch_one(&self.pool) + .await + .map_sql_err()?; + row_u64(&row, "total") + } + + pub async fn list_monitoring_usage_errors( + &self, + query: &UsageMonitoringErrorListQuery, + ) -> Result, DataLayerError> { + if query.created_from_unix_secs >= query.created_until_unix_secs { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder + .push(" WHERE `usage`.created_at_unix_ms >= ") + .push_bind(to_i64( + query.created_from_unix_secs, + "usage.created_at_unix_ms", + )?) + .push(" AND `usage`.created_at_unix_ms < ") + .push_bind(to_i64( + query.created_until_unix_secs, + "usage.created_at_unix_ms", + )?) + .push(" AND ") + .push(MONITORING_ERROR_PREDICATE) + .push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC"); + if let Some(limit) = query.limit { + builder + .push(" LIMIT ") + .push_bind(usize_to_i64(limit, "usage monitoring limit")?); + } + self.fetch_usage_items(builder).await + } + + pub async fn list_recent_usage_audits( + &self, + user_id: Option<&str>, + limit: usize, + ) -> Result, DataLayerError> { + if limit == 0 { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + if let Some(user_id) = user_id { + builder + .push(" WHERE `usage`.user_id = ") + .push_bind(user_id.to_string()); + } + builder + .push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC LIMIT ") + .push_bind(usize_to_i64(limit, "recent usage limit")?); + self.fetch_usage_items(builder).await + } + + pub async fn load_usage_records_by_api_key_ids( + &self, + api_key_ids: &[String], + ) -> Result, DataLayerError> { + if api_key_ids.is_empty() { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder.push(" WHERE `usage`.api_key_id IN ("); + push_string_list(&mut builder, api_key_ids); + builder.push(") ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC"); + self.fetch_usage_items(builder).await + } + + pub async fn load_usage_records_by_provider_api_key_ids( + &self, + provider_api_key_ids: &[String], + ) -> Result, DataLayerError> { + if provider_api_key_ids.is_empty() { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder + .push(" WHERE (") + .push(EFFECTIVE_PROVIDER_API_KEY_ID_EXPR) + .push(") IN ("); + push_string_list(&mut builder, provider_api_key_ids); + builder.push(") ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC"); + self.fetch_usage_items(builder).await + } + + pub async fn load_usage_records_by_provider_api_key_windows( + &self, + requests: &[ProviderApiKeyWindowUsageRequest], + ) -> Result, DataLayerError> { + if requests.is_empty() { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder.push(" WHERE "); + for request in requests { + let provider_api_key_id = request.provider_api_key_id.trim(); + if provider_api_key_id.is_empty() { + return Err(DataLayerError::InvalidInput( + "provider api key window usage provider_api_key_id cannot be empty".to_string(), + )); + } + let window_code = request.window_code.trim(); + if window_code.is_empty() { + return Err(DataLayerError::InvalidInput( + "provider api key window usage window_code cannot be empty".to_string(), + )); + } + if request.start_unix_secs >= request.end_unix_secs { + return Err(DataLayerError::InvalidInput( + "provider api key window usage range must be non-empty".to_string(), + )); + } + } + { + let mut separated = builder.separated(" OR "); + for request in requests { + separated + .push("((") + .push(EFFECTIVE_PROVIDER_API_KEY_ID_EXPR) + .push(") = ") + .push_bind(request.provider_api_key_id.trim().to_string()) + .push(" AND `usage`.created_at_unix_ms >= ") + .push_bind(to_i64(request.start_unix_secs, "usage.created_at_unix_ms")?) + .push(" AND `usage`.created_at_unix_ms < ") + .push_bind(to_i64(request.end_unix_secs, "usage.created_at_unix_ms")?) + .push(")"); + } + } + builder.push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC"); + self.fetch_usage_items(builder).await + } + + pub async fn load_usage_records_for_provider_since( + &self, + provider_id: &str, + since_unix_secs: u64, + ) -> Result, DataLayerError> { + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder + .push(" WHERE (") + .push(EFFECTIVE_PROVIDER_ID_EXPR) + .push(") = ") + .push_bind(provider_id.to_string()) + .push(" AND `usage`.created_at_unix_ms >= ") + .push_bind(to_i64(since_unix_secs, "usage.created_at_unix_ms")?) + .push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC"); + self.fetch_usage_items(builder).await + } + + async fn fetch_usage_items( + &self, + mut builder: QueryBuilder<'_, MySql>, + ) -> Result, DataLayerError> { + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(|row| map_usage_row(row, false)).collect() + } +} + +fn build_range_query( + filter: &MysqlUsageReadFilter, +) -> Result, DataLayerError> { + let mut builder = QueryBuilder::::new(USAGE_COLUMNS); + builder + .push(" WHERE `usage`.created_at_unix_ms >= ") + .push_bind(to_i64( + filter.created_from_unix_secs, + "usage.created_at_unix_ms", + )?) + .push(" AND `usage`.created_at_unix_ms < ") + .push_bind(to_i64( + filter.created_until_unix_secs, + "usage.created_at_unix_ms", + )?); + push_optional_text_filter(&mut builder, "`usage`.user_id", filter.user_id.as_deref()); + push_optional_text_filter( + &mut builder, + "`usage`.api_key_id", + filter.api_key_id.as_deref(), + ); + push_optional_text_filter( + &mut builder, + "`usage`.provider_name", + filter.provider_name.as_deref(), + ); + if let Some(provider_id) = filter.provider_id.as_deref() { + builder + .push(" AND (") + .push(EFFECTIVE_PROVIDER_ID_EXPR) + .push(") = ") + .push_bind(provider_id.to_string()); + } + push_optional_text_filter(&mut builder, "`usage`.model", filter.model.as_deref()); + push_optional_text_filter( + &mut builder, + "`usage`.api_format", + filter.api_format.as_deref(), + ); + push_optional_text_filter( + &mut builder, + "`usage`.endpoint_kind", + filter.endpoint_kind.as_deref(), + ); + if let Some(is_stream) = filter.is_stream { + builder + .push(" AND `usage`.is_stream = ") + .push_bind(is_stream); + } + if let Some(has_format_conversion) = filter.has_format_conversion { + builder + .push(" AND CASE WHEN usage_routing_snapshots.request_id IS NOT NULL ") + .push("THEN COALESCE(usage_routing_snapshots.has_format_conversion, FALSE) ") + .push("ELSE COALESCE(`usage`.has_format_conversion, FALSE) END = ") + .push_bind(has_format_conversion); + } + if filter.finalized_only { + builder.push( + " AND `usage`.status NOT IN ('pending', 'streaming') \ +AND `usage`.provider_name NOT IN ('unknown', 'pending')", + ); + } + if filter.completed_only { + builder.push(" AND `usage`.status = 'completed'"); + } + builder.push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC"); + Ok(builder) +} + +fn push_list_filters( + builder: &mut QueryBuilder<'_, MySql>, + query: &UsageAuditListQuery, + has_where: &mut bool, +) -> Result<(), DataLayerError> { + if let Some(value) = query.created_from_unix_secs { + push_where(builder, has_where); + builder + .push("`usage`.created_at_unix_ms >= ") + .push_bind(to_i64(value, "usage.created_at_unix_ms")?); + } + if let Some(value) = query.created_until_unix_secs { + push_where(builder, has_where); + builder + .push("`usage`.created_at_unix_ms < ") + .push_bind(to_i64(value, "usage.created_at_unix_ms")?); + } + for (column, value) in [ + ("`usage`.user_id", query.user_id.as_deref()), + ("`usage`.provider_name", query.provider_name.as_deref()), + ("`usage`.model", query.model.as_deref()), + ("`usage`.api_format", query.api_format.as_deref()), + ] { + if let Some(value) = value { + push_where(builder, has_where); + builder + .push(column) + .push(" = ") + .push_bind(value.to_string()); + } + } + if let Some(client_family) = query + .client_family + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + push_where(builder, has_where); + builder + .push("LOWER(COALESCE(NULLIF(TRIM(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.client_session_affinity.client_family'))), ''), NULLIF(TRIM(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.client_family'))), ''))) = ") + .push_bind(client_family.to_ascii_lowercase()); + } + if query.exclude_unknown_model_or_provider { + push_where(builder, has_where); + builder.push( + "(LOWER(TRIM(COALESCE(`usage`.model, ''))) NOT IN ('unknown', 'unknow') \ +AND LOWER(TRIM(COALESCE(`usage`.provider_name, ''))) NOT IN ('unknown', 'unknow'))", + ); + } + if let Some(statuses) = query + .statuses + .as_deref() + .filter(|values| !values.is_empty()) + { + push_where(builder, has_where); + builder.push("`usage`.status IN ("); + push_string_list(builder, statuses); + builder.push(")"); + } + if !query.exclude_status_codes.is_empty() { + push_where(builder, has_where); + builder.push("(`usage`.status_code IS NULL OR `usage`.status_code NOT IN ("); + { + let mut separated = builder.separated(", "); + for status_code in &query.exclude_status_codes { + separated.push_bind(i64::from(*status_code)); + } + } + builder.push("))"); + } + if let Some(is_stream) = query.is_stream { + push_where(builder, has_where); + builder.push("`usage`.is_stream = ").push_bind(is_stream); + } + if query.error_only { + push_where(builder, has_where); + builder.push( + "(`usage`.status = 'failed' \ +OR COALESCE(`usage`.status_code, 0) >= 400 \ +OR (`usage`.error_message IS NOT NULL AND TRIM(`usage`.error_message) <> ''))", + ); + } + Ok(()) +} + +fn push_keyword_filters( + builder: &mut QueryBuilder<'_, MySql>, + query: &UsageAuditKeywordSearchQuery, + has_where: &mut bool, +) -> Result<(), DataLayerError> { + push_list_filters( + builder, + &UsageAuditListQuery { + created_from_unix_secs: query.created_from_unix_secs, + created_until_unix_secs: query.created_until_unix_secs, + user_id: query.user_id.clone(), + provider_name: query.provider_name.clone(), + model: query.model.clone(), + api_format: query.api_format.clone(), + client_family: query.client_family.clone(), + exclude_unknown_model_or_provider: query.exclude_unknown_model_or_provider, + statuses: query.statuses.clone(), + exclude_status_codes: query.exclude_status_codes.clone(), + is_stream: query.is_stream, + error_only: query.error_only, + limit: None, + offset: None, + newest_first: query.newest_first, + }, + has_where, + )?; + + for (index, keyword) in query.keywords.iter().enumerate() { + let keyword = keyword.trim(); + if keyword.is_empty() { + continue; + } + let pattern = format!("%{}%", keyword.to_ascii_lowercase()); + push_where(builder, has_where); + builder + .push("(LOWER(COALESCE(`usage`.model, '')) LIKE ") + .push_bind(pattern.clone()) + .push(" OR LOWER(COALESCE(`usage`.provider_name, '')) LIKE ") + .push_bind(pattern.clone()); + if query.auth_user_reader_available { + if let Some(ids) = query + .matched_user_ids_by_keyword + .get(index) + .filter(|ids| !ids.is_empty()) + { + builder.push(" OR `usage`.user_id IN ("); + push_string_list(builder, ids); + builder.push(")"); + } + } else { + builder + .push(" OR LOWER(COALESCE(`usage`.username, '')) LIKE ") + .push_bind(pattern.clone()); + } + if query.auth_api_key_reader_available { + if let Some(ids) = query + .matched_api_key_ids_by_keyword + .get(index) + .filter(|ids| !ids.is_empty()) + { + builder.push(" OR `usage`.api_key_id IN ("); + push_string_list(builder, ids); + builder.push(")"); + } + } else { + builder + .push(" OR LOWER(COALESCE(`usage`.api_key_name, '')) LIKE ") + .push_bind(pattern); + } + builder.push(")"); + } + + if let Some(username_keyword) = query + .username_keyword + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + push_where(builder, has_where); + if query.auth_user_reader_available { + if query.matched_user_ids_for_username.is_empty() { + builder.push("FALSE"); + } else { + builder.push("`usage`.user_id IN ("); + push_string_list(builder, &query.matched_user_ids_for_username); + builder.push(")"); + } + } else { + builder + .push("LOWER(COALESCE(`usage`.username, '')) LIKE ") + .push_bind(format!("%{}%", username_keyword.to_ascii_lowercase())); + } + } + Ok(()) +} + +fn push_order_limit_offset( + builder: &mut QueryBuilder<'_, MySql>, + newest_first: bool, + limit: Option, + offset: Option, +) -> Result<(), DataLayerError> { + if newest_first { + builder.push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC"); + } else { + builder.push(" ORDER BY created_at_unix_ms ASC, `usage`.request_id ASC"); + } + match (limit, offset) { + (Some(limit), offset) => { + builder + .push(" LIMIT ") + .push_bind(usize_to_i64(limit, "usage list limit")?); + if let Some(offset) = offset { + builder + .push(" OFFSET ") + .push_bind(usize_to_i64(offset, "usage list offset")?); + } + } + (None, Some(offset)) => { + builder + .push(" LIMIT 18446744073709551615 OFFSET ") + .push_bind(usize_to_i64(offset, "usage list offset")?); + } + (None, None) => {} + } + Ok(()) +} + +fn push_where(builder: &mut QueryBuilder<'_, MySql>, has_where: &mut bool) { + builder.push(if *has_where { " AND " } else { " WHERE " }); + *has_where = true; +} + +fn push_optional_text_filter( + builder: &mut QueryBuilder<'_, MySql>, + column: &'static str, + value: Option<&str>, +) { + if let Some(value) = value { + builder + .push(" AND ") + .push(column) + .push(" = ") + .push_bind(value.to_string()); + } +} + +fn push_string_list(builder: &mut QueryBuilder<'_, MySql>, values: &[String]) { + let mut separated = builder.separated(", "); + for value in values { + separated.push_bind(value.clone()); + } +} + +fn usize_to_i64(value: usize, field: &str) -> Result { + i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow"))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn range_reads_always_have_both_time_bounds() { + let filter = MysqlUsageReadFilter::new(100, 200) + .with_user_id(Some("user-1")) + .finalized_only(); + let query = build_range_query(&filter).expect("range query should build"); + let sql = query.sql(); + assert!(sql.contains("created_at_unix_ms >= ?")); + assert!(sql.contains("created_at_unix_ms < ?")); + assert!(sql.contains("`usage`.user_id = ?")); + assert!(sql.contains("status NOT IN ('pending', 'streaming')")); + } + + #[test] + fn audit_reads_keep_pagination_in_mysql() { + let mut query = QueryBuilder::::new(USAGE_COLUMNS); + push_order_limit_offset(&mut query, true, Some(25), Some(50)) + .expect("pagination should build"); + assert!(query.sql().contains("LIMIT ? OFFSET ?")); + } + + #[test] + fn offset_only_uses_mysql_unbounded_limit_syntax() { + let mut query = QueryBuilder::::new(USAGE_COLUMNS); + push_order_limit_offset(&mut query, false, None, Some(50)) + .expect("pagination should build"); + assert!(query.sql().contains("LIMIT 18446744073709551615 OFFSET ?")); + } + + #[test] + fn usage_projection_and_legacy_keyword_search_keep_snapshot_names() { + assert!(USAGE_COLUMNS.contains("`usage`.username")); + assert!(USAGE_COLUMNS.contains("`usage`.api_key_name")); + + let mut query = QueryBuilder::::new(USAGE_COLUMNS); + let mut has_where = false; + push_keyword_filters( + &mut query, + &UsageAuditKeywordSearchQuery { + keywords: vec!["legacy".to_string()], + ..UsageAuditKeywordSearchQuery::default() + }, + &mut has_where, + ) + .expect("keyword query should build"); + assert!(query + .sql() + .contains("LOWER(COALESCE(`usage`.username, '')) LIKE ?")); + assert!(query + .sql() + .contains("LOWER(COALESCE(`usage`.api_key_name, '')) LIKE ?")); + } +} diff --git a/crates/aether-data/adapters/mysql/src/usage/snapshots.rs b/crates/aether-data/adapters/mysql/src/usage/snapshots.rs new file mode 100644 index 000000000..a929c089a --- /dev/null +++ b/crates/aether-data/adapters/mysql/src/usage/snapshots.rs @@ -0,0 +1,1043 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord}; +use aether_data_contracts::DataLayerError; +use serde_json::{Map, Value}; +use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row, Transaction}; + +use crate::error::SqlResultExt; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct UsageRoutingSnapshot { + candidate_id: Option, + candidate_index: Option, + key_name: Option, + planner_kind: Option, + route_family: Option, + route_kind: Option, + execution_path: Option, + local_execution_runtime_miss_reason: Option, + selected_provider_id: Option, + selected_endpoint_id: Option, + selected_provider_api_key_id: Option, + has_format_conversion: Option, +} + +impl UsageRoutingSnapshot { + fn has_metadata_fields(&self) -> bool { + self.candidate_id.is_some() + || self.candidate_index.is_some() + || self.key_name.is_some() + || self.planner_kind.is_some() + || self.route_family.is_some() + || self.route_kind.is_some() + || self.execution_path.is_some() + || self.local_execution_runtime_miss_reason.is_some() + } + + fn any_present(&self) -> bool { + self.has_metadata_fields() + || self.selected_provider_id.is_some() + || self.selected_endpoint_id.is_some() + || self.selected_provider_api_key_id.is_some() + || self.has_format_conversion.is_some() + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub(crate) struct UsageSettlementPricingSnapshot { + billing_status: Option, + billing_snapshot_schema_version: Option, + billing_snapshot_status: Option, + settlement_snapshot_schema_version: Option, + settlement_snapshot: Option, + billing_dimensions: Option, + billing_input_tokens: Option, + billing_effective_input_tokens: Option, + billing_output_tokens: Option, + billing_cache_creation_tokens: Option, + billing_cache_creation_5m_tokens: Option, + billing_cache_creation_1h_tokens: Option, + billing_cache_read_tokens: Option, + billing_total_input_context: Option, + billing_cache_creation_cost_usd: Option, + billing_cache_read_cost_usd: Option, + billing_total_cost_usd: Option, + billing_actual_total_cost_usd: Option, + billing_pricing_source: Option, + billing_rule_id: Option, + billing_rule_version: Option, + rate_multiplier: Option, + is_free_tier: Option, + input_price_per_1m: Option, + output_price_per_1m: Option, + cache_creation_price_per_1m: Option, + cache_read_price_per_1m: Option, + price_per_request: Option, +} + +impl UsageSettlementPricingSnapshot { + fn any_present(&self) -> bool { + self.billing_status.is_some() + || self.billing_snapshot_schema_version.is_some() + || self.billing_snapshot_status.is_some() + || self.settlement_snapshot_schema_version.is_some() + || self.settlement_snapshot.is_some() + || self.billing_dimensions.is_some() + || self.billing_input_tokens.is_some() + || self.billing_effective_input_tokens.is_some() + || self.billing_output_tokens.is_some() + || self.billing_cache_creation_tokens.is_some() + || self.billing_cache_creation_5m_tokens.is_some() + || self.billing_cache_creation_1h_tokens.is_some() + || self.billing_cache_read_tokens.is_some() + || self.billing_total_input_context.is_some() + || self.billing_cache_creation_cost_usd.is_some() + || self.billing_cache_read_cost_usd.is_some() + || self.billing_total_cost_usd.is_some() + || self.billing_actual_total_cost_usd.is_some() + || self.billing_pricing_source.is_some() + || self.billing_rule_id.is_some() + || self.billing_rule_version.is_some() + || self.rate_multiplier.is_some() + || self.is_free_tier.is_some() + || self.input_price_per_1m.is_some() + || self.output_price_per_1m.is_some() + || self.cache_creation_price_per_1m.is_some() + || self.cache_read_price_per_1m.is_some() + || self.price_per_request.is_some() + } +} + +pub(crate) fn from_usage( + usage: &UpsertUsageRecord, +) -> Result<(UsageRoutingSnapshot, UsageSettlementPricingSnapshot), DataLayerError> { + Ok(( + routing_snapshot_from_usage(usage), + settlement_snapshot_from_usage(usage)?, + )) +} + +fn routing_snapshot_from_usage(usage: &UpsertUsageRecord) -> UsageRoutingSnapshot { + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + let mut snapshot = UsageRoutingSnapshot { + candidate_id: usage + .candidate_id + .clone() + .or_else(|| metadata_string(metadata, "candidate_id")), + candidate_index: usage + .candidate_index + .or_else(|| metadata_u64(metadata, "candidate_index")), + key_name: usage + .key_name + .clone() + .or_else(|| metadata_string(metadata, "key_name")), + planner_kind: usage + .planner_kind + .clone() + .or_else(|| metadata_string(metadata, "planner_kind")), + route_family: usage + .route_family + .clone() + .or_else(|| metadata_string(metadata, "route_family")), + route_kind: usage + .route_kind + .clone() + .or_else(|| metadata_string(metadata, "route_kind")), + execution_path: usage + .execution_path + .clone() + .or_else(|| metadata_string(metadata, "execution_path")), + local_execution_runtime_miss_reason: usage + .local_execution_runtime_miss_reason + .clone() + .or_else(|| metadata_string(metadata, "local_execution_runtime_miss_reason")), + selected_provider_id: None, + selected_endpoint_id: None, + selected_provider_api_key_id: None, + has_format_conversion: None, + }; + if snapshot.has_metadata_fields() { + snapshot.selected_provider_id = usage.provider_id.clone(); + snapshot.selected_endpoint_id = usage.provider_endpoint_id.clone(); + snapshot.selected_provider_api_key_id = usage.provider_api_key_id.clone(); + snapshot.has_format_conversion = usage.has_format_conversion; + } + snapshot +} + +fn settlement_snapshot_from_usage( + usage: &UpsertUsageRecord, +) -> Result { + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + let billing_dimensions = metadata_or_snapshot_dimensions(metadata); + let has_billing_dimensions = billing_dimensions.is_some(); + let usage_input_tokens = optional_i64(usage.input_tokens, "input_tokens")?; + let usage_output_tokens = optional_i64(usage.output_tokens, "output_tokens")?; + let usage_cache_creation_uncategorized_tokens = optional_i64( + usage.cache_creation_input_tokens, + "cache_creation_input_tokens", + )?; + let usage_cache_creation_5m_tokens = optional_i64( + usage.cache_creation_ephemeral_5m_input_tokens, + "cache_creation_ephemeral_5m_input_tokens", + )?; + let usage_cache_creation_1h_tokens = optional_i64( + usage.cache_creation_ephemeral_1h_input_tokens, + "cache_creation_ephemeral_1h_input_tokens", + )?; + let usage_cache_read_tokens = + optional_i64(usage.cache_read_input_tokens, "cache_read_input_tokens")?; + let usage_cache_creation_tokens = cache_creation_tokens_from_parts( + usage_cache_creation_uncategorized_tokens, + usage_cache_creation_5m_tokens, + usage_cache_creation_1h_tokens, + ); + let billing_cache_creation_tokens = billing_dimension_i64(metadata, "cache_creation_tokens") + .or_else(|| { + cache_creation_tokens_from_parts( + billing_dimension_i64(metadata, "cache_creation_uncategorized_tokens"), + billing_dimension_i64(metadata, "cache_creation_ephemeral_5m_tokens"), + billing_dimension_i64(metadata, "cache_creation_ephemeral_1h_tokens"), + ) + }) + .or(usage_cache_creation_tokens); + let billing_cache_creation_5m_tokens = + billing_dimension_i64(metadata, "cache_creation_ephemeral_5m_tokens") + .or(usage_cache_creation_5m_tokens); + let billing_cache_creation_1h_tokens = + billing_dimension_i64(metadata, "cache_creation_ephemeral_1h_tokens") + .or(usage_cache_creation_1h_tokens); + let billing_input_tokens = + billing_dimension_i64(metadata, "input_tokens").or(usage_input_tokens); + let billing_output_tokens = + billing_dimension_i64(metadata, "output_tokens").or(usage_output_tokens); + let billing_cache_read_tokens = + billing_dimension_i64(metadata, "cache_read_tokens").or(usage_cache_read_tokens); + let api_family = normalized_api_family(usage); + let billing_effective_input_tokens = billing_dimension_i64(metadata, "effective_input_tokens") + .or_else(|| { + has_billing_dimensions + .then(|| billing_dimension_i64(metadata, "input_tokens")) + .flatten() + }) + .or_else(|| { + effective_input_tokens( + billing_input_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens, + &api_family, + ) + }); + let billing_total_input_context = billing_dimension_i64(metadata, "total_input_context") + .or_else(|| { + total_input_context( + billing_input_tokens, + billing_effective_input_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens, + &api_family, + ) + }); + + Ok(UsageSettlementPricingSnapshot { + billing_status: Some(usage.billing_status.clone()), + billing_snapshot_schema_version: metadata_string( + metadata, + "billing_snapshot_schema_version", + ) + .or_else(|| billing_snapshot_string(metadata, "schema_version")), + billing_snapshot_status: metadata_string(metadata, "billing_snapshot_status") + .or_else(|| billing_snapshot_string(metadata, "status")), + settlement_snapshot_schema_version: settlement_snapshot_schema_version(metadata), + settlement_snapshot: settlement_snapshot_value(metadata), + billing_dimensions, + billing_input_tokens, + billing_effective_input_tokens, + billing_output_tokens, + billing_cache_creation_tokens, + billing_cache_creation_5m_tokens, + billing_cache_creation_1h_tokens, + billing_cache_read_tokens, + billing_total_input_context, + billing_cache_creation_cost_usd: settlement_cache_creation_cost(metadata) + .or(usage.cache_creation_cost_usd), + billing_cache_read_cost_usd: settlement_cost_breakdown_number(metadata, "cache_read_cost") + .or(usage.cache_read_cost_usd), + billing_total_cost_usd: settlement_snapshot_number(metadata, "total_cost") + .or_else(|| billing_snapshot_number(metadata, "total_cost")) + .or(usage.total_cost_usd), + billing_actual_total_cost_usd: settlement_snapshot_number(metadata, "actual_total_cost") + .or(usage.actual_total_cost_usd), + billing_pricing_source: settlement_nested_string( + metadata, + "pricing_snapshot", + "pricing_source", + ), + billing_rule_id: settlement_nested_string(metadata, "billing_plan_snapshot", "rule_id") + .or_else(|| billing_snapshot_string_field(metadata, "rule_id")), + billing_rule_version: settlement_nested_string( + metadata, + "billing_plan_snapshot", + "rule_version", + ), + rate_multiplier: metadata_number(metadata, "rate_multiplier"), + is_free_tier: metadata_bool(metadata, "is_free_tier"), + input_price_per_1m: metadata_number(metadata, "input_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "input_price_per_1m")), + output_price_per_1m: metadata_number(metadata, "output_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "output_price_per_1m")) + .or(usage.output_price_per_1m), + cache_creation_price_per_1m: metadata_number(metadata, "cache_creation_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "cache_creation_price_per_1m")), + cache_read_price_per_1m: metadata_number(metadata, "cache_read_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "cache_read_price_per_1m")), + price_per_request: metadata_number(metadata, "price_per_request") + .or_else(|| billing_snapshot_resolved_number(metadata, "price_per_request")), + }) +} + +pub(crate) async fn sync( + tx: &mut Transaction<'_, MySql>, + request_id: &str, + routing: &UsageRoutingSnapshot, + settlement: &UsageSettlementPricingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + sync_routing(tx, request_id, routing, replace_existing).await?; + sync_settlement(tx, request_id, settlement, replace_existing).await +} + +async fn sync_routing( + tx: &mut Transaction<'_, MySql>, + request_id: &str, + snapshot: &UsageRoutingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + if !snapshot.any_present() && !replace_existing { + return Ok(()); + } + let now = unix_now()?; + let mut query = QueryBuilder::::new( + "INSERT INTO usage_routing_snapshots (request_id, candidate_id, candidate_index, \ + key_name, planner_kind, route_family, route_kind, execution_path, \ + local_execution_runtime_miss_reason, selected_provider_id, selected_endpoint_id, \ + selected_provider_api_key_id, has_format_conversion, created_at, updated_at) VALUES (", + ); + { + let mut values = query.separated(", "); + values + .push_bind(request_id) + .push_bind(snapshot.candidate_id.as_deref()) + .push_bind(optional_i64(snapshot.candidate_index, "candidate_index")?) + .push_bind(snapshot.key_name.as_deref()) + .push_bind(snapshot.planner_kind.as_deref()) + .push_bind(snapshot.route_family.as_deref()) + .push_bind(snapshot.route_kind.as_deref()) + .push_bind(snapshot.execution_path.as_deref()) + .push_bind(snapshot.local_execution_runtime_miss_reason.as_deref()) + .push_bind(snapshot.selected_provider_id.as_deref()) + .push_bind(snapshot.selected_endpoint_id.as_deref()) + .push_bind(snapshot.selected_provider_api_key_id.as_deref()) + .push_bind(snapshot.has_format_conversion) + .push_bind(now) + .push_bind(now); + } + query.push(") ON DUPLICATE KEY UPDATE "); + push_mysql_updates( + &mut query, + &[ + "candidate_id", + "candidate_index", + "key_name", + "planner_kind", + "route_family", + "route_kind", + "execution_path", + "local_execution_runtime_miss_reason", + "selected_provider_id", + "selected_endpoint_id", + "selected_provider_api_key_id", + "has_format_conversion", + ], + "usage_routing_snapshots", + replace_existing, + ); + query.push(", updated_at = VALUES(updated_at)"); + query.build().execute(&mut **tx).await.map_sql_err()?; + Ok(()) +} + +async fn sync_settlement( + tx: &mut Transaction<'_, MySql>, + request_id: &str, + snapshot: &UsageSettlementPricingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + if !snapshot.any_present() && !replace_existing { + return Ok(()); + } + let now = unix_now()?; + let settlement_json = json_text(snapshot.settlement_snapshot.as_ref())?; + let dimensions_json = json_text(snapshot.billing_dimensions.as_ref())?; + let mut query = QueryBuilder::::new( + "INSERT INTO usage_settlement_snapshots (request_id, billing_status, \ + billing_snapshot_schema_version, billing_snapshot_status, \ + settlement_snapshot_schema_version, settlement_snapshot, billing_dimensions, \ + billing_input_tokens, billing_effective_input_tokens, billing_output_tokens, \ + billing_cache_creation_tokens, billing_cache_creation_5m_tokens, \ + billing_cache_creation_1h_tokens, billing_cache_read_tokens, \ + billing_total_input_context, billing_cache_creation_cost_usd, \ + billing_cache_read_cost_usd, billing_total_cost_usd, \ + billing_actual_total_cost_usd, billing_pricing_source, billing_rule_id, \ + billing_rule_version, rate_multiplier, is_free_tier, input_price_per_1m, \ + output_price_per_1m, cache_creation_price_per_1m, cache_read_price_per_1m, \ + price_per_request, created_at, updated_at) VALUES (", + ); + { + let mut values = query.separated(", "); + values + .push_bind(request_id) + .push_bind(snapshot.billing_status.as_deref().unwrap_or("pending")) + .push_bind(snapshot.billing_snapshot_schema_version.as_deref()) + .push_bind(snapshot.billing_snapshot_status.as_deref()) + .push_bind(snapshot.settlement_snapshot_schema_version.as_deref()) + .push_bind(settlement_json.as_deref()) + .push_bind(dimensions_json.as_deref()) + .push_bind(snapshot.billing_input_tokens) + .push_bind(snapshot.billing_effective_input_tokens) + .push_bind(snapshot.billing_output_tokens) + .push_bind(snapshot.billing_cache_creation_tokens) + .push_bind(snapshot.billing_cache_creation_5m_tokens) + .push_bind(snapshot.billing_cache_creation_1h_tokens) + .push_bind(snapshot.billing_cache_read_tokens) + .push_bind(snapshot.billing_total_input_context) + .push_bind(snapshot.billing_cache_creation_cost_usd) + .push_bind(snapshot.billing_cache_read_cost_usd) + .push_bind(snapshot.billing_total_cost_usd) + .push_bind(snapshot.billing_actual_total_cost_usd) + .push_bind(snapshot.billing_pricing_source.as_deref()) + .push_bind(snapshot.billing_rule_id.as_deref()) + .push_bind(snapshot.billing_rule_version.as_deref()) + .push_bind(snapshot.rate_multiplier) + .push_bind(snapshot.is_free_tier) + .push_bind(snapshot.input_price_per_1m) + .push_bind(snapshot.output_price_per_1m) + .push_bind(snapshot.cache_creation_price_per_1m) + .push_bind(snapshot.cache_read_price_per_1m) + .push_bind(snapshot.price_per_request) + .push_bind(now) + .push_bind(now); + } + query.push(") ON DUPLICATE KEY UPDATE "); + if replace_existing { + query.push("billing_status = VALUES(billing_status), "); + } + push_mysql_updates( + &mut query, + &[ + "billing_snapshot_schema_version", + "billing_snapshot_status", + "settlement_snapshot_schema_version", + "settlement_snapshot", + "billing_dimensions", + "billing_input_tokens", + "billing_effective_input_tokens", + "billing_output_tokens", + "billing_cache_creation_tokens", + "billing_cache_creation_5m_tokens", + "billing_cache_creation_1h_tokens", + "billing_cache_read_tokens", + "billing_total_input_context", + "billing_cache_creation_cost_usd", + "billing_cache_read_cost_usd", + "billing_total_cost_usd", + "billing_actual_total_cost_usd", + "billing_pricing_source", + "billing_rule_id", + "billing_rule_version", + "rate_multiplier", + "is_free_tier", + "input_price_per_1m", + "output_price_per_1m", + "cache_creation_price_per_1m", + "cache_read_price_per_1m", + "price_per_request", + ], + "usage_settlement_snapshots", + replace_existing, + ); + query.push(", updated_at = VALUES(updated_at)"); + query.build().execute(&mut **tx).await.map_sql_err()?; + Ok(()) +} + +fn push_mysql_updates( + query: &mut QueryBuilder<'_, MySql>, + fields: &[&str], + table: &str, + replace_existing: bool, +) { + for (index, field) in fields.iter().enumerate() { + if index > 0 { + query.push(", "); + } + query.push(*field).push(" = "); + if replace_existing { + query.push("VALUES(").push(*field).push(")"); + } else { + query + .push("COALESCE(VALUES(") + .push(*field) + .push("), ") + .push(table) + .push(".") + .push(*field) + .push(")"); + } + } +} + +pub(crate) fn hydrate_row( + row: &MySqlRow, + audit: &mut StoredRequestUsageAudit, +) -> Result<(), DataLayerError> { + audit.candidate_id = row.try_get("routing_candidate_id").map_sql_err()?; + audit.candidate_index = row + .try_get::, _>("routing_candidate_index") + .map_sql_err()? + .map(|value| { + u64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!( + "usage routing candidate_index is negative: {value}" + )) + }) + }) + .transpose()?; + audit.key_name = row.try_get("routing_key_name").map_sql_err()?; + audit.planner_kind = row.try_get("routing_planner_kind").map_sql_err()?; + audit.route_family = row.try_get("routing_route_family").map_sql_err()?; + audit.route_kind = row.try_get("routing_route_kind").map_sql_err()?; + audit.execution_path = row.try_get("routing_execution_path").map_sql_err()?; + audit.local_execution_runtime_miss_reason = row + .try_get("routing_local_execution_runtime_miss_reason") + .map_sql_err()?; + + let snapshot = settlement_snapshot_from_row(row)?; + if let Some(effective) = nonnegative_u64(snapshot.billing_effective_input_tokens) { + audit.total_tokens = effective + .saturating_add(audit.output_tokens) + .saturating_add(audit.cache_creation_input_tokens) + .saturating_add(audit.cache_read_input_tokens); + } else if let Some(context) = nonnegative_u64(snapshot.billing_total_input_context) { + audit.total_tokens = context.saturating_add(audit.output_tokens); + } + audit.request_metadata = attach_settlement_metadata(audit.request_metadata.take(), &snapshot); + Ok(()) +} + +fn settlement_snapshot_from_row( + row: &MySqlRow, +) -> Result { + Ok(UsageSettlementPricingSnapshot { + billing_status: None, + billing_snapshot_schema_version: row + .try_get("settlement_billing_snapshot_schema_version") + .map_sql_err()?, + billing_snapshot_status: row + .try_get("settlement_billing_snapshot_status") + .map_sql_err()?, + settlement_snapshot_schema_version: row + .try_get("settlement_snapshot_schema_version") + .map_sql_err()?, + settlement_snapshot: json_value_from_row(row, "settlement_snapshot")?, + billing_dimensions: json_value_from_row(row, "settlement_billing_dimensions")?, + billing_input_tokens: row + .try_get("settlement_billing_input_tokens") + .map_sql_err()?, + billing_effective_input_tokens: row + .try_get("settlement_billing_effective_input_tokens") + .map_sql_err()?, + billing_output_tokens: row + .try_get("settlement_billing_output_tokens") + .map_sql_err()?, + billing_cache_creation_tokens: row + .try_get("settlement_billing_cache_creation_tokens") + .map_sql_err()?, + billing_cache_creation_5m_tokens: row + .try_get("settlement_billing_cache_creation_5m_tokens") + .map_sql_err()?, + billing_cache_creation_1h_tokens: row + .try_get("settlement_billing_cache_creation_1h_tokens") + .map_sql_err()?, + billing_cache_read_tokens: row + .try_get("settlement_billing_cache_read_tokens") + .map_sql_err()?, + billing_total_input_context: row + .try_get("settlement_billing_total_input_context") + .map_sql_err()?, + billing_cache_creation_cost_usd: row + .try_get("settlement_billing_cache_creation_cost_usd") + .map_sql_err()?, + billing_cache_read_cost_usd: row + .try_get("settlement_billing_cache_read_cost_usd") + .map_sql_err()?, + billing_total_cost_usd: row + .try_get("settlement_billing_total_cost_usd") + .map_sql_err()?, + billing_actual_total_cost_usd: row + .try_get("settlement_billing_actual_total_cost_usd") + .map_sql_err()?, + billing_pricing_source: row + .try_get("settlement_billing_pricing_source") + .map_sql_err()?, + billing_rule_id: row.try_get("settlement_billing_rule_id").map_sql_err()?, + billing_rule_version: row + .try_get("settlement_billing_rule_version") + .map_sql_err()?, + rate_multiplier: row.try_get("settlement_rate_multiplier").map_sql_err()?, + is_free_tier: row + .try_get::, _>("settlement_is_free_tier") + .map_sql_err()?, + input_price_per_1m: row.try_get("settlement_input_price_per_1m").map_sql_err()?, + output_price_per_1m: row + .try_get("settlement_output_price_per_1m") + .map_sql_err()?, + cache_creation_price_per_1m: row + .try_get("settlement_cache_creation_price_per_1m") + .map_sql_err()?, + cache_read_price_per_1m: row + .try_get("settlement_cache_read_price_per_1m") + .map_sql_err()?, + price_per_request: row.try_get("settlement_price_per_request").map_sql_err()?, + }) +} + +fn attach_settlement_metadata( + metadata: Option, + snapshot: &UsageSettlementPricingSnapshot, +) -> Option { + if !snapshot.any_present() { + return metadata; + } + let mut metadata = match metadata { + Some(Value::Object(object)) => object, + Some(value) => return Some(value), + None => Map::new(), + }; + insert_string( + &mut metadata, + "billing_snapshot_schema_version", + snapshot.billing_snapshot_schema_version.as_deref(), + ); + insert_string( + &mut metadata, + "billing_snapshot_status", + snapshot.billing_snapshot_status.as_deref(), + ); + insert_string( + &mut metadata, + "settlement_snapshot_schema_version", + snapshot.settlement_snapshot_schema_version.as_deref(), + ); + insert_value( + &mut metadata, + "settlement_snapshot", + snapshot.settlement_snapshot.as_ref(), + ); + insert_value( + &mut metadata, + "billing_dimensions", + snapshot.billing_dimensions.as_ref(), + ); + insert_number(&mut metadata, "rate_multiplier", snapshot.rate_multiplier); + insert_bool(&mut metadata, "is_free_tier", snapshot.is_free_tier); + insert_number( + &mut metadata, + "input_price_per_1m", + snapshot.input_price_per_1m, + ); + insert_number( + &mut metadata, + "output_price_per_1m", + snapshot.output_price_per_1m, + ); + insert_number( + &mut metadata, + "cache_creation_price_per_1m", + snapshot.cache_creation_price_per_1m, + ); + insert_number( + &mut metadata, + "cache_read_price_per_1m", + snapshot.cache_read_price_per_1m, + ); + insert_number( + &mut metadata, + "price_per_request", + snapshot.price_per_request, + ); + (!metadata.is_empty()).then_some(Value::Object(metadata)) +} + +fn insert_string(metadata: &mut Map, key: &str, value: Option<&str>) { + if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) { + metadata.insert(key.to_string(), Value::String(value.to_string())); + } +} + +fn insert_number(metadata: &mut Map, key: &str, value: Option) { + if let Some(number) = value + .filter(|value| value.is_finite()) + .and_then(serde_json::Number::from_f64) + { + metadata.insert(key.to_string(), Value::Number(number)); + } +} + +fn insert_bool(metadata: &mut Map, key: &str, value: Option) { + if let Some(value) = value { + metadata.insert(key.to_string(), Value::Bool(value)); + } +} + +fn insert_value(metadata: &mut Map, key: &str, value: Option<&Value>) { + if let Some(value) = value { + metadata.insert(key.to_string(), value.clone()); + } +} + +fn metadata_string(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn metadata_number(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn metadata_u64(metadata: Option<&Map>, key: &str) -> Option { + metadata.and_then(|object| { + object.get(key).and_then(|value| { + value + .as_u64() + .or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok())) + }) + }) +} + +fn metadata_bool(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_bool) +} + +fn billing_snapshot_object(metadata: Option<&Map>) -> Option<&Map> { + metadata + .and_then(|object| object.get("billing_snapshot")) + .and_then(Value::as_object) +} + +fn billing_snapshot_string(metadata: Option<&Map>, key: &str) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn billing_snapshot_resolved_number( + metadata: Option<&Map>, + key: &str, +) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("resolved_variables")) + .and_then(Value::as_object) + .and_then(|variables| variables.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_snapshot_object( + metadata: Option<&Map>, +) -> Option<&Map> { + metadata + .and_then(|object| object.get("settlement_snapshot")) + .and_then(Value::as_object) +} + +fn settlement_snapshot_schema_version(metadata: Option<&Map>) -> Option { + metadata_string(metadata, "settlement_snapshot_schema_version").or_else(|| { + settlement_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("schema_version")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + }) +} + +fn settlement_snapshot_value(metadata: Option<&Map>) -> Option { + metadata + .and_then(|object| object.get("settlement_snapshot")) + .cloned() +} + +fn settlement_child_value<'a>( + metadata: Option<&'a Map>, + child: &str, +) -> Option<&'a Value> { + settlement_snapshot_object(metadata).and_then(|snapshot| snapshot.get(child)) +} + +fn settlement_child_object<'a>( + metadata: Option<&'a Map>, + child: &str, +) -> Option<&'a Map> { + settlement_child_value(metadata, child).and_then(Value::as_object) +} + +fn metadata_or_snapshot_dimensions(metadata: Option<&Map>) -> Option { + metadata + .and_then(|object| object.get("billing_dimensions")) + .cloned() + .or_else(|| settlement_child_value(metadata, "resolved_dimensions").cloned()) + .or_else(|| { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("resolved_dimensions")) + .cloned() + }) +} + +fn billing_dimension_i64(metadata: Option<&Map>, key: &str) -> Option { + metadata_or_snapshot_dimensions(metadata) + .and_then(|dimensions| dimensions.get(key).and_then(json_i64)) + .filter(|value| *value >= 0) +} + +fn json_i64(value: &Value) -> Option { + value + .as_i64() + .or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok())) +} + +fn settlement_snapshot_number(metadata: Option<&Map>, key: &str) -> Option { + settlement_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn billing_snapshot_number(metadata: Option<&Map>, key: &str) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_cost_breakdown_number( + metadata: Option<&Map>, + key: &str, +) -> Option { + settlement_child_object(metadata, "cost_breakdown") + .or_else(|| { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("cost_breakdown")) + .and_then(Value::as_object) + }) + .and_then(|breakdown| breakdown.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_cache_creation_cost(metadata: Option<&Map>) -> Option { + let mut found = false; + let total = [ + "cache_creation_uncategorized_cost", + "cache_creation_ephemeral_5m_cost", + "cache_creation_ephemeral_1h_cost", + "cache_creation_cost", + ] + .into_iter() + .fold(0.0, |sum, key| { + if let Some(value) = settlement_cost_breakdown_number(metadata, key) { + found = true; + sum + value + } else { + sum + } + }); + found.then_some(total) +} + +fn settlement_nested_string( + metadata: Option<&Map>, + child: &str, + key: &str, +) -> Option { + settlement_child_object(metadata, child) + .and_then(|object| object.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn billing_snapshot_string_field( + metadata: Option<&Map>, + key: &str, +) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn optional_i64(value: Option, field: &str) -> Result, DataLayerError> { + value + .map(|value| { + i64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!("usage {field} exceeds bigint: {value}")) + }) + }) + .transpose() +} + +fn cache_creation_tokens_from_parts( + uncategorized: Option, + ephemeral_5m: Option, + ephemeral_1h: Option, +) -> Option { + let categorized = ephemeral_5m + .unwrap_or_default() + .saturating_add(ephemeral_1h.unwrap_or_default()); + match uncategorized { + Some(0) if categorized > 0 => Some(categorized), + Some(value) => Some(value), + None if categorized > 0 => Some(categorized), + None => None, + } +} + +fn normalized_api_family(usage: &UpsertUsageRecord) -> String { + usage + .endpoint_api_format + .as_deref() + .or(usage.api_format.as_deref()) + .unwrap_or_default() + .split(':') + .next() + .unwrap_or_default() + .trim() + .to_ascii_lowercase() +} + +fn effective_input_tokens( + input_tokens: Option, + cache_creation_tokens: Option, + cache_read_tokens: Option, + api_family: &str, +) -> Option { + let input_tokens = input_tokens?; + let cache_creation_tokens = cache_creation_tokens.unwrap_or_default(); + let cache_read_tokens = cache_read_tokens.unwrap_or_default(); + if input_tokens > 0 { + if api_family == "openai" && (cache_creation_tokens > 0 || cache_read_tokens > 0) { + return Some( + input_tokens + .saturating_sub(cache_creation_tokens) + .saturating_sub(cache_read_tokens), + ); + } + if matches!(api_family, "gemini" | "google") && cache_read_tokens > 0 { + return Some(input_tokens.saturating_sub(cache_read_tokens)); + } + } + Some(input_tokens) +} + +fn total_input_context( + input_tokens: Option, + effective_input_tokens: Option, + cache_creation_tokens: Option, + cache_read_tokens: Option, + api_family: &str, +) -> Option { + if input_tokens.is_none() + && effective_input_tokens.is_none() + && cache_creation_tokens.is_none() + && cache_read_tokens.is_none() + { + return None; + } + let input_tokens = input_tokens.unwrap_or_default(); + let effective_input_tokens = effective_input_tokens.unwrap_or(input_tokens); + let cache_creation_tokens = cache_creation_tokens.unwrap_or_default(); + let cache_read_tokens = cache_read_tokens.unwrap_or_default(); + match api_family { + "claude" | "anthropic" => Some( + input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + "openai" => Some( + effective_input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + "gemini" | "google" => Some(effective_input_tokens.saturating_add(cache_read_tokens)), + _ => Some( + input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + } +} + +fn json_text(value: Option<&Value>) -> Result, DataLayerError> { + value + .map(|value| { + serde_json::to_string(value).map_err(|error| { + DataLayerError::UnexpectedValue(format!( + "failed to serialize usage settlement snapshot: {error}" + )) + }) + }) + .transpose() +} + +fn json_value_from_row(row: &MySqlRow, column: &str) -> Result, DataLayerError> { + row.try_get::, _>(column) + .map_sql_err()? + .map(|value| { + serde_json::from_str(&value).map_err(|error| { + DataLayerError::UnexpectedValue(format!( + "invalid usage settlement JSON in {column}: {error}" + )) + }) + }) + .transpose() +} + +fn nonnegative_u64(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + +fn unix_now() -> Result { + let seconds = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))? + .as_secs(); + i64::try_from(seconds) + .map_err(|_| DataLayerError::UnexpectedValue("unix timestamp overflow".to_string())) +} diff --git a/crates/aether-data/adapters/mysql/src/usage/tests.rs b/crates/aether-data/adapters/mysql/src/usage/tests.rs index 9d60fd604..77f64af8a 100644 --- a/crates/aether-data/adapters/mysql/src/usage/tests.rs +++ b/crates/aether-data/adapters/mysql/src/usage/tests.rs @@ -1,6 +1,10 @@ use super::{MysqlUsageStorage, MysqlUsageWriteRepository}; use crate::run_migrations; -use aether_data_contracts::repository::usage::{UpsertUsageRecord, UsageWriteRepository}; +use aether_data_contracts::repository::usage::{ + UpsertUsageRecord, UsageAuditListQuery, UsageBodyCaptureState, UsageCleanupExecutionMode, + UsageCleanupTargets, UsageCleanupWindow, UsageWriteRepository, +}; +use chrono::DateTime; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -48,6 +52,45 @@ fn mysql_dashboard_reads_imported_daily_aggregates() { assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests")); } +#[test] +fn mysql_usage_stat_rebuilds_aggregate_in_sql() { + let source = include_str!("../usage.rs"); + assert!(source.contains("UPDATE api_keys\nJOIN (")); + assert!(source.contains("AND status NOT IN ('pending', 'streaming')")); + assert!(source.contains("MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR")); + assert!(source.contains("MAX(created_at_unix_ms) AS last_used_at")); + assert!(source.contains("UPDATE provider_api_keys\nJOIN (")); + assert!(source.contains("GROUP BY provider_api_key_id")); + assert!(!source.contains("struct ProviderKeyStats")); + assert!(super::MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR + .contains("status IN ('completed', 'success', 'ok', 'billed', 'settled')")); + assert!(super::MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR + .contains("error_message IS NULL OR TRIM(error_message) = ''")); + assert!(super::MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR + .contains("status NOT IN ('pending', 'streaming')")); + assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR.contains( + "COALESCE(`usage`.input_tokens, 0) - COALESCE(`usage`.cache_read_input_tokens, 0)" + )); + assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + .contains("GREATEST(COALESCE(`usage`.output_tokens, 0), 0)")); + assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + .contains("NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0)")); + let snapshot_position = super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + .find("settlement.billing_effective_input_tokens") + .expect("canonical total must read settlement snapshots"); + let raw_position = super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + .find("NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0)") + .expect("canonical total must preserve non-zero legacy raw totals"); + assert!(snapshot_position < raw_position); + assert!(source.matches("canonical_total_tokens_expr =").count() >= 4); + assert_eq!( + source + .matches("LEFT JOIN usage_settlement_snapshots AS settlement") + .count(), + 4 + ); +} + #[test] fn mysql_usage_upsert_keeps_terminal_state_when_streaming_arrives_late() { assert!(super::UPSERT_USAGE_SQL.contains( @@ -110,7 +153,7 @@ fn mysql_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_l } #[tokio::test] -async fn mysql_usage_write_repository_upserts_when_url_is_set() { +async fn mysql_usage_write_repository_upserts_and_flushes_counters_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") .ok() .filter(|value| !value.trim().is_empty()) @@ -133,29 +176,35 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() { let api_key_id = format!("api-key-{suffix}"); let provider_id = format!("provider-{suffix}"); let provider_key_id = format!("provider-key-{suffix}"); + let model_name = format!("model-{suffix}"); seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + sqlx::query("INSERT INTO global_models (id, name, created_at, updated_at) VALUES (?, ?, 1, 1)") + .bind(format!("global-model-{suffix}")) + .bind(&model_name) + .execute(&pool) + .await + .expect("global model should seed"); let repository = MysqlUsageWriteRepository::new(pool.clone()); - let record = repository - .upsert(sample_usage( - &format!("request-{suffix}"), - &user_id, - &api_key_id, - &provider_id, - &provider_key_id, - "completed", - "pending", - 1_000, - )) - .await - .expect("usage should upsert"); + let mut usage = sample_usage( + &format!("request-{suffix}"), + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "pending", + 1_000, + ); + usage.model.clone_from(&model_name); + let record = repository.upsert(usage).await.expect("usage should upsert"); assert_eq!(record.api_key_id.as_deref(), Some(api_key_id.as_str())); assert_eq!( record.provider_api_key_id.as_deref(), Some(provider_key_id.as_str()) ); - assert_eq!(record.total_tokens, 7); + assert_eq!(record.total_tokens, 5); assert_eq!( record.request_metadata.as_ref().unwrap()["upstream_is_stream"], true @@ -168,6 +217,25 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() { .expect("usage stream mode should load"); assert_eq!(upstream_is_stream, Some(true)); + repository + .flush_usage_counter_deltas(100) + .await + .expect("usage counter deltas should flush"); + assert!( + repository + .rebuild_api_key_usage_stats() + .await + .expect("api key stats should rebuild") + >= 1 + ); + assert!( + repository + .rebuild_provider_api_key_usage_stats() + .await + .expect("provider api key stats should rebuild") + >= 1 + ); + let stats = sqlx::query_as::<_, (i64, i64, f64, Option)>( "SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = ?", ) @@ -175,7 +243,7 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() { .fetch_one(&pool) .await .expect("api key stats should load"); - assert_eq!(stats, (1, 7, 0.5, Some(1_000))); + assert_eq!(stats, (1, 5, 0.5, Some(1_000))); let provider_stats = sqlx::query_as::<_, (i64, i64, i64, i64, f64, i64, Option)>( "SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, total_response_time_ms, last_used_at FROM provider_api_keys WHERE id = ?", @@ -184,7 +252,575 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() { .fetch_one(&pool) .await .expect("provider key stats should load"); - assert_eq!(provider_stats, (1, 1, 0, 7, 0.5, 42, Some(1_000))); + assert_eq!(provider_stats, (1, 1, 0, 5, 0.5, 42, Some(1_000))); + let model_usage_count: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = ?") + .bind(&model_name) + .fetch_one(&pool) + .await + .expect("global model usage count should load"); + assert_eq!(model_usage_count, 1); +} + +#[tokio::test] +async fn mysql_canonical_totals_preserve_legacy_total_tokens_only_rows_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!( + "skipping mysql canonical total_tokens test because AETHER_TEST_MYSQL_URL is unset" + ); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + + let suffix = unique_suffix(); + let user_id = format!("total-only-user-{suffix}"); + let api_key_id = format!("total-only-api-key-{suffix}"); + let provider_id = format!("total-only-provider-{suffix}"); + let provider_key_id = format!("total-only-provider-key-{suffix}"); + seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + + let repository = MysqlUsageWriteRepository::new(pool.clone()); + let mut usage = sample_usage( + &format!("total-only-request-{suffix}"), + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "pending", + 1_000, + ); + usage.input_tokens = Some(0); + usage.output_tokens = Some(0); + usage.cache_creation_input_tokens = Some(0); + usage.cache_creation_ephemeral_5m_input_tokens = Some(0); + usage.cache_creation_ephemeral_1h_input_tokens = Some(0); + usage.cache_read_input_tokens = Some(0); + usage.total_tokens = Some(77); + + let request_id = format!("total-only-request-{suffix}"); + usage.request_id.clone_from(&request_id); + repository + .upsert(usage) + .await + .expect("legacy total-only usage should upsert"); + sqlx::query("DELETE FROM usage_settlement_snapshots WHERE request_id = ?") + .bind(&request_id) + .execute(&pool) + .await + .expect("legacy fixture must not have a settlement snapshot"); + let stored = repository + .find_by_request_id(&request_id) + .await + .expect("legacy total-only usage should load") + .expect("legacy total-only usage should exist"); + assert_eq!(stored.total_tokens, 77); + + repository + .rebuild_api_key_usage_stats() + .await + .expect("api key stats should rebuild"); + repository + .rebuild_provider_api_key_usage_stats() + .await + .expect("provider api key stats should rebuild"); + + let api_key_total: i64 = sqlx::query_scalar("SELECT total_tokens FROM api_keys WHERE id = ?") + .bind(&api_key_id) + .fetch_one(&pool) + .await + .expect("api key total should load"); + let provider_key_total: i64 = + sqlx::query_scalar("SELECT total_tokens FROM provider_api_keys WHERE id = ?") + .bind(&provider_key_id) + .fetch_one(&pool) + .await + .expect("provider key total should load"); + assert_eq!(api_key_total, 77); + assert_eq!(provider_key_total, 77); + + let totals = MysqlUsageStorage::new(pool) + .summarize_usage_totals_by_user_ids(std::slice::from_ref(&user_id)) + .await + .expect("user usage totals should load"); + assert_eq!(totals.len(), 1); + assert_eq!(totals[0].total_tokens, 77); +} + +#[tokio::test] +async fn mysql_concurrent_same_request_upserts_enqueue_counters_once_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!( + "skipping mysql usage counter concurrency test because AETHER_TEST_MYSQL_URL is unset" + ); + return; + }; + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(4) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + + let suffix = unique_suffix(); + let user_id = format!("counter-user-{suffix}"); + let api_key_id = format!("counter-api-key-{suffix}"); + let provider_id = format!("counter-provider-{suffix}"); + let provider_key_id = format!("counter-provider-key-{suffix}"); + let model_name = format!("counter-model-{suffix}"); + let request_id = format!("counter-request-{suffix}"); + seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + sqlx::query("INSERT INTO global_models (id, name, created_at, updated_at) VALUES (?, ?, 1, 1)") + .bind(format!("counter-global-model-{suffix}")) + .bind(&model_name) + .execute(&pool) + .await + .expect("global model should seed"); + + let repository = MysqlUsageWriteRepository::new(pool.clone()); + let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let repository = repository.clone(); + let barrier = barrier.clone(); + let mut usage = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "pending", + 1_000, + ); + usage.model.clone_from(&model_name); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + repository.upsert(usage).await + })); + } + for task in tasks { + task.await + .expect("concurrent usage writer should join") + .expect("concurrent usage should persist"); + } + repository + .flush_usage_counter_deltas(1_000) + .await + .expect("usage counter deltas should flush"); + + let api_key_requests: i64 = + sqlx::query_scalar("SELECT total_requests FROM api_keys WHERE id = ?") + .bind(&api_key_id) + .fetch_one(&pool) + .await + .expect("api key counter should load"); + let provider_key_requests: i64 = + sqlx::query_scalar("SELECT request_count FROM provider_api_keys WHERE id = ?") + .bind(&provider_key_id) + .fetch_one(&pool) + .await + .expect("provider key counter should load"); + let model_requests: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = ?") + .bind(&model_name) + .fetch_one(&pool) + .await + .expect("model counter should load"); + let outbox_rows: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("usage counter outbox should load"); + let routing_rows: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("routing snapshot should count"); + let settlement_rows: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("settlement snapshot should count"); + assert_eq!(api_key_requests, 1); + assert_eq!(provider_key_requests, 1); + assert_eq!(model_requests, 1); + assert_eq!(outbox_rows, 3); + assert_eq!(routing_rows, 1); + assert_eq!(settlement_rows, 1); +} + +#[tokio::test] +async fn mysql_usage_http_capture_round_trips_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql usage HTTP capture test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + let suffix = unique_suffix(); + let user_id = format!("capture-user-{suffix}"); + let api_key_id = format!("capture-api-key-{suffix}"); + let provider_id = format!("capture-provider-{suffix}"); + let provider_key_id = format!("capture-provider-key-{suffix}"); + let request_id = format!("capture-request-{suffix}"); + seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + let writer = MysqlUsageWriteRepository::new(pool.clone()); + + let mut rich = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "pending", + "pending", + 1_000, + ); + rich.request_headers = Some(serde_json::json!({"x-client": "one"})); + rich.provider_request_headers = Some(serde_json::json!({"x-provider": "two"})); + rich.request_body = Some(serde_json::json!({"request": true})); + rich.provider_request_body = Some(serde_json::json!({"provider": true})); + rich.request_body_state = Some(UsageBodyCaptureState::Inline); + rich.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + let stored = writer + .upsert(rich) + .await + .expect("MySQL canonical capture should upsert"); + assert_eq!( + stored.request_headers, + Some(serde_json::json!({"x-client": "one"})) + ); + assert_eq!( + stored.request_body, + Some(serde_json::json!({"request": true})) + ); + assert_eq!( + stored.request_body_state, + Some(UsageBodyCaptureState::Reference) + ); + assert_eq!( + stored.request_body_ref.as_deref(), + Some(format!("usage://request/{request_id}/request_body").as_str()) + ); + let blob_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("MySQL canonical blobs should count"); + assert_eq!(blob_count, 2); + let legacy_body: Option = + sqlx::query_scalar("SELECT CAST(request_body AS CHAR) FROM `usage` WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("legacy body should load"); + assert!(legacy_body.is_none()); + + let sparse = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "streaming", + "pending", + 1_001, + ); + let sparse_stored = writer + .upsert(sparse) + .await + .expect("MySQL sparse capture should upsert"); + assert_eq!(sparse_stored.request_headers, stored.request_headers); + assert_eq!(sparse_stored.request_body, stored.request_body); + + let mut clear = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "streaming", + "pending", + 1_002, + ); + clear.request_body = Some(serde_json::json!({"residual": true})); + clear.request_body_state = Some(UsageBodyCaptureState::None); + let cleared = writer + .upsert(clear) + .await + .expect("MySQL explicit none should clear"); + assert!(cleared.request_body.is_none()); + assert!(cleared.request_body_ref.is_none()); + assert_eq!( + cleared.request_body_state, + Some(UsageBodyCaptureState::None) + ); + assert_eq!(cleared.provider_request_body, stored.provider_request_body); +} + +#[tokio::test] +async fn mysql_usage_canonical_snapshots_round_trip_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql usage snapshot test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + let suffix = unique_suffix(); + let user_id = format!("snapshot-user-{suffix}"); + let api_key_id = format!("snapshot-api-key-{suffix}"); + let provider_id = format!("snapshot-provider-{suffix}"); + let provider_key_id = format!("snapshot-provider-key-{suffix}"); + let request_id = format!("snapshot-request-{suffix}"); + seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + let writer = MysqlUsageWriteRepository::new(pool.clone()); + + let mut rich = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "streaming", + "pending", + 1_000, + ); + rich.candidate_id = None; + rich.candidate_index = None; + rich.key_name = None; + rich.planner_kind = None; + rich.route_family = None; + rich.route_kind = None; + rich.execution_path = None; + rich.input_tokens = Some(100); + rich.output_tokens = Some(20); + rich.cache_creation_input_tokens = Some(10); + rich.cache_creation_ephemeral_5m_input_tokens = Some(10); + rich.cache_read_input_tokens = Some(30); + rich.request_metadata = Some(serde_json::json!({ + "trace_id": "rich", + "candidate_id": "candidate-canonical", + "candidate_index": 2, + "key_name": "key-canonical", + "planner_kind": "fallback", + "route_family": "chat", + "route_kind": "remote", + "execution_path": "converted", + "billing_snapshot_schema_version": "v3", + "billing_snapshot_status": "resolved", + "rate_multiplier": 0.5, + "input_price_per_1m": 1.1, + "output_price_per_1m": 2.2, + "billing_dimensions": { + "input_tokens": 100, + "effective_input_tokens": 60, + "output_tokens": 20, + "cache_creation_tokens": 10, + "cache_read_tokens": 30, + "total_input_context": 100 + }, + "settlement_snapshot": { + "schema_version": "v3", + "total_cost": 1.25, + "actual_total_cost": 1.0, + "pricing_snapshot": {"pricing_source": "catalog"}, + "billing_plan_snapshot": {"rule_id": "rule-1", "rule_version": "7"} + } + })); + let stored = writer + .upsert(rich) + .await + .expect("MySQL canonical snapshots should upsert"); + assert_eq!(stored.routing_candidate_id(), Some("candidate-canonical")); + assert_eq!(stored.routing_candidate_index(), Some(2)); + assert_eq!(stored.provider_id.as_deref(), Some(provider_id.as_str())); + assert_eq!(stored.output_tokens, 20); + assert_eq!(stored.cache_creation_input_tokens, 10); + assert_eq!(stored.cache_read_input_tokens, 30); + assert_eq!(stored.total_tokens, 120); + assert_eq!(stored.total_cost_usd, 1.25); + assert_eq!(stored.actual_total_cost_usd, 1.0); + assert_eq!(stored.settlement_rate_multiplier(), Some(0.5)); + assert_eq!(stored.settlement_input_price_per_1m(), Some(1.1)); + + let mut sparse = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "streaming", + "pending", + 1_001, + ); + sparse.provider_id = None; + sparse.provider_endpoint_id = None; + sparse.provider_api_key_id = None; + sparse.has_format_conversion = None; + sparse.candidate_id = None; + sparse.candidate_index = None; + sparse.key_name = None; + sparse.planner_kind = None; + sparse.route_family = None; + sparse.route_kind = None; + sparse.execution_path = None; + sparse.local_execution_runtime_miss_reason = None; + sparse.input_tokens = None; + sparse.output_tokens = None; + sparse.total_tokens = None; + sparse.cache_creation_input_tokens = None; + sparse.cache_creation_ephemeral_5m_input_tokens = None; + sparse.cache_creation_ephemeral_1h_input_tokens = None; + sparse.cache_read_input_tokens = None; + sparse.cache_creation_cost_usd = None; + sparse.cache_read_cost_usd = None; + sparse.output_price_per_1m = None; + sparse.total_cost_usd = None; + sparse.actual_total_cost_usd = None; + sparse.request_metadata = Some(serde_json::json!({"trace_id": "sparse"})); + let sparse_stored = writer + .upsert(sparse) + .await + .expect("MySQL sparse snapshots should merge"); + assert_eq!( + sparse_stored.routing_candidate_id(), + Some("candidate-canonical") + ); + assert_eq!( + sparse_stored.provider_id.as_deref(), + Some(provider_id.as_str()) + ); + assert_eq!(sparse_stored.output_tokens, 20); + assert_eq!(sparse_stored.total_tokens, 120); + assert_eq!(sparse_stored.total_cost_usd, 1.25); + assert_eq!(sparse_stored.settlement_rate_multiplier(), Some(0.5)); + + let mut terminal = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "settled", + 1_002, + ); + terminal.provider_id = None; + terminal.provider_endpoint_id = None; + terminal.provider_api_key_id = None; + terminal.has_format_conversion = None; + terminal.candidate_id = None; + terminal.candidate_index = None; + terminal.key_name = None; + terminal.planner_kind = None; + terminal.route_family = None; + terminal.route_kind = None; + terminal.execution_path = None; + terminal.local_execution_runtime_miss_reason = None; + terminal.input_tokens = None; + terminal.output_tokens = None; + terminal.total_tokens = None; + terminal.cache_creation_input_tokens = None; + terminal.cache_creation_ephemeral_5m_input_tokens = None; + terminal.cache_creation_ephemeral_1h_input_tokens = None; + terminal.cache_read_input_tokens = None; + terminal.cache_creation_cost_usd = None; + terminal.cache_read_cost_usd = None; + terminal.output_price_per_1m = None; + terminal.total_cost_usd = None; + terminal.actual_total_cost_usd = None; + terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal"})); + let terminal_stored = writer + .upsert(terminal) + .await + .expect("MySQL terminal snapshots should replace"); + assert_eq!(terminal_stored.status, "completed"); + assert_eq!(terminal_stored.billing_status, "settled"); + assert!(terminal_stored.candidate_id.is_none()); + assert!(terminal_stored.provider_id.is_none()); + assert_eq!(terminal_stored.total_tokens, 0); + assert_eq!(terminal_stored.total_cost_usd, 0.0); + assert_eq!(terminal_stored.settlement_rate_multiplier(), None); + + let mut late = sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "pending", + "pending", + 1_003, + ); + late.candidate_id = Some("late-candidate".to_string()); + late.request_metadata = Some(serde_json::json!({ + "trace_id": "late", + "rate_multiplier": 9, + "settlement_snapshot": {"schema_version": "late", "total_cost": 99} + })); + let after_late = writer + .upsert(late) + .await + .expect("late MySQL pending usage should return terminal record"); + assert_eq!(after_late.status, "completed"); + assert_eq!(after_late.billing_status, "settled"); + assert!(after_late.candidate_id.is_none()); + assert_eq!(after_late.total_cost_usd, 0.0); + + let snapshot_counts: (i64, i64) = sqlx::query_as( + r#" +SELECT + (SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id = ?), + (SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = ?) +"#, + ) + .bind(&request_id) + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("MySQL canonical snapshots should count"); + assert_eq!(snapshot_counts, (1, 1)); } #[tokio::test] @@ -211,12 +847,14 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set( let api_key_id = format!("api-key-read-{suffix}"); let provider_id = format!("provider-read-{suffix}"); let provider_key_id = format!("provider-key-read-{suffix}"); + let first_request_id = format!("request-read-1-{suffix}"); + let second_request_id = format!("request-read-2-{suffix}"); seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; let writer = MysqlUsageWriteRepository::new(pool.clone()); writer .upsert(sample_usage( - &format!("request-read-1-{suffix}"), + &first_request_id, &user_id, &api_key_id, &provider_id, @@ -229,7 +867,7 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set( .expect("usage should upsert"); writer .upsert(sample_usage( - &format!("request-read-2-{suffix}"), + &second_request_id, &user_id, &api_key_id, &provider_id, @@ -241,17 +879,33 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set( .await .expect("usage should upsert"); + sqlx::query("UPDATE `usage` SET username = ?, api_key_name = ? WHERE request_id = ?") + .bind("legacy-user") + .bind("legacy-key") + .bind(&first_request_id) + .execute(&pool) + .await + .expect("legacy usage names should update"); + let reader = MysqlUsageStorage::new(pool); let records = reader - .load_usage_records() + .list_usage_audits(&UsageAuditListQuery { + created_from_unix_secs: Some(900), + created_until_unix_secs: Some(1_100), + user_id: Some(user_id.clone()), + newest_first: true, + ..UsageAuditListQuery::default() + }) .await .expect("usage records should load"); let loaded = records .iter() - .find(|item| item.request_id == format!("request-read-1-{suffix}")) + .find(|item| item.request_id == first_request_id) .expect("usage should exist"); - assert_eq!(loaded.total_tokens, 7); + assert_eq!(loaded.total_tokens, 5); assert_eq!(loaded.billing_status, "settled"); + assert_eq!(loaded.username.as_deref(), Some("legacy-user")); + assert_eq!(loaded.api_key_name.as_deref(), Some("legacy-key")); assert_eq!( records .iter() @@ -261,6 +915,164 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set( ); } +#[tokio::test] +async fn mysql_usage_cleanup_executes_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql usage cleanup test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + + let suffix = unique_suffix(); + let user_id = format!("cleanup-user-{suffix}"); + let api_key_id = format!("cleanup-api-key-{suffix}"); + let provider_id = format!("cleanup-provider-{suffix}"); + let provider_key_id = format!("cleanup-provider-key-{suffix}"); + let request_id = format!("cleanup-request-{suffix}"); + seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; + let repository = MysqlUsageWriteRepository::new(pool.clone()); + repository + .upsert(sample_usage( + &request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "settled", + 10, + )) + .await + .expect("cleanup usage should seed"); + sqlx::query( + "UPDATE `usage` SET request_headers = '{\"keep\":true}', request_body = '{\"body\":true}' WHERE request_id = ?", + ) + .bind(&request_id) + .execute(&pool) + .await + .expect("cleanup fields should seed"); + + let window = UsageCleanupWindow { + detail_cutoff: DateTime::from_timestamp(20, 0).expect("valid detail cutoff"), + compressed_cutoff: DateTime::from_timestamp(5, 0).expect("valid compressed cutoff"), + header_cutoff: DateTime::from_timestamp(20, 0).expect("valid header cutoff"), + log_cutoff: DateTime::from_timestamp(5, 0).expect("valid log cutoff"), + }; + let detail_only = UsageCleanupTargets { + detail_body: true, + compressed_body: false, + headers: false, + records: false, + expired_keys: false, + }; + let preview = repository + .preview_usage_cleanup(&window, detail_only, UsageCleanupExecutionMode::Policy) + .await + .expect("MySQL cleanup preview should load"); + assert!(preview.detail >= 1); + let summary = repository + .cleanup_usage( + &window, + 1, + false, + detail_only, + UsageCleanupExecutionMode::Policy, + ) + .await + .expect("MySQL detail cleanup should succeed"); + assert!(summary.body_externalized >= 1); + let body_ref: String = + sqlx::query_scalar("SELECT request_body_ref FROM usage_http_audits WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("externalized body ref should load"); + assert_eq!( + body_ref, + format!("usage://request/{request_id}/request_body") + ); + + let headers_only = UsageCleanupTargets { + detail_body: false, + compressed_body: false, + headers: true, + records: false, + expired_keys: false, + }; + let summary = repository + .cleanup_usage( + &window, + 1, + false, + headers_only, + UsageCleanupExecutionMode::Policy, + ) + .await + .expect("MySQL header cleanup should succeed"); + assert!(summary.header_cleaned >= 1); + let request_headers: Option = + sqlx::query_scalar("SELECT request_headers FROM `usage` WHERE request_id = ?") + .bind(&request_id) + .fetch_one(&pool) + .await + .expect("cleaned headers should load"); + assert!(request_headers.is_none()); + + let before_now_request_id = format!("cleanup-before-now-{suffix}"); + repository + .upsert(sample_usage( + &before_now_request_id, + &user_id, + &api_key_id, + &provider_id, + &provider_key_id, + "completed", + "settled", + 10, + )) + .await + .expect("before-now usage should seed"); + sqlx::query( + "UPDATE `usage` SET request_body = '{\"raw\":true}', request_body_compressed = ? WHERE request_id = ?", + ) + .bind(vec![0x1f_u8, 0x8b]) + .bind(&before_now_request_id) + .execute(&pool) + .await + .expect("before-now fields should seed"); + let summary = repository + .cleanup_usage( + &window, + 1, + false, + UsageCleanupTargets::body_targets(), + UsageCleanupExecutionMode::BeforeNowBodyFields, + ) + .await + .expect("MySQL before-now cleanup should succeed"); + assert!(summary.body_externalized >= 1); + assert!(summary.body_cleaned >= 1); + let body_fields: (Option, Option>) = sqlx::query_as( + "SELECT CAST(request_body AS CHAR), request_body_compressed FROM `usage` WHERE request_id = ?", + ) + .bind(&before_now_request_id) + .fetch_one(&pool) + .await + .expect("before-now fields should load"); + assert_eq!(body_fields, (None, None)); +} + async fn seed_stats_targets( pool: &sqlx::MySqlPool, user_id: &str, diff --git a/crates/aether-data/adapters/mysql/src/video_tasks.rs b/crates/aether-data/adapters/mysql/src/video_tasks.rs index a3b1e7db9..9996d6403 100644 --- a/crates/aether-data/adapters/mysql/src/video_tasks.rs +++ b/crates/aether-data/adapters/mysql/src/video_tasks.rs @@ -99,32 +99,6 @@ impl MysqlVideoTaskRepository { .map_sql_err()?; row.as_ref().map(map_video_task_row).transpose() } - - async fn reload_ids(&self, ids: &[String]) -> Result, DataLayerError> { - if ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(VIDEO_TASK_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for id in ids { - separated.push_bind(id); - } - } - builder.push(")"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut tasks = rows - .iter() - .map(map_video_task_row) - .collect::, _>>()?; - tasks.sort_by(|left, right| { - left.next_poll_at_unix_secs - .cmp(&right.next_poll_at_unix_secs) - .then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs)) - }); - Ok(tasks) - } } #[async_trait] @@ -320,20 +294,77 @@ impl VideoTaskWriteRepository for MysqlVideoTaskRepository { if limit == 0 { return Ok(Vec::new()); } - let due = self.list_due(now_unix_secs, limit).await?; - let ids = due.iter().map(|task| task.id.clone()).collect::>(); - for id in &ids { - sqlx::query( - "UPDATE video_tasks SET next_poll_at = ?, updated_at = GREATEST(updated_at, ?) WHERE id = ?", - ) - .bind(u64_to_i64(claim_until_unix_secs, "video task claim_until")?) - .bind(u64_to_i64(now_unix_secs, "video task now")?) - .bind(id) - .execute(&self.pool) - .await - .map_sql_err()?; + let now = u64_to_i64(now_unix_secs, "video task now")?; + let claim_until = u64_to_i64(claim_until_unix_secs, "video task claim_until")?; + let limit = limit_i64(limit, "due video task claim limit")?; + let mut tx = self.pool.begin().await.map_sql_err()?; + let due_rows = sqlx::query( + r#" +SELECT id +FROM video_tasks +WHERE status IN ('submitted', 'queued', 'processing') + AND next_poll_at IS NOT NULL + AND next_poll_at <= ? + AND poll_count < max_poll_count +ORDER BY next_poll_at ASC, updated_at ASC +LIMIT ? +FOR UPDATE SKIP LOCKED +"#, + ) + .bind(now) + .bind(limit) + .fetch_all(&mut *tx) + .await + .map_sql_err()?; + let ids = due_rows + .iter() + .map(|row| row.try_get::("id").map_sql_err()) + .collect::, _>>()?; + if ids.is_empty() { + tx.commit().await.map_sql_err()?; + return Ok(Vec::new()); } - self.reload_ids(&ids).await + + let mut update = QueryBuilder::::new("UPDATE video_tasks SET next_poll_at = "); + update + .push_bind(claim_until) + .push(", updated_at = GREATEST(updated_at, ") + .push_bind(now) + .push(") WHERE id IN ("); + { + let mut separated = update.separated(", "); + for id in &ids { + separated.push_bind(id); + } + } + update + .push(") AND status IN ('submitted', 'queued', 'processing')") + .push(" AND next_poll_at IS NOT NULL AND next_poll_at <= ") + .push_bind(now) + .push(" AND poll_count < max_poll_count"); + update.build().execute(&mut *tx).await.map_sql_err()?; + + let mut reload = QueryBuilder::::new(VIDEO_TASK_COLUMNS); + reload.push(" WHERE id IN ("); + { + let mut separated = reload.separated(", "); + for id in &ids { + separated.push_bind(id); + } + } + reload.push(")"); + let rows = reload.build().fetch_all(&mut *tx).await.map_sql_err()?; + let mut tasks = rows + .iter() + .map(map_video_task_row) + .collect::, _>>()?; + tasks.sort_by(|left, right| { + left.next_poll_at_unix_secs + .cmp(&right.next_poll_at_unix_secs) + .then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs)) + }); + tx.commit().await.map_sql_err()?; + Ok(tasks) } } @@ -676,6 +707,11 @@ fn optional_u32_to_i32(value: Option, name: &str) -> Result, Da #[cfg(test)] mod tests { use super::MysqlVideoTaskRepository; + use crate::run_migrations; + use aether_data_contracts::repository::video_tasks::{ + UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository, + }; + use std::sync::Arc; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -687,4 +723,116 @@ mod tests { let _repository = MysqlVideoTaskRepository::new(pool); } + + #[tokio::test] + async fn mysql_claim_due_does_not_return_one_task_to_multiple_workers_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!( + "skipping mysql claim_due concurrency test because AETHER_TEST_MYSQL_URL is unset" + ); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(4) + .connect(&database_url) + .await + .expect("mysql test pool should connect"); + run_migrations(&pool) + .await + .expect("mysql migrations should run"); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let task_id = format!("claim-{}", &suffix[..20]); + let repository = MysqlVideoTaskRepository::new(pool.clone()); + repository + .upsert(claimable_task(&task_id)) + .await + .expect("claimable mysql task should insert"); + + let barrier = Arc::new(tokio::sync::Barrier::new(3)); + let first_repository = repository.clone(); + let first_barrier = barrier.clone(); + let first = tokio::spawn(async move { + first_barrier.wait().await; + first_repository.claim_due(100, 130, 1).await + }); + let second_repository = repository.clone(); + let second_barrier = barrier.clone(); + let second = tokio::spawn(async move { + second_barrier.wait().await; + second_repository.claim_due(100, 130, 1).await + }); + barrier.wait().await; + + let first_result = first.await; + let second_result = second.await; + let followup_result = repository.claim_due(100, 130, 1).await; + let cleanup_result = sqlx::query("DELETE FROM video_tasks WHERE id = ?") + .bind(&task_id) + .execute(&pool) + .await; + + let first_claimed = first_result + .expect("first mysql claim worker should join") + .expect("first mysql claim should execute"); + let second_claimed = second_result + .expect("second mysql claim worker should join") + .expect("second mysql claim should execute"); + let followup_claimed = followup_result.expect("follow-up mysql claim should execute"); + cleanup_result.expect("mysql claim fixture should clean up"); + + let claimed_ids = first_claimed + .iter() + .chain(&second_claimed) + .map(|task| task.id.as_str()) + .collect::>(); + assert_eq!(claimed_ids, vec![task_id.as_str()]); + assert!(followup_claimed.is_empty()); + } + + fn claimable_task(id: &str) -> UpsertVideoTask { + UpsertVideoTask { + id: id.to_string(), + short_id: None, + request_id: format!("request-{id}"), + user_id: None, + api_key_id: None, + username: None, + api_key_name: None, + external_task_id: Some(format!("external-{id}")), + provider_id: None, + endpoint_id: None, + key_id: None, + client_api_format: Some("openai:video".to_string()), + provider_api_format: Some("openai:video".to_string()), + format_converted: false, + model: Some("sora-2".to_string()), + prompt: Some("claim test".to_string()), + original_request_body: None, + duration_seconds: None, + resolution: None, + aspect_ratio: None, + size: None, + status: VideoTaskStatus::Submitted, + progress_percent: 0, + progress_message: None, + retry_count: 0, + poll_interval_seconds: 10, + next_poll_at_unix_secs: Some(100), + poll_count: 0, + max_poll_count: 360, + created_at_unix_ms: 90, + submitted_at_unix_secs: Some(90), + completed_at_unix_secs: None, + updated_at_unix_secs: 100, + error_code: None, + error_message: None, + video_url: None, + request_metadata: None, + } + } } diff --git a/crates/aether-data/adapters/mysql/src/wallet.rs b/crates/aether-data/adapters/mysql/src/wallet.rs index 4fbf9b310..8b3cbf935 100644 --- a/crates/aether-data/adapters/mysql/src/wallet.rs +++ b/crates/aether-data/adapters/mysql/src/wallet.rs @@ -1,6 +1,6 @@ use async_trait::async_trait; use chrono::Utc; -use sqlx::{mysql::MySqlRow, Row}; +use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use aether_data_contracts::repository::wallet::{ redeem_code_credits_recharge_balance, redeem_code_payment_method, @@ -17,11 +17,13 @@ use aether_data_contracts::repository::wallet::{ RedeemWalletCodeInput, RedeemWalletCodeOutcome, StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminRedeemCode, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, - StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage, - StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage, - StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger, + StoredAdminRedeemCodePage, StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage, + StoredAdminWalletListItem, StoredAdminWalletListPage, StoredAdminWalletRefund, + StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem, + StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction, + StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome, - WalletReadRepository, WalletReadSeed, WalletReadSnapshot, WalletWriteRepository, + WalletReadRepository, WalletWriteRepository, }; use aether_data_contracts::DataLayerError; @@ -37,38 +39,60 @@ impl MysqlWalletReadRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } } +} - async fn load_snapshot(&self) -> Result { - Ok(WalletReadSnapshot::new(WalletReadSeed { - wallets: self.load_wallets().await?, - payment_orders: self.load_payment_orders().await?, - payment_callbacks: self.load_payment_callbacks().await?, - wallet_transactions: self.load_wallet_transactions().await?, - refunds: self.load_refunds().await?, - redeem_batches: self.load_redeem_batches().await?, - redeem_codes: self.load_redeem_codes().await?, - })) - } - - async fn load_wallets(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" +const ADMIN_WALLET_LIST_SELECT_SQL: &str = r#" SELECT - id, user_id, api_key_id, balance, gift_balance, limit_mode, currency, - status, total_recharged, total_consumed, total_refunded, total_adjusted, - updated_at AS updated_at_unix_secs -FROM wallets -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_wallet_row).collect() - } + w.id, w.user_id, w.api_key_id, w.balance, w.gift_balance, w.limit_mode, + w.currency, w.status, w.total_recharged, w.total_consumed, w.total_refunded, + w.total_adjusted, users.username AS user_name, api_keys.name AS api_key_name, + w.created_at AS created_at_unix_ms, w.updated_at AS updated_at_unix_secs +FROM wallets w +LEFT JOIN users ON users.id = w.user_id +LEFT JOIN api_keys ON api_keys.id = w.api_key_id +WHERE 1 = 1 +"#; - async fn load_payment_orders(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" +const ADMIN_WALLET_LEDGER_SELECT_SQL: &str = r#" +SELECT + tx.id, tx.wallet_id, tx.category, tx.reason_code, tx.amount, + tx.balance_before, tx.balance_after, tx.recharge_balance_before, + tx.recharge_balance_after, tx.gift_balance_before, tx.gift_balance_after, + tx.link_type, tx.link_id, tx.operator_id, tx.description, + w.user_id, w.api_key_id, w.status AS wallet_status, + wallet_users.username AS wallet_user_name, + api_keys.name AS api_key_name, + operator_users.username AS operator_name, + operator_users.email AS operator_email, + tx.created_at AS created_at_unix_ms +FROM wallet_transactions tx +JOIN wallets w ON w.id = tx.wallet_id +LEFT JOIN users wallet_users ON wallet_users.id = w.user_id +LEFT JOIN api_keys ON api_keys.id = w.api_key_id +LEFT JOIN users operator_users ON operator_users.id = tx.operator_id +WHERE 1 = 1 +"#; + +const ADMIN_WALLET_REFUND_REQUEST_SELECT_SQL: &str = r#" +SELECT + rr.id, rr.refund_no, rr.wallet_id, rr.user_id, rr.payment_order_id, + rr.source_type, rr.source_id, rr.refund_mode, rr.amount_usd, rr.status, + rr.reason, rr.failure_reason, rr.gateway_refund_id, rr.payout_method, + rr.payout_reference, rr.payout_proof, rr.requested_by, rr.approved_by, + rr.processed_by, w.user_id AS wallet_user_id, w.api_key_id AS wallet_api_key_id, + w.status AS wallet_status, wallet_users.username AS wallet_user_name, + api_keys.name AS api_key_name, rr.created_at AS created_at_unix_ms, + rr.updated_at AS updated_at_unix_secs, + rr.processed_at AS processed_at_unix_secs, + rr.completed_at AS completed_at_unix_secs +FROM refund_requests rr +JOIN wallets w ON w.id = rr.wallet_id +LEFT JOIN users wallet_users ON wallet_users.id = w.user_id +LEFT JOIN api_keys ON api_keys.id = w.api_key_id +WHERE w.user_id IS NOT NULL +"#; + +const ADMIN_PAYMENT_ORDER_SELECT_SQL: &str = r#" SELECT id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, @@ -79,129 +103,325 @@ SELECT credited_at AS credited_at_unix_secs, expires_at AS expires_at_unix_secs FROM payment_orders -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_payment_order_row).collect() - } +WHERE 1 = 1 +"#; - async fn load_payment_callbacks( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" +const ADMIN_PAYMENT_CALLBACK_SELECT_SQL: &str = r#" SELECT id, payment_order_id, payment_method, callback_key, order_no, gateway_order_id, payload_hash, signature_valid, status, payload, error_message, created_at AS created_at_unix_ms, processed_at AS processed_at_unix_secs FROM payment_callbacks -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_payment_callback_row).collect() - } +WHERE 1 = 1 +"#; - async fn load_wallet_transactions( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - tx.id, tx.wallet_id, tx.category, tx.reason_code, tx.amount, - tx.balance_before, tx.balance_after, tx.recharge_balance_before, - tx.recharge_balance_after, tx.gift_balance_before, tx.gift_balance_after, - tx.link_type, tx.link_id, tx.operator_id, tx.description, - operator_users.username AS operator_name, - operator_users.email AS operator_email, - tx.created_at AS created_at_unix_ms -FROM wallet_transactions tx -LEFT JOIN users operator_users ON operator_users.id = tx.operator_id -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_wallet_transaction_row).collect() - } - - async fn load_refunds(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, refund_no, wallet_id, user_id, payment_order_id, source_type, - source_id, refund_mode, amount_usd, status, reason, failure_reason, - gateway_refund_id, payout_method, payout_reference, payout_proof, - requested_by, approved_by, processed_by, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs, - processed_at AS processed_at_unix_secs, - completed_at AS completed_at_unix_secs -FROM refund_requests -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_refund_row).collect() - } - - async fn load_redeem_batches(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" +const ADMIN_REDEEM_BATCH_SELECT_SQL: &str = r#" SELECT batches.id, batches.name, batches.amount_usd, batches.currency, batches.balance_bucket, batches.total_count, - CAST(COALESCE(SUM(CASE WHEN codes.status = 'redeemed' THEN 1 ELSE 0 END), 0) AS SIGNED) AS redeemed_count, - CAST(COALESCE(SUM(CASE WHEN codes.status = 'active' THEN 1 ELSE 0 END), 0) AS SIGNED) AS active_count, + CAST(COALESCE(stats.redeemed_count, 0) AS SIGNED) AS redeemed_count, + CAST(COALESCE(stats.active_count, 0) AS SIGNED) AS active_count, batches.status, batches.description, batches.created_by, batches.expires_at AS expires_at_unix_secs, batches.created_at AS created_at_unix_ms, batches.updated_at AS updated_at_unix_secs FROM redeem_code_batches AS batches -LEFT JOIN redeem_codes AS codes ON codes.batch_id = batches.id -GROUP BY - batches.id, batches.name, batches.amount_usd, batches.currency, - batches.balance_bucket, batches.total_count, batches.status, - batches.description, batches.created_by, batches.expires_at, - batches.created_at, batches.updated_at -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_redeem_batch_row).collect() - } +LEFT JOIN ( + SELECT + batch_id, + SUM(CASE WHEN status = 'redeemed' THEN 1 ELSE 0 END) AS redeemed_count, + SUM(CASE WHEN status = 'active' THEN 1 ELSE 0 END) AS active_count + FROM redeem_codes + GROUP BY batch_id +) AS stats ON stats.batch_id = batches.id +WHERE 1 = 1 +"#; - async fn load_redeem_codes(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - codes.id, codes.batch_id, batches.name AS batch_name, codes.code_prefix, - codes.code_suffix, codes.status, codes.redeemed_by_user_id, - redeemed_users.username AS redeemed_by_user_name, - codes.redeemed_wallet_id, codes.redeemed_payment_order_id, - orders.order_no AS redeemed_order_no, - codes.redeemed_at AS redeemed_at_unix_secs, - codes.disabled_by, - batches.expires_at AS expires_at_unix_secs, - codes.created_at AS created_at_unix_ms, - codes.updated_at AS updated_at_unix_secs -FROM redeem_codes AS codes -JOIN redeem_code_batches AS batches ON batches.id = codes.batch_id -LEFT JOIN users AS redeemed_users ON redeemed_users.id = codes.redeemed_by_user_id -LEFT JOIN payment_orders AS orders ON orders.id = codes.redeemed_payment_order_id -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_redeem_code_row).collect() +fn wallets_by_owner_ids_builder<'a>( + owner_column: &'static str, + owner_ids: &'a [String], +) -> QueryBuilder<'a, MySql> { + assert!(matches!(owner_column, "user_id" | "api_key_id")); + let mut builder = QueryBuilder::::new(wallet_select_sql("")); + builder.push("WHERE ").push(owner_column).push(" IN ("); + let mut separated = builder.separated(", "); + for owner_id in owner_ids { + separated.push_bind(owner_id); } + separated.push_unseparated(") ORDER BY id ASC"); + builder +} + +fn push_admin_wallet_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminWalletListQuery, +) { + if let Some(status) = query.status.as_deref() { + builder.push(" AND w.status = ").push_bind(status); + } + match query.owner_type.as_deref() { + Some("user") => { + builder.push(" AND w.user_id IS NOT NULL"); + } + Some("api_key") => { + builder.push(" AND w.api_key_id IS NOT NULL"); + } + _ => {} + } +} + +fn admin_wallet_count_builder<'a>(query: &'a AdminWalletListQuery) -> QueryBuilder<'a, MySql> { + let mut builder = + QueryBuilder::::new("SELECT COUNT(*) AS total FROM wallets w WHERE 1 = 1"); + push_admin_wallet_filters(&mut builder, query); + builder +} + +fn admin_wallet_list_builder<'a>( + query: &'a AdminWalletListQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_WALLET_LIST_SELECT_SQL); + push_admin_wallet_filters(&mut builder, query); + builder + .push(" ORDER BY w.updated_at DESC, w.id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_wallet_ledger_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminWalletLedgerQuery, +) { + if let Some(category) = query.category.as_deref() { + builder.push(" AND tx.category = ").push_bind(category); + } + if let Some(reason_code) = query.reason_code.as_deref() { + builder + .push(" AND tx.reason_code = ") + .push_bind(reason_code); + } + match query.owner_type.as_deref() { + Some("user") => { + builder.push(" AND w.user_id IS NOT NULL"); + } + Some("api_key") => { + builder.push(" AND w.api_key_id IS NOT NULL"); + } + _ => {} + } +} + +fn admin_wallet_ledger_count_builder<'a>( + query: &'a AdminWalletLedgerQuery, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new( + "SELECT COUNT(*) AS total FROM wallet_transactions tx JOIN wallets w ON w.id = tx.wallet_id WHERE 1 = 1", + ); + push_admin_wallet_ledger_filters(&mut builder, query); + builder +} + +fn admin_wallet_ledger_list_builder<'a>( + query: &'a AdminWalletLedgerQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_WALLET_LEDGER_SELECT_SQL); + push_admin_wallet_ledger_filters(&mut builder, query); + builder + .push(" ORDER BY tx.created_at DESC, tx.id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_wallet_refund_request_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminWalletRefundRequestListQuery, +) { + if let Some(status) = query.status.as_deref() { + builder.push(" AND rr.status = ").push_bind(status); + } +} + +fn admin_wallet_refund_request_count_builder<'a>( + query: &'a AdminWalletRefundRequestListQuery, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new( + "SELECT COUNT(*) AS total FROM refund_requests rr JOIN wallets w ON w.id = rr.wallet_id WHERE w.user_id IS NOT NULL", + ); + push_admin_wallet_refund_request_filters(&mut builder, query); + builder +} + +fn admin_wallet_refund_request_list_builder<'a>( + query: &'a AdminWalletRefundRequestListQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_WALLET_REFUND_REQUEST_SELECT_SQL); + push_admin_wallet_refund_request_filters(&mut builder, query); + builder + .push(" ORDER BY rr.created_at DESC, rr.id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_payment_order_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminPaymentOrderListQuery, + now: i64, +) { + if let Some(payment_method) = query.payment_method.as_deref() { + builder + .push(" AND payment_method = ") + .push_bind(payment_method); + } + if let Some(status) = query.status.as_deref() { + builder + .push(" AND (CASE WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at < ") + .push_bind(now) + .push(" THEN 'expired' ELSE status END) = ") + .push_bind(status); + } +} + +fn admin_payment_order_count_builder<'a>( + query: &'a AdminPaymentOrderListQuery, + now: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = + QueryBuilder::::new("SELECT COUNT(*) AS total FROM payment_orders WHERE 1 = 1"); + push_admin_payment_order_filters(&mut builder, query, now); + builder +} + +fn admin_payment_order_list_builder<'a>( + query: &'a AdminPaymentOrderListQuery, + now: i64, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_PAYMENT_ORDER_SELECT_SQL); + push_admin_payment_order_filters(&mut builder, query, now); + builder + .push(" ORDER BY created_at DESC, id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_payment_callback_filter<'a>( + builder: &mut QueryBuilder<'a, MySql>, + payment_method: Option<&'a str>, +) { + if let Some(payment_method) = payment_method { + builder + .push(" AND payment_method = ") + .push_bind(payment_method); + } +} + +fn admin_payment_callback_count_builder(payment_method: Option<&str>) -> QueryBuilder<'_, MySql> { + let mut builder = + QueryBuilder::::new("SELECT COUNT(*) AS total FROM payment_callbacks WHERE 1 = 1"); + push_admin_payment_callback_filter(&mut builder, payment_method); + builder +} + +fn admin_payment_callback_list_builder( + payment_method: Option<&str>, + limit: i64, + offset: i64, +) -> QueryBuilder<'_, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_PAYMENT_CALLBACK_SELECT_SQL); + push_admin_payment_callback_filter(&mut builder, payment_method); + builder + .push(" ORDER BY created_at DESC, id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_redeem_batch_filter<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminRedeemCodeBatchListQuery, +) { + if let Some(status) = query.status.as_deref() { + builder.push(" AND batches.status = ").push_bind(status); + } +} + +fn admin_redeem_batch_count_builder<'a>( + query: &'a AdminRedeemCodeBatchListQuery, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new( + "SELECT COUNT(*) AS total FROM redeem_code_batches AS batches WHERE 1 = 1", + ); + push_admin_redeem_batch_filter(&mut builder, query); + builder +} + +fn admin_redeem_batch_list_builder<'a>( + query: &'a AdminRedeemCodeBatchListQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(ADMIN_REDEEM_BATCH_SELECT_SQL); + push_admin_redeem_batch_filter(&mut builder, query); + builder + .push(" ORDER BY batches.created_at DESC, batches.id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder +} + +fn push_admin_redeem_code_filters<'a>( + builder: &mut QueryBuilder<'a, MySql>, + query: &'a AdminRedeemCodeListQuery, +) { + builder + .push(" AND codes.batch_id = ") + .push_bind(&query.batch_id); + if let Some(status) = query.status.as_deref() { + builder.push(" AND codes.status = ").push_bind(status); + } +} + +fn admin_redeem_code_count_builder<'a>( + query: &'a AdminRedeemCodeListQuery, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new( + "SELECT COUNT(*) AS total FROM redeem_codes AS codes WHERE 1 = 1", + ); + push_admin_redeem_code_filters(&mut builder, query); + builder +} + +fn admin_redeem_code_list_builder<'a>( + query: &'a AdminRedeemCodeListQuery, + limit: i64, + offset: i64, +) -> QueryBuilder<'a, MySql> { + let mut builder = QueryBuilder::::new(redeem_code_select_sql("WHERE 1 = 1")); + push_admin_redeem_code_filters(&mut builder, query); + builder + .push(" ORDER BY codes.created_at DESC, codes.id DESC LIMIT ") + .push_bind(limit) + .push(" OFFSET ") + .push_bind(offset); + builder } #[async_trait] @@ -210,7 +430,18 @@ impl WalletReadRepository for MysqlWalletReadRepository { &self, key: WalletLookupKey<'_>, ) -> Result, DataLayerError> { - Ok(self.load_snapshot().await?.find(key)) + let (where_clause, bind) = match key { + WalletLookupKey::WalletId(value) => ("WHERE id = ? LIMIT 1", value), + WalletLookupKey::UserId(value) => ("WHERE user_id = ? LIMIT 1", value), + WalletLookupKey::ApiKeyId(value) => ("WHERE api_key_id = ? LIMIT 1", value), + }; + let sql = wallet_select_sql(where_clause); + let row = sqlx::query(&sql) + .bind(bind) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_wallet_row).transpose() } async fn update_auth_user_wallet_limit_mode( @@ -347,44 +578,111 @@ impl WalletReadRepository for MysqlWalletReadRepository { &self, user_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_wallets_by_user_ids(user_ids)) + if user_ids.is_empty() { + return Ok(Vec::new()); + } + let mut builder = wallets_by_owner_ids_builder("user_id", user_ids); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_wallet_row).collect() } async fn list_wallets_by_api_key_ids( &self, api_key_ids: &[String], ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .list_wallets_by_api_key_ids(api_key_ids)) + if api_key_ids.is_empty() { + return Ok(Vec::new()); + } + let mut builder = wallets_by_owner_ids_builder("api_key_id", api_key_ids); + let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; + rows.iter().map(map_wallet_row).collect() } async fn list_admin_wallets( &self, query: &AdminWalletListQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_admin_wallets(query)) + let mut count_builder = admin_wallet_count_builder(query); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_wallet_list_builder( + query, + i64_from_usize(query.limit, "wallet limit")?, + i64_from_usize(query.offset, "wallet offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_admin_wallet_list_item_row) + .collect::, _>>()?; + Ok(StoredAdminWalletListPage { items, total }) } async fn list_admin_wallet_ledger( &self, query: &AdminWalletLedgerQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_admin_wallet_ledger(query)) + let mut count_builder = admin_wallet_ledger_count_builder(query); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_wallet_ledger_list_builder( + query, + i64_from_usize(query.limit, "wallet ledger limit")?, + i64_from_usize(query.offset, "wallet ledger offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_admin_wallet_ledger_item_row) + .collect::, _>>()?; + Ok(StoredAdminWalletLedgerPage { items, total }) } async fn list_admin_wallet_refund_requests( &self, query: &AdminWalletRefundRequestListQuery, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_wallet_refund_requests(query)) + let mut count_builder = admin_wallet_refund_request_count_builder(query); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_wallet_refund_request_list_builder( + query, + i64_from_usize(query.limit, "wallet refund request limit")?, + i64_from_usize(query.offset, "wallet refund request offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_admin_wallet_refund_request_item_row) + .collect::, _>>()?; + Ok(StoredAdminWalletRefundRequestPage { items, total }) } async fn list_admin_wallet_transactions( @@ -393,10 +691,41 @@ impl WalletReadRepository for MysqlWalletReadRepository { limit: usize, offset: usize, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_wallet_transactions(wallet_id, limit, offset)) + let total = read_count_row( + sqlx::query("SELECT COUNT(*) AS total FROM wallet_transactions WHERE wallet_id = ?") + .bind(wallet_id) + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let rows = sqlx::query( + r#" +SELECT + tx.id, tx.wallet_id, tx.category, tx.reason_code, tx.amount, + tx.balance_before, tx.balance_after, tx.recharge_balance_before, + tx.recharge_balance_after, tx.gift_balance_before, tx.gift_balance_after, + tx.link_type, tx.link_id, tx.operator_id, tx.description, + operator_users.username AS operator_name, + operator_users.email AS operator_email, + tx.created_at AS created_at_unix_ms +FROM wallet_transactions tx +LEFT JOIN users operator_users ON operator_users.id = tx.operator_id +WHERE tx.wallet_id = ? +ORDER BY tx.created_at DESC, tx.id DESC +LIMIT ? OFFSET ? +"#, + ) + .bind(wallet_id) + .bind(i64_from_usize(limit, "wallet transaction limit")?) + .bind(i64_from_usize(offset, "wallet transaction offset")?) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_wallet_transaction_row) + .collect::, _>>()?; + Ok(StoredAdminWalletTransactionPage { items, total }) } async fn find_wallet_today_usage( @@ -444,7 +773,7 @@ WHERE wallet_id = ? .bind(wallet_id) .bind(billing_timezone) .bind(billing_date) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) + .bind(i64_from_usize(limit, "wallet daily usage history limit")?) .fetch_all(&self.pool) .await .map_sql_err()?; @@ -464,30 +793,72 @@ WHERE wallet_id = ? limit: usize, offset: usize, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_wallet_refunds(wallet_id, limit, offset)) + let total = read_count_row( + sqlx::query("SELECT COUNT(*) AS total FROM refund_requests WHERE wallet_id = ?") + .bind(wallet_id) + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let sql = refund_select_sql( + "WHERE wallet_id = ? ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?", + ); + let rows = sqlx::query(&sql) + .bind(wallet_id) + .bind(i64_from_usize(limit, "wallet refund limit")?) + .bind(i64_from_usize(offset, "wallet refund offset")?) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_refund_row) + .collect::, _>>()?; + Ok(StoredAdminWalletRefundPage { items, total }) } async fn list_admin_payment_orders( &self, query: &AdminPaymentOrderListQuery, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_payment_orders(query, current_unix_secs_i64().max(0) as u64)) + let now = current_unix_secs_i64(); + let mut count_builder = admin_payment_order_count_builder(query, now); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_payment_order_list_builder( + query, + now, + i64_from_usize(query.limit, "payment order limit")?, + i64_from_usize(query.offset, "payment order offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_payment_order_row) + .collect::, _>>()?; + Ok(StoredAdminPaymentOrderPage { items, total }) } async fn find_admin_payment_order( &self, order_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .find_admin_payment_order(order_id)) + let sql = payment_order_select_sql("WHERE id = ? LIMIT 1"); + let row = sqlx::query(&sql) + .bind(order_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_payment_order_row).transpose() } async fn list_wallet_payment_orders_by_user_id( @@ -496,27 +867,83 @@ WHERE wallet_id = ? limit: usize, offset: usize, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_wallet_payment_orders_by_user_id(user_id, limit, offset)) + let total = read_count_row( + sqlx::query("SELECT COUNT(*) AS total FROM payment_orders WHERE user_id = ?") + .bind(user_id) + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let rows = sqlx::query( + r#" +SELECT + id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, + exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, + payment_provider, payment_channel, order_kind, product_id, product_snapshot, + gateway_order_id, gateway_response, + CASE + WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at < ? THEN 'expired' + ELSE status + END AS status, + created_at AS created_at_unix_ms, + paid_at AS paid_at_unix_secs, + credited_at AS credited_at_unix_secs, + expires_at AS expires_at_unix_secs +FROM payment_orders +WHERE user_id = ? +ORDER BY created_at DESC, id DESC +LIMIT ? OFFSET ? +"#, + ) + .bind(current_unix_secs_i64()) + .bind(user_id) + .bind(i64_from_usize(limit, "wallet payment order limit")?) + .bind(i64_from_usize(offset, "wallet payment order offset")?) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_payment_order_row) + .collect::, _>>()?; + Ok(StoredAdminPaymentOrderPage { items, total }) } async fn count_pending_refunds_by_user_id(&self, user_id: &str) -> Result { - Ok(self - .load_snapshot() - .await? - .count_pending_refunds_by_user_id(user_id)) + read_count_row( + sqlx::query( + r#" +SELECT COUNT(*) AS total +FROM refund_requests +WHERE user_id = ? + AND status IN ('pending_approval', 'approved', 'processing') +"#, + ) + .bind(user_id) + .fetch_one(&self.pool) + .await + .map_sql_err()?, + ) } async fn count_pending_payment_orders_by_user_id( &self, user_id: &str, ) -> Result { - Ok(self - .load_snapshot() - .await? - .count_pending_payment_orders_by_user_id(user_id)) + read_count_row( + sqlx::query( + r#" +SELECT COUNT(*) AS total +FROM payment_orders +WHERE user_id = ? + AND status IN ('pending', 'paid') +"#, + ) + .bind(user_id) + .fetch_one(&self.pool) + .await + .map_sql_err()?, + ) } async fn find_wallet_payment_order_by_user_id( @@ -524,10 +951,34 @@ WHERE wallet_id = ? user_id: &str, order_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .find_wallet_payment_order_by_user_id(user_id, order_id)) + let row = sqlx::query( + r#" +SELECT + id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, + exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, + payment_provider, payment_channel, order_kind, product_id, product_snapshot, + gateway_order_id, gateway_response, + CASE + WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at < ? THEN 'expired' + ELSE status + END AS status, + created_at AS created_at_unix_ms, + paid_at AS paid_at_unix_secs, + credited_at AS credited_at_unix_secs, + expires_at AS expires_at_unix_secs +FROM payment_orders +WHERE user_id = ? + AND id = ? +LIMIT 1 +"#, + ) + .bind(current_unix_secs_i64()) + .bind(user_id) + .bind(order_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_payment_order_row).transpose() } async fn find_pending_plan_purchase_order_by_user_id( @@ -561,10 +1012,14 @@ LIMIT 1 wallet_id: &str, refund_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .find_wallet_refund(wallet_id, refund_id)) + let sql = refund_select_sql("WHERE wallet_id = ? AND id = ? LIMIT 1"); + let row = sqlx::query(&sql) + .bind(wallet_id) + .bind(refund_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_refund_row).transpose() } async fn list_admin_payment_callbacks( @@ -573,37 +1028,100 @@ LIMIT 1 limit: usize, offset: usize, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_payment_callbacks(payment_method, limit, offset)) + let mut count_builder = admin_payment_callback_count_builder(payment_method); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_payment_callback_list_builder( + payment_method, + i64_from_usize(limit, "payment callback limit")?, + i64_from_usize(offset, "payment callback offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_payment_callback_row) + .collect::, _>>()?; + Ok(StoredAdminPaymentCallbackPage { items, total }) } async fn list_admin_redeem_code_batches( &self, query: &AdminRedeemCodeBatchListQuery, ) -> Result { - Ok(self - .load_snapshot() - .await? - .list_admin_redeem_code_batches(query)) + let mut count_builder = admin_redeem_batch_count_builder(query); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_redeem_batch_list_builder( + query, + i64_from_usize(query.limit, "redeem code batch limit")?, + i64_from_usize(query.offset, "redeem code batch offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_redeem_batch_row) + .collect::, _>>()?; + Ok(StoredAdminRedeemCodeBatchPage { items, total }) } async fn find_admin_redeem_code_batch( &self, batch_id: &str, ) -> Result, DataLayerError> { - Ok(self - .load_snapshot() - .await? - .find_admin_redeem_code_batch(batch_id)) + let sql = redeem_batch_select_sql("WHERE batches.id = ?"); + let row = sqlx::query(&sql) + .bind(batch_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_redeem_batch_row).transpose() } async fn list_admin_redeem_codes( &self, query: &AdminRedeemCodeListQuery, ) -> Result { - Ok(self.load_snapshot().await?.list_admin_redeem_codes(query)) + let mut count_builder = admin_redeem_code_count_builder(query); + let total = read_count_row( + count_builder + .build() + .fetch_one(&self.pool) + .await + .map_sql_err()?, + )?; + let mut list_builder = admin_redeem_code_list_builder( + query, + i64_from_usize(query.limit, "redeem code limit")?, + i64_from_usize(query.offset, "redeem code offset")?, + ); + let rows = list_builder + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; + let items = rows + .iter() + .map(map_redeem_code_row) + .collect::, _>>()?; + Ok(StoredAdminRedeemCodePage { items, total }) } } @@ -3157,10 +3675,123 @@ fn map_wallet_row(row: &MySqlRow) -> Result Result { + Ok(StoredAdminWalletListItem { + id: get(row, "id")?, + user_id: get(row, "user_id")?, + api_key_id: get(row, "api_key_id")?, + balance: get(row, "balance")?, + gift_balance: get(row, "gift_balance")?, + limit_mode: get(row, "limit_mode")?, + currency: get(row, "currency")?, + status: get(row, "status")?, + total_recharged: get(row, "total_recharged")?, + total_consumed: get(row, "total_consumed")?, + total_refunded: get(row, "total_refunded")?, + total_adjusted: get(row, "total_adjusted")?, + user_name: get(row, "user_name")?, + api_key_name: get(row, "api_key_name")?, + created_at_unix_ms: optional_timestamp( + get(row, "created_at_unix_ms")?, + "wallets.created_at", + )?, + updated_at_unix_secs: optional_timestamp( + get(row, "updated_at_unix_secs")?, + "wallets.updated_at", + )?, + }) +} + +fn map_admin_wallet_ledger_item_row( + row: &MySqlRow, +) -> Result { + Ok(StoredAdminWalletLedgerItem { + id: get(row, "id")?, + wallet_id: get(row, "wallet_id")?, + category: get(row, "category")?, + reason_code: get(row, "reason_code")?, + amount: get(row, "amount")?, + balance_before: get(row, "balance_before")?, + balance_after: get(row, "balance_after")?, + recharge_balance_before: get(row, "recharge_balance_before")?, + recharge_balance_after: get(row, "recharge_balance_after")?, + gift_balance_before: get(row, "gift_balance_before")?, + gift_balance_after: get(row, "gift_balance_after")?, + link_type: get(row, "link_type")?, + link_id: get(row, "link_id")?, + operator_id: get(row, "operator_id")?, + operator_name: get(row, "operator_name")?, + operator_email: get(row, "operator_email")?, + description: get(row, "description")?, + wallet_user_id: get(row, "user_id")?, + wallet_user_name: get(row, "wallet_user_name")?, + wallet_api_key_id: get(row, "api_key_id")?, + api_key_name: get(row, "api_key_name")?, + wallet_status: get(row, "wallet_status")?, + created_at_unix_ms: optional_timestamp( + get(row, "created_at_unix_ms")?, + "wallet_transactions.created_at", + )?, + }) +} + +fn map_admin_wallet_refund_request_item_row( + row: &MySqlRow, +) -> Result { + Ok(StoredAdminWalletRefundRequestItem { + id: get(row, "id")?, + refund_no: get(row, "refund_no")?, + wallet_id: get(row, "wallet_id")?, + user_id: get(row, "user_id")?, + payment_order_id: get(row, "payment_order_id")?, + source_type: get(row, "source_type")?, + source_id: get(row, "source_id")?, + refund_mode: get(row, "refund_mode")?, + amount_usd: get(row, "amount_usd")?, + status: get(row, "status")?, + reason: get(row, "reason")?, + failure_reason: get(row, "failure_reason")?, + gateway_refund_id: get(row, "gateway_refund_id")?, + payout_method: get(row, "payout_method")?, + payout_reference: get(row, "payout_reference")?, + payout_proof: optional_json(get(row, "payout_proof")?, "refund_requests.payout_proof")?, + requested_by: get(row, "requested_by")?, + approved_by: get(row, "approved_by")?, + processed_by: get(row, "processed_by")?, + wallet_user_id: get(row, "wallet_user_id")?, + wallet_user_name: get(row, "wallet_user_name")?, + wallet_api_key_id: get(row, "wallet_api_key_id")?, + api_key_name: get(row, "api_key_name")?, + wallet_status: get(row, "wallet_status")?, + created_at_unix_ms: optional_timestamp( + get(row, "created_at_unix_ms")?, + "refund_requests.created_at", + )?, + updated_at_unix_secs: optional_timestamp( + get(row, "updated_at_unix_secs")?, + "refund_requests.updated_at", + )?, + processed_at_unix_secs: optional_timestamp( + get(row, "processed_at_unix_secs")?, + "refund_requests.processed_at", + )?, + completed_at_unix_secs: optional_timestamp( + get(row, "completed_at_unix_secs")?, + "refund_requests.completed_at", + )?, + }) +} + fn current_unix_secs_i64() -> i64 { Utc::now().timestamp().max(0) } +fn i64_from_usize(value: usize, field_name: &str) -> Result { + i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field_name} overflow"))) +} + fn json_string(value: &serde_json::Value, field_name: &str) -> Result { serde_json::to_string(value).map_err(|err| { DataLayerError::UnexpectedValue(format!("{field_name} could not be encoded: {err}")) @@ -4209,6 +4840,10 @@ where row.try_get(field).map_sql_err() } +fn read_count_row(row: MySqlRow) -> Result { + nonnegative_u64(get(&row, "total")?, "count total") +} + fn optional_json( value: Option, field_name: &str, diff --git a/crates/aether-data/adapters/mysql/src/wallet/tests.rs b/crates/aether-data/adapters/mysql/src/wallet/tests.rs index 2177f8fed..3ed08cd1d 100644 --- a/crates/aether-data/adapters/mysql/src/wallet/tests.rs +++ b/crates/aether-data/adapters/mysql/src/wallet/tests.rs @@ -1,10 +1,119 @@ -use super::MysqlWalletReadRepository; +use super::{ + admin_payment_callback_list_builder, admin_payment_order_list_builder, + admin_redeem_batch_list_builder, admin_redeem_code_list_builder, + admin_wallet_ledger_list_builder, admin_wallet_list_builder, + admin_wallet_refund_request_list_builder, wallets_by_owner_ids_builder, + MysqlWalletReadRepository, +}; use crate::run_migrations; use aether_data_contracts::repository::wallet::{ - AdminPaymentOrderListQuery, AdminRedeemCodeListQuery, AdminWalletListQuery, WalletLookupKey, - WalletReadRepository, + AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery, + AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletRefundRequestListQuery, + WalletLookupKey, WalletReadRepository, }; +#[test] +fn mysql_wallet_builders_bind_filters_and_page_bounds() { + let injected_status = "active' OR 1 = 1 --".to_string(); + let query = AdminWalletListQuery { + status: Some(injected_status.clone()), + owner_type: Some("user".to_string()), + limit: 7, + offset: 3, + }; + let builder = admin_wallet_list_builder(&query, 7, 3); + let sql = compact_sql(builder.sql()); + + assert!(sql.contains("AND w.status = ?")); + assert!(sql.contains("AND w.user_id IS NOT NULL")); + assert!(sql.contains("ORDER BY w.updated_at DESC, w.id DESC LIMIT ? OFFSET ?")); + assert!(!sql.contains(&injected_status)); + assert_eq!(sql.matches('?').count(), 3); + + let ids = vec!["user-1".to_string(), "user-2' OR 1 = 1 --".to_string()]; + let builder = wallets_by_owner_ids_builder("user_id", &ids); + let sql = compact_sql(builder.sql()); + assert!(sql.contains("WHERE user_id IN (?, ?) ORDER BY id ASC")); + assert!(!sql.contains(&ids[0])); + assert!(!sql.contains(&ids[1])); +} + +#[test] +fn mysql_wallet_admin_builders_cover_filters_ordering_and_mapping_columns() { + let ledger_query = AdminWalletLedgerQuery { + category: Some("credit".to_string()), + reason_code: Some("manual_adjustment".to_string()), + owner_type: Some("api_key".to_string()), + limit: 5, + offset: 2, + }; + let ledger_sql = compact_sql(admin_wallet_ledger_list_builder(&ledger_query, 5, 2).sql()); + assert!(ledger_sql.contains("tx.category = ?")); + assert!(ledger_sql.contains("tx.reason_code = ?")); + assert!(ledger_sql.contains("w.api_key_id IS NOT NULL")); + assert!(ledger_sql.contains("wallet_users.username AS wallet_user_name")); + assert!(ledger_sql.contains("ORDER BY tx.created_at DESC, tx.id DESC LIMIT ? OFFSET ?")); + + let refund_query = AdminWalletRefundRequestListQuery { + status: Some("pending_approval".to_string()), + limit: 4, + offset: 1, + }; + let refund_sql = + compact_sql(admin_wallet_refund_request_list_builder(&refund_query, 4, 1).sql()); + assert!(refund_sql.contains("WHERE w.user_id IS NOT NULL")); + assert!(refund_sql.contains("rr.status = ?")); + assert!(refund_sql.contains("rr.payout_proof")); + assert!(refund_sql.contains("ORDER BY rr.created_at DESC, rr.id DESC LIMIT ? OFFSET ?")); + + let order_query = AdminPaymentOrderListQuery { + status: Some("expired".to_string()), + payment_method: Some("card".to_string()), + limit: 8, + offset: 6, + }; + let order_sql = compact_sql(admin_payment_order_list_builder(&order_query, 100, 8, 6).sql()); + assert!(order_sql.contains("payment_method = ?")); + assert!(order_sql.contains( + "CASE WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at < ? THEN 'expired'" + )); + assert!(order_sql.contains("ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?")); + + let callback_sql = + compact_sql(admin_payment_callback_list_builder(Some("redeem_code"), 9, 4).sql()); + assert!(callback_sql.contains("payment_method = ?")); + assert!(callback_sql.contains("ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?")); + + let batch_query = AdminRedeemCodeBatchListQuery { + status: Some("active".to_string()), + limit: 2, + offset: 1, + }; + let batch_sql = compact_sql(admin_redeem_batch_list_builder(&batch_query, 2, 1).sql()); + assert!(batch_sql.contains("batches.status = ?")); + assert!(batch_sql.contains("AS redeemed_count")); + assert!( + batch_sql.contains("ORDER BY batches.created_at DESC, batches.id DESC LIMIT ? OFFSET ?") + ); + + let code_query = AdminRedeemCodeListQuery { + batch_id: "batch-1".to_string(), + status: Some("redeemed".to_string()), + limit: 3, + offset: 2, + }; + let code_sql = compact_sql(admin_redeem_code_list_builder(&code_query, 3, 2).sql()); + assert!(code_sql.contains("codes.batch_id = ?")); + assert!(code_sql.contains("codes.status = ?")); + assert!(code_sql.contains("orders.order_no AS redeemed_order_no")); + assert!(code_sql.contains("ORDER BY codes.created_at DESC, codes.id DESC LIMIT ? OFFSET ?")); + assert!(!code_sql.contains("batch-1")); +} + +fn compact_sql(sql: &str) -> String { + sql.split_whitespace().collect::>().join(" ") +} + #[tokio::test] async fn mysql_wallet_read_repository_reads_wallet_contract_views() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") @@ -34,22 +143,93 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() { .expect("wallet should exist"); assert_eq!(wallet.total_adjusted, 3.0); + let user_wallets = repository + .list_wallets_by_user_ids(&[ + "user-2".to_string(), + "missing-user".to_string(), + "user-1".to_string(), + ]) + .await + .expect("user wallets should list"); + assert_eq!( + user_wallets + .iter() + .map(|wallet| wallet.id.as_str()) + .collect::>(), + vec!["wallet-1", "wallet-2"] + ); + let api_key_wallets = repository + .list_wallets_by_api_key_ids(&["api-key-1".to_string()]) + .await + .expect("api key wallets should list"); + assert_eq!(api_key_wallets[0].id, "wallet-api-1"); + assert!(repository + .list_wallets_by_user_ids(&[]) + .await + .expect("empty user ids should short circuit") + .is_empty()); + let page = repository .list_admin_wallets(&AdminWalletListQuery { status: Some("active".to_string()), owner_type: Some("user".to_string()), + limit: 1, + offset: 1, + }) + .await + .expect("admin wallets should list"); + assert_eq!(page.total, 2); + assert_eq!(page.items.len(), 1); + let wallet_item = &page.items[0]; + assert_eq!(wallet_item.id, "wallet-1"); + assert_eq!(wallet_item.total_adjusted, 3.0); + + let unknown_owner = repository + .list_admin_wallets(&AdminWalletListQuery { + status: Some("active".to_string()), + owner_type: Some("unknown".to_string()), limit: 10, offset: 0, }) .await - .expect("admin wallets should list"); - let wallet_item = page - .items - .iter() - .find(|item| item.id == "wallet-1") - .expect("seeded wallet should be listed"); - assert!(page.total >= 1); - assert_eq!(wallet_item.total_adjusted, 3.0); + .expect("unknown owner type should preserve snapshot semantics"); + assert_eq!(unknown_owner.total, 3); + + let ledger = repository + .list_admin_wallet_ledger(&AdminWalletLedgerQuery { + category: Some("credit".to_string()), + reason_code: Some("manual_adjustment".to_string()), + owner_type: Some("user".to_string()), + limit: 1, + offset: 1, + }) + .await + .expect("admin ledger should list"); + assert_eq!(ledger.total, 2); + assert_eq!(ledger.items[0].id, "tx-1"); + assert_eq!(ledger.items[0].wallet_user_name.as_deref(), Some("Alice")); + + let refund_requests = repository + .list_admin_wallet_refund_requests(&AdminWalletRefundRequestListQuery { + status: Some("pending_approval".to_string()), + limit: 10, + offset: 0, + }) + .await + .expect("admin refund requests should list"); + assert_eq!(refund_requests.total, 1); + assert_eq!(refund_requests.items[0].id, "refund-2"); + assert_eq!( + refund_requests.items[0].wallet_user_name.as_deref(), + Some("Alice") + ); + + let transactions = repository + .list_admin_wallet_transactions("wallet-1", 1, 1) + .await + .expect("wallet transactions should page"); + assert_eq!(transactions.total, 2); + assert_eq!(transactions.items[0].id, "tx-1"); let orders = repository .list_admin_payment_orders(&AdminPaymentOrderListQuery { @@ -66,23 +246,99 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() { true ); + let expired_orders = repository + .list_admin_payment_orders(&AdminPaymentOrderListQuery { + status: Some("expired".to_string()), + payment_method: Some("card".to_string()), + limit: 10, + offset: 0, + }) + .await + .expect("expired payment orders should filter by effective status"); + assert_eq!(expired_orders.total, 1); + assert_eq!(expired_orders.items[0].id, "order-2"); + + let user_orders = repository + .list_wallet_payment_orders_by_user_id("user-1", 2, 0) + .await + .expect("user payment orders should page"); + assert_eq!(user_orders.total, 3); + assert_eq!(user_orders.items.len(), 2); + assert_eq!(user_orders.items[0].id, "order-3"); + assert_eq!(user_orders.items[1].id, "order-2"); + assert_eq!(user_orders.items[1].status, "expired"); + assert_eq!( + repository + .count_pending_payment_orders_by_user_id("user-1") + .await + .expect("pending payment orders should count"), + 2 + ); + assert_eq!( + repository + .find_admin_payment_order("order-1") + .await + .expect("admin payment order should query") + .expect("admin payment order should exist") + .id, + "order-1" + ); + assert!(repository + .find_wallet_payment_order_by_user_id("user-2", "order-1") + .await + .expect("cross-user payment order lookup should query") + .is_none()); + let refunds = repository - .list_admin_wallet_refunds("wallet-1", 10, 0) + .list_admin_wallet_refunds("wallet-1", 1, 1) .await .expect("refunds should list"); - assert_eq!(refunds.total, 1); + assert_eq!(refunds.total, 2); + assert_eq!(refunds.items[0].id, "refund-1"); assert_eq!( refunds.items[0].payout_proof.as_ref().unwrap()["proof"], "ok" ); + assert_eq!( + repository + .count_pending_refunds_by_user_id("user-1") + .await + .expect("pending refunds should count"), + 1 + ); + assert!(repository + .find_wallet_refund("wallet-2", "refund-1") + .await + .expect("cross-wallet refund lookup should query") + .is_none()); let callbacks = repository - .list_admin_payment_callbacks(Some("redeem_code"), 10, 0) + .list_admin_payment_callbacks(Some("redeem_code"), 1, 1) .await .expect("callbacks should list"); - assert_eq!(callbacks.total, 1); + assert_eq!(callbacks.total, 2); + assert_eq!(callbacks.items[0].id, "callback-1"); assert!(callbacks.items[0].signature_valid); + let batches = repository + .list_admin_redeem_code_batches(&AdminRedeemCodeBatchListQuery { + status: Some("active".to_string()), + limit: 1, + offset: 0, + }) + .await + .expect("redeem batches should filter"); + assert_eq!(batches.total, 1); + assert_eq!(batches.items[0].id, "batch-1"); + assert_eq!(batches.items[0].redeemed_count, 1); + assert_eq!(batches.items[0].active_count, 1); + let batch = repository + .find_admin_redeem_code_batch("batch-2") + .await + .expect("redeem batch should query") + .expect("redeem batch should exist"); + assert_eq!(batch.active_count, 1); + let codes = repository .list_admin_redeem_codes(&AdminRedeemCodeListQuery { batch_id: "batch-1".to_string(), @@ -94,6 +350,17 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() { .expect("redeem codes should list"); assert_eq!(codes.total, 1); assert_eq!(codes.items[0].masked_code, "ABCD****WXYZ"); + let paged_codes = repository + .list_admin_redeem_codes(&AdminRedeemCodeListQuery { + batch_id: "batch-1".to_string(), + status: None, + limit: 1, + offset: 1, + }) + .await + .expect("redeem codes should page"); + assert_eq!(paged_codes.total, 2); + assert_eq!(paged_codes.items[0].id, "code-1"); let today = super::current_billing_date("UTC").expect("UTC should parse"); sqlx::query("UPDATE wallet_daily_usage_ledgers SET billing_date = ? WHERE id = 'daily-1'") @@ -107,6 +374,8 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() { .expect("daily usage should query") .expect("daily usage should exist"); assert_eq!(daily.total_requests, 2); + + cleanup_rows(repository.pool()).await; } impl MysqlWalletReadRepository { @@ -118,14 +387,15 @@ impl MysqlWalletReadRepository { async fn cleanup_rows(pool: &sqlx::MySqlPool) { for sql in [ "DELETE FROM wallet_daily_usage_ledgers WHERE id = 'daily-1'", - "DELETE FROM redeem_codes WHERE id = 'code-1'", - "DELETE FROM redeem_code_batches WHERE id = 'batch-1'", - "DELETE FROM wallet_transactions WHERE id = 'tx-1'", - "DELETE FROM refund_requests WHERE id = 'refund-1'", - "DELETE FROM payment_callbacks WHERE id = 'callback-1'", - "DELETE FROM payment_orders WHERE id = 'order-1'", - "DELETE FROM wallets WHERE id = 'wallet-1'", - "DELETE FROM users WHERE id = 'user-1'", + "DELETE FROM redeem_codes WHERE id IN ('code-1', 'code-2', 'code-3')", + "DELETE FROM redeem_code_batches WHERE id IN ('batch-1', 'batch-2')", + "DELETE FROM wallet_transactions WHERE id IN ('tx-1', 'tx-2', 'tx-api-1')", + "DELETE FROM refund_requests WHERE id IN ('refund-1', 'refund-2', 'refund-api-1')", + "DELETE FROM payment_callbacks WHERE id IN ('callback-1', 'callback-2', 'callback-3')", + "DELETE FROM payment_orders WHERE id IN ('order-1', 'order-2', 'order-3', 'order-4')", + "DELETE FROM wallets WHERE id IN ('wallet-1', 'wallet-2', 'wallet-api-1')", + "DELETE FROM api_keys WHERE id = 'api-key-1'", + "DELETE FROM users WHERE id IN ('user-1', 'user-2')", ] { sqlx::query(sql) .execute(pool) @@ -138,72 +408,99 @@ async fn seed_rows(pool: &sqlx::MySqlPool) { sqlx::query( r#" INSERT INTO users (id, username, email, auth_source, created_at, updated_at) -VALUES ('user-1', 'Alice', 'alice@example.com', 'local', 1, 1) +VALUES + ('user-1', 'Alice', 'alice@example.com', 'local', 1, 1), + ('user-2', 'Bob', 'bob@example.com', 'local', 1, 1) "#, ) .execute(pool) .await - .expect("user should seed"); + .expect("users should seed"); + + sqlx::query( + r#" +INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at) +VALUES ('api-key-1', 'user-1', 'wallet-test-api-key-hash-1', 'Standalone Key', 1, 1) +"#, + ) + .execute(pool) + .await + .expect("api key should seed"); sqlx::query( r#" INSERT INTO wallets ( - id, user_id, balance, gift_balance, total_recharged, total_consumed, + id, user_id, api_key_id, balance, gift_balance, total_recharged, total_consumed, total_refunded, total_adjusted, created_at, updated_at -) VALUES ( - 'wallet-1', 'user-1', 10.0, 2.0, 20.0, 4.0, 1.0, 3.0, 1, 2 -) +) VALUES + ('wallet-1', 'user-1', NULL, 10.0, 2.0, 20.0, 4.0, 1.0, 3.0, 1, 2), + ('wallet-2', 'user-2', NULL, 4.0, 1.0, 5.0, 2.0, 0.0, 0.0, 1, 3), + ('wallet-api-1', NULL, 'api-key-1', 7.0, 0.0, 7.0, 0.0, 0.0, 0.0, 1, 4) "#, ) .execute(pool) .await - .expect("wallet should seed"); + .expect("wallets should seed"); sqlx::query( r#" INSERT INTO payment_orders ( id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, gateway_response, status, created_at -) VALUES ( - 'order-1', 'order-no-1', 'wallet-1', 'user-1', 5.0, 1.0, 4.0, - 'redeem_code', '{"ok":true}', 'credited', 3 -) + refundable_amount_usd, payment_method, gateway_response, status, created_at, + expires_at +) VALUES + ('order-1', 'order-no-1', 'wallet-1', 'user-1', 5.0, 1.0, 4.0, + 'redeem_code', '{"ok":true}', 'credited', 3, NULL), + ('order-2', 'order-no-2', 'wallet-1', 'user-1', 2.0, 0.0, 2.0, + 'card', NULL, 'pending', 12, 1), + ('order-3', 'order-no-3', 'wallet-1', 'user-1', 3.0, 0.0, 3.0, + 'card', NULL, 'pending', 13, 4102444800), + ('order-4', 'order-no-4', 'wallet-2', 'user-2', 4.0, 0.0, 4.0, + 'bank', NULL, 'paid', 14, 4102444800) "#, ) .execute(pool) .await - .expect("payment order should seed"); + .expect("payment orders should seed"); sqlx::query( r#" INSERT INTO payment_callbacks ( id, payment_order_id, payment_method, callback_key, order_no, signature_valid, payload, created_at -) VALUES ( - 'callback-1', 'order-1', 'redeem_code', 'callback-key-1', - 'order-no-1', 1, '{"event":"paid"}', 4 -) +) VALUES + ('callback-1', 'order-1', 'redeem_code', 'callback-key-1', + 'order-no-1', 1, '{"event":"paid"}', 4), + ('callback-2', 'order-1', 'redeem_code', 'callback-key-2', + 'order-no-1', 1, '{"event":"duplicate"}', 5), + ('callback-3', 'order-2', 'card', 'callback-key-3', + 'order-no-2', 0, '{"event":"created"}', 6) "#, ) .execute(pool) .await - .expect("callback should seed"); + .expect("callbacks should seed"); sqlx::query( r#" INSERT INTO refund_requests ( id, refund_no, wallet_id, user_id, payment_order_id, source_type, refund_mode, amount_usd, status, payout_proof, created_at, updated_at -) VALUES ( - 'refund-1', 'refund-no-1', 'wallet-1', 'user-1', 'order-1', - 'payment_order', 'offline_payout', 1.0, 'completed', - '{"proof":"ok"}', 5, 6 -) +) VALUES + ('refund-1', 'refund-no-1', 'wallet-1', 'user-1', 'order-1', + 'payment_order', 'offline_payout', 1.0, 'completed', + '{"proof":"ok"}', 5, 6), + ('refund-2', 'refund-no-2', 'wallet-1', 'user-1', 'order-2', + 'payment_order', 'original_channel', 0.5, 'pending_approval', + NULL, 6, 7), + ('refund-api-1', 'refund-no-api-1', 'wallet-api-1', NULL, NULL, + 'manual', 'offline_payout', 0.5, 'pending_approval', + NULL, 7, 8) "#, ) .execute(pool) .await - .expect("refund should seed"); + .expect("refunds should seed"); sqlx::query( r#" @@ -211,28 +508,31 @@ INSERT INTO wallet_transactions ( id, wallet_id, category, reason_code, amount, balance_before, balance_after, recharge_balance_before, recharge_balance_after, gift_balance_before, gift_balance_after, created_at -) VALUES ( - 'tx-1', 'wallet-1', 'credit', 'manual_adjustment', 3.0, 7.0, 10.0, - 5.0, 8.0, 2.0, 2.0, 7 -) +) VALUES + ('tx-1', 'wallet-1', 'credit', 'manual_adjustment', 3.0, 7.0, 10.0, + 5.0, 8.0, 2.0, 2.0, 7), + ('tx-2', 'wallet-1', 'credit', 'manual_adjustment', 1.0, 11.0, 12.0, + 9.0, 10.0, 2.0, 2.0, 8), + ('tx-api-1', 'wallet-api-1', 'credit', 'manual_adjustment', 1.0, 6.0, 7.0, + 6.0, 7.0, 0.0, 0.0, 9) "#, ) .execute(pool) .await - .expect("transaction should seed"); + .expect("transactions should seed"); sqlx::query( r#" INSERT INTO redeem_code_batches ( - id, name, amount_usd, total_count, created_at, updated_at -) VALUES ( - 'batch-1', 'Batch One', 5.0, 1, 8, 9 -) + id, name, amount_usd, total_count, status, created_at, updated_at +) VALUES + ('batch-1', 'Batch One', 5.0, 2, 'active', 8, 9), + ('batch-2', 'Batch Two', 8.0, 1, 'disabled', 9, 10) "#, ) .execute(pool) .await - .expect("redeem batch should seed"); + .expect("redeem batches should seed"); sqlx::query( r#" @@ -240,15 +540,18 @@ INSERT INTO redeem_codes ( id, batch_id, code_hash, code_prefix, code_suffix, status, redeemed_by_user_id, redeemed_wallet_id, redeemed_payment_order_id, redeemed_at, created_at, updated_at -) VALUES ( - 'code-1', 'batch-1', 'hash-1', 'ABCD', 'WXYZ', 'redeemed', - 'user-1', 'wallet-1', 'order-1', 10, 8, 10 -) +) VALUES + ('code-1', 'batch-1', 'hash-1', 'ABCD', 'WXYZ', 'redeemed', + 'user-1', 'wallet-1', 'order-1', 10, 8, 10), + ('code-2', 'batch-1', 'hash-2', 'EFGH', 'QRST', 'active', + NULL, NULL, NULL, NULL, 9, 10), + ('code-3', 'batch-2', 'hash-3', 'IJKL', 'MNOP', 'active', + NULL, NULL, NULL, NULL, 10, 11) "#, ) .execute(pool) .await - .expect("redeem code should seed"); + .expect("redeem codes should seed"); sqlx::query( r#" diff --git a/crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql b/crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql new file mode 100644 index 000000000..7e321201d --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql @@ -0,0 +1,784 @@ +-- PostgreSQL and legacy SQLite exports encode timestamps as ISO/SQL datetime text. +-- Normalize imported and repository-facing timestamps back to the INTEGER Unix-second contract. +-- Numeric strings already receive INTEGER affinity on insert; unparseable text is preserved. + +UPDATE "users" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "users" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "users" +SET "last_login_at" = CAST(strftime('%s', trim("last_login_at")) AS INTEGER) +WHERE typeof("last_login_at") = 'text' + AND strftime('%s', trim("last_login_at")) IS NOT NULL; + +UPDATE "users" +SET "privacy_policy_accepted_at" = CAST(strftime('%s', trim("privacy_policy_accepted_at")) AS INTEGER) +WHERE typeof("privacy_policy_accepted_at") = 'text' + AND strftime('%s', trim("privacy_policy_accepted_at")) IS NOT NULL; + +UPDATE "api_keys" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "api_keys" +SET "last_used_at" = CAST(strftime('%s', trim("last_used_at")) AS INTEGER) +WHERE typeof("last_used_at") = 'text' + AND strftime('%s', trim("last_used_at")) IS NOT NULL; + +UPDATE "api_keys" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "api_keys" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "providers" +SET "quota_last_reset_at" = CAST(strftime('%s', trim("quota_last_reset_at")) AS INTEGER) +WHERE typeof("quota_last_reset_at") = 'text' + AND strftime('%s', trim("quota_last_reset_at")) IS NOT NULL; + +UPDATE "providers" +SET "quota_expires_at" = CAST(strftime('%s', trim("quota_expires_at")) AS INTEGER) +WHERE typeof("quota_expires_at") = 'text' + AND strftime('%s', trim("quota_expires_at")) IS NOT NULL; + +UPDATE "providers" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "providers" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "last_429_at" = CAST(strftime('%s', trim("last_429_at")) AS INTEGER) +WHERE typeof("last_429_at") = 'text' + AND strftime('%s', trim("last_429_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "last_probe_increase_at" = CAST(strftime('%s', trim("last_probe_increase_at")) AS INTEGER) +WHERE typeof("last_probe_increase_at") = 'text' + AND strftime('%s', trim("last_probe_increase_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "last_used_at" = CAST(strftime('%s', trim("last_used_at")) AS INTEGER) +WHERE typeof("last_used_at") = 'text' + AND strftime('%s', trim("last_used_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "last_models_fetch_at" = CAST(strftime('%s', trim("last_models_fetch_at")) AS INTEGER) +WHERE typeof("last_models_fetch_at") = 'text' + AND strftime('%s', trim("last_models_fetch_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "oauth_invalid_at" = CAST(strftime('%s', trim("oauth_invalid_at")) AS INTEGER) +WHERE typeof("oauth_invalid_at") = 'text' + AND strftime('%s', trim("oauth_invalid_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "provider_api_keys" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "provider_endpoints" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "provider_endpoints" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "global_models" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "global_models" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "models" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "models" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "auth_modules" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "auth_modules" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "oauth_providers" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "oauth_providers" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "user_oauth_links" +SET "linked_at" = CAST(strftime('%s', trim("linked_at")) AS INTEGER) +WHERE typeof("linked_at") = 'text' + AND strftime('%s', trim("linked_at")) IS NOT NULL; + +UPDATE "user_oauth_links" +SET "last_login_at" = CAST(strftime('%s', trim("last_login_at")) AS INTEGER) +WHERE typeof("last_login_at") = 'text' + AND strftime('%s', trim("last_login_at")) IS NOT NULL; + +UPDATE "user_groups" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "user_groups" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "user_group_members" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "proxy_nodes" +SET "last_heartbeat_at" = CAST(strftime('%s', trim("last_heartbeat_at")) AS INTEGER) +WHERE typeof("last_heartbeat_at") = 'text' + AND strftime('%s', trim("last_heartbeat_at")) IS NOT NULL; + +UPDATE "proxy_nodes" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "proxy_nodes" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "proxy_nodes" +SET "tunnel_connected_at" = CAST(strftime('%s', trim("tunnel_connected_at")) AS INTEGER) +WHERE typeof("tunnel_connected_at") = 'text' + AND strftime('%s', trim("tunnel_connected_at")) IS NOT NULL; + +UPDATE "system_configs" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "system_configs" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "wallets" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "wallets" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "wallet_transactions" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "wallet_daily_usage_ledgers" +SET "first_finalized_at" = CAST(strftime('%s', trim("first_finalized_at")) AS INTEGER) +WHERE typeof("first_finalized_at") = 'text' + AND strftime('%s', trim("first_finalized_at")) IS NOT NULL; + +UPDATE "wallet_daily_usage_ledgers" +SET "last_finalized_at" = CAST(strftime('%s', trim("last_finalized_at")) AS INTEGER) +WHERE typeof("last_finalized_at") = 'text' + AND strftime('%s', trim("last_finalized_at")) IS NOT NULL; + +UPDATE "wallet_daily_usage_ledgers" +SET "aggregated_at" = CAST(strftime('%s', trim("aggregated_at")) AS INTEGER) +WHERE typeof("aggregated_at") = 'text' + AND strftime('%s', trim("aggregated_at")) IS NOT NULL; + +UPDATE "wallet_daily_usage_ledgers" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "wallet_daily_usage_ledgers" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "payment_orders" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "payment_orders" +SET "paid_at" = CAST(strftime('%s', trim("paid_at")) AS INTEGER) +WHERE typeof("paid_at") = 'text' + AND strftime('%s', trim("paid_at")) IS NOT NULL; + +UPDATE "payment_orders" +SET "credited_at" = CAST(strftime('%s', trim("credited_at")) AS INTEGER) +WHERE typeof("credited_at") = 'text' + AND strftime('%s', trim("credited_at")) IS NOT NULL; + +UPDATE "payment_orders" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "payment_callbacks" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "payment_callbacks" +SET "processed_at" = CAST(strftime('%s', trim("processed_at")) AS INTEGER) +WHERE typeof("processed_at") = 'text' + AND strftime('%s', trim("processed_at")) IS NOT NULL; + +UPDATE "refund_requests" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "refund_requests" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "refund_requests" +SET "processed_at" = CAST(strftime('%s', trim("processed_at")) AS INTEGER) +WHERE typeof("processed_at") = 'text' + AND strftime('%s', trim("processed_at")) IS NOT NULL; + +UPDATE "refund_requests" +SET "completed_at" = CAST(strftime('%s', trim("completed_at")) AS INTEGER) +WHERE typeof("completed_at") = 'text' + AND strftime('%s', trim("completed_at")) IS NOT NULL; + +UPDATE "redeem_code_batches" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "redeem_code_batches" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "redeem_code_batches" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "redeem_codes" +SET "redeemed_at" = CAST(strftime('%s', trim("redeemed_at")) AS INTEGER) +WHERE typeof("redeemed_at") = 'text' + AND strftime('%s', trim("redeemed_at")) IS NOT NULL; + +UPDATE "redeem_codes" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "redeem_codes" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "usage" +SET "created_at_unix_ms" = CAST(strftime('%s', trim("created_at_unix_ms")) AS INTEGER) +WHERE typeof("created_at_unix_ms") = 'text' + AND strftime('%s', trim("created_at_unix_ms")) IS NOT NULL; + +UPDATE "usage" +SET "updated_at_unix_secs" = CAST(strftime('%s', trim("updated_at_unix_secs")) AS INTEGER) +WHERE typeof("updated_at_unix_secs") = 'text' + AND strftime('%s', trim("updated_at_unix_secs")) IS NOT NULL; + +UPDATE "usage" +SET "finalized_at" = CAST(strftime('%s', trim("finalized_at")) AS INTEGER) +WHERE typeof("finalized_at") = 'text' + AND strftime('%s', trim("finalized_at")) IS NOT NULL; + +UPDATE "billing_rules" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "billing_rules" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "dimension_collectors" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "dimension_collectors" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "usage_settlement_snapshots" +SET "finalized_at" = CAST(strftime('%s', trim("finalized_at")) AS INTEGER) +WHERE typeof("finalized_at") = 'text' + AND strftime('%s', trim("finalized_at")) IS NOT NULL; + +UPDATE "usage_settlement_snapshots" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "usage_settlement_snapshots" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "announcements" +SET "start_time" = CAST(strftime('%s', trim("start_time")) AS INTEGER) +WHERE typeof("start_time") = 'text' + AND strftime('%s', trim("start_time")) IS NOT NULL; + +UPDATE "announcements" +SET "end_time" = CAST(strftime('%s', trim("end_time")) AS INTEGER) +WHERE typeof("end_time") = 'text' + AND strftime('%s', trim("end_time")) IS NOT NULL; + +UPDATE "announcements" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "announcements" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "management_tokens" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "management_tokens" +SET "last_used_at" = CAST(strftime('%s', trim("last_used_at")) AS INTEGER) +WHERE typeof("last_used_at") = 'text' + AND strftime('%s', trim("last_used_at")) IS NOT NULL; + +UPDATE "management_tokens" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "management_tokens" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "video_tasks" +SET "next_poll_at" = CAST(strftime('%s', trim("next_poll_at")) AS INTEGER) +WHERE typeof("next_poll_at") = 'text' + AND strftime('%s', trim("next_poll_at")) IS NOT NULL; + +UPDATE "video_tasks" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "video_tasks" +SET "submitted_at" = CAST(strftime('%s', trim("submitted_at")) AS INTEGER) +WHERE typeof("submitted_at") = 'text' + AND strftime('%s', trim("submitted_at")) IS NOT NULL; + +UPDATE "video_tasks" +SET "completed_at" = CAST(strftime('%s', trim("completed_at")) AS INTEGER) +WHERE typeof("completed_at") = 'text' + AND strftime('%s', trim("completed_at")) IS NOT NULL; + +UPDATE "video_tasks" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "payment_gateway_configs" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "payment_gateway_configs" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "billing_plans" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "billing_plans" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +UPDATE "user_plan_entitlements" +SET "starts_at" = CAST(strftime('%s', trim("starts_at")) AS INTEGER) +WHERE typeof("starts_at") = 'text' + AND strftime('%s', trim("starts_at")) IS NOT NULL; + +UPDATE "user_plan_entitlements" +SET "expires_at" = CAST(strftime('%s', trim("expires_at")) AS INTEGER) +WHERE typeof("expires_at") = 'text' + AND strftime('%s', trim("expires_at")) IS NOT NULL; + +UPDATE "user_plan_entitlements" +SET "created_at" = CAST(strftime('%s', trim("created_at")) AS INTEGER) +WHERE typeof("created_at") = 'text' + AND strftime('%s', trim("created_at")) IS NOT NULL; + +UPDATE "user_plan_entitlements" +SET "updated_at" = CAST(strftime('%s', trim("updated_at")) AS INTEGER) +WHERE typeof("updated_at") = 'text' + AND strftime('%s', trim("updated_at")) IS NOT NULL; + +-- Do not mark the migration successful while an incompatible timestamp can still reach SQLx. +DROP TABLE IF EXISTS temp._aether_timestamp_storage_guard; +CREATE TEMP TABLE _aether_timestamp_storage_guard ( + invalid_count INTEGER NOT NULL, + CONSTRAINT imported_timestamp_storage_must_be_integer CHECK (invalid_count = 0) +); + +INSERT INTO _aether_timestamp_storage_guard (invalid_count) +SELECT COUNT(*) +FROM ( + SELECT 1 FROM "users" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "users" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "users" + WHERE "last_login_at" IS NOT NULL AND typeof("last_login_at") <> 'integer' + UNION ALL + SELECT 1 FROM "users" + WHERE "privacy_policy_accepted_at" IS NOT NULL AND typeof("privacy_policy_accepted_at") <> 'integer' + UNION ALL + SELECT 1 FROM "api_keys" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "api_keys" + WHERE "last_used_at" IS NOT NULL AND typeof("last_used_at") <> 'integer' + UNION ALL + SELECT 1 FROM "api_keys" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "api_keys" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "providers" + WHERE "quota_last_reset_at" IS NOT NULL AND typeof("quota_last_reset_at") <> 'integer' + UNION ALL + SELECT 1 FROM "providers" + WHERE "quota_expires_at" IS NOT NULL AND typeof("quota_expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "providers" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "providers" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "last_429_at" IS NOT NULL AND typeof("last_429_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "last_probe_increase_at" IS NOT NULL AND typeof("last_probe_increase_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "last_used_at" IS NOT NULL AND typeof("last_used_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "last_models_fetch_at" IS NOT NULL AND typeof("last_models_fetch_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "oauth_invalid_at" IS NOT NULL AND typeof("oauth_invalid_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_api_keys" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_endpoints" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "provider_endpoints" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "global_models" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "global_models" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "models" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "models" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "auth_modules" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "auth_modules" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "oauth_providers" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "oauth_providers" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_oauth_links" + WHERE "linked_at" IS NOT NULL AND typeof("linked_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_oauth_links" + WHERE "last_login_at" IS NOT NULL AND typeof("last_login_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_groups" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_groups" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_group_members" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "proxy_nodes" + WHERE "last_heartbeat_at" IS NOT NULL AND typeof("last_heartbeat_at") <> 'integer' + UNION ALL + SELECT 1 FROM "proxy_nodes" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "proxy_nodes" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "proxy_nodes" + WHERE "tunnel_connected_at" IS NOT NULL AND typeof("tunnel_connected_at") <> 'integer' + UNION ALL + SELECT 1 FROM "system_configs" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "system_configs" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallets" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallets" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_transactions" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_daily_usage_ledgers" + WHERE "first_finalized_at" IS NOT NULL AND typeof("first_finalized_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_daily_usage_ledgers" + WHERE "last_finalized_at" IS NOT NULL AND typeof("last_finalized_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_daily_usage_ledgers" + WHERE "aggregated_at" IS NOT NULL AND typeof("aggregated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_daily_usage_ledgers" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "wallet_daily_usage_ledgers" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_orders" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_orders" + WHERE "paid_at" IS NOT NULL AND typeof("paid_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_orders" + WHERE "credited_at" IS NOT NULL AND typeof("credited_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_orders" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_callbacks" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_callbacks" + WHERE "processed_at" IS NOT NULL AND typeof("processed_at") <> 'integer' + UNION ALL + SELECT 1 FROM "refund_requests" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "refund_requests" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "refund_requests" + WHERE "processed_at" IS NOT NULL AND typeof("processed_at") <> 'integer' + UNION ALL + SELECT 1 FROM "refund_requests" + WHERE "completed_at" IS NOT NULL AND typeof("completed_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_code_batches" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_code_batches" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_code_batches" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_codes" + WHERE "redeemed_at" IS NOT NULL AND typeof("redeemed_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_codes" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "redeem_codes" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "usage" + WHERE "created_at_unix_ms" IS NOT NULL AND typeof("created_at_unix_ms") <> 'integer' + UNION ALL + SELECT 1 FROM "usage" + WHERE "updated_at_unix_secs" IS NOT NULL AND typeof("updated_at_unix_secs") <> 'integer' + UNION ALL + SELECT 1 FROM "usage" + WHERE "finalized_at" IS NOT NULL AND typeof("finalized_at") <> 'integer' + UNION ALL + SELECT 1 FROM "billing_rules" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "billing_rules" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "dimension_collectors" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "dimension_collectors" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "usage_settlement_snapshots" + WHERE "finalized_at" IS NOT NULL AND typeof("finalized_at") <> 'integer' + UNION ALL + SELECT 1 FROM "usage_settlement_snapshots" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "usage_settlement_snapshots" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "announcements" + WHERE "start_time" IS NOT NULL AND typeof("start_time") <> 'integer' + UNION ALL + SELECT 1 FROM "announcements" + WHERE "end_time" IS NOT NULL AND typeof("end_time") <> 'integer' + UNION ALL + SELECT 1 FROM "announcements" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "announcements" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "management_tokens" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "management_tokens" + WHERE "last_used_at" IS NOT NULL AND typeof("last_used_at") <> 'integer' + UNION ALL + SELECT 1 FROM "management_tokens" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "management_tokens" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "video_tasks" + WHERE "next_poll_at" IS NOT NULL AND typeof("next_poll_at") <> 'integer' + UNION ALL + SELECT 1 FROM "video_tasks" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "video_tasks" + WHERE "submitted_at" IS NOT NULL AND typeof("submitted_at") <> 'integer' + UNION ALL + SELECT 1 FROM "video_tasks" + WHERE "completed_at" IS NOT NULL AND typeof("completed_at") <> 'integer' + UNION ALL + SELECT 1 FROM "video_tasks" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_gateway_configs" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "payment_gateway_configs" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "billing_plans" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "billing_plans" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_plan_entitlements" + WHERE "starts_at" IS NOT NULL AND typeof("starts_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_plan_entitlements" + WHERE "expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_plan_entitlements" + WHERE "created_at" IS NOT NULL AND typeof("created_at") <> 'integer' + UNION ALL + SELECT 1 FROM "user_plan_entitlements" + WHERE "updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer' +) AS invalid_timestamp_values; + +DROP TABLE temp._aether_timestamp_storage_guard; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql b/crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql new file mode 100644 index 000000000..ecfcf2621 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql @@ -0,0 +1,270 @@ +-- Complete the portable schema contract that is already represented by the +-- logical/generated schema and PostgreSQL usage capture tables. + +ALTER TABLE provider_api_keys ADD COLUMN last_error_at INTEGER; +ALTER TABLE provider_api_keys ADD COLUMN last_error_msg TEXT; + +CREATE TABLE IF NOT EXISTS api_key_provider_mappings ( + id TEXT PRIMARY KEY NOT NULL, + api_key_id TEXT NOT NULL, + provider_id TEXT NOT NULL, + priority_adjustment INTEGER NOT NULL DEFAULT 0, + weight_multiplier REAL NOT NULL DEFAULT 1, + is_enabled INTEGER NOT NULL DEFAULT 1, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (api_key_id, provider_id) +); + +CREATE INDEX IF NOT EXISTS api_key_provider_mappings_api_key_id_idx + ON api_key_provider_mappings (api_key_id); +CREATE INDEX IF NOT EXISTS api_key_provider_mappings_provider_id_idx + ON api_key_provider_mappings (provider_id); +CREATE INDEX IF NOT EXISTS idx_apikey_provider_enabled + ON api_key_provider_mappings (api_key_id, is_enabled); + +CREATE TABLE IF NOT EXISTS provider_usage_tracking ( + id TEXT PRIMARY KEY NOT NULL, + provider_id TEXT NOT NULL, + window_start INTEGER NOT NULL, + window_end INTEGER NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + successful_requests INTEGER NOT NULL DEFAULT 0, + failed_requests INTEGER NOT NULL DEFAULT 0, + avg_response_time_ms REAL NOT NULL DEFAULT 0, + total_response_time_ms REAL NOT NULL DEFAULT 0, + total_cost_usd REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS provider_usage_tracking_provider_id_idx + ON provider_usage_tracking (provider_id); +CREATE INDEX IF NOT EXISTS provider_usage_tracking_window_start_idx + ON provider_usage_tracking (window_start); +CREATE INDEX IF NOT EXISTS idx_provider_window + ON provider_usage_tracking (provider_id, window_start); +CREATE INDEX IF NOT EXISTS idx_window_time + ON provider_usage_tracking (window_start, window_end); + +-- The baseline already has partial uniqueness guards for enabled billing +-- configuration. These full indexes cover reads that also inspect disabled rows. +CREATE INDEX IF NOT EXISTS billing_rules_global_model_task_idx + ON billing_rules (global_model_id, task_type, is_enabled); +CREATE INDEX IF NOT EXISTS billing_rules_model_task_idx + ON billing_rules (model_id, task_type, is_enabled); +CREATE INDEX IF NOT EXISTS dimension_collectors_enabled_idx + ON dimension_collectors (api_format, task_type, dimension_name, priority, is_enabled); + +ALTER TABLE video_tasks ADD COLUMN converted_request_body TEXT; +ALTER TABLE video_tasks ADD COLUMN max_retries INTEGER NOT NULL DEFAULT 3; +ALTER TABLE video_tasks ADD COLUMN video_urls TEXT; +ALTER TABLE video_tasks ADD COLUMN thumbnail_url TEXT; +ALTER TABLE video_tasks ADD COLUMN video_size_bytes INTEGER; +ALTER TABLE video_tasks ADD COLUMN video_expires_at INTEGER; +ALTER TABLE video_tasks ADD COLUMN stored_video_path TEXT; +ALTER TABLE video_tasks ADD COLUMN storage_provider TEXT; +ALTER TABLE video_tasks ADD COLUMN remixed_from_task_id TEXT; +ALTER TABLE video_tasks ADD COLUMN webhook_url TEXT; +ALTER TABLE video_tasks ADD COLUMN webhook_sent INTEGER NOT NULL DEFAULT 0; +ALTER TABLE video_tasks ADD COLUMN webhook_sent_at INTEGER; +ALTER TABLE video_tasks ADD COLUMN video_duration_seconds REAL; + +-- Portable compatibility columns. New canonical HTTP payload writes use the +-- normalized usage_http_audits and usage_body_blobs tables below. +ALTER TABLE "usage" ADD COLUMN input_output_total_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN cache_creation_input_tokens_5m INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN cache_creation_input_tokens_1h INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN input_context_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN input_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN output_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN cache_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN cache_creation_cost_usd_5m REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN cache_creation_cost_usd_1h REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN request_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_input_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_output_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_cache_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_cache_creation_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_cache_creation_cost_usd_5m REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_cache_creation_cost_usd_1h REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_cache_read_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN actual_request_cost_usd REAL NOT NULL DEFAULT 0; +ALTER TABLE "usage" ADD COLUMN rate_multiplier REAL NOT NULL DEFAULT 1; +ALTER TABLE "usage" ADD COLUMN input_price_per_1m REAL; +ALTER TABLE "usage" ADD COLUMN cache_creation_price_per_1m REAL; +ALTER TABLE "usage" ADD COLUMN cache_creation_price_per_1m_5m REAL; +ALTER TABLE "usage" ADD COLUMN cache_creation_price_per_1m_1h REAL; +ALTER TABLE "usage" ADD COLUMN cache_read_price_per_1m REAL; +ALTER TABLE "usage" ADD COLUMN price_per_request REAL; +ALTER TABLE "usage" ADD COLUMN request_headers TEXT; +ALTER TABLE "usage" ADD COLUMN request_body TEXT; +ALTER TABLE "usage" ADD COLUMN provider_request_headers TEXT; +ALTER TABLE "usage" ADD COLUMN provider_request_body TEXT; +ALTER TABLE "usage" ADD COLUMN response_headers TEXT; +ALTER TABLE "usage" ADD COLUMN response_body TEXT; +ALTER TABLE "usage" ADD COLUMN client_response_headers TEXT; +ALTER TABLE "usage" ADD COLUMN client_response_body TEXT; +ALTER TABLE "usage" ADD COLUMN request_body_compressed BLOB; +ALTER TABLE "usage" ADD COLUMN provider_request_body_compressed BLOB; +ALTER TABLE "usage" ADD COLUMN response_body_compressed BLOB; +ALTER TABLE "usage" ADD COLUMN client_response_body_compressed BLOB; +ALTER TABLE "usage" ADD COLUMN created_at INTEGER; +ALTER TABLE "usage" ADD COLUMN username TEXT; +ALTER TABLE "usage" ADD COLUMN api_key_name TEXT; + +CREATE TABLE IF NOT EXISTS usage_body_blobs ( + body_ref TEXT PRIMARY KEY NOT NULL, + request_id TEXT NOT NULL, + body_field TEXT NOT NULL, + payload_gzip BLOB NOT NULL, + created_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + updated_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + UNIQUE (request_id, body_field), + FOREIGN KEY (request_id) REFERENCES "usage" (request_id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS ix_usage_body_blobs_request_id + ON usage_body_blobs (request_id); + +CREATE TABLE IF NOT EXISTS usage_http_audits ( + request_id TEXT PRIMARY KEY NOT NULL, + request_headers TEXT, + provider_request_headers TEXT, + response_headers TEXT, + client_response_headers TEXT, + request_body_ref TEXT, + provider_request_body_ref TEXT, + response_body_ref TEXT, + client_response_body_ref TEXT, + request_body_state TEXT, + provider_request_body_state TEXT, + response_body_state TEXT, + client_response_body_state TEXT, + body_capture_mode TEXT NOT NULL DEFAULT 'none', + created_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + updated_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + FOREIGN KEY (request_id) REFERENCES "usage" (request_id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS ix_usage_http_audits_updated_at + ON usage_http_audits (updated_at); + +-- Billing V3 keeps the immutable pricing/token snapshot separate from the +-- mutable compatibility columns on usage. +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_snapshot_schema_version TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_snapshot_status TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN rate_multiplier REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN is_free_tier INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN input_price_per_1m REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN output_price_per_1m REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN cache_creation_price_per_1m REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN cache_read_price_per_1m REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN price_per_request REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN settlement_snapshot_schema_version TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN settlement_snapshot TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_dimensions TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_input_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_effective_input_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_output_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_creation_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_creation_5m_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_creation_1h_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_read_tokens INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_total_input_context INTEGER; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_creation_cost_usd REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_cache_read_cost_usd REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_total_cost_usd REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_actual_total_cost_usd REAL; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_pricing_source TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_rule_id TEXT; +ALTER TABLE usage_settlement_snapshots + ADD COLUMN billing_rule_version TEXT; + +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_schema_version + ON usage_settlement_snapshots (settlement_snapshot_schema_version); +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_pricing_source + ON usage_settlement_snapshots (billing_pricing_source); + +CREATE TABLE IF NOT EXISTS stats_summary ( + id TEXT PRIMARY KEY NOT NULL, + cutoff_date INTEGER NOT NULL, + all_time_requests INTEGER NOT NULL DEFAULT 0, + all_time_success_requests INTEGER NOT NULL DEFAULT 0, + all_time_error_requests INTEGER NOT NULL DEFAULT 0, + all_time_input_tokens INTEGER NOT NULL DEFAULT 0, + all_time_output_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cost REAL NOT NULL DEFAULT 0, + all_time_actual_cost REAL NOT NULL DEFAULT 0, + total_users INTEGER NOT NULL DEFAULT 0, + active_users INTEGER NOT NULL DEFAULT 0, + total_api_keys INTEGER NOT NULL DEFAULT 0, + active_api_keys INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +CREATE TABLE IF NOT EXISTS user_model_usage_counts ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + model TEXT NOT NULL, + usage_count INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, model) +); + +CREATE INDEX IF NOT EXISTS idx_user_model_usage_user + ON user_model_usage_counts (user_id); +CREATE INDEX IF NOT EXISTS idx_user_model_usage_model + ON user_model_usage_counts (model); + +ALTER TABLE stats_daily ADD COLUMN input_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily ADD COLUMN output_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily ADD COLUMN cache_creation_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily ADD COLUMN cache_read_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily ADD COLUMN p50_response_time_ms INTEGER; +ALTER TABLE stats_daily ADD COLUMN p90_response_time_ms INTEGER; +ALTER TABLE stats_daily ADD COLUMN p99_response_time_ms INTEGER; +ALTER TABLE stats_daily ADD COLUMN p50_first_byte_time_ms INTEGER; +ALTER TABLE stats_daily ADD COLUMN p90_first_byte_time_ms INTEGER; +ALTER TABLE stats_daily ADD COLUMN p99_first_byte_time_ms INTEGER; + +-- PostgreSQL uses a partial index for this active-row cleanup path. SQLite can +-- preserve the same selectivity and ordering. +CREATE INDEX IF NOT EXISTS idx_usage_stale_pending_created_request + ON "usage" (created_at_unix_ms, request_id) + WHERE status IN ('pending', 'streaming'); + +CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_created_at_desc + ON provider_api_keys (provider_id, created_at DESC, name, id); +CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_last_used_at_desc + ON provider_api_keys (provider_id, last_used_at DESC, name, id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql b/crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql new file mode 100644 index 000000000..0ed6b4983 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql @@ -0,0 +1,447 @@ +ALTER TABLE stats_user_daily + ADD COLUMN actual_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN effective_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN total_input_context INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN cache_creation_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN cache_read_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_output_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_user_daily + ADD COLUMN settled_first_finalized_at_unix_secs INTEGER; +ALTER TABLE stats_user_daily + ADD COLUMN settled_last_finalized_at_unix_secs INTEGER; + +ALTER TABLE stats_hourly_user + ADD COLUMN cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN actual_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_output_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_first_finalized_at_unix_secs INTEGER; +ALTER TABLE stats_hourly_user + ADD COLUMN settled_last_finalized_at_unix_secs INTEGER; + +ALTER TABLE stats_daily + ADD COLUMN effective_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN total_input_context INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN cache_hit_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN cache_hit_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_cache_hit_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_total_input_context INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_cache_creation_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN completed_cache_read_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_output_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily + ADD COLUMN settled_first_finalized_at_unix_secs INTEGER; +ALTER TABLE stats_daily + ADD COLUMN settled_last_finalized_at_unix_secs INTEGER; + +ALTER TABLE stats_hourly + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN cache_hit_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN cache_hit_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_cache_hit_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_total_input_context INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_cache_creation_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN completed_cache_read_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_total_cost REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_total_requests INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_output_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly + ADD COLUMN settled_first_finalized_at_unix_secs INTEGER; +ALTER TABLE stats_hourly + ADD COLUMN settled_last_finalized_at_unix_secs INTEGER; + +ALTER TABLE stats_daily_model + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_daily_model + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily_model + ADD COLUMN cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_daily_model + ADD COLUMN cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_model + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_model + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user_model + ADD COLUMN response_time_sum_ms REAL NOT NULL DEFAULT 0; +ALTER TABLE stats_hourly_user_model + ADD COLUMN response_time_samples INTEGER NOT NULL DEFAULT 0; + +CREATE TABLE IF NOT EXISTS stats_user_summary ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL UNIQUE, + username TEXT, + cutoff_date INTEGER NOT NULL, + all_time_requests INTEGER NOT NULL DEFAULT 0, + all_time_success_requests INTEGER NOT NULL DEFAULT 0, + all_time_error_requests INTEGER NOT NULL DEFAULT 0, + all_time_input_tokens INTEGER NOT NULL DEFAULT 0, + all_time_output_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cost REAL NOT NULL DEFAULT 0, + all_time_actual_cost REAL NOT NULL DEFAULT 0, + active_days INTEGER NOT NULL DEFAULT 0, + first_active_date INTEGER, + last_active_date INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); +CREATE INDEX idx_stats_user_summary_cutoff_date ON stats_user_summary (cutoff_date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", model) +); +CREATE INDEX idx_stats_user_daily_model_date ON stats_user_daily_model ("date"); +CREATE INDEX idx_stats_user_daily_model_user_id ON stats_user_daily_model (user_id); + +CREATE TABLE IF NOT EXISTS stats_user_daily_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", provider_name) +); +CREATE INDEX idx_stats_user_daily_provider_date ON stats_user_daily_provider ("date"); +CREATE INDEX idx_stats_user_daily_provider_user_id ON stats_user_daily_provider (user_id); + +CREATE TABLE IF NOT EXISTS stats_user_daily_api_format ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + api_format TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", api_format) +); +CREATE INDEX idx_stats_user_daily_api_format_date ON stats_user_daily_api_format ("date"); +CREATE INDEX idx_stats_user_daily_api_format_user_id ON stats_user_daily_api_format (user_id); + +CREATE TABLE IF NOT EXISTS stats_daily_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE ("date", model, provider_name) +); +CREATE INDEX idx_stats_daily_model_provider_date ON stats_daily_model_provider ("date"); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", model, provider_name) +); +CREATE INDEX idx_stats_user_daily_model_provider_date + ON stats_user_daily_model_provider ("date"); +CREATE INDEX idx_stats_user_daily_model_provider_user_date + ON stats_user_daily_model_provider (user_id, "date"); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings ( + id TEXT PRIMARY KEY NOT NULL, + "date" INTEGER NOT NULL UNIQUE, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_provider ( + id TEXT PRIMARY KEY NOT NULL, + "date" INTEGER NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE ("date", provider_name) +); +CREATE INDEX idx_stats_daily_cost_savings_provider_date + ON stats_daily_cost_savings_provider ("date"); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model ( + id TEXT PRIMARY KEY NOT NULL, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE ("date", model) +); +CREATE INDEX idx_stats_daily_cost_savings_model_date + ON stats_daily_cost_savings_model ("date"); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE ("date", model, provider_name) +); +CREATE INDEX idx_stats_daily_cost_savings_model_provider_date + ON stats_daily_cost_savings_model_provider ("date"); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date") +); +CREATE INDEX idx_stats_user_daily_cost_savings_date + ON stats_user_daily_cost_savings ("date"); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", provider_name) +); +CREATE INDEX idx_stats_user_daily_cost_savings_provider_date + ON stats_user_daily_cost_savings_provider ("date"); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", model) +); +CREATE INDEX idx_stats_user_daily_cost_savings_model_date + ON stats_user_daily_cost_savings_model ("date"); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + "date" INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, "date", model, provider_name) +); +CREATE INDEX idx_stats_user_daily_cost_savings_model_provider_date + ON stats_user_daily_cost_savings_model_provider ("date"); + +-- Existing completed buckets predate the enriched dimensions above. Preserve the rows for +-- reads, but make the bounded aggregation worker replay every historical bucket. +UPDATE stats_hourly SET is_complete = 0 WHERE is_complete <> 0; +UPDATE stats_daily SET is_complete = 0 WHERE is_complete <> 0; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql b/crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql new file mode 100644 index 000000000..c0201bc17 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql @@ -0,0 +1,70 @@ +CREATE TABLE IF NOT EXISTS usage_routing_snapshots ( + request_id TEXT PRIMARY KEY NOT NULL, + candidate_id TEXT, + candidate_index INTEGER, + key_name TEXT, + planner_kind TEXT, + route_family TEXT, + route_kind TEXT, + execution_path TEXT, + local_execution_runtime_miss_reason TEXT, + selected_provider_id TEXT, + selected_endpoint_id TEXT, + selected_provider_api_key_id TEXT, + has_format_conversion INTEGER, + created_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + updated_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)), + FOREIGN KEY (request_id) REFERENCES "usage" (request_id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_route_family_kind + ON usage_routing_snapshots (route_family, route_kind); +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_candidate_id + ON usage_routing_snapshots (candidate_id); + +INSERT INTO usage_routing_snapshots ( + request_id, + candidate_id, + candidate_index, + key_name, + planner_kind, + route_family, + route_kind, + execution_path, + local_execution_runtime_miss_reason, + selected_provider_id, + selected_endpoint_id, + selected_provider_api_key_id, + has_format_conversion, + created_at, + updated_at +) +SELECT + request_id, + candidate_id, + candidate_index, + key_name, + planner_kind, + route_family, + route_kind, + execution_path, + local_execution_runtime_miss_reason, + provider_id, + provider_endpoint_id, + provider_api_key_id, + has_format_conversion, + COALESCE(NULLIF(created_at_unix_ms, 0), NULLIF(updated_at_unix_secs, 0), CAST(strftime('%s', 'now') AS INTEGER)), + COALESCE(NULLIF(updated_at_unix_secs, 0), NULLIF(created_at_unix_ms, 0), CAST(strftime('%s', 'now') AS INTEGER)) +FROM "usage" +WHERE candidate_id IS NOT NULL + OR candidate_index IS NOT NULL + OR key_name IS NOT NULL + OR planner_kind IS NOT NULL + OR route_family IS NOT NULL + OR route_kind IS NOT NULL + OR execution_path IS NOT NULL + OR local_execution_runtime_miss_reason IS NOT NULL + OR provider_id IS NOT NULL + OR provider_endpoint_id IS NOT NULL + OR provider_api_key_id IS NOT NULL + OR has_format_conversion <> 0; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql b/crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql new file mode 100644 index 000000000..a7ad4e434 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql @@ -0,0 +1,821 @@ +-- Complete the legacy import repair for every INTEGER-affinity timestamp column +-- that existed before the 20260725 parity migrations. Invalid values are left +-- untouched so the guard at the end fails the migration instead of coercing data. + +UPDATE "announcement_reads" +SET "read_at" = CASE + WHEN typeof("read_at") = 'text' AND strftime('%s', trim("read_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("read_at")) AS INTEGER) + ELSE "read_at" + END +WHERE (typeof("read_at") = 'text' AND strftime('%s', trim("read_at")) IS NOT NULL); + +UPDATE "audit_logs" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL); + +UPDATE "background_task_events" +SET "created_at_unix_secs" = CASE + WHEN typeof("created_at_unix_secs") = 'text' AND strftime('%s', trim("created_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at_unix_secs")) AS INTEGER) + ELSE "created_at_unix_secs" + END +WHERE (typeof("created_at_unix_secs") = 'text' AND strftime('%s', trim("created_at_unix_secs")) IS NOT NULL); + +UPDATE "background_task_runs" +SET "created_at_unix_secs" = CASE + WHEN typeof("created_at_unix_secs") = 'text' AND strftime('%s', trim("created_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at_unix_secs")) AS INTEGER) + ELSE "created_at_unix_secs" + END, + "finished_at_unix_secs" = CASE + WHEN typeof("finished_at_unix_secs") = 'text' AND strftime('%s', trim("finished_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("finished_at_unix_secs")) AS INTEGER) + ELSE "finished_at_unix_secs" + END, + "started_at_unix_secs" = CASE + WHEN typeof("started_at_unix_secs") = 'text' AND strftime('%s', trim("started_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("started_at_unix_secs")) AS INTEGER) + ELSE "started_at_unix_secs" + END, + "updated_at_unix_secs" = CASE + WHEN typeof("updated_at_unix_secs") = 'text' AND strftime('%s', trim("updated_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at_unix_secs")) AS INTEGER) + ELSE "updated_at_unix_secs" + END +WHERE (typeof("created_at_unix_secs") = 'text' AND strftime('%s', trim("created_at_unix_secs")) IS NOT NULL) + OR (typeof("finished_at_unix_secs") = 'text' AND strftime('%s', trim("finished_at_unix_secs")) IS NOT NULL) + OR (typeof("started_at_unix_secs") = 'text' AND strftime('%s', trim("started_at_unix_secs")) IS NOT NULL) + OR (typeof("updated_at_unix_secs") = 'text' AND strftime('%s', trim("updated_at_unix_secs")) IS NOT NULL); + +UPDATE "entitlement_usage_ledgers" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL); + +UPDATE "gemini_file_mappings" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "expires_at" = CASE + WHEN typeof("expires_at") = 'text' AND strftime('%s', trim("expires_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("expires_at")) AS INTEGER) + ELSE "expires_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("expires_at") = 'text' AND strftime('%s', trim("expires_at")) IS NOT NULL); + +UPDATE "ldap_configs" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "pool_member_scores" +SET "last_failure_at" = CASE + WHEN typeof("last_failure_at") = 'text' AND strftime('%s', trim("last_failure_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_failure_at")) AS INTEGER) + ELSE "last_failure_at" + END, + "last_probe_attempt_at" = CASE + WHEN typeof("last_probe_attempt_at") = 'text' AND strftime('%s', trim("last_probe_attempt_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_probe_attempt_at")) AS INTEGER) + ELSE "last_probe_attempt_at" + END, + "last_probe_failure_at" = CASE + WHEN typeof("last_probe_failure_at") = 'text' AND strftime('%s', trim("last_probe_failure_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_probe_failure_at")) AS INTEGER) + ELSE "last_probe_failure_at" + END, + "last_probe_success_at" = CASE + WHEN typeof("last_probe_success_at") = 'text' AND strftime('%s', trim("last_probe_success_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_probe_success_at")) AS INTEGER) + ELSE "last_probe_success_at" + END, + "last_ranked_at" = CASE + WHEN typeof("last_ranked_at") = 'text' AND strftime('%s', trim("last_ranked_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_ranked_at")) AS INTEGER) + ELSE "last_ranked_at" + END, + "last_scheduled_at" = CASE + WHEN typeof("last_scheduled_at") = 'text' AND strftime('%s', trim("last_scheduled_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_scheduled_at")) AS INTEGER) + ELSE "last_scheduled_at" + END, + "last_success_at" = CASE + WHEN typeof("last_success_at") = 'text' AND strftime('%s', trim("last_success_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_success_at")) AS INTEGER) + ELSE "last_success_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("last_failure_at") = 'text' AND strftime('%s', trim("last_failure_at")) IS NOT NULL) + OR (typeof("last_probe_attempt_at") = 'text' AND strftime('%s', trim("last_probe_attempt_at")) IS NOT NULL) + OR (typeof("last_probe_failure_at") = 'text' AND strftime('%s', trim("last_probe_failure_at")) IS NOT NULL) + OR (typeof("last_probe_success_at") = 'text' AND strftime('%s', trim("last_probe_success_at")) IS NOT NULL) + OR (typeof("last_ranked_at") = 'text' AND strftime('%s', trim("last_ranked_at")) IS NOT NULL) + OR (typeof("last_scheduled_at") = 'text' AND strftime('%s', trim("last_scheduled_at")) IS NOT NULL) + OR (typeof("last_success_at") = 'text' AND strftime('%s', trim("last_success_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "proxy_node_events" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL); + +UPDATE "proxy_node_metrics_1h" +SET "bucket_start_unix_secs" = CASE + WHEN typeof("bucket_start_unix_secs") = 'text' AND strftime('%s', trim("bucket_start_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("bucket_start_unix_secs")) AS INTEGER) + ELSE "bucket_start_unix_secs" + END +WHERE (typeof("bucket_start_unix_secs") = 'text' AND strftime('%s', trim("bucket_start_unix_secs")) IS NOT NULL); + +UPDATE "proxy_node_metrics_1m" +SET "bucket_start_unix_secs" = CASE + WHEN typeof("bucket_start_unix_secs") = 'text' AND strftime('%s', trim("bucket_start_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("bucket_start_unix_secs")) AS INTEGER) + ELSE "bucket_start_unix_secs" + END +WHERE (typeof("bucket_start_unix_secs") = 'text' AND strftime('%s', trim("bucket_start_unix_secs")) IS NOT NULL); + +UPDATE "referral_rewards" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "request_candidates" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "finished_at" = CASE + WHEN typeof("finished_at") = 'text' AND strftime('%s', trim("finished_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("finished_at")) AS INTEGER) + ELSE "finished_at" + END, + "started_at" = CASE + WHEN typeof("started_at") = 'text' AND strftime('%s', trim("started_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("started_at")) AS INTEGER) + ELSE "started_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("finished_at") = 'text' AND strftime('%s', trim("finished_at")) IS NOT NULL) + OR (typeof("started_at") = 'text' AND strftime('%s', trim("started_at")) IS NOT NULL); + +UPDATE "routing_group_bindings" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "routing_group_versions" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL); + +UPDATE "routing_groups" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "published_at" = CASE + WHEN typeof("published_at") = 'text' AND strftime('%s', trim("published_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("published_at")) AS INTEGER) + ELSE "published_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("published_at") = 'text' AND strftime('%s', trim("published_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_daily_api_key" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_daily_error" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_daily_model" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_daily_provider" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_daily" +SET "aggregated_at" = CASE + WHEN typeof("aggregated_at") = 'text' AND strftime('%s', trim("aggregated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("aggregated_at")) AS INTEGER) + ELSE "aggregated_at" + END, + "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("aggregated_at") = 'text' AND strftime('%s', trim("aggregated_at")) IS NOT NULL) + OR (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_hourly_model" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "hour_utc" = CASE + WHEN typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL + THEN CAST(strftime('%s', trim("hour_utc")) AS INTEGER) + ELSE "hour_utc" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_hourly_provider" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "hour_utc" = CASE + WHEN typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL + THEN CAST(strftime('%s', trim("hour_utc")) AS INTEGER) + ELSE "hour_utc" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_hourly_user_model" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "hour_utc" = CASE + WHEN typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL + THEN CAST(strftime('%s', trim("hour_utc")) AS INTEGER) + ELSE "hour_utc" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_hourly_user" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "hour_utc" = CASE + WHEN typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL + THEN CAST(strftime('%s', trim("hour_utc")) AS INTEGER) + ELSE "hour_utc" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_hourly" +SET "aggregated_at" = CASE + WHEN typeof("aggregated_at") = 'text' AND strftime('%s', trim("aggregated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("aggregated_at")) AS INTEGER) + ELSE "aggregated_at" + END, + "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "hour_utc" = CASE + WHEN typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL + THEN CAST(strftime('%s', trim("hour_utc")) AS INTEGER) + ELSE "hour_utc" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("aggregated_at") = 'text' AND strftime('%s', trim("aggregated_at")) IS NOT NULL) + OR (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("hour_utc") = 'text' AND strftime('%s', trim("hour_utc")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "stats_user_daily" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "date" = CASE + WHEN typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL + THEN CAST(strftime('%s', trim("date")) AS INTEGER) + ELSE "date" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("date") = 'text' AND strftime('%s', trim("date")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "usage_counter_deltas" +SET "candidate_last_used_at_unix_secs" = CASE + WHEN typeof("candidate_last_used_at_unix_secs") = 'text' AND strftime('%s', trim("candidate_last_used_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("candidate_last_used_at_unix_secs")) AS INTEGER) + ELSE "candidate_last_used_at_unix_secs" + END, + "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "last_used_at_unix_secs" = CASE + WHEN typeof("last_used_at_unix_secs") = 'text' AND strftime('%s', trim("last_used_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_used_at_unix_secs")) AS INTEGER) + ELSE "last_used_at_unix_secs" + END, + "processed_at" = CASE + WHEN typeof("processed_at") = 'text' AND strftime('%s', trim("processed_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("processed_at")) AS INTEGER) + ELSE "processed_at" + END, + "removed_last_used_at_unix_secs" = CASE + WHEN typeof("removed_last_used_at_unix_secs") = 'text' AND strftime('%s', trim("removed_last_used_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("removed_last_used_at_unix_secs")) AS INTEGER) + ELSE "removed_last_used_at_unix_secs" + END, + "usage_created_at_unix_secs" = CASE + WHEN typeof("usage_created_at_unix_secs") = 'text' AND strftime('%s', trim("usage_created_at_unix_secs")) IS NOT NULL + THEN CAST(strftime('%s', trim("usage_created_at_unix_secs")) AS INTEGER) + ELSE "usage_created_at_unix_secs" + END +WHERE (typeof("candidate_last_used_at_unix_secs") = 'text' AND strftime('%s', trim("candidate_last_used_at_unix_secs")) IS NOT NULL) + OR (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("last_used_at_unix_secs") = 'text' AND strftime('%s', trim("last_used_at_unix_secs")) IS NOT NULL) + OR (typeof("processed_at") = 'text' AND strftime('%s', trim("processed_at")) IS NOT NULL) + OR (typeof("removed_last_used_at_unix_secs") = 'text' AND strftime('%s', trim("removed_last_used_at_unix_secs")) IS NOT NULL) + OR (typeof("usage_created_at_unix_secs") = 'text' AND strftime('%s', trim("usage_created_at_unix_secs")) IS NOT NULL); + +UPDATE "user_invite_codes" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "user_preferences" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "user_referrals" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "first_paid_at" = CASE + WHEN typeof("first_paid_at") = 'text' AND strftime('%s', trim("first_paid_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("first_paid_at")) AS INTEGER) + ELSE "first_paid_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("first_paid_at") = 'text' AND strftime('%s', trim("first_paid_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +UPDATE "user_sessions" +SET "created_at" = CASE + WHEN typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("created_at")) AS INTEGER) + ELSE "created_at" + END, + "expires_at" = CASE + WHEN typeof("expires_at") = 'text' AND strftime('%s', trim("expires_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("expires_at")) AS INTEGER) + ELSE "expires_at" + END, + "last_seen_at" = CASE + WHEN typeof("last_seen_at") = 'text' AND strftime('%s', trim("last_seen_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("last_seen_at")) AS INTEGER) + ELSE "last_seen_at" + END, + "revoked_at" = CASE + WHEN typeof("revoked_at") = 'text' AND strftime('%s', trim("revoked_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("revoked_at")) AS INTEGER) + ELSE "revoked_at" + END, + "rotated_at" = CASE + WHEN typeof("rotated_at") = 'text' AND strftime('%s', trim("rotated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("rotated_at")) AS INTEGER) + ELSE "rotated_at" + END, + "updated_at" = CASE + WHEN typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL + THEN CAST(strftime('%s', trim("updated_at")) AS INTEGER) + ELSE "updated_at" + END +WHERE (typeof("created_at") = 'text' AND strftime('%s', trim("created_at")) IS NOT NULL) + OR (typeof("expires_at") = 'text' AND strftime('%s', trim("expires_at")) IS NOT NULL) + OR (typeof("last_seen_at") = 'text' AND strftime('%s', trim("last_seen_at")) IS NOT NULL) + OR (typeof("revoked_at") = 'text' AND strftime('%s', trim("revoked_at")) IS NOT NULL) + OR (typeof("rotated_at") = 'text' AND strftime('%s', trim("rotated_at")) IS NOT NULL) + OR (typeof("updated_at") = 'text' AND strftime('%s', trim("updated_at")) IS NOT NULL); + +DROP TABLE IF EXISTS temp._aether_remaining_timestamp_storage_guard; +CREATE TEMP TABLE _aether_remaining_timestamp_storage_guard ( + invalid_count INTEGER NOT NULL CHECK (invalid_count = 0) +); + +INSERT INTO _aether_remaining_timestamp_storage_guard (invalid_count) +SELECT + EXISTS ( + SELECT 1 + FROM "announcement_reads" + WHERE ("read_at" IS NOT NULL AND typeof("read_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "audit_logs" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "background_task_events" + WHERE ("created_at_unix_secs" IS NOT NULL AND typeof("created_at_unix_secs") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "background_task_runs" + WHERE ("created_at_unix_secs" IS NOT NULL AND typeof("created_at_unix_secs") <> 'integer') + OR ("finished_at_unix_secs" IS NOT NULL AND typeof("finished_at_unix_secs") <> 'integer') + OR ("started_at_unix_secs" IS NOT NULL AND typeof("started_at_unix_secs") <> 'integer') + OR ("updated_at_unix_secs" IS NOT NULL AND typeof("updated_at_unix_secs") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "entitlement_usage_ledgers" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "gemini_file_mappings" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "ldap_configs" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "pool_member_scores" + WHERE ("last_failure_at" IS NOT NULL AND typeof("last_failure_at") <> 'integer') + OR ("last_probe_attempt_at" IS NOT NULL AND typeof("last_probe_attempt_at") <> 'integer') + OR ("last_probe_failure_at" IS NOT NULL AND typeof("last_probe_failure_at") <> 'integer') + OR ("last_probe_success_at" IS NOT NULL AND typeof("last_probe_success_at") <> 'integer') + OR ("last_ranked_at" IS NOT NULL AND typeof("last_ranked_at") <> 'integer') + OR ("last_scheduled_at" IS NOT NULL AND typeof("last_scheduled_at") <> 'integer') + OR ("last_success_at" IS NOT NULL AND typeof("last_success_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "proxy_node_events" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "proxy_node_metrics_1h" + WHERE ("bucket_start_unix_secs" IS NOT NULL AND typeof("bucket_start_unix_secs") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "proxy_node_metrics_1m" + WHERE ("bucket_start_unix_secs" IS NOT NULL AND typeof("bucket_start_unix_secs") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "referral_rewards" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "request_candidates" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("finished_at" IS NOT NULL AND typeof("finished_at") <> 'integer') + OR ("started_at" IS NOT NULL AND typeof("started_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "routing_group_bindings" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "routing_group_versions" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "routing_groups" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("published_at" IS NOT NULL AND typeof("published_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_daily_api_key" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_daily_error" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_daily_model" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_daily_provider" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_daily" + WHERE ("aggregated_at" IS NOT NULL AND typeof("aggregated_at") <> 'integer') + OR ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_hourly_model" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("hour_utc" IS NOT NULL AND typeof("hour_utc") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_hourly_provider" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("hour_utc" IS NOT NULL AND typeof("hour_utc") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_hourly_user_model" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("hour_utc" IS NOT NULL AND typeof("hour_utc") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_hourly_user" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("hour_utc" IS NOT NULL AND typeof("hour_utc") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_hourly" + WHERE ("aggregated_at" IS NOT NULL AND typeof("aggregated_at") <> 'integer') + OR ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("hour_utc" IS NOT NULL AND typeof("hour_utc") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "stats_user_daily" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("date" IS NOT NULL AND typeof("date") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "usage_counter_deltas" + WHERE ("candidate_last_used_at_unix_secs" IS NOT NULL AND typeof("candidate_last_used_at_unix_secs") <> 'integer') + OR ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("last_used_at_unix_secs" IS NOT NULL AND typeof("last_used_at_unix_secs") <> 'integer') + OR ("processed_at" IS NOT NULL AND typeof("processed_at") <> 'integer') + OR ("removed_last_used_at_unix_secs" IS NOT NULL AND typeof("removed_last_used_at_unix_secs") <> 'integer') + OR ("usage_created_at_unix_secs" IS NOT NULL AND typeof("usage_created_at_unix_secs") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "user_invite_codes" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "user_preferences" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "user_referrals" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("first_paid_at" IS NOT NULL AND typeof("first_paid_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ) + + EXISTS ( + SELECT 1 + FROM "user_sessions" + WHERE ("created_at" IS NOT NULL AND typeof("created_at") <> 'integer') + OR ("expires_at" IS NOT NULL AND typeof("expires_at") <> 'integer') + OR ("last_seen_at" IS NOT NULL AND typeof("last_seen_at") <> 'integer') + OR ("revoked_at" IS NOT NULL AND typeof("revoked_at") <> 'integer') + OR ("rotated_at" IS NOT NULL AND typeof("rotated_at") <> 'integer') + OR ("updated_at" IS NOT NULL AND typeof("updated_at") <> 'integer') + ); + +DROP TABLE temp._aether_remaining_timestamp_storage_guard; diff --git a/crates/aether-data/adapters/sqlite/src/billing.rs b/crates/aether-data/adapters/sqlite/src/billing.rs index c53b19d8d..f5242fb84 100644 --- a/crates/aether-data/adapters/sqlite/src/billing.rs +++ b/crates/aether-data/adapters/sqlite/src/billing.rs @@ -642,23 +642,8 @@ LIMIT 1 input: &PaymentGatewayConfigWriteInput, ) -> Result, DataLayerError> { let provider = input.provider.trim().to_ascii_lowercase(); - let existing_secret = if input.preserve_existing_secret { - sqlx::query_scalar::<_, String>( - "SELECT merchant_key_encrypted FROM payment_gateway_configs WHERE provider = ?", - ) - .bind(&provider) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - } else { - None - }; - let secret = if input.preserve_existing_secret { - existing_secret - } else { - input.merchant_key_encrypted.clone() - }; let now = current_unix_secs_i64(); + let mut tx = self.pool.begin().await.map_sql_err()?; sqlx::query( r#" INSERT INTO payment_gateway_configs ( @@ -672,7 +657,10 @@ ON CONFLICT(provider) DO UPDATE SET endpoint_url = excluded.endpoint_url, callback_base_url = excluded.callback_base_url, merchant_id = excluded.merchant_id, - merchant_key_encrypted = excluded.merchant_key_encrypted, + merchant_key_encrypted = CASE + WHEN ? THEN payment_gateway_configs.merchant_key_encrypted + ELSE excluded.merchant_key_encrypted + END, pay_currency = excluded.pay_currency, usd_exchange_rate = excluded.usd_exchange_rate, min_recharge_usd = excluded.min_recharge_usd, @@ -685,22 +673,35 @@ ON CONFLICT(provider) DO UPDATE SET .bind(&input.endpoint_url) .bind(input.callback_base_url.as_deref()) .bind(&input.merchant_id) - .bind(secret.as_deref()) + .bind(input.merchant_key_encrypted.as_deref()) .bind(&input.pay_currency) .bind(input.usd_exchange_rate) .bind(input.min_recharge_usd) .bind(json_to_string(&input.channels_json)?) .bind(now) .bind(now) - .execute(&self.pool) + .bind(input.preserve_existing_secret) + .execute(&mut *tx) .await .map_sql_err()?; - match self.find_payment_gateway_config(&provider).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "upserted payment gateway config missing".to_string(), - )), - } + let row = sqlx::query( + r#" +SELECT + provider, enabled, endpoint_url, callback_base_url, merchant_id, + merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, + channels_json, created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs +FROM payment_gateway_configs +WHERE provider = ? +LIMIT 1 +"#, + ) + .bind(&provider) + .fetch_one(&mut *tx) + .await + .map_sql_err()?; + let record = map_payment_gateway_config_sqlite(&row)?; + tx.commit().await.map_sql_err()?; + Ok(AdminBillingMutationOutcome::Applied(record)) } async fn list_billing_plans( @@ -1404,7 +1405,7 @@ mod tests { use crate::run_migrations; use aether_data_contracts::repository::billing::{ AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, AdminBillingRuleWriteInput, - BillingPlanWriteInput, BillingReadRepository, + BillingPlanWriteInput, BillingReadRepository, PaymentGatewayConfigWriteInput, }; #[tokio::test] @@ -1647,6 +1648,77 @@ VALUES ('order-1', 'order-no-1', 'wallet-1', 0, 'epay', 'plan_purchase', .is_some()); } + #[tokio::test] + async fn sqlite_gateway_upsert_preserves_secret_atomically_without_dropping_insert_secret() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteBillingReadRepository::new(pool); + + let mut input = PaymentGatewayConfigWriteInput { + provider: " ParityPay ".to_string(), + enabled: true, + endpoint_url: "https://pay.example/first".to_string(), + callback_base_url: Some("https://api.example/callback".to_string()), + merchant_id: "merchant-1".to_string(), + merchant_key_encrypted: Some("secret-first".to_string()), + preserve_existing_secret: true, + pay_currency: "USD".to_string(), + usd_exchange_rate: 1.0, + min_recharge_usd: 2.0, + channels_json: json!([{"id": "card"}]), + }; + let inserted = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway insert should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway insert outcome: {other:?}"), + }; + assert_eq!(inserted.provider, "paritypay"); + assert_eq!( + inserted.merchant_key_encrypted.as_deref(), + Some("secret-first") + ); + + input.endpoint_url = "https://pay.example/preserved".to_string(); + input.merchant_key_encrypted = Some("secret-ignored".to_string()); + let preserved = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway preserve update should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway preserve outcome: {other:?}"), + }; + assert_eq!(preserved.endpoint_url, "https://pay.example/preserved"); + assert_eq!( + preserved.merchant_key_encrypted.as_deref(), + Some("secret-first") + ); + + input.preserve_existing_secret = false; + input.merchant_key_encrypted = Some("secret-replaced".to_string()); + let replaced = match repository + .upsert_payment_gateway_config(&input) + .await + .expect("gateway replacement update should run") + { + AdminBillingMutationOutcome::Applied(record) => record, + other => panic!("unexpected gateway replace outcome: {other:?}"), + }; + assert_eq!( + replaced.merchant_key_encrypted.as_deref(), + Some("secret-replaced") + ); + } + async fn seed_billing_context(pool: &sqlx::SqlitePool) { sqlx::query( r#" diff --git a/crates/aether-data/adapters/sqlite/src/candidates.rs b/crates/aether-data/adapters/sqlite/src/candidates.rs index bc4285ee5..065ea216f 100644 --- a/crates/aether-data/adapters/sqlite/src/candidates.rs +++ b/crates/aether-data/adapters/sqlite/src/candidates.rs @@ -1,7 +1,7 @@ use std::collections::{BTreeMap, BTreeSet}; use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; +use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqliteConnection}; use aether_data_contracts::repository::candidates::{ request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket, @@ -52,24 +52,6 @@ impl SqliteRequestCandidateRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } - - async fn find_by_unique( - &self, - request_id: &str, - candidate_index: u32, - retry_index: u32, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1" - )) - .bind(request_id) - .bind(to_i32(candidate_index)?) - .bind(to_i32(retry_index)?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_candidate_row).transpose() - } } #[async_trait] @@ -241,16 +223,50 @@ impl RequestCandidateWriteRepository for SqliteRequestCandidateRepository { candidate: UpsertRequestCandidateRecord, ) -> Result { candidate.validate()?; - let existing = self - .find_by_unique( - &candidate.request_id, - candidate.candidate_index, - candidate.retry_index, - ) - .await?; - let merged = merge_candidate(candidate, existing)?; - upsert_merged_candidate(&self.pool, &merged).await?; - Ok(merged) + let mut tx = self.pool.begin().await.map_sql_err()?; + match upsert_candidate_in_transaction(&mut tx, candidate).await { + Ok(candidate) => { + tx.commit().await.map_sql_err()?; + Ok(candidate) + } + Err(err) => { + tx.rollback().await.map_sql_err()?; + Err(err) + } + } + } + + async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + if candidates.is_empty() { + return Ok(0); + } + for candidate in &candidates { + candidate.validate()?; + } + + let mut tx = self.pool.begin().await.map_sql_err()?; + let result: Result = async { + let mut persisted = 0usize; + for candidate in candidates { + upsert_candidate_in_transaction(&mut tx, candidate).await?; + persisted = persisted.saturating_add(1); + } + Ok(persisted) + } + .await; + match result { + Ok(persisted) => { + tx.commit().await.map_sql_err()?; + Ok(persisted) + } + Err(err) => { + tx.rollback().await.map_sql_err()?; + Err(err) + } + } } async fn delete_created_before( @@ -283,8 +299,90 @@ WHERE id IN ( } } +async fn upsert_candidate_in_transaction( + tx: &mut sqlx::Transaction<'_, Sqlite>, + candidate: UpsertRequestCandidateRecord, +) -> Result { + // Write first so SQLite upgrades the deferred transaction before the Rust + // merge reads the latest state for this unique key. + let insert_candidate = merge_candidate(candidate.clone(), None)?; + insert_candidate_if_absent(tx, &insert_candidate).await?; + let existing = find_by_unique( + tx, + &candidate.request_id, + candidate.candidate_index, + candidate.retry_index, + ) + .await? + .ok_or_else(|| { + DataLayerError::UnexpectedValue( + "request candidate row was not locked after insert-if-absent".to_string(), + ) + })?; + let merged = merge_candidate(candidate, Some(existing))?; + upsert_merged_candidate(tx, &merged).await?; + find_by_unique( + tx, + &merged.request_id, + merged.candidate_index, + merged.retry_index, + ) + .await? + .ok_or_else(|| { + DataLayerError::UnexpectedValue( + "request candidate row disappeared after atomic upsert".to_string(), + ) + }) +} + +async fn insert_candidate_if_absent( + connection: &mut SqliteConnection, + candidate: &StoredRequestCandidate, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO request_candidates ( + id, request_id, candidate_index, retry_index, status, created_at +) +VALUES (?, ?, ?, ?, ?, ?) +ON CONFLICT(request_id, candidate_index, retry_index) DO NOTHING +"#, + ) + .bind(&candidate.id) + .bind(&candidate.request_id) + .bind(to_i32(candidate.candidate_index)?) + .bind(to_i32(candidate.retry_index)?) + .bind(status_to_database(candidate.status)) + .bind(u64_to_i64( + candidate.created_at_unix_ms, + "request candidate created_at", + )?) + .execute(connection) + .await + .map_sql_err()?; + Ok(()) +} + +async fn find_by_unique( + connection: &mut SqliteConnection, + request_id: &str, + candidate_index: u32, + retry_index: u32, +) -> Result, DataLayerError> { + let row = sqlx::query(&format!( + "{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1" + )) + .bind(request_id) + .bind(to_i32(candidate_index)?) + .bind(to_i32(retry_index)?) + .fetch_optional(connection) + .await + .map_sql_err()?; + row.as_ref().map(map_candidate_row).transpose() +} + async fn upsert_merged_candidate( - pool: &SqlitePool, + connection: &mut SqliteConnection, candidate: &StoredRequestCandidate, ) -> Result<(), DataLayerError> { sqlx::query( @@ -418,7 +516,7 @@ ON CONFLICT(request_id, candidate_index, retry_index) DO UPDATE SET candidate.finished_at_unix_ms, "request candidate finished_at", )?) - .execute(pool) + .execute(connection) .await .map_sql_err()?; Ok(()) @@ -871,7 +969,107 @@ mod tests { } #[tokio::test] - async fn sqlite_atomic_conflict_keeps_candidate_lifecycle_monotonic() { + async fn sqlite_concurrent_upserts_merge_without_lifecycle_regression() { + let database_path = std::env::temp_dir().join(format!( + "aether-candidate-concurrency-{}.sqlite", + uuid::Uuid::new_v4() + )); + let options = sqlx::sqlite::SqliteConnectOptions::new() + .filename(&database_path) + .create_if_missing(true) + .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) + .busy_timeout(std::time::Duration::from_secs(30)); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(12) + .connect_with(options) + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteRequestCandidateRepository::new(pool.clone()); + let request_id = format!("candidate-concurrency-{}", uuid::Uuid::new_v4()); + + let mut initial = sample_upsert( + "initial", + RequestCandidateStatus::Pending, + Some(json!({"initial": true})), + 3_000_000, + ); + initial.request_id = request_id.clone(); + repository + .upsert(initial) + .await + .expect("initial candidate should insert"); + + const WRITERS: usize = 10; + let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(WRITERS)); + let mut tasks = Vec::new(); + for writer in 0..WRITERS { + let repository = repository.clone(); + let request_id = request_id.clone(); + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + let status = if writer == 0 { + RequestCandidateStatus::Success + } else { + RequestCandidateStatus::Streaming + }; + let mut extra_data = serde_json::Map::new(); + extra_data.insert(format!("writer_{writer}"), json!(writer)); + let mut candidate = sample_upsert( + format!("writer-{writer}").as_str(), + status, + Some(serde_json::Value::Object(extra_data)), + 3_100_000 + u64::try_from(writer).expect("writer index should fit") * 10, + ); + candidate.request_id = request_id; + if writer != 0 { + candidate.latency_ms = Some(9_000 + writer as u64); + candidate.finished_at_unix_ms = Some(9_000_000 + writer as u64); + } + barrier.wait().await; + repository.upsert(candidate).await + })); + } + for task in tasks { + task.await + .expect("candidate writer should join") + .expect("candidate writer should persist"); + } + + let candidates = repository + .list_by_request_id(&request_id) + .await + .expect("concurrent candidate should load"); + assert_eq!(candidates.len(), 1); + let candidate = &candidates[0]; + assert_eq!(candidate.id, "initial"); + assert_eq!(candidate.status, RequestCandidateStatus::Success); + assert_eq!(candidate.latency_ms, Some(123)); + assert_eq!(candidate.finished_at_unix_ms, Some(3_100_002)); + let extra_data = candidate + .extra_data + .as_ref() + .and_then(serde_json::Value::as_object) + .expect("merged extra data should be an object"); + assert_eq!(extra_data.get("initial"), Some(&json!(true))); + for writer in 0..WRITERS { + assert_eq!( + extra_data.get(format!("writer_{writer}").as_str()), + Some(&json!(writer)) + ); + } + + drop(repository); + pool.close().await; + let _ = std::fs::remove_file(&database_path); + let _ = std::fs::remove_file(format!("{}-wal", database_path.display())); + let _ = std::fs::remove_file(format!("{}-shm", database_path.display())); + } + + #[tokio::test] + async fn sqlite_batch_upsert_preserves_order_and_rolls_back_on_error() { let pool = sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") @@ -880,111 +1078,98 @@ mod tests { run_migrations(&pool) .await .expect("sqlite migrations should run"); - let repository = SqliteRequestCandidateRepository::new(pool.clone()); + let repository = SqliteRequestCandidateRepository::new(pool); + let request_id = "request-batch"; - let terminal = super::merge_candidate( - sample_upsert( - "terminal", - RequestCandidateStatus::Success, - Some(json!({"terminal": true})), - 2_000_000, - ), - None, - ) - .expect("terminal candidate should build"); - super::upsert_merged_candidate(&pool, &terminal) - .await - .expect("terminal candidate should insert"); - - let mut stale_streaming = sample_upsert( - "stale-streaming", - RequestCandidateStatus::Streaming, - Some(json!({"stale": true})), - 1_999_000, + let mut pending = sample_upsert( + "batch-first", + RequestCandidateStatus::Pending, + Some(json!({"pending": true})), + 4_000_000, ); - stale_streaming.latency_ms = Some(9_999); - stale_streaming.finished_at_unix_ms = Some(9_999_999); - let stale_streaming = super::merge_candidate(stale_streaming, None) - .expect("stale streaming candidate should build"); - super::upsert_merged_candidate(&pool, &stale_streaming) - .await - .expect("stale streaming conflict should execute"); - - let mut streaming_input = sample_upsert( - "streaming", + pending.request_id = request_id.to_string(); + let mut streaming = sample_upsert( + "batch-second", RequestCandidateStatus::Streaming, - None, - 2_100_000, + Some(json!({"streaming": true})), + 4_000_100, ); - streaming_input.candidate_index = 1; - let streaming = super::merge_candidate(streaming_input, None) - .expect("streaming candidate should build"); - super::upsert_merged_candidate(&pool, &streaming) - .await - .expect("streaming candidate should insert"); + streaming.request_id = request_id.to_string(); + streaming.is_cached = None; + let mut success = sample_upsert( + "batch-third", + RequestCandidateStatus::Success, + Some(json!({"success": true})), + 4_000_200, + ); + success.request_id = request_id.to_string(); + success.is_cached = Some(true); + let mut late_pending = sample_upsert( + "batch-fourth", + RequestCandidateStatus::Pending, + Some(json!({"late": true})), + 4_000_300, + ); + late_pending.request_id = request_id.to_string(); + late_pending.is_cached = None; + late_pending.latency_ms = Some(9_999); + late_pending.finished_at_unix_ms = Some(9_999_999); - let mut stale_pending = sample_upsert( - "stale-pending", + assert_eq!( + repository + .upsert_many(vec![pending, streaming, success, late_pending]) + .await + .expect("ordered batch should persist"), + 4 + ); + let candidates = repository + .list_by_request_id(request_id) + .await + .expect("batch candidate should load"); + assert_eq!(candidates.len(), 1); + let candidate = &candidates[0]; + assert_eq!(candidate.id, "batch-first"); + assert_eq!(candidate.status, RequestCandidateStatus::Success); + assert!(candidate.is_cached); + assert_eq!(candidate.latency_ms, Some(123)); + assert_eq!(candidate.finished_at_unix_ms, Some(4_000_202)); + assert_eq!( + candidate.extra_data, + Some(json!({ + "pending": true, + "streaming": true, + "success": true, + "late": true + })) + ); + + let rollback_request_id = "request-batch-rollback"; + let mut valid = sample_upsert( + "rollback-valid", RequestCandidateStatus::Pending, None, - 2_099_000, + 5_000_000, ); - stale_pending.candidate_index = 1; - let stale_pending = - super::merge_candidate(stale_pending, None).expect("pending candidate should build"); - super::upsert_merged_candidate(&pool, &stale_pending) - .await - .expect("stale pending conflict should execute"); - - let mut pending_input = - sample_upsert("pending", RequestCandidateStatus::Pending, None, 2_200_000); - pending_input.candidate_index = 2; - pending_input.latency_ms = Some(321); - pending_input.finished_at_unix_ms = None; - let pending = - super::merge_candidate(pending_input, None).expect("pending candidate should build"); - super::upsert_merged_candidate(&pool, &pending) - .await - .expect("pending candidate should insert"); - - let mut stale_available = sample_upsert( - "stale-available", - RequestCandidateStatus::Available, + valid.request_id = rollback_request_id.to_string(); + let mut invalid = sample_upsert( + "rollback-invalid", + RequestCandidateStatus::Success, None, - 2_199_000, + 5_000_100, ); - stale_available.candidate_index = 2; - stale_available.latency_ms = Some(9_999); - stale_available.finished_at_unix_ms = Some(9_999_999); - let stale_available = super::merge_candidate(stale_available, None) - .expect("stale available candidate should build"); - super::upsert_merged_candidate(&pool, &stale_available) - .await - .expect("stale available conflict should execute"); + invalid.request_id = rollback_request_id.to_string(); + invalid.candidate_index = 1; + invalid.latency_ms = Some(u64::MAX); - let candidates = repository - .list_by_request_id("request-1") + repository + .upsert_many(vec![valid, invalid]) .await - .expect("request candidates should load"); - let terminal = candidates - .iter() - .find(|candidate| candidate.candidate_index == 0) - .expect("terminal candidate should remain"); - assert_eq!(terminal.status, RequestCandidateStatus::Success); - assert_eq!(terminal.latency_ms, Some(123)); - assert_eq!(terminal.finished_at_unix_ms, Some(2_000_002)); - let streaming = candidates - .iter() - .find(|candidate| candidate.candidate_index == 1) - .expect("streaming candidate should remain"); - assert_eq!(streaming.status, RequestCandidateStatus::Streaming); - let pending = candidates - .iter() - .find(|candidate| candidate.candidate_index == 2) - .expect("pending candidate should remain"); - assert_eq!(pending.status, RequestCandidateStatus::Pending); - assert_eq!(pending.latency_ms, Some(321)); - assert_eq!(pending.finished_at_unix_ms, None); + .expect_err("invalid later row should roll back the batch"); + assert!(repository + .list_by_request_id(rollback_request_id) + .await + .expect("rolled-back batch should be readable") + .is_empty()); } fn sample_upsert( diff --git a/crates/aether-data/adapters/sqlite/src/global_models.rs b/crates/aether-data/adapters/sqlite/src/global_models.rs index a43f481b3..fb74a5507 100644 --- a/crates/aether-data/adapters/sqlite/src/global_models.rs +++ b/crates/aether-data/adapters/sqlite/src/global_models.rs @@ -1094,8 +1094,8 @@ fn map_public_catalog_model_row( .and_then(|value| value.get("icon_url")) .and_then(serde_json::Value::as_str) .map(ToString::to_string), - first_tier_price(pricing, "input_price_per_1m"), - first_tier_price(pricing, "output_price_per_1m"), + Some(first_tier_price(pricing, "input_price_per_1m").unwrap_or(0.0)), + Some(first_tier_price(pricing, "output_price_per_1m").unwrap_or(0.0)), first_tier_price(pricing, "cache_creation_price_per_1m"), first_tier_price(pricing, "cache_read_price_per_1m"), row.try_get("supports_vision").map_sql_err()?, @@ -1375,6 +1375,44 @@ mod tests { .is_none()); } + #[tokio::test] + async fn sqlite_public_catalog_defaults_missing_input_and_output_prices_to_zero() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_provider(&pool).await; + sqlx::query( + r#" +INSERT INTO models ( + id, provider_id, provider_model_name, is_active, is_available, created_at, updated_at +) VALUES ('model-no-price', 'provider-1', 'provider-no-price', 1, 1, 1, 1) +"#, + ) + .execute(&pool) + .await + .expect("unpriced provider model should seed"); + + let items = SqliteGlobalModelReadRepository::new(pool) + .list_public_catalog_models(&PublicCatalogModelListQuery { + provider_id: Some("provider-1".to_string()), + offset: 0, + limit: 10, + }) + .await + .expect("public catalog should load"); + + assert_eq!(items.len(), 1); + assert_eq!(items[0].input_price_per_1m, Some(0.0)); + assert_eq!(items[0].output_price_per_1m, Some(0.0)); + assert_eq!(items[0].cache_creation_price_per_1m, None); + assert_eq!(items[0].cache_read_price_per_1m, None); + } + async fn seed_rows(pool: &sqlx::SqlitePool) { seed_provider(pool).await; sqlx::query( diff --git a/crates/aether-data/adapters/sqlite/src/migrations.rs b/crates/aether-data/adapters/sqlite/src/migrations.rs index 7a4e1b480..76acd9d16 100644 --- a/crates/aether-data/adapters/sqlite/src/migrations.rs +++ b/crates/aether-data/adapters/sqlite/src/migrations.rs @@ -1,5 +1,5 @@ use sqlx::{ - migrate::{Migrate, MigrateError, Migrator}, + migrate::{AppliedMigration, Migrate, MigrateError, Migrator}, SqlitePool, }; @@ -17,9 +17,15 @@ pub async fn pending_migrations( let mut conn = pool.acquire().await?; let applied_migrations = match conn.list_applied_migrations().await { Ok(applied_migrations) => applied_migrations, - Err(err) if is_missing_sqlx_migrations_table_error(&err) => Vec::new(), + Err(err) if is_missing_sqlx_migrations_table_error(&err) => { + return Ok(pending_migrations_from_applied(&[])); + } Err(err) => return Err(err), }; + if let Some(version) = conn.dirty_version().await? { + return Err(MigrateError::Dirty(version)); + } + validate_applied_migrations(&applied_migrations)?; Ok(pending_migrations_from_applied(&applied_migrations)) } @@ -56,9 +62,34 @@ fn pending_migrations_from_applied( .collect() } +fn validate_applied_migrations( + applied_migrations: &[AppliedMigration], +) -> Result<(), MigrateError> { + if MIGRATOR.ignore_missing { + return Ok(()); + } + let known_versions = MIGRATOR + .iter() + .map(|migration| migration.version) + .collect::>(); + if let Some(migration) = applied_migrations + .iter() + .find(|migration| !known_versions.contains(&migration.version)) + { + return Err(MigrateError::VersionMissing(migration.version)); + } + Ok(()) +} + #[cfg(test)] mod tests { - use super::{pending_migrations, run_migrations, MIGRATOR}; + use std::borrow::Cow; + + use super::{ + pending_migrations, prepare_database_for_startup, run_migrations, + validate_applied_migrations, MIGRATOR, + }; + use sqlx::migrate::{AppliedMigration, MigrateError}; #[tokio::test] async fn migrates_empty_database_and_clears_pending_set() { @@ -78,4 +109,468 @@ mod tests { .expect("pending migrations after run") .is_empty()); } + + #[test] + fn rejects_applied_migration_versions_unknown_to_this_binary() { + let version = MIGRATOR + .iter() + .map(|migration| migration.version) + .max() + .expect("sqlite migrations should not be empty") + + 1; + let error = validate_applied_migrations(&[AppliedMigration { + version, + checksum: Cow::Borrowed(&[]), + }]) + .expect_err("unknown applied migration should block startup"); + + assert!(matches!(error, MigrateError::VersionMissing(found) if found == version)); + } + + #[tokio::test] + async fn migrates_cross_driver_schema_parity_contract() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + + run_migrations(&pool).await.expect("run sqlite migrations"); + assert!(MIGRATOR + .iter() + .any(|migration| migration.version == 20260725010000)); + + let parity_table_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM sqlite_master +WHERE type = 'table' + AND name IN ( + 'api_key_provider_mappings', + 'provider_usage_tracking', + 'stats_summary', + 'user_model_usage_counts', + 'usage_body_blobs', + 'usage_http_audits' + ) +"#, + ) + .fetch_one(&pool) + .await + .expect("parity tables should be inspectable"); + assert_eq!(parity_table_count, 6); + + let usage_column_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM pragma_table_info('usage') +WHERE name IN ( + 'input_output_total_tokens', + 'cache_creation_input_tokens_5m', + 'cache_creation_input_tokens_1h', + 'input_context_tokens', + 'input_cost_usd', + 'output_cost_usd', + 'cache_cost_usd', + 'cache_creation_cost_usd_5m', + 'cache_creation_cost_usd_1h', + 'request_cost_usd', + 'actual_input_cost_usd', + 'actual_output_cost_usd', + 'actual_cache_cost_usd', + 'actual_cache_creation_cost_usd', + 'actual_cache_creation_cost_usd_5m', + 'actual_cache_creation_cost_usd_1h', + 'actual_cache_read_cost_usd', + 'actual_request_cost_usd', + 'rate_multiplier', + 'input_price_per_1m', + 'cache_creation_price_per_1m', + 'cache_creation_price_per_1m_5m', + 'cache_creation_price_per_1m_1h', + 'cache_read_price_per_1m', + 'price_per_request', + 'request_headers', + 'request_body', + 'provider_request_headers', + 'provider_request_body', + 'response_headers', + 'response_body', + 'client_response_headers', + 'client_response_body', + 'request_body_compressed', + 'provider_request_body_compressed', + 'response_body_compressed', + 'client_response_body_compressed', + 'created_at', + 'username', + 'api_key_name' +) +"#, + ) + .fetch_one(&pool) + .await + .expect("usage parity columns should be inspectable"); + assert_eq!(usage_column_count, 40); + + let settlement_column_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM pragma_table_info('usage_settlement_snapshots') +WHERE name IN ( + 'billing_snapshot_schema_version', + 'billing_snapshot_status', + 'rate_multiplier', + 'is_free_tier', + 'input_price_per_1m', + 'output_price_per_1m', + 'cache_creation_price_per_1m', + 'cache_read_price_per_1m', + 'price_per_request', + 'settlement_snapshot_schema_version', + 'settlement_snapshot', + 'billing_dimensions', + 'billing_input_tokens', + 'billing_effective_input_tokens', + 'billing_output_tokens', + 'billing_cache_creation_tokens', + 'billing_cache_creation_5m_tokens', + 'billing_cache_creation_1h_tokens', + 'billing_cache_read_tokens', + 'billing_total_input_context', + 'billing_cache_creation_cost_usd', + 'billing_cache_read_cost_usd', + 'billing_total_cost_usd', + 'billing_actual_total_cost_usd', + 'billing_pricing_source', + 'billing_rule_id', + 'billing_rule_version' +) +"#, + ) + .fetch_one(&pool) + .await + .expect("Billing V3 columns should be inspectable"); + assert_eq!(settlement_column_count, 27); + + let catalog_video_stats_column_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM ( + SELECT name FROM pragma_table_info('provider_api_keys') + WHERE name IN ('last_error_at', 'last_error_msg') + UNION ALL + SELECT name FROM pragma_table_info('video_tasks') + WHERE name IN ( + 'converted_request_body', 'max_retries', 'video_urls', 'thumbnail_url', + 'video_size_bytes', 'video_expires_at', 'stored_video_path', 'storage_provider', + 'remixed_from_task_id', 'webhook_url', 'webhook_sent', 'webhook_sent_at', + 'video_duration_seconds' + ) + UNION ALL + SELECT name FROM pragma_table_info('stats_daily') + WHERE name IN ( + 'input_cost', 'output_cost', 'cache_creation_cost', 'cache_read_cost', + 'p50_response_time_ms', 'p90_response_time_ms', 'p99_response_time_ms', + 'p50_first_byte_time_ms', 'p90_first_byte_time_ms', 'p99_first_byte_time_ms' + ) +) +"#, + ) + .fetch_one(&pool) + .await + .expect("catalog, video, and stats parity columns should be inspectable"); + assert_eq!(catalog_video_stats_column_count, 25); + + let parity_index_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM sqlite_master +WHERE type = 'index' + AND name IN ( + 'ix_usage_body_blobs_request_id', + 'ix_usage_http_audits_updated_at', + 'ix_usage_settlement_snapshots_schema_version', + 'ix_usage_settlement_snapshots_pricing_source', + 'idx_usage_stale_pending_created_request', + 'idx_provider_api_keys_provider_created_at_desc', + 'idx_provider_api_keys_provider_last_used_at_desc' + ) +"#, + ) + .fetch_one(&pool) + .await + .expect("parity indexes should be inspectable"); + assert_eq!(parity_index_count, 7); + } + + #[tokio::test] + async fn migrates_advanced_stats_parity_contract() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + + run_migrations(&pool).await.expect("run sqlite migrations"); + assert!(MIGRATOR + .iter() + .any(|migration| migration.version == 20260725020000)); + + let table_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM sqlite_master +WHERE type = 'table' + AND name IN ( + 'stats_user_summary', + 'stats_user_daily_api_format', + 'stats_user_daily_model', + 'stats_user_daily_provider', + 'stats_user_daily_model_provider', + 'stats_daily_model_provider', + 'stats_daily_cost_savings', + 'stats_daily_cost_savings_provider', + 'stats_daily_cost_savings_model', + 'stats_daily_cost_savings_model_provider', + 'stats_user_daily_cost_savings', + 'stats_user_daily_cost_savings_provider', + 'stats_user_daily_cost_savings_model', + 'stats_user_daily_cost_savings_model_provider' + ) +"#, + ) + .fetch_one(&pool) + .await + .expect("advanced stats tables should be inspectable"); + assert_eq!(table_count, 14); + + let enriched_column_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM ( + SELECT name FROM pragma_table_info('stats_daily') + WHERE name IN ( + 'effective_input_tokens', 'total_input_context', 'response_time_sum_ms', + 'response_time_samples', 'cache_hit_total_requests', 'cache_hit_requests', + 'completed_total_input_context', 'settled_total_cost', + 'settled_first_finalized_at_unix_secs' + ) + UNION ALL + SELECT name FROM pragma_table_info('stats_hourly') + WHERE name IN ( + 'response_time_sum_ms', 'response_time_samples', 'cache_hit_total_requests', + 'completed_total_input_context', 'settled_total_cost' + ) + UNION ALL + SELECT name FROM pragma_table_info('stats_user_daily') + WHERE name IN ( + 'effective_input_tokens', 'total_input_context', 'actual_total_cost', + 'response_time_samples', 'settled_total_cost' + ) +) +"#, + ) + .fetch_one(&pool) + .await + .expect("advanced stats columns should be inspectable"); + assert_eq!(enriched_column_count, 19); + } + + #[tokio::test] + async fn advanced_stats_migration_invalidates_completed_legacy_buckets() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + + for migration in MIGRATOR + .iter() + .filter(|migration| migration.version < 20260725020000) + { + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .unwrap_or_else(|err| panic!("migration {} should run: {err}", migration.version)); + } + + sqlx::query( + r#" +INSERT INTO stats_hourly (id, hour_utc, is_complete, created_at, updated_at) +VALUES ('legacy-hour', 3600, 1, 1, 1) +"#, + ) + .execute(&pool) + .await + .expect("legacy hourly bucket should seed"); + sqlx::query( + r#" +INSERT INTO stats_daily (id, "date", is_complete, created_at, updated_at) +VALUES ('legacy-day', 86400, 1, 1, 1) +"#, + ) + .execute(&pool) + .await + .expect("legacy daily bucket should seed"); + + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == 20260725020000) + .expect("advanced stats migration should be embedded"); + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect("advanced stats migration should run"); + + let hourly_complete: i64 = + sqlx::query_scalar("SELECT is_complete FROM stats_hourly WHERE id = 'legacy-hour'") + .fetch_one(&pool) + .await + .expect("hourly completion state should load"); + let daily_complete: i64 = + sqlx::query_scalar("SELECT is_complete FROM stats_daily WHERE id = 'legacy-day'") + .fetch_one(&pool) + .await + .expect("daily completion state should load"); + assert_eq!((hourly_complete, daily_complete), (0, 0)); + } + + #[tokio::test] + async fn routing_snapshot_migration_backfills_legacy_usage_columns() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + + for migration in MIGRATOR + .iter() + .filter(|migration| migration.version < 20260725030000) + { + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .unwrap_or_else(|err| panic!("migration {} should run: {err}", migration.version)); + } + + sqlx::query( + r#" +INSERT INTO "usage" ( + request_id, provider_name, model, provider_id, provider_endpoint_id, + provider_api_key_id, candidate_id, candidate_index, route_family, + has_format_conversion, created_at_unix_ms, updated_at_unix_secs +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +"#, + ) + .bind("routing-backfill-request") + .bind("provider") + .bind("model") + .bind("provider-1") + .bind("endpoint-1") + .bind("provider-key-1") + .bind("candidate-1") + .bind(7_i64) + .bind("direct") + .bind(true) + .bind(111_i64) + .bind(222_i64) + .execute(&pool) + .await + .expect("legacy routing usage should insert"); + + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == 20260725030000) + .expect("routing snapshot migration should be embedded"); + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect("routing snapshot migration should run"); + + let row = sqlx::query_as::< + _, + ( + String, + Option, + Option, + Option, + Option, + Option, + i64, + i64, + ), + >( + r#" +SELECT candidate_id, candidate_index, route_family, selected_provider_id, + selected_provider_api_key_id, has_format_conversion, created_at, updated_at +FROM usage_routing_snapshots +WHERE request_id = ? +"#, + ) + .bind("routing-backfill-request") + .fetch_one(&pool) + .await + .expect("backfilled routing snapshot should load"); + assert_eq!( + row, + ( + "candidate-1".to_string(), + Some(7), + Some("direct".to_string()), + Some("provider-1".to_string()), + Some("provider-key-1".to_string()), + Some(true), + 111, + 222, + ) + ); + + sqlx::query("DELETE FROM \"usage\" WHERE request_id = ?") + .bind("routing-backfill-request") + .execute(&pool) + .await + .expect("usage row should delete"); + let remaining: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_routing_snapshots") + .fetch_one(&pool) + .await + .expect("routing snapshot count should load"); + assert_eq!(remaining, 0); + } + + #[tokio::test] + async fn pending_and_startup_preparation_reject_dirty_migration_state() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + run_migrations(&pool).await.expect("run sqlite migrations"); + + let dirty_version: i64 = sqlx::query_scalar("SELECT MAX(version) FROM _sqlx_migrations") + .fetch_one(&pool) + .await + .expect("latest sqlite migration version should load"); + sqlx::query("UPDATE _sqlx_migrations SET success = FALSE WHERE version = ?") + .bind(dirty_version) + .execute(&pool) + .await + .expect("sqlite migration should be marked dirty"); + + let pending_error = pending_migrations(&pool) + .await + .expect_err("dirty sqlite migration should fail pending inspection"); + assert!( + matches!(&pending_error, MigrateError::Dirty(version) if *version == dirty_version), + "unexpected pending migration error: {pending_error}" + ); + + let preparation_error = prepare_database_for_startup(&pool) + .await + .expect_err("dirty sqlite migration should fail startup preparation"); + assert!( + matches!(&preparation_error, MigrateError::Dirty(version) if *version == dirty_version), + "unexpected startup preparation error: {preparation_error}" + ); + } } diff --git a/crates/aether-data/adapters/sqlite/src/settlement.rs b/crates/aether-data/adapters/sqlite/src/settlement.rs index 171e36a8b..dd00095f2 100644 --- a/crates/aether-data/adapters/sqlite/src/settlement.rs +++ b/crates/aether-data/adapters/sqlite/src/settlement.rs @@ -109,6 +109,13 @@ DO UPDATE SET updated_at = excluded.updated_at "#; +const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#" +INSERT INTO usage_counter_deltas ( + id, request_id, kind, target_id, total_cost_usd_delta, created_at +) +VALUES (?, ?, 'provider_monthly', ?, ?, ?) +"#; + #[derive(Debug, Clone)] pub struct SqliteSettlementRepository { pool: SqlitePool, @@ -152,6 +159,36 @@ fn now_unix_secs() -> Result { .map_err(|_| DataLayerError::InvalidInput("timestamp overflow".to_string())) } +async fn enqueue_provider_monthly_usage_delta_sqlite( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, + provider_id: &str, + total_cost_usd_delta: f64, + created_at: i64, +) -> Result<(), DataLayerError> { + let request_id = request_id.trim(); + let provider_id = provider_id.trim(); + if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 { + return Ok(()); + } + if !total_cost_usd_delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "provider monthly usage delta is not finite for {provider_id}" + ))); + } + + sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL) + .bind(uuid::Uuid::new_v4().to_string()) + .bind(request_id) + .bind(provider_id) + .bind(total_cost_usd_delta) + .bind(created_at) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + #[derive(Debug, Default)] struct DailyQuotaDebitResult { debited_usd: f64, @@ -360,6 +397,12 @@ impl SettlementWriteRepository for SqliteSettlementRepository { let updated_at = now_unix_secs()?; let mut tx = self.pool.begin().await.map_sql_err()?; + // SQLite transactions are deferred. Acquire the single writer slot before reading the + // billing status so concurrent settlement attempts cannot both observe `pending`. + sqlx::query("UPDATE \"usage\" SET billing_status = billing_status WHERE 0") + .execute(&mut *tx) + .await + .map_sql_err()?; let row = sqlx::query(FIND_USAGE_FOR_SETTLEMENT_SQL) .bind(&input.request_id) .fetch_optional(&mut *tx) @@ -633,31 +676,14 @@ WHERE id = ? .as_deref() .filter(|value| !value.is_empty()) { - sqlx::query( - r#" -UPDATE providers -SET - monthly_used_usd = CAST(COALESCE(monthly_used_usd, 0) AS REAL) + ?, - updated_at = ? -WHERE id = ? -"#, + enqueue_provider_monthly_usage_delta_sqlite( + &mut tx, + &input.request_id, + provider_id, + input.actual_total_cost_usd, + updated_at, ) - .bind(input.actual_total_cost_usd) - .bind(updated_at) - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - settlement.provider_monthly_used_usd = sqlx::query( - "SELECT CAST(monthly_used_usd AS REAL) AS monthly_used_usd FROM providers WHERE id = ? LIMIT 1", - ) - .bind(provider_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .map(|row| sqlite_real(&row, "monthly_used_usd")) - .transpose()?; + .await?; } } @@ -700,6 +726,7 @@ mod tests { SettlementWriteRepository, UsageSettlementInput, }; use sqlx::Row; + use std::time::Duration; #[tokio::test] async fn sqlite_repository_settles_usage_once() { @@ -737,7 +764,7 @@ mod tests { assert_eq!(settlement.wallet_balance_after, Some(6.0)); assert_eq!(settlement.wallet_recharge_balance_after, Some(4.0)); assert_eq!(settlement.wallet_gift_balance_after, Some(2.0)); - assert_eq!(settlement.provider_monthly_used_usd, Some(11.0)); + assert_eq!(settlement.provider_monthly_used_usd, None); let wallet = sqlx::query( "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'", @@ -772,7 +799,40 @@ mod tests { .fetch_one(&pool) .await .expect("provider should load"); - assert_eq!(provider_used, 11.0); + assert_eq!(provider_used, 5.0); + let provider_delta: (i64, f64) = sqlx::query_as( + r#" +SELECT COUNT(*), CAST(COALESCE(SUM(total_cost_usd_delta), 0) AS REAL) +FROM usage_counter_deltas +WHERE request_id = 'request-1' + AND kind = 'provider_monthly' + AND target_id = 'provider-1' +"#, + ) + .fetch_one(&pool) + .await + .expect("provider delta should load"); + assert_eq!(provider_delta, (1, 6.0)); + + let snapshot: (String, Option, Option, Option) = sqlx::query_as( + r#" +SELECT billing_status, wallet_id, wallet_balance_after, finalized_at +FROM usage_settlement_snapshots +WHERE request_id = 'request-1' +"#, + ) + .fetch_one(&pool) + .await + .expect("canonical settlement snapshot should load"); + assert_eq!( + snapshot, + ( + "settled".to_string(), + Some("wallet-1".to_string()), + Some(6.0), + Some(1_234), + ) + ); } #[tokio::test] @@ -851,7 +911,7 @@ mod tests { assert_eq!(settlement.wallet_balance_after, Some(-3.0)); assert_eq!(settlement.wallet_recharge_balance_after, Some(-3.0)); assert_eq!(settlement.wallet_gift_balance_after, Some(0.0)); - assert_eq!(settlement.provider_monthly_used_usd, Some(20.0)); + assert_eq!(settlement.provider_monthly_used_usd, None); let wallet = sqlx::query( "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'", @@ -862,6 +922,13 @@ mod tests { assert_eq!(wallet.try_get::("balance").unwrap(), -3.0); assert_eq!(wallet.try_get::("gift_balance").unwrap(), 0.0); assert_eq!(wallet.try_get::("total_consumed").unwrap(), 15.0); + let provider_delta: f64 = sqlx::query_scalar( + "SELECT total_cost_usd_delta FROM usage_counter_deltas WHERE request_id = 'request-overdraw' AND kind = 'provider_monthly'", + ) + .fetch_one(&pool) + .await + .expect("provider delta should load"); + assert_eq!(provider_delta, 15.0); } #[tokio::test] @@ -916,6 +983,74 @@ mod tests { assert_eq!(quota_used, 6.0); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn sqlite_repository_serializes_concurrent_settlement_attempts() { + let database_path = std::env::temp_dir().join(format!( + "aether-settlement-parity-{}.db", + uuid::Uuid::new_v4() + )); + let options = sqlx::sqlite::SqliteConnectOptions::new() + .filename(&database_path) + .create_if_missing(true) + .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(2) + .connect_with(options) + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_settlement_rows(&pool).await; + + let repository = SqliteSettlementRepository::new(pool.clone()); + let input = UsageSettlementInput { + request_id: "request-1".to_string(), + user_id: Some("user-1".to_string()), + api_key_id: None, + api_key_is_standalone: false, + provider_id: Some("provider-1".to_string()), + status: "completed".to_string(), + billing_status: "pending".to_string(), + total_cost_usd: 3.0, + actual_total_cost_usd: 6.0, + finalized_at_unix_secs: Some(1_234), + }; + let (first, second) = tokio::join!( + repository.settle_usage(input.clone()), + repository.settle_usage(input) + ); + let first = first + .expect("first settlement should succeed") + .expect("usage should exist"); + let second = second + .expect("second settlement should succeed") + .expect("usage should exist"); + assert_eq!(first.billing_status, "settled"); + assert_eq!(second.billing_status, "settled"); + + let wallet: (f64, f64, f64) = sqlx::query_as( + "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'", + ) + .fetch_one(&pool) + .await + .expect("wallet should load"); + assert_eq!(wallet, (4.0, 2.0, 6.0)); + let delta_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = 'request-1' AND kind = 'provider_monthly'", + ) + .fetch_one(&pool) + .await + .expect("provider deltas should count"); + assert_eq!(delta_count, 1); + + pool.close().await; + let _ = std::fs::remove_file(&database_path); + let _ = std::fs::remove_file(format!("{}-wal", database_path.display())); + let _ = std::fs::remove_file(format!("{}-shm", database_path.display())); + } + async fn seed_settlement_rows(pool: &sqlx::SqlitePool) { sqlx::query( r#" diff --git a/crates/aether-data/adapters/sqlite/src/usage.rs b/crates/aether-data/adapters/sqlite/src/usage.rs index 49e0930d3..6030a894d 100644 --- a/crates/aether-data/adapters/sqlite/src/usage.rs +++ b/crates/aether-data/adapters/sqlite/src/usage.rs @@ -3,7 +3,7 @@ use std::io::Read; use std::time::{SystemTime, UNIX_EPOCH}; use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; -use aether_data_contracts::repository::usage::{parse_usage_body_ref, UsageBodyField}; +use aether_data_contracts::repository::usage::UsageBodyField; use async_trait::async_trait; use flate2::read::GzDecoder; use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; @@ -27,28 +27,46 @@ use aether_data_contracts::repository::usage::{ UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, - UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery, - UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, - UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, - UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, - UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, - UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery, - UsageWriteRepository, + UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupExecutionMode, + UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, + UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, + UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, + UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, + UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, + UsageReadRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, + UsageTimeSeriesQuery, UsageWriteRepository, }; use aether_data_contracts::DataLayerError; +mod cleanup; +mod counters; +mod http_capture; +mod snapshots; + const USAGE_COLUMNS: &str = r#" SELECT id, - request_id, + "usage".request_id, user_id, api_key_id, provider_name, model, target_model, - provider_id, - provider_endpoint_id, - provider_api_key_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_id + ELSE "usage".provider_id + END AS provider_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_endpoint_id + ELSE "usage".provider_endpoint_id + END AS provider_endpoint_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.selected_provider_api_key_id + ELSE "usage".provider_api_key_id + END AS provider_api_key_id, request_type, api_format, api_family, @@ -56,43 +74,199 @@ SELECT endpoint_api_format, provider_api_family, provider_endpoint_kind, - has_format_conversion, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN COALESCE(usage_routing_snapshots.has_format_conversion, 0) + ELSE COALESCE("usage".has_format_conversion, 0) + END AS has_format_conversion, is_stream, upstream_is_stream, input_tokens, - output_tokens, + COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0) + AS output_tokens, total_tokens, - cache_creation_input_tokens, - cache_creation_ephemeral_5m_input_tokens, - cache_creation_ephemeral_1h_input_tokens, - cache_read_input_tokens, - CAST(cache_creation_cost_usd AS REAL) AS cache_creation_cost_usd, - CAST(cache_read_cost_usd AS REAL) AS cache_read_cost_usd, - CAST(output_price_per_1m AS REAL) AS output_price_per_1m, - CAST(total_cost_usd AS REAL) AS total_cost_usd, - CAST(actual_total_cost_usd AS REAL) AS actual_total_cost_usd, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_tokens, + CASE + WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL + OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0) + + COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0) + END, + "usage".cache_creation_input_tokens, + 0 + ) AS cache_creation_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_5m_tokens, + "usage".cache_creation_ephemeral_5m_input_tokens, + 0 + ) AS cache_creation_ephemeral_5m_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_creation_1h_tokens, + "usage".cache_creation_ephemeral_1h_input_tokens, + 0 + ) AS cache_creation_ephemeral_1h_input_tokens, + COALESCE( + usage_settlement_snapshots.billing_cache_read_tokens, + "usage".cache_read_input_tokens, + 0 + ) AS cache_read_input_tokens, + CAST(COALESCE( + usage_settlement_snapshots.billing_cache_creation_cost_usd, + "usage".cache_creation_cost_usd, + 0 + ) AS REAL) AS cache_creation_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.billing_cache_read_cost_usd, + "usage".cache_read_cost_usd, + 0 + ) AS REAL) AS cache_read_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.output_price_per_1m, + "usage".output_price_per_1m + ) AS REAL) AS output_price_per_1m, + CAST(COALESCE( + usage_settlement_snapshots.billing_total_cost_usd, + "usage".total_cost_usd, + 0 + ) AS REAL) AS total_cost_usd, + CAST(COALESCE( + usage_settlement_snapshots.billing_actual_total_cost_usd, + "usage".actual_total_cost_usd, + 0 + ) AS REAL) AS actual_total_cost_usd, status_code, error_message, error_category, response_time_ms, first_byte_time_ms, status, - billing_status, + COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) + AS billing_status, + COALESCE(usage_http_audits.request_headers, "usage".request_headers) AS request_headers, + "usage".request_body, + "usage".request_body_compressed, + COALESCE( + usage_http_audits.provider_request_headers, + "usage".provider_request_headers + ) AS provider_request_headers, + "usage".provider_request_body, + "usage".provider_request_body_compressed, + COALESCE(usage_http_audits.response_headers, "usage".response_headers) AS response_headers, + "usage".response_body, + "usage".response_body_compressed, + COALESCE( + usage_http_audits.client_response_headers, + "usage".client_response_headers + ) AS client_response_headers, + "usage".client_response_body, + "usage".client_response_body_compressed, + usage_http_audits.request_body_ref AS http_request_body_ref, + usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref, + usage_http_audits.response_body_ref AS http_response_body_ref, + usage_http_audits.client_response_body_ref AS http_client_response_body_ref, + usage_http_audits.request_body_state AS http_request_body_state, + usage_http_audits.provider_request_body_state AS http_provider_request_body_state, + usage_http_audits.response_body_state AS http_response_body_state, + usage_http_audits.client_response_body_state AS http_client_response_body_state, request_metadata, - candidate_id, - candidate_index, - NULL AS username, - NULL AS api_key_name, - key_name, - planner_kind, - route_family, - route_kind, - execution_path, - local_execution_runtime_miss_reason, - finalized_at AS finalized_at_unix_secs, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.candidate_id + ELSE "usage".candidate_id + END AS routing_candidate_id, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.candidate_index + ELSE "usage".candidate_index + END AS routing_candidate_index, + "usage".username AS username, + "usage".api_key_name AS api_key_name, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.key_name + ELSE "usage".key_name + END AS routing_key_name, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.planner_kind + ELSE "usage".planner_kind + END AS routing_planner_kind, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.route_family + ELSE "usage".route_family + END AS routing_route_family, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.route_kind + ELSE "usage".route_kind + END AS routing_route_kind, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.execution_path + ELSE "usage".execution_path + END AS routing_execution_path, + CASE + WHEN usage_routing_snapshots.request_id IS NOT NULL + THEN usage_routing_snapshots.local_execution_runtime_miss_reason + ELSE "usage".local_execution_runtime_miss_reason + END AS routing_local_execution_runtime_miss_reason, + usage_settlement_snapshots.billing_snapshot_schema_version + AS settlement_billing_snapshot_schema_version, + usage_settlement_snapshots.billing_snapshot_status AS settlement_billing_snapshot_status, + CAST(usage_settlement_snapshots.rate_multiplier AS REAL) AS settlement_rate_multiplier, + usage_settlement_snapshots.is_free_tier AS settlement_is_free_tier, + CAST(usage_settlement_snapshots.input_price_per_1m AS REAL) + AS settlement_input_price_per_1m, + CAST(usage_settlement_snapshots.output_price_per_1m AS REAL) + AS settlement_output_price_per_1m, + CAST(usage_settlement_snapshots.cache_creation_price_per_1m AS REAL) + AS settlement_cache_creation_price_per_1m, + CAST(usage_settlement_snapshots.cache_read_price_per_1m AS REAL) + AS settlement_cache_read_price_per_1m, + CAST(usage_settlement_snapshots.price_per_request AS REAL) + AS settlement_price_per_request, + usage_settlement_snapshots.settlement_snapshot_schema_version + AS settlement_snapshot_schema_version, + usage_settlement_snapshots.settlement_snapshot AS settlement_snapshot, + usage_settlement_snapshots.billing_dimensions AS settlement_billing_dimensions, + usage_settlement_snapshots.billing_input_tokens AS settlement_billing_input_tokens, + usage_settlement_snapshots.billing_effective_input_tokens + AS settlement_billing_effective_input_tokens, + usage_settlement_snapshots.billing_output_tokens AS settlement_billing_output_tokens, + usage_settlement_snapshots.billing_cache_creation_tokens + AS settlement_billing_cache_creation_tokens, + usage_settlement_snapshots.billing_cache_creation_5m_tokens + AS settlement_billing_cache_creation_5m_tokens, + usage_settlement_snapshots.billing_cache_creation_1h_tokens + AS settlement_billing_cache_creation_1h_tokens, + usage_settlement_snapshots.billing_cache_read_tokens + AS settlement_billing_cache_read_tokens, + usage_settlement_snapshots.billing_total_input_context + AS settlement_billing_total_input_context, + CAST(usage_settlement_snapshots.billing_cache_creation_cost_usd AS REAL) + AS settlement_billing_cache_creation_cost_usd, + CAST(usage_settlement_snapshots.billing_cache_read_cost_usd AS REAL) + AS settlement_billing_cache_read_cost_usd, + CAST(usage_settlement_snapshots.billing_total_cost_usd AS REAL) + AS settlement_billing_total_cost_usd, + CAST(usage_settlement_snapshots.billing_actual_total_cost_usd AS REAL) + AS settlement_billing_actual_total_cost_usd, + usage_settlement_snapshots.billing_pricing_source AS settlement_billing_pricing_source, + usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id, + usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version, + COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at) + AS finalized_at_unix_secs, created_at_unix_ms, updated_at_unix_secs FROM "usage" +LEFT JOIN usage_http_audits + ON usage_http_audits.request_id = "usage".request_id +LEFT JOIN usage_routing_snapshots + ON usage_routing_snapshots.request_id = "usage".request_id +LEFT JOIN usage_settlement_snapshots + ON usage_settlement_snapshots.request_id = "usage".request_id "#; const UPSERT_USAGE_SQL: &str = r#" @@ -274,6 +448,102 @@ ON CONFLICT (request_id) DO UPDATE SET END "#; +const SQLITE_FIRST_BYTE_BATCH_SIZE: usize = 32; +const UPSERT_FIRST_BYTE_BATCH_PREFIX_SQL: &str = r#" +INSERT INTO "usage" ( + id, + request_id, + user_id, + api_key_id, + provider_name, + model, + target_model, + provider_id, + provider_endpoint_id, + provider_api_key_id, + request_type, + api_format, + api_family, + endpoint_kind, + endpoint_api_format, + provider_api_family, + provider_endpoint_kind, + has_format_conversion, + is_stream, + upstream_is_stream, + status_code, + response_time_ms, + first_byte_time_ms, + status, + billing_status, + request_metadata, + created_at, + created_at_unix_ms, + updated_at_unix_secs +) +"#; + +const UPSERT_FIRST_BYTE_BATCH_UPDATE_PREFIX_SQL: &str = r#" +ON CONFLICT (request_id) +DO UPDATE SET + user_id = COALESCE(excluded.user_id, "usage".user_id), + api_key_id = COALESCE(excluded.api_key_id, "usage".api_key_id), + provider_name = excluded.provider_name, + model = excluded.model, + target_model = COALESCE(excluded.target_model, "usage".target_model), + provider_id = COALESCE(excluded.provider_id, "usage".provider_id), + provider_endpoint_id = COALESCE(excluded.provider_endpoint_id, "usage".provider_endpoint_id), + provider_api_key_id = COALESCE(excluded.provider_api_key_id, "usage".provider_api_key_id), + request_type = COALESCE(excluded.request_type, "usage".request_type), + api_format = COALESCE(excluded.api_format, "usage".api_format), + api_family = COALESCE(excluded.api_family, "usage".api_family), + endpoint_kind = COALESCE(excluded.endpoint_kind, "usage".endpoint_kind), + endpoint_api_format = COALESCE(excluded.endpoint_api_format, "usage".endpoint_api_format), + provider_api_family = COALESCE(excluded.provider_api_family, "usage".provider_api_family), + provider_endpoint_kind = COALESCE( + excluded.provider_endpoint_kind, + "usage".provider_endpoint_kind + ), + has_format_conversion = +"#; + +const UPSERT_FIRST_BYTE_BATCH_UPDATE_SUFFIX_SQL: &str = r#", + is_stream = 1, + upstream_is_stream = COALESCE( + CASE json_type(excluded.request_metadata, '$.upstream_is_stream') + WHEN 'true' THEN 1 + WHEN 'false' THEN 0 + ELSE NULL + END, + "usage".upstream_is_stream, + "usage".is_stream, + 1 + ), + status_code = COALESCE(excluded.status_code, "usage".status_code), + response_time_ms = CASE + WHEN excluded.response_time_ms IS NULL OR excluded.response_time_ms = 0 + THEN "usage".response_time_ms + ELSE excluded.response_time_ms + END, + first_byte_time_ms = CASE + WHEN "usage".first_byte_time_ms IS NOT NULL AND "usage".first_byte_time_ms <> 0 + THEN "usage".first_byte_time_ms + WHEN excluded.first_byte_time_ms IS NULL OR excluded.first_byte_time_ms = 0 + THEN "usage".first_byte_time_ms + ELSE excluded.first_byte_time_ms + END, + status = 'streaming', + request_metadata = COALESCE("usage".request_metadata, excluded.request_metadata), + updated_at_unix_secs = MAX( + COALESCE(NULLIF("usage".updated_at_unix_secs, 0), 0), + COALESCE(NULLIF(excluded.updated_at_unix_secs, 0), 0), + COALESCE("usage".created_at_unix_ms, 0) + ) +WHERE "usage".billing_status = 'pending' + AND "usage".status IN ('pending', 'streaming') + AND "usage".finalized_at IS NULL +"#; + const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#" SELECT "usage".request_id, @@ -407,35 +677,84 @@ END "#; const SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR: &str = r#" -( - CASE - WHEN ( - LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'openai' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'openai:%' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'gemini' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'gemini:%' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'google' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'google:%' - ) - AND COALESCE(input_tokens, 0) > 0 - AND COALESCE(cache_read_input_tokens, 0) > 0 - THEN MAX(COALESCE(input_tokens, 0) - COALESCE(cache_read_input_tokens, 0), 0) - ELSE MAX(COALESCE(input_tokens, 0), 0) - END - + MAX(COALESCE(output_tokens, 0), 0) - + ( +MAX( + COALESCE( CASE - WHEN COALESCE(cache_creation_input_tokens, 0) = 0 - AND ( - COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) - + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0) - ) > 0 - THEN COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) - + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0) - ELSE MAX(COALESCE(cache_creation_input_tokens, 0), 0) - END - ) - + MAX(COALESCE(cache_read_input_tokens, 0), 0) + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + MAX(settlement.billing_effective_input_tokens, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + + MAX( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) > 0 + THEN COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ELSE COALESCE("usage".cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + MAX( + COALESCE( + settlement.billing_cache_read_tokens, + "usage".cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + MAX(settlement.billing_total_input_context, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + END, + NULLIF(MAX(COALESCE("usage".total_tokens, 0), 0), 0), + ( + CASE + WHEN ( + LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) = 'openai' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) LIKE 'openai:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) = 'gemini' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) LIKE 'gemini:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) = 'google' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) LIKE 'google:%' + ) + AND COALESCE("usage".input_tokens, 0) > 0 + AND COALESCE("usage".cache_read_input_tokens, 0) > 0 + THEN MAX( + COALESCE("usage".input_tokens, 0) - COALESCE("usage".cache_read_input_tokens, 0), + 0 + ) + ELSE MAX(COALESCE("usage".input_tokens, 0), 0) + END + + MAX(COALESCE("usage".output_tokens, 0), 0) + + ( + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) > 0 + THEN COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ELSE MAX(COALESCE("usage".cache_creation_input_tokens, 0), 0) + END + ) + + MAX(COALESCE("usage".cache_read_input_tokens, 0), 0) + ), + 0 + ), + 0 ) "#; @@ -742,7 +1061,7 @@ fn push_sqlite_usage_order_limit_offset( if newest_first { builder.push(" ORDER BY created_at_unix_ms DESC, id ASC"); } else { - builder.push(" ORDER BY created_at_unix_ms ASC, request_id ASC"); + builder.push(" ORDER BY created_at_unix_ms ASC, \"usage\".request_id ASC"); } if let Some(limit) = limit { builder.push(" LIMIT ").push_bind(limit as i64); @@ -1057,6 +1376,61 @@ pub struct SqliteUsageWriteRepository { pool: SqlitePool, } +#[derive(Debug)] +struct PreparedFirstByteUsage { + usage: UpsertUsageRecord, + request_metadata_json: Option, + response_time_ms: Option, + first_byte_time_ms: Option, + created_at_unix_secs: i64, + updated_at_unix_secs: i64, +} + +impl PreparedFirstByteUsage { + fn try_from_usage(usage: UpsertUsageRecord) -> Result { + usage.validate()?; + if usage.status != "streaming" || usage.billing_status != "pending" { + return Err(DataLayerError::InvalidInput( + "first-byte usage upsert requires streaming status with pending billing" + .to_string(), + )); + } + + let usage = strip_deprecated_usage_display_fields(usage); + let request_metadata_json = usage + .request_metadata + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; + let response_time_ms = usage + .response_time_ms + .map(|value| to_i64(value, "usage.response_time_ms")) + .transpose()?; + let first_byte_time_ms = usage + .first_byte_time_ms + .map(|value| to_i64(value, "usage.first_byte_time_ms")) + .transpose()?; + let created_at_unix_secs = to_i64( + usage + .created_at_unix_ms + .unwrap_or(usage.updated_at_unix_secs), + "usage.created_at_unix_ms", + )?; + let updated_at_unix_secs = + to_i64(usage.updated_at_unix_secs, "usage.updated_at_unix_secs")?; + + Ok(Self { + usage, + request_metadata_json, + response_time_ms, + first_byte_time_ms, + created_at_unix_secs, + updated_at_unix_secs, + }) + } +} + #[derive(Debug, Clone)] pub struct SqliteUsageReadRepository { pool: SqlitePool, @@ -1466,7 +1840,7 @@ LEFT JOIN first_byte_percentiles ON first_byte_percentiles.provider_id = provide mut builder: QueryBuilder<'_, Sqlite>, ) -> Result, DataLayerError> { let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_usage_row).collect() + rows.iter().map(|row| map_usage_row(row, false)).collect() } pub async fn list_usage_audits( @@ -1484,7 +1858,7 @@ LEFT JOIN first_byte_percentiles ON first_byte_percentiles.provider_id = provide ); let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_usage_row).collect() + rows.iter().map(|row| map_usage_row(row, false)).collect() } pub async fn count_usage_audits( @@ -1514,7 +1888,7 @@ LEFT JOIN first_byte_percentiles ON first_byte_percentiles.provider_id = provide ); let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_usage_row).collect() + rows.iter().map(|row| map_usage_row(row, false)).collect() } pub async fn count_usage_audits_by_keyword_search( @@ -1543,7 +1917,7 @@ SELECT COUNT(*) AS total_requests, COALESCE(SUM(MAX(COALESCE(input_tokens, 0), 0)), 0) AS input_tokens, COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS recorded_total_tokens, + COALESCE(SUM({total_tokens_expr}), 0) AS recorded_total_tokens, COALESCE(SUM({cache_creation_expr}), 0) AS cache_creation_tokens, COALESCE(SUM(MAX(COALESCE(cache_creation_ephemeral_5m_input_tokens, 0), 0)), 0) AS cache_creation_ephemeral_5m_tokens, @@ -1566,8 +1940,11 @@ SELECT END ), 0) AS error_requests FROM "usage" +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id "#, - cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR + cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR, + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )); let mut has_where = false; push_sqlite_usage_summary_filters(&mut builder, query, &mut has_where); @@ -1613,7 +1990,7 @@ SELECT {display_expr} AS display_name, {secondary_expr} AS secondary_name, COUNT(*) AS request_count, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens, + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, COALESCE(SUM({effective_input_expr}), 0) AS effective_input_tokens, COALESCE(SUM({total_input_context_expr}), 0) AS total_input_context, @@ -1629,11 +2006,14 @@ SELECT {avg_response_expr} AS avg_response_time_ms, {success_count_expr} AS success_count FROM "usage" +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id "#, effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR, total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR, secondary_expr = secondary_expr, - cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR + cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR, + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )); let mut has_where = false; push_sqlite_usage_where(&mut builder, &mut has_where); @@ -1722,16 +2102,17 @@ WHERE user_id IN ( ); } - let mut builder = QueryBuilder::::new( + let mut builder = QueryBuilder::::new(format!( r#" SELECT "usage".user_id, COUNT(*) AS request_count, - COALESCE(SUM(MAX(COALESCE("usage".total_tokens, 0), 0)), 0) AS total_tokens + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens FROM "usage" JOIN ( "#, - ); + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + )); for (index, user_id) in unique_user_ids.iter().enumerate() { if index > 0 { builder.push(" UNION ALL "); @@ -1747,6 +2128,8 @@ JOIN ( builder.push( r#" ) AS requested ON requested.user_id = "usage".user_id +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id WHERE "usage".created_at_unix_ms >= requested.cutoff_unix_secs AND "usage".status NOT IN ('pending', 'streaming') AND "usage".provider_name NOT IN ('unknown', 'pending') @@ -1946,12 +2329,14 @@ impl UsageReadRepository for SqliteUsageReadRepository { &self, id: &str, ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE id = ? LIMIT 1")) + let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE \"usage\".id = ? LIMIT 1")) .bind(id) .fetch_optional(&self.pool) .await .map_sql_err()?; - row.as_ref().map(map_usage_row).transpose() + row.as_ref() + .map(|row| map_usage_row(row, false)) + .transpose() } async fn list_by_ids( @@ -1962,7 +2347,7 @@ impl UsageReadRepository for SqliteUsageReadRepository { return Ok(Vec::new()); } let mut builder = QueryBuilder::::new(USAGE_COLUMNS); - builder.push(" WHERE id IN ("); + builder.push(" WHERE \"usage\".id IN ("); { let mut separated = builder.separated(", "); for id in ids { @@ -1977,61 +2362,30 @@ impl UsageReadRepository for SqliteUsageReadRepository { &self, request_id: &str, ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE request_id = ? LIMIT 1")) - .bind(request_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_usage_row).transpose() + let row = sqlx::query(&format!( + "{USAGE_COLUMNS} WHERE \"usage\".request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + let usage = row + .as_ref() + .map(|row| map_usage_row(row, true)) + .transpose()?; + match usage { + Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage) + .await + .map(Some), + None => Ok(None), + } } async fn resolve_body_ref( &self, body_ref: &str, ) -> Result, DataLayerError> { - let has_blob_table: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'usage_body_blobs'", - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - if has_blob_table > 0 { - let blob_row = - sqlx::query("SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? LIMIT 1") - .bind(body_ref) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - if let Some(row) = blob_row.as_ref() { - let payload_gzip = row.try_get::, _>("payload_gzip").map_sql_err()?; - return inflate_usage_json_value(&payload_gzip).map(Some); - } - } - - let Some((request_id, field)) = parse_usage_body_ref(body_ref) else { - return Ok(None); - }; - let (inline_column, compressed_column) = sqlite_usage_body_sql_columns(field); - let row = sqlx::query(&format!( - "SELECT {inline_column} AS inline_body, {compressed_column} AS compressed_body FROM \"usage\" WHERE request_id = ? LIMIT 1" - )) - .bind(request_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - let Some(row) = row.as_ref() else { - return Ok(None); - }; - if let Some(raw) = row - .try_get::, _>("inline_body") - .map_sql_err()? - { - return parse_usage_json_text(&raw).map(Some); - } - row.try_get::>, _>("compressed_body") - .map_sql_err()? - .map(|bytes| inflate_usage_json_value(&bytes)) - .transpose() + http_capture::resolve_body_ref(&self.pool, body_ref).await } async fn list_usage_audits( @@ -2362,6 +2716,8 @@ SELECT COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples FROM "usage" +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id "#, effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR, total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, @@ -2426,14 +2782,17 @@ SELECT model, provider_name AS provider, COUNT(*) AS requests, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens, + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL THEN MAX(COALESCE(response_time_ms, 0), 0) ELSE 0 END), 0) AS response_time_sum_ms, COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples FROM "usage" -"# +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +"#, + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )); let mut has_where = false; push_sqlite_usage_range( @@ -2536,7 +2895,7 @@ SELECT {group_expr} AS group_key, COUNT(*) AS request_count, COALESCE(SUM(MAX(COALESCE(input_tokens, 0), 0)), 0) AS input_tokens, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens, + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, COALESCE(SUM({effective_input_expr}), 0) AS effective_input_tokens, COALESCE(SUM({total_input_context_expr}), 0) AS total_input_context, @@ -2559,11 +2918,14 @@ SELECT COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS overall_response_time_samples FROM "usage" +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id "#, effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR, total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR, cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR, - success_flag_expr = SQLITE_USAGE_SUCCESS_FLAG_EXPR + success_flag_expr = SQLITE_USAGE_SUCCESS_FLAG_EXPR, + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )); let mut has_where = false; push_sqlite_usage_range( @@ -3459,6 +3821,8 @@ SELECT COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd FROM "usage" +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id "#, total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )); @@ -3531,22 +3895,25 @@ FROM "usage" return Ok(BTreeMap::new()); } - let mut builder = QueryBuilder::::new( + let mut builder = QueryBuilder::::new(format!( r#" SELECT - api_key_id, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens + "usage".api_key_id, + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens FROM "usage" -WHERE api_key_id IN ( +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +WHERE "usage".api_key_id IN ( "#, - ); + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + )); { let mut separated = builder.separated(", "); for api_key_id in api_key_ids { separated.push_bind(api_key_id.clone()); } } - builder.push(") GROUP BY api_key_id ORDER BY api_key_id ASC"); + builder.push(") GROUP BY \"usage\".api_key_id ORDER BY \"usage\".api_key_id ASC"); let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; let mut totals = BTreeMap::new(); @@ -3567,25 +3934,30 @@ WHERE api_key_id IN ( return Ok(BTreeMap::new()); } - let mut builder = QueryBuilder::::new( + let mut builder = QueryBuilder::::new(format!( r#" SELECT - provider_api_key_id, + "usage".provider_api_key_id, COUNT(*) AS request_count, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - MAX(created_at_unix_ms) AS last_used_at_unix_secs + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, + COALESCE(SUM(COALESCE(CAST("usage".total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, + MAX("usage".created_at_unix_ms) AS last_used_at_unix_secs FROM "usage" -WHERE provider_api_key_id IN ( +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +WHERE "usage".provider_api_key_id IN ( "#, - ); + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + )); { let mut separated = builder.separated(", "); for provider_api_key_id in provider_api_key_ids { separated.push_bind(provider_api_key_id.clone()); } } - builder.push(") GROUP BY provider_api_key_id ORDER BY provider_api_key_id ASC"); + builder.push( + ") GROUP BY \"usage\".provider_api_key_id ORDER BY \"usage\".provider_api_key_id ASC", + ); let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; let mut summaries = BTreeMap::new(); @@ -3632,18 +4004,21 @@ WHERE provider_api_key_id IN ( )); } - let row = sqlx::query( + let row = sqlx::query(&format!( r#" SELECT COUNT(*) AS request_count, - COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, + COALESCE(SUM(COALESCE(CAST("usage".total_cost_usd AS REAL), 0)), 0) AS total_cost_usd FROM "usage" -WHERE provider_api_key_id = ? - AND created_at_unix_ms >= ? - AND created_at_unix_ms < ? +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +WHERE "usage".provider_api_key_id = ? + AND "usage".created_at_unix_ms >= ? + AND "usage".created_at_unix_ms < ? "#, - ) + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + )) .bind(provider_api_key_id) .bind(request.start_unix_secs as i64) .bind(request.end_unix_secs as i64) @@ -3736,6 +4111,22 @@ WHERE provider_id = ? Ok(summaries.into_values().collect()) } + + async fn read_usage_counter_health( + &self, + ) -> Result + { + counters::read_health(&self.pool).await + } + + async fn read_usage_counter_pending_health( + &self, + ) -> Result< + aether_data_contracts::repository::usage::UsageCounterPendingHealthSnapshot, + DataLayerError, + > { + counters::read_pending_health(&self.pool).await + } } fn map_sqlite_usage_daily_summary( @@ -3757,6 +4148,12 @@ fn usage_current_unix_secs() -> u64 { .unwrap_or_default() } +fn first_byte_transition_allowed(existing: &StoredRequestUsageAudit) -> bool { + existing.billing_status == "pending" + && matches!(existing.status.as_str(), "pending" | "streaming") + && existing.finalized_at_unix_secs.is_none() +} + impl SqliteUsageWriteRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } @@ -3766,48 +4163,299 @@ impl SqliteUsageWriteRepository { &self, request_id: &str, ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE request_id = ? LIMIT 1")) - .bind(request_id) - .fetch_optional(&self.pool) + let row = sqlx::query(&format!( + "{USAGE_COLUMNS} WHERE \"usage\".request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + let usage = row + .as_ref() + .map(|row| map_usage_row(row, true)) + .transpose()?; + match usage { + Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage) + .await + .map(Some), + None => Ok(None), + } + } + + async fn upsert_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + usage: UpsertUsageRecord, + ) -> Result<(), DataLayerError> { + let mut usage = strip_deprecated_usage_display_fields(usage); + usage.validate()?; + let prepared_capture = http_capture::prepare_usage_http_capture(&mut usage)?; + let existing = counters::lock_and_load_usage(tx, &usage.request_id).await?; + let recovers_terminal_failure = existing.as_ref().is_some_and(|existing| { + usage_can_recover_terminal_failure( + &existing.status, + &existing.billing_status, + &usage.status, + &usage.billing_status, + ) + }); + if existing.as_ref().is_some_and(|existing| { + (existing.billing_status == "settled" || existing.billing_status == "void") + && !recovers_terminal_failure + }) { + return Ok(()); + } + + let capture_update_allowed = recovers_terminal_failure + || http_capture::capture_update_allowed(existing.as_ref(), &usage.status); + if capture_update_allowed { + http_capture::apply_previous_metadata_tombstones(&mut usage, existing.as_ref()); + } + let prepared_snapshots = capture_update_allowed + .then(|| snapshots::from_usage(&usage)) + .transpose()?; + bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? + .execute(&mut **tx) .await .map_sql_err()?; - row.as_ref().map(map_usage_row).transpose() + if capture_update_allowed { + http_capture::sync_usage_http_capture(tx, &usage.request_id, &prepared_capture).await?; + let (routing_snapshot, settlement_snapshot) = prepared_snapshots + .as_ref() + .expect("capture-allowed usage has prepared snapshots"); + snapshots::sync( + tx, + &usage.request_id, + routing_snapshot, + settlement_snapshot, + matches!(usage.status.as_str(), "completed" | "failed" | "cancelled"), + ) + .await?; + } + counters::enqueue_usage_transition_for_request(tx, &usage.request_id, existing.as_ref()) + .await + } + + async fn execute_first_byte_batch( + tx: &mut sqlx::Transaction<'_, Sqlite>, + rows: &[&PreparedFirstByteUsage], + preserve_existing_format_conversion: bool, + ) -> Result<(), DataLayerError> { + if rows.is_empty() { + return Ok(()); + } + + let mut builder = QueryBuilder::::new(UPSERT_FIRST_BYTE_BATCH_PREFIX_SQL); + builder.push_values(rows, |mut values, row| { + values + .push_bind(row.usage.request_id.clone()) + .push_bind(row.usage.request_id.clone()) + .push_bind(row.usage.user_id.clone()) + .push_bind(row.usage.api_key_id.clone()) + .push_bind(row.usage.provider_name.clone()) + .push_bind(row.usage.model.clone()) + .push_bind(row.usage.target_model.clone()) + .push_bind(row.usage.provider_id.clone()) + .push_bind(row.usage.provider_endpoint_id.clone()) + .push_bind(row.usage.provider_api_key_id.clone()) + .push_bind(row.usage.request_type.clone()) + .push_bind(row.usage.api_format.clone()) + .push_bind(row.usage.api_family.clone()) + .push_bind(row.usage.endpoint_kind.clone()) + .push_bind(row.usage.endpoint_api_format.clone()) + .push_bind(row.usage.provider_api_family.clone()) + .push_bind(row.usage.provider_endpoint_kind.clone()) + .push_bind(row.usage.has_format_conversion.unwrap_or(false)) + .push("1") + .push_bind( + row.usage + .request_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY)) + .and_then(serde_json::Value::as_bool) + .unwrap_or(true), + ) + .push_bind(row.usage.status_code.map(i64::from)) + .push_bind(row.response_time_ms) + .push_bind(row.first_byte_time_ms) + .push("'streaming'") + .push("'pending'") + .push_bind(row.request_metadata_json.clone()) + .push_bind(row.created_at_unix_secs) + .push_bind(row.created_at_unix_secs) + .push_bind(row.updated_at_unix_secs); + }); + builder.push(UPSERT_FIRST_BYTE_BATCH_UPDATE_PREFIX_SQL); + if preserve_existing_format_conversion { + builder.push("COALESCE(\"usage\".has_format_conversion, 0)"); + } else { + builder.push( + "COALESCE(excluded.has_format_conversion, \"usage\".has_format_conversion, 0)", + ); + } + builder.push(UPSERT_FIRST_BYTE_BATCH_UPDATE_SUFFIX_SQL); + builder.build().execute(&mut **tx).await.map_sql_err()?; + Ok(()) + } + + async fn upsert_first_byte_many_native( + &self, + usages: Vec, + ) -> Result<(), DataLayerError> { + if usages.is_empty() { + return Ok(()); + } + + let mut request_id_counts = BTreeMap::::new(); + for usage in &usages { + *request_id_counts + .entry(usage.request_id.clone()) + .or_default() += 1; + } + let mut batch_rows = Vec::new(); + let mut fallback_rows = Vec::new(); + for (sequence, usage) in usages.into_iter().enumerate() { + let prepared = PreparedFirstByteUsage::try_from_usage(usage)?; + if request_id_counts + .get(&prepared.usage.request_id) + .copied() + .unwrap_or_default() + == 1 + { + batch_rows.push(prepared); + } else { + fallback_rows.push((sequence, prepared)); + } + } + + let mut tx = self.pool.begin().await.map_sql_err()?; + let mut before = BTreeMap::>::new(); + let request_ids = batch_rows + .iter() + .map(|row| row.usage.request_id.clone()) + .collect::>(); + for request_id in request_ids { + let existing = counters::lock_and_load_usage(&mut tx, &request_id).await?; + before.insert(request_id, existing); + } + let eligible_rows = batch_rows + .iter() + .filter(|row| { + before + .get(&row.usage.request_id) + .and_then(Option::as_ref) + .is_none_or(first_byte_transition_allowed) + }) + .collect::>(); + for preserve_existing_format_conversion in [true, false] { + let matching = eligible_rows + .iter() + .copied() + .filter(|row| { + row.usage.has_format_conversion.is_none() == preserve_existing_format_conversion + }) + .collect::>(); + for chunk in matching.chunks(SQLITE_FIRST_BYTE_BATCH_SIZE) { + Self::execute_first_byte_batch(&mut tx, chunk, preserve_existing_format_conversion) + .await?; + } + } + for row in eligible_rows { + counters::enqueue_usage_transition_for_request( + &mut tx, + &row.usage.request_id, + before.get(&row.usage.request_id).and_then(Option::as_ref), + ) + .await?; + } + + fallback_rows.sort_by_key(|(sequence, _)| *sequence); + for (_, row) in fallback_rows { + let existing = counters::lock_and_load_usage(&mut tx, &row.usage.request_id).await?; + if existing + .as_ref() + .is_some_and(|existing| !first_byte_transition_allowed(existing)) + { + continue; + } + Self::execute_first_byte_batch( + &mut tx, + &[&row], + row.usage.has_format_conversion.is_none(), + ) + .await?; + counters::enqueue_usage_transition_for_request( + &mut tx, + &row.usage.request_id, + existing.as_ref(), + ) + .await?; + } + + tx.commit().await.map_sql_err() } } #[async_trait] impl UsageWriteRepository for SqliteUsageWriteRepository { + fn supports_first_byte_usage_fast_path(&self) -> bool { + true + } + + fn supports_first_byte_usage_batch(&self) -> bool { + true + } + + fn supports_pending_usage_batch(&self) -> bool { + true + } + async fn upsert( &self, usage: UpsertUsageRecord, ) -> Result { - let usage = strip_deprecated_usage_display_fields(usage); usage.validate()?; + let request_id = usage.request_id.clone(); + let mut tx = self.pool.begin().await.map_sql_err()?; + Self::upsert_in_tx(&mut tx, usage).await?; + tx.commit().await.map_sql_err()?; + self.find_by_request_id(&request_id).await?.ok_or_else(|| { + DataLayerError::UnexpectedValue("usage upsert returned no row".to_string()) + }) + } - if let Some(existing) = self.find_by_request_id(&usage.request_id).await? { - if (existing.billing_status == "settled" || existing.billing_status == "void") - && !usage_can_recover_terminal_failure( - &existing.status, - &existing.billing_status, - &usage.status, - &usage.billing_status, - ) - { - return Ok(existing); + async fn upsert_first_byte(&self, usage: UpsertUsageRecord) -> Result<(), DataLayerError> { + self.upsert_first_byte_many_native(vec![usage]).await + } + + async fn upsert_first_byte_many( + &self, + usages: Vec, + ) -> Result<(), DataLayerError> { + self.upsert_first_byte_many_native(usages).await + } + + async fn upsert_pending_many( + &self, + usages: Vec, + ) -> Result<(), DataLayerError> { + if usages.is_empty() { + return Ok(()); + } + for usage in &usages { + usage.validate()?; + if usage.status != "pending" || usage.billing_status != "pending" { + return Err(DataLayerError::InvalidInput( + "pending usage batch requires pending status with pending billing".to_string(), + )); } } - bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? - .execute(&self.pool) - .await - .map_sql_err()?; - self.rebuild_api_key_usage_stats().await?; - self.rebuild_provider_api_key_usage_stats().await?; - self.find_by_request_id(&usage.request_id) - .await? - .ok_or_else(|| { - DataLayerError::UnexpectedValue("usage upsert returned no row".to_string()) - }) + let mut tx = self.pool.begin().await.map_sql_err()?; + for usage in usages { + Self::upsert_in_tx(&mut tx, usage).await?; + } + tx.commit().await.map_sql_err() } async fn rebuild_api_key_usage_stats(&self) -> Result { @@ -3824,19 +4472,28 @@ SET total_requests = 0, .await .map_sql_err()?; - let rows = sqlx::query( + let rows = sqlx::query(&format!( r#" SELECT - api_key_id, + "usage".api_key_id, COUNT(*) AS total_requests, - COALESCE(SUM(total_tokens), 0) AS total_tokens, - CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL) AS total_cost_usd, - MAX(updated_at_unix_secs) AS last_used_at + COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, + CAST(COALESCE(SUM("usage".total_cost_usd), 0) AS REAL) AS total_cost_usd, + MAX(COALESCE( + "usage".created_at, + "usage".created_at_unix_ms, + "usage".updated_at_unix_secs + )) AS last_used_at FROM "usage" -WHERE api_key_id IS NOT NULL AND api_key_id <> '' -GROUP BY api_key_id +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +WHERE "usage".api_key_id IS NOT NULL + AND TRIM("usage".api_key_id) <> '' + AND "usage".status NOT IN ('pending', 'streaming') +GROUP BY "usage".api_key_id "#, - ) + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR + )) .fetch_all(&self.pool) .await .map_sql_err()?; @@ -3888,31 +4545,34 @@ SET request_count = 0, let rows = sqlx::query(&format!( r#" SELECT - provider_api_key_id, + "usage".provider_api_key_id, COUNT(*) AS request_count, COALESCE(SUM({success_flag_expr}), 0) AS success_count, COALESCE(SUM({error_flag_expr}), 0) AS error_count, COALESCE(SUM(CASE - WHEN status IN ('pending', 'streaming') THEN 0 - ELSE MAX(COALESCE(total_tokens, 0), 0) + WHEN "usage".status IN ('pending', 'streaming') THEN 0 + ELSE {total_tokens_expr} END), 0) AS total_tokens, COALESCE(SUM(CASE - WHEN status IN ('pending', 'streaming') THEN 0 - ELSE COALESCE(CAST(total_cost_usd AS REAL), 0) + WHEN "usage".status IN ('pending', 'streaming') THEN 0 + ELSE COALESCE(CAST("usage".total_cost_usd AS REAL), 0) END), 0) AS total_cost_usd, COALESCE(SUM(CASE - WHEN {success_flag_expr} = 1 AND response_time_ms IS NOT NULL - THEN MAX(COALESCE(response_time_ms, 0), 0) + WHEN {success_flag_expr} = 1 AND "usage".response_time_ms IS NOT NULL + THEN MAX(COALESCE("usage".response_time_ms, 0), 0) ELSE 0 END), 0) AS total_response_time_ms, - MAX(created_at_unix_ms) AS last_used_at + MAX("usage".created_at_unix_ms) AS last_used_at FROM "usage" -WHERE provider_api_key_id IS NOT NULL - AND TRIM(provider_api_key_id) <> '' -GROUP BY provider_api_key_id +LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id +WHERE "usage".provider_api_key_id IS NOT NULL + AND TRIM("usage".provider_api_key_id) <> '' +GROUP BY "usage".provider_api_key_id "#, success_flag_expr = SQLITE_PROVIDER_KEY_SUCCESS_FLAG_EXPR, - error_flag_expr = SQLITE_PROVIDER_KEY_ERROR_FLAG_EXPR + error_flag_expr = SQLITE_PROVIDER_KEY_ERROR_FLAG_EXPR, + total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR )) .fetch_all(&self.pool) .await @@ -3968,7 +4628,6 @@ WHERE id = ? return Ok(PendingUsageCleanupSummary::default()); } - let cutoff_unix_ms = cutoff_unix_secs.saturating_mul(1000); let now_unix_ms = now_unix_secs.saturating_mul(1000); let mut summary = PendingUsageCleanupSummary::default(); let batch_size_u64 = u64::try_from(batch_size).map_err(|_| { @@ -3980,7 +4639,7 @@ WHERE id = ? loop { let mut tx = self.pool.begin().await.map_sql_err()?; let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL) - .bind(to_i64(cutoff_unix_ms, "stale pending usage cutoff")?) + .bind(to_i64(cutoff_unix_secs, "stale pending usage cutoff")?) .bind(to_i64(batch_size_u64, "stale pending usage batch size")?) .fetch_all(&mut *tx) .await @@ -4114,6 +4773,71 @@ WHERE request_id = ? Ok(summary) } + + async fn flush_usage_counter_deltas( + &self, + batch_size: usize, + ) -> Result + { + counters::flush(&self.pool, batch_size).await + } + + async fn enqueue_proxy_node_counter_delta( + &self, + delta: aether_data_contracts::repository::usage::ProxyNodeCounterDelta, + ) -> Result { + counters::enqueue_proxy_node(&self.pool, delta).await + } + + async fn enqueue_management_token_counter_delta( + &self, + delta: aether_data_contracts::repository::usage::ManagementTokenCounterDelta, + ) -> Result { + counters::enqueue_management_token(&self.pool, delta).await + } + + async fn enqueue_api_key_last_used_delta( + &self, + delta: aether_data_contracts::repository::usage::ApiKeyLastUsedDelta, + ) -> Result { + counters::enqueue_api_key_last_used(&self.pool, delta).await + } + + async fn cleanup_processed_usage_counter_deltas( + &self, + cutoff_unix_secs: u64, + batch_size: usize, + ) -> Result { + counters::cleanup_processed(&self.pool, cutoff_unix_secs, batch_size).await + } + + async fn cleanup_usage( + &self, + window: &UsageCleanupWindow, + batch_size: usize, + auto_delete_expired_keys: bool, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, + ) -> Result { + cleanup::cleanup_usage( + &self.pool, + window, + batch_size, + auto_delete_expired_keys, + targets, + mode, + ) + .await + } + + async fn preview_usage_cleanup( + &self, + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, + ) -> Result { + cleanup::preview_usage_cleanup(&self.pool, window, targets, mode).await + } } struct StalePendingUsageRow { @@ -4188,6 +4912,33 @@ INSERT INTO usage_settlement_snapshots ( ON CONFLICT (request_id) DO UPDATE SET billing_status = excluded.billing_status, + billing_snapshot_schema_version = NULL, + billing_snapshot_status = NULL, + settlement_snapshot_schema_version = NULL, + settlement_snapshot = NULL, + billing_dimensions = NULL, + billing_input_tokens = NULL, + billing_effective_input_tokens = NULL, + billing_output_tokens = NULL, + billing_cache_creation_tokens = NULL, + billing_cache_creation_5m_tokens = NULL, + billing_cache_creation_1h_tokens = NULL, + billing_cache_read_tokens = NULL, + billing_total_input_context = NULL, + billing_cache_creation_cost_usd = NULL, + billing_cache_read_cost_usd = NULL, + billing_total_cost_usd = NULL, + billing_actual_total_cost_usd = NULL, + billing_pricing_source = NULL, + billing_rule_id = NULL, + billing_rule_version = NULL, + rate_multiplier = NULL, + is_free_tier = NULL, + input_price_per_1m = NULL, + output_price_per_1m = NULL, + cache_creation_price_per_1m = NULL, + cache_read_price_per_1m = NULL, + price_per_request = NULL, finalized_at = COALESCE(usage_settlement_snapshots.finalized_at, excluded.finalized_at), updated_at = excluded.updated_at "#, @@ -4293,7 +5044,7 @@ fn bind_upsert<'q>( .unwrap_or(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens); let created_at = usage .created_at_unix_ms - .unwrap_or(usage.updated_at_unix_secs.saturating_mul(1000)); + .unwrap_or(usage.updated_at_unix_secs); let request_metadata = usage .request_metadata .as_ref() @@ -4369,7 +5120,10 @@ fn bind_upsert<'q>( Ok(query) } -fn map_usage_row(row: &SqliteRow) -> Result { +fn map_usage_row( + row: &SqliteRow, + resolve_legacy_compressed: bool, +) -> Result { let mut audit = StoredRequestUsageAudit::new( row.try_get("id").map_sql_err()?, row.try_get("request_id").map_sql_err()?, @@ -4428,24 +5182,13 @@ fn map_usage_row(row: &SqliteRow) -> Result, _>("upstream_is_stream") .map_sql_err()? .map(|value| value != 0); merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream); - audit.candidate_id = row.try_get("candidate_id").map_sql_err()?; - audit.candidate_index = row - .try_get::, _>("candidate_index") - .map_sql_err()? - .map(|value| value as u64); - audit.key_name = row.try_get("key_name").map_sql_err()?; - audit.planner_kind = row.try_get("planner_kind").map_sql_err()?; - audit.route_family = row.try_get("route_family").map_sql_err()?; - audit.route_kind = row.try_get("route_kind").map_sql_err()?; - audit.execution_path = row.try_get("execution_path").map_sql_err()?; - audit.local_execution_runtime_miss_reason = row - .try_get("local_execution_runtime_miss_reason") - .map_sql_err()?; + snapshots::hydrate_row(row, &mut audit)?; Ok(audit) } diff --git a/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs b/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs new file mode 100644 index 000000000..b3299f545 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs @@ -0,0 +1,1142 @@ +use std::io::Write; + +use aether_data_contracts::repository::usage::{ + parse_usage_body_ref, usage_body_ref, UsageBodyField, UsageCleanupExecutionMode, + UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, +}; +use chrono::{DateTime, Utc}; +use flate2::{write::GzEncoder, Compression}; +use serde_json::Value; +use sqlx::Row; +use tracing::warn; + +use crate::error::SqlResultExt; +use crate::{DataLayerError, SqlitePool}; + +const RAW_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +"#; + +const COMPRESSED_BODY_PREDICATE: &str = r#" +request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_body_blobs + WHERE usage_body_blobs.request_id = "usage".request_id +) +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = "usage".request_id + AND ( + usage_http_audits.request_body_ref IS NOT NULL + OR usage_http_audits.provider_request_body_ref IS NOT NULL + OR usage_http_audits.response_body_ref IS NOT NULL + OR usage_http_audits.client_response_body_ref IS NOT NULL + ) +) +"#; + +const ALL_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_body_blobs + WHERE usage_body_blobs.request_id = "usage".request_id +) +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = "usage".request_id + AND ( + usage_http_audits.request_body_ref IS NOT NULL + OR usage_http_audits.provider_request_body_ref IS NOT NULL + OR usage_http_audits.response_body_ref IS NOT NULL + OR usage_http_audits.client_response_body_ref IS NOT NULL + ) +) +"#; + +const INLINE_OR_COMPRESSED_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +"#; + +const HEADER_PREDICATE: &str = r#" +request_headers IS NOT NULL +OR response_headers IS NOT NULL +OR provider_request_headers IS NOT NULL +OR client_response_headers IS NOT NULL +OR EXISTS ( + SELECT 1 FROM usage_http_audits + WHERE usage_http_audits.request_id = "usage".request_id + AND ( + usage_http_audits.request_headers IS NOT NULL + OR usage_http_audits.response_headers IS NOT NULL + OR usage_http_audits.provider_request_headers IS NOT NULL + OR usage_http_audits.client_response_headers IS NOT NULL + ) +) +"#; + +const LEGACY_BODY_REF_PREDICATE: &str = r#" +request_metadata IS NOT NULL +AND json_valid(request_metadata) +AND ( + json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.request_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.provider_request_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.response_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.client_response_body_ref') IS NOT NULL +) +"#; + +const DETAIL_BODY_PREDICATE: &str = r#" +request_body IS NOT NULL +OR response_body IS NOT NULL +OR provider_request_body IS NOT NULL +OR client_response_body IS NOT NULL +OR request_body_compressed IS NOT NULL +OR response_body_compressed IS NOT NULL +OR provider_request_body_compressed IS NOT NULL +OR client_response_body_compressed IS NOT NULL +OR ( + request_metadata IS NOT NULL + AND json_valid(request_metadata) + AND ( + json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.request_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.provider_request_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.response_body_ref') IS NOT NULL + OR json_type(CASE WHEN json_valid(request_metadata) THEN request_metadata ELSE '{}' END, '$.client_response_body_ref') IS NOT NULL + ) +) +"#; + +#[derive(Debug)] +struct CleanupRow { + id: String, + request_id: String, +} + +#[derive(Debug)] +struct BodyRow { + id: String, + request_id: String, + request_body: Option, + request_body_compressed: Option>, + provider_request_body: Option, + provider_request_body_compressed: Option>, + response_body: Option, + response_body_compressed: Option>, + client_response_body: Option, + client_response_body_compressed: Option>, +} + +#[derive(Debug, Default)] +struct DetachedRefs { + request_body_ref: Option, + provider_request_body_ref: Option, + response_body_ref: Option, + client_response_body_ref: Option, +} + +impl DetachedRefs { + fn any_present(&self) -> bool { + self.request_body_ref.is_some() + || self.provider_request_body_ref.is_some() + || self.response_body_ref.is_some() + || self.client_response_body_ref.is_some() + } +} + +struct DetachedBlob { + body_ref: String, + body_field: &'static str, + payload_gzip: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BodyCleanupKind { + Raw, + Compressed, + All, +} + +impl BodyCleanupKind { + fn predicate(self) -> &'static str { + match self { + Self::Raw => RAW_BODY_PREDICATE, + Self::Compressed => COMPRESSED_BODY_PREDICATE, + Self::All => ALL_BODY_PREDICATE, + } + } + + fn clears_detached(self) -> bool { + self != Self::Raw + } +} + +pub(crate) async fn cleanup_usage( + pool: &SqlitePool, + window: &UsageCleanupWindow, + batch_size: usize, + auto_delete_expired_keys: bool, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, +) -> Result { + if batch_size == 0 || !targets.any_selected() { + return Ok(UsageCleanupSummary::default()); + } + + if mode == UsageCleanupExecutionMode::BeforeNowBodyFields { + let body_externalized = if targets.detail_body { + cleanup_body_fields( + pool, + window.detail_cutoff, + None, + batch_size, + BodyCleanupKind::Raw, + ) + .await? + } else { + 0 + }; + let body_cleaned = if targets.compressed_body { + cleanup_body_fields( + pool, + window.compressed_cutoff, + None, + batch_size, + BodyCleanupKind::Compressed, + ) + .await? + } else { + 0 + }; + return Ok(UsageCleanupSummary { + body_externalized, + body_cleaned, + ..UsageCleanupSummary::default() + }); + } + + let records_deleted = if targets.records { + delete_old_usage_records(pool, window.log_cutoff, batch_size).await? + } else { + 0 + }; + let record_cutoff = targets.records.then_some(window.log_cutoff); + let header_cleaned = if targets.headers { + cleanup_headers(pool, window.header_cutoff, record_cutoff, batch_size).await? + } else { + 0 + }; + let body_cleaned = if targets.compressed_body { + cleanup_body_fields( + pool, + window.compressed_cutoff, + record_cutoff, + batch_size, + BodyCleanupKind::All, + ) + .await? + } else { + 0 + }; + let detail_newer_than = detail_body_newer_than(window, targets); + let legacy_body_refs_migrated = if targets.detail_body { + migrate_legacy_body_refs(pool, window.detail_cutoff, detail_newer_than, batch_size).await? + } else { + 0 + }; + let body_externalized = if targets.detail_body { + externalize_detail_bodies(pool, window.detail_cutoff, detail_newer_than, batch_size).await? + } else { + 0 + }; + let keys_cleaned = if targets.expired_keys { + match cleanup_expired_api_keys(pool, auto_delete_expired_keys).await { + Ok(count) => count, + Err(err) => { + warn!(error = %err, "SQLite usage cleanup expired api key sweep failed"); + 0 + } + } + } else { + 0 + }; + + Ok(UsageCleanupSummary { + body_externalized, + legacy_body_refs_migrated, + body_cleaned, + header_cleaned, + keys_cleaned, + records_deleted, + }) +} + +pub(crate) async fn preview_usage_cleanup( + pool: &SqlitePool, + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, + mode: UsageCleanupExecutionMode, +) -> Result { + if mode == UsageCleanupExecutionMode::BeforeNowBodyFields { + let detail = if targets.detail_body { + count_candidates(pool, RAW_BODY_PREDICATE, window.detail_cutoff, None).await? + } else { + 0 + }; + let compressed = if targets.compressed_body { + count_candidates( + pool, + COMPRESSED_BODY_PREDICATE, + window.compressed_cutoff, + None, + ) + .await? + } else { + 0 + }; + return Ok(UsageCleanupPreviewCounts { + detail, + compressed, + header: 0, + log: 0, + }); + } + + let record_cutoff = targets.records.then_some(window.log_cutoff); + let detail = if targets.detail_body { + count_candidates( + pool, + DETAIL_BODY_PREDICATE, + window.detail_cutoff, + detail_body_newer_than(window, targets), + ) + .await? + } else { + 0 + }; + let compressed = if targets.compressed_body { + count_candidates( + pool, + ALL_BODY_PREDICATE, + window.compressed_cutoff, + record_cutoff, + ) + .await? + } else { + 0 + }; + let header = if targets.headers { + count_candidates(pool, HEADER_PREDICATE, window.header_cutoff, record_cutoff).await? + } else { + 0 + }; + let log = if targets.records { + let count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM \"usage\" WHERE created_at_unix_ms < ?") + .bind(window.log_cutoff.timestamp()) + .fetch_one(pool) + .await + .map_sql_err()?; + u64::try_from(count).unwrap_or(0) + } else { + 0 + }; + + Ok(UsageCleanupPreviewCounts { + detail, + compressed, + header, + log, + }) +} + +fn detail_body_newer_than( + window: &UsageCleanupWindow, + targets: UsageCleanupTargets, +) -> Option> { + [ + targets.compressed_body.then_some(window.compressed_cutoff), + targets.records.then_some(window.log_cutoff), + ] + .into_iter() + .flatten() + .max() +} + +async fn count_candidates( + pool: &SqlitePool, + predicate: &str, + cutoff: DateTime, + newer_than: Option>, +) -> Result { + if invalid_window(cutoff, newer_than) { + return Ok(0); + } + let sql = format!( + r#" +SELECT COUNT(*) +FROM "usage" +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({predicate}) +"# + ); + let newer_than = newer_than.map(|value| value.timestamp()); + let count: i64 = sqlx::query_scalar(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .fetch_one(pool) + .await + .map_sql_err()?; + Ok(u64::try_from(count).unwrap_or(0)) +} + +async fn fetch_cleanup_rows( + pool: &SqlitePool, + predicate: &str, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result, DataLayerError> { + if invalid_window(cutoff, newer_than) { + return Ok(Vec::new()); + } + let sql = format!( + r#" +SELECT id, request_id +FROM "usage" +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({predicate}) +ORDER BY created_at_unix_ms ASC, id ASC +LIMIT ? +"# + ); + let newer_than = newer_than.map(|value| value.timestamp()); + sqlx::query(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .bind(i64::try_from(batch_size).unwrap_or(i64::MAX)) + .fetch_all(pool) + .await + .map_sql_err()? + .into_iter() + .map(|row| { + Ok(CleanupRow { + id: row.try_get("id").map_sql_err()?, + request_id: row.try_get("request_id").map_sql_err()?, + }) + }) + .collect() +} + +fn invalid_window(cutoff: DateTime, newer_than: Option>) -> bool { + matches!(newer_than, Some(value) if value >= cutoff) +} + +async fn delete_old_usage_records( + pool: &SqlitePool, + cutoff: DateTime, + batch_size: usize, +) -> Result { + let mut total = 0usize; + loop { + let rows = fetch_cleanup_rows(pool, "1 = 1", cutoff, None, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + let mut deleted = 0usize; + for row in rows { + deleted += usize::try_from( + sqlx::query("DELETE FROM \"usage\" WHERE id = ?") + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(), + ) + .unwrap_or(usize::MAX); + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(deleted); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn cleanup_headers( + pool: &SqlitePool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "SQLite usage header cleanup skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = + fetch_cleanup_rows(pool, HEADER_PREDICATE, cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + for row in rows { + sqlx::query( + r#" +UPDATE "usage" +SET request_headers = NULL, + response_headers = NULL, + provider_request_headers = NULL, + client_response_headers = NULL +WHERE id = ? +"#, + ) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query( + r#" +UPDATE usage_http_audits +SET request_headers = NULL, + response_headers = NULL, + provider_request_headers = NULL, + client_response_headers = NULL, + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE request_id = ? +"#, + ) + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + delete_empty_http_audit(&mut tx, &row.request_id).await?; + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(row_count); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn cleanup_body_fields( + pool: &SqlitePool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, + kind: BodyCleanupKind, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "SQLite usage body cleanup skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = + fetch_cleanup_rows(pool, kind.predicate(), cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + for row in rows { + if kind == BodyCleanupKind::All { + sqlx::query( + r#" +UPDATE "usage" +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL, + request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"#, + ) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + } else if kind == BodyCleanupKind::Compressed { + sqlx::query( + r#" +UPDATE "usage" +SET request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"#, + ) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + } else { + sqlx::query( + r#" +UPDATE "usage" +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL +WHERE id = ? +"#, + ) + .bind(&row.id) + .execute(&mut *tx) + .await + .map_sql_err()?; + } + + if kind.clears_detached() { + sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query( + r#" +UPDATE usage_http_audits +SET request_body_ref = NULL, + provider_request_body_ref = NULL, + response_body_ref = NULL, + client_response_body_ref = NULL, + body_capture_mode = 'none', + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE request_id = ? +"#, + ) + .bind(&row.request_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + delete_empty_http_audit(&mut tx, &row.request_id).await?; + } + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(row_count); + if row_count < batch_size { + break; + } + } + Ok(total) +} + +async fn delete_empty_http_audit( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +DELETE FROM usage_http_audits +WHERE request_id = ? + AND request_headers IS NULL + AND response_headers IS NULL + AND provider_request_headers IS NULL + AND client_response_headers IS NULL + AND request_body_ref IS NULL + AND provider_request_body_ref IS NULL + AND response_body_ref IS NULL + AND client_response_body_ref IS NULL +"#, + ) + .bind(request_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn migrate_legacy_body_refs( + pool: &SqlitePool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "SQLite usage legacy body-ref migration skipped due to invalid window"); + return Ok(0); + } + let mut total = 0usize; + loop { + let rows = fetch_cleanup_rows( + pool, + LEGACY_BODY_REF_PREDICATE, + cutoff, + newer_than, + batch_size, + ) + .await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut tx = pool.begin().await.map_sql_err()?; + let mut migrated = 0usize; + for row in rows { + let metadata: Option = + sqlx::query_scalar("SELECT request_metadata FROM \"usage\" WHERE id = ? LIMIT 1") + .bind(&row.id) + .fetch_optional(&mut *tx) + .await + .map_sql_err()? + .flatten(); + let Some((refs, metadata)) = + legacy_body_ref_plan(&row.request_id, metadata.as_deref())? + else { + continue; + }; + if refs.any_present() { + upsert_http_audit_refs(&mut tx, &row.request_id, &refs).await?; + } + let updated = sqlx::query( + r#" +UPDATE "usage" +SET request_metadata = ?, + updated_at_unix_secs = CAST(strftime('%s', 'now') AS INTEGER) +WHERE id = ? +"#, + ) + .bind(metadata) + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + if updated > 0 { + migrated += 1; + } + } + tx.commit().await.map_sql_err()?; + total = total.saturating_add(migrated); + if row_count < batch_size || migrated == 0 { + break; + } + } + Ok(total) +} + +fn legacy_body_ref_plan( + request_id: &str, + metadata: Option<&str>, +) -> Result)>, DataLayerError> { + let Some(metadata) = metadata else { + return Ok(None); + }; + let value: Value = serde_json::from_str(metadata).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid usage request_metadata JSON: {err}")) + })?; + let Value::Object(mut object) = value else { + return Ok(None); + }; + let mut refs = DetachedRefs::default(); + let mut removed = false; + for field in [ + UsageBodyField::RequestBody, + UsageBodyField::ProviderRequestBody, + UsageBodyField::ResponseBody, + UsageBodyField::ClientResponseBody, + ] { + let Some(value) = object.remove(field.as_ref_key()) else { + continue; + }; + removed = true; + let parsed = value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(parse_usage_body_ref) + .filter(|(parsed_request_id, parsed_field)| { + parsed_request_id == request_id && *parsed_field == field + }) + .map(|(parsed_request_id, parsed_field)| { + usage_body_ref(&parsed_request_id, parsed_field) + }); + set_ref(&mut refs, field, parsed); + } + if !removed { + return Ok(None); + } + let metadata = if object.is_empty() { + None + } else { + Some( + serde_json::to_string(&Value::Object(object)).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "failed to serialize request_metadata: {err}" + )) + })?, + ) + }; + Ok(Some((refs, metadata))) +} + +async fn externalize_detail_bodies( + pool: &SqlitePool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result { + if invalid_window(cutoff, newer_than) { + warn!(%cutoff, ?newer_than, "SQLite usage body externalization skipped due to invalid window"); + return Ok(0); + } + let batch_size = batch_size.clamp(1, 25); + let mut total = 0usize; + loop { + let rows = fetch_body_rows(pool, cutoff, newer_than, batch_size).await?; + if rows.is_empty() { + break; + } + let row_count = rows.len(); + let mut externalized = 0usize; + for row in rows { + let (blobs, refs) = build_detached_bodies(&row)?; + let mut tx = pool.begin().await.map_sql_err()?; + for blob in blobs { + sqlx::query( + r#" +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES (?, ?, ?, ?) +ON CONFLICT(body_ref) DO UPDATE SET + request_id = excluded.request_id, + body_field = excluded.body_field, + payload_gzip = excluded.payload_gzip, + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +"#, + ) + .bind(blob.body_ref) + .bind(&row.request_id) + .bind(blob.body_field) + .bind(blob.payload_gzip) + .execute(&mut *tx) + .await + .map_sql_err()?; + } + if refs.any_present() { + upsert_http_audit_refs(&mut tx, &row.request_id, &refs).await?; + } + let updated = sqlx::query( + r#" +UPDATE "usage" +SET request_body = NULL, + response_body = NULL, + provider_request_body = NULL, + client_response_body = NULL, + request_body_compressed = NULL, + response_body_compressed = NULL, + provider_request_body_compressed = NULL, + client_response_body_compressed = NULL +WHERE id = ? +"#, + ) + .bind(row.id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + tx.commit().await.map_sql_err()?; + if updated > 0 { + externalized += 1; + } + } + total = total.saturating_add(externalized); + if row_count < batch_size || externalized == 0 { + break; + } + } + Ok(total) +} + +async fn fetch_body_rows( + pool: &SqlitePool, + cutoff: DateTime, + newer_than: Option>, + batch_size: usize, +) -> Result, DataLayerError> { + let newer_than = newer_than.map(|value| value.timestamp()); + let sql = format!( + r#" +SELECT id, + request_id, + request_body, + request_body_compressed, + provider_request_body, + provider_request_body_compressed, + response_body, + response_body_compressed, + client_response_body, + client_response_body_compressed +FROM "usage" +WHERE created_at_unix_ms < ? + AND (? IS NULL OR created_at_unix_ms >= ?) + AND ({INLINE_OR_COMPRESSED_BODY_PREDICATE}) +ORDER BY created_at_unix_ms ASC, id ASC +LIMIT ? +"# + ); + sqlx::query(&sql) + .bind(cutoff.timestamp()) + .bind(newer_than) + .bind(newer_than) + .bind(i64::try_from(batch_size).unwrap_or(i64::MAX)) + .fetch_all(pool) + .await + .map_sql_err()? + .into_iter() + .map(|row| { + Ok(BodyRow { + id: row.try_get("id").map_sql_err()?, + request_id: row.try_get("request_id").map_sql_err()?, + request_body: parse_optional_json(row.try_get("request_body").map_sql_err()?)?, + request_body_compressed: row.try_get("request_body_compressed").map_sql_err()?, + provider_request_body: parse_optional_json( + row.try_get("provider_request_body").map_sql_err()?, + )?, + provider_request_body_compressed: row + .try_get("provider_request_body_compressed") + .map_sql_err()?, + response_body: parse_optional_json(row.try_get("response_body").map_sql_err()?)?, + response_body_compressed: row.try_get("response_body_compressed").map_sql_err()?, + client_response_body: parse_optional_json( + row.try_get("client_response_body").map_sql_err()?, + )?, + client_response_body_compressed: row + .try_get("client_response_body_compressed") + .map_sql_err()?, + }) + }) + .collect() +} + +fn parse_optional_json(raw: Option) -> Result, DataLayerError> { + raw.map(|raw| { + serde_json::from_str(&raw).map_err(|err| { + DataLayerError::UnexpectedValue(format!("invalid inline usage body JSON: {err}")) + }) + }) + .transpose() +} + +fn build_detached_bodies( + row: &BodyRow, +) -> Result<(Vec, DetachedRefs), DataLayerError> { + let mut blobs = Vec::new(); + let mut refs = DetachedRefs::default(); + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::RequestBody, + row.request_body.as_ref(), + row.request_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ProviderRequestBody, + row.provider_request_body.as_ref(), + row.provider_request_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ResponseBody, + row.response_body.as_ref(), + row.response_body_compressed.as_deref(), + )?; + add_detached_body( + &mut blobs, + &mut refs, + &row.request_id, + UsageBodyField::ClientResponseBody, + row.client_response_body.as_ref(), + row.client_response_body_compressed.as_deref(), + )?; + Ok((blobs, refs)) +} + +fn add_detached_body( + blobs: &mut Vec, + refs: &mut DetachedRefs, + request_id: &str, + field: UsageBodyField, + raw: Option<&Value>, + compressed: Option<&[u8]>, +) -> Result<(), DataLayerError> { + let payload_gzip = match raw { + Some(value) => Some(compress_json(value)?), + None => compressed.map(ToOwned::to_owned), + }; + let Some(payload_gzip) = payload_gzip else { + return Ok(()); + }; + let body_ref = usage_body_ref(request_id, field); + blobs.push(DetachedBlob { + body_ref: body_ref.clone(), + body_field: field.as_storage_field(), + payload_gzip, + }); + set_ref(refs, field, Some(body_ref)); + Ok(()) +} + +fn compress_json(value: &Value) -> Result, DataLayerError> { + let bytes = serde_json::to_vec(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage body: {err}")) + })?; + let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6)); + encoder.write_all(&bytes).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to gzip usage body: {err}")) + })?; + encoder.finish().map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to finish usage body gzip: {err}")) + }) +} + +fn set_ref(refs: &mut DetachedRefs, field: UsageBodyField, value: Option) { + match field { + UsageBodyField::RequestBody => refs.request_body_ref = value, + UsageBodyField::ProviderRequestBody => refs.provider_request_body_ref = value, + UsageBodyField::ResponseBody => refs.response_body_ref = value, + UsageBodyField::ClientResponseBody => refs.client_response_body_ref = value, + } +} + +async fn upsert_http_audit_refs( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, + refs: &DetachedRefs, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO usage_http_audits ( + request_id, + request_body_ref, + provider_request_body_ref, + response_body_ref, + client_response_body_ref, + body_capture_mode +) +VALUES (?, ?, ?, ?, ?, 'ref_backed') +ON CONFLICT(request_id) DO UPDATE SET + request_body_ref = COALESCE(excluded.request_body_ref, usage_http_audits.request_body_ref), + provider_request_body_ref = COALESCE( + excluded.provider_request_body_ref, + usage_http_audits.provider_request_body_ref + ), + response_body_ref = COALESCE(excluded.response_body_ref, usage_http_audits.response_body_ref), + client_response_body_ref = COALESCE( + excluded.client_response_body_ref, + usage_http_audits.client_response_body_ref + ), + body_capture_mode = 'ref_backed', + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +"#, + ) + .bind(request_id) + .bind(refs.request_body_ref.as_deref()) + .bind(refs.provider_request_body_ref.as_deref()) + .bind(refs.response_body_ref.as_deref()) + .bind(refs.client_response_body_ref.as_deref()) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn cleanup_expired_api_keys( + pool: &SqlitePool, + auto_delete_expired_keys: bool, +) -> Result { + let now = Utc::now().timestamp(); + let rows = sqlx::query( + r#" +SELECT id, auto_delete_on_expiry +FROM api_keys +WHERE expires_at <= ? + AND is_active = 1 +ORDER BY expires_at ASC, id ASC +"#, + ) + .bind(now) + .fetch_all(pool) + .await + .map_sql_err()?; + let mut cleaned = 0usize; + for row in rows { + let id: String = row.try_get("id").map_sql_err()?; + let auto_delete = row + .try_get::, _>("auto_delete_on_expiry") + .map_sql_err()? + .map(|value| value != 0) + .unwrap_or(auto_delete_expired_keys); + let mut tx = pool.begin().await.map_sql_err()?; + let affected = if auto_delete { + sqlx::query( + "UPDATE wallets SET status = 'disabled', updated_at = ? WHERE api_key_id = ? AND status <> 'disabled'", + ) + .bind(now) + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()?; + sqlx::query("DELETE FROM api_keys WHERE id = ?") + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected() + } else { + sqlx::query( + "UPDATE api_keys SET is_active = 0, updated_at = ? WHERE id = ? AND is_active = 1", + ) + .bind(now) + .bind(&id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected() + }; + tx.commit().await.map_sql_err()?; + if affected > 0 { + cleaned += 1; + } + } + Ok(cleaned) +} diff --git a/crates/aether-data/adapters/sqlite/src/usage/counters.rs b/crates/aether-data/adapters/sqlite/src/usage/counters.rs new file mode 100644 index 000000000..4e265b5f8 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/src/usage/counters.rs @@ -0,0 +1,1283 @@ +use std::collections::BTreeMap; + +use aether_data_contracts::repository::usage::{ + api_key_usage_contribution, model_usage_contribution, provider_api_key_usage_contribution, + ApiKeyLastUsedDelta, ApiKeyUsageDelta, ManagementTokenCounterDelta, ModelUsageDelta, + ProviderApiKeyUsageDelta, ProxyNodeCounterDelta, StoredRequestUsageAudit, + UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, +}; +use aether_data_contracts::DataLayerError; +use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; + +use crate::error::SqlResultExt; +use crate::sqlite_real; + +const KIND_API_KEY: &str = "api_key"; +const KIND_PROVIDER_API_KEY: &str = "provider_api_key"; +const KIND_MODEL: &str = "model"; +const KIND_PROVIDER_MONTHLY: &str = "provider_monthly"; +const KIND_PROXY_NODE: &str = "proxy_node"; +const KIND_MANAGEMENT_TOKEN: &str = "management_token"; +const KIND_API_KEY_LAST_USED: &str = "api_key_last_used"; + +const CLAIM_SQL: &str = r#" +SELECT + id, + kind, + target_id, + request_count_delta, + total_requests_delta, + success_count_delta, + error_count_delta, + dns_failures_delta, + stream_errors_delta, + total_tokens_delta, + total_cost_usd_delta, + total_response_time_ms_delta, + last_used_at_unix_secs, + last_used_ip, + candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs, + usage_created_at_unix_secs +FROM usage_counter_deltas +WHERE processed_at IS NULL +ORDER BY created_at ASC, id ASC +LIMIT ? +"#; + +struct DeltaRow { + id: String, + kind: String, + target_id: String, + request_count_delta: i64, + total_requests_delta: i64, + success_count_delta: i64, + error_count_delta: i64, + dns_failures_delta: i64, + stream_errors_delta: i64, + total_tokens_delta: i64, + total_cost_usd_delta: f64, + total_response_time_ms_delta: i64, + last_used_at_unix_secs: Option, + last_used_ip: Option, + candidate_last_used_at_unix_secs: Option, + removed_last_used_at_unix_secs: Option, + usage_created_at_unix_secs: Option, +} + +#[derive(Default)] +struct Aggregates { + api_keys: BTreeMap, + provider_api_keys: BTreeMap, + models: BTreeMap, + provider_monthly: BTreeMap, + proxy_nodes: BTreeMap, + management_tokens: BTreeMap, + api_key_last_used: BTreeMap, +} + +impl Aggregates { + fn from_rows(rows: &[DeltaRow]) -> Result { + let mut aggregates = Self::default(); + for row in rows { + if !row.total_cost_usd_delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "usage_counter_deltas.total_cost_usd_delta is not finite for {}", + row.id + ))); + } + match row.kind.as_str() { + KIND_API_KEY => { + let entry = aggregates + .api_keys + .entry(row.target_id.clone()) + .or_default(); + entry.total_requests += row.total_requests_delta; + entry.total_tokens += row.total_tokens_delta; + entry.total_cost_usd += row.total_cost_usd_delta; + merge_optional_max( + &mut entry.candidate_last_used_at_unix_secs, + row.candidate_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.removed_last_used_at_unix_secs, + row.removed_last_used_at_unix_secs, + ); + } + KIND_PROVIDER_API_KEY => { + let entry = aggregates + .provider_api_keys + .entry(row.target_id.clone()) + .or_default(); + entry.request_count += row.request_count_delta; + entry.success_count += row.success_count_delta; + entry.error_count += row.error_count_delta; + entry.total_tokens += row.total_tokens_delta; + entry.total_cost_usd += row.total_cost_usd_delta; + entry.total_response_time_ms += row.total_response_time_ms_delta; + merge_optional_max( + &mut entry.candidate_last_used_at_unix_secs, + row.candidate_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.removed_last_used_at_unix_secs, + row.removed_last_used_at_unix_secs, + ); + merge_optional_max( + &mut entry.usage_created_at_unix_secs, + row.usage_created_at_unix_secs, + ); + } + KIND_MODEL => { + aggregates + .models + .entry(row.target_id.clone()) + .or_default() + .request_count += row.request_count_delta; + } + KIND_PROVIDER_MONTHLY => { + *aggregates + .provider_monthly + .entry(row.target_id.clone()) + .or_default() += row.total_cost_usd_delta; + } + KIND_PROXY_NODE => { + let entry = aggregates + .proxy_nodes + .entry(row.target_id.clone()) + .or_insert(ProxyNodeCounterDelta { + node_id: row.target_id.clone(), + total_requests_delta: 0, + failed_requests_delta: 0, + dns_failures_delta: 0, + stream_errors_delta: 0, + }); + entry.total_requests_delta += row.total_requests_delta; + entry.failed_requests_delta += row.error_count_delta; + entry.dns_failures_delta += row.dns_failures_delta; + entry.stream_errors_delta += row.stream_errors_delta; + } + KIND_MANAGEMENT_TOKEN => { + let entry = aggregates + .management_tokens + .entry(row.target_id.clone()) + .or_insert(ManagementTokenCounterDelta { + token_id: row.target_id.clone(), + usage_count_delta: 0, + last_used_at_unix_secs: None, + last_used_ip: None, + }); + entry.usage_count_delta += row.request_count_delta; + merge_latest_timestamp_with_value( + &mut entry.last_used_at_unix_secs, + &mut entry.last_used_ip, + row.last_used_at_unix_secs, + row.last_used_ip.clone(), + ); + } + KIND_API_KEY_LAST_USED => { + let Some(last_used_at_unix_secs) = row.last_used_at_unix_secs else { + continue; + }; + let entry = aggregates + .api_key_last_used + .entry(row.target_id.clone()) + .or_insert(ApiKeyLastUsedDelta { + api_key_id: row.target_id.clone(), + last_used_at_unix_secs, + }); + if last_used_at_unix_secs > entry.last_used_at_unix_secs { + entry.last_used_at_unix_secs = last_used_at_unix_secs; + } + } + other => { + return Err(DataLayerError::UnexpectedValue(format!( + "unknown usage counter delta kind: {other}" + ))); + } + } + } + Ok(aggregates) + } +} + +pub(super) async fn flush( + pool: &SqlitePool, + batch_size: usize, +) -> Result { + if batch_size == 0 { + return Ok(UsageCounterFlushSummary::default()); + } + let limit = i64::try_from(batch_size).map_err(|_| { + DataLayerError::InvalidInput(format!( + "usage counter flush batch size is out of range: {batch_size}" + )) + })?; + + let mut tx = pool.begin().await.map_sql_err()?; + // Force a RESERVED write lock before reading the outbox. This serializes SQLite flushers so + // two deferred transactions cannot claim and apply the same rows. + sqlx::query("UPDATE usage_counter_deltas SET processed_at = processed_at WHERE 0") + .execute(&mut *tx) + .await + .map_sql_err()?; + let rows = sqlx::query(CLAIM_SQL) + .bind(limit) + .fetch_all(&mut *tx) + .await + .map_sql_err()? + .iter() + .map(map_row) + .collect::, _>>()?; + if rows.is_empty() { + tx.rollback().await.map_sql_err()?; + return Ok(UsageCounterFlushSummary::default()); + } + + let aggregates = Aggregates::from_rows(&rows)?; + for (target_id, delta) in &aggregates.api_keys { + apply_api_key(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.models { + apply_model(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.provider_api_keys { + apply_provider_api_key(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.provider_monthly { + apply_provider_monthly(&mut tx, target_id, *delta).await?; + } + for (target_id, delta) in &aggregates.proxy_nodes { + apply_proxy_node(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.management_tokens { + apply_management_token(&mut tx, target_id, delta).await?; + } + for (target_id, delta) in &aggregates.api_key_last_used { + apply_api_key_last_used(&mut tx, target_id, delta).await?; + } + + let now = current_unix_secs(); + let mut mark = QueryBuilder::::new("UPDATE usage_counter_deltas SET processed_at = "); + mark.push_bind(now).push(" WHERE id IN ("); + { + let mut ids = mark.separated(", "); + for row in &rows { + ids.push_bind(&row.id); + } + } + mark.push(")"); + mark.build().execute(&mut *tx).await.map_sql_err()?; + tx.commit().await.map_sql_err()?; + + Ok(UsageCounterFlushSummary { + rows_claimed: rows.len(), + api_key_targets: aggregates.api_keys.len(), + provider_api_key_targets: aggregates.provider_api_keys.len(), + model_targets: aggregates.models.len(), + provider_monthly_targets: aggregates.provider_monthly.len(), + proxy_node_targets: aggregates.proxy_nodes.len(), + management_token_targets: aggregates.management_tokens.len(), + api_key_last_used_targets: aggregates.api_key_last_used.len(), + }) +} + +pub(super) async fn enqueue_proxy_node( + pool: &SqlitePool, + delta: ProxyNodeCounterDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let node_id = delta.node_id.trim().to_string(); + let request_id = format!("proxy_node:{node_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_PROXY_NODE, + target_id: &node_id, + total_requests_delta: delta.total_requests_delta, + error_count_delta: delta.failed_requests_delta, + dns_failures_delta: delta.dns_failures_delta, + stream_errors_delta: delta.stream_errors_delta, + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn enqueue_management_token( + pool: &SqlitePool, + delta: ManagementTokenCounterDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let token_id = delta.token_id.trim().to_string(); + let last_used_ip = delta + .last_used_ip + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let last_used_at = delta + .last_used_at_unix_secs + .unwrap_or_else(|| current_unix_secs().max(0) as u64); + let request_id = format!("management_token:{token_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_MANAGEMENT_TOKEN, + target_id: &token_id, + request_count_delta: delta.usage_count_delta, + last_used_at_unix_secs: Some(last_used_at), + last_used_ip: last_used_ip.as_deref(), + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn enqueue_api_key_last_used( + pool: &SqlitePool, + delta: ApiKeyLastUsedDelta, +) -> Result { + if delta.is_noop() { + return Ok(false); + } + let api_key_id = delta.api_key_id.trim().to_string(); + let request_id = format!("api_key_last_used:{api_key_id}:{}", uuid::Uuid::new_v4()); + let mut tx = pool.begin().await.map_sql_err()?; + insert_delta( + &mut tx, + DeltaInsert { + request_id: &request_id, + kind: KIND_API_KEY_LAST_USED, + target_id: &api_key_id, + last_used_at_unix_secs: Some(delta.last_used_at_unix_secs), + ..DeltaInsert::default() + }, + ) + .await?; + tx.commit().await.map_sql_err()?; + Ok(true) +} + +pub(super) async fn cleanup_processed( + pool: &SqlitePool, + cutoff_unix_secs: u64, + batch_size: usize, +) -> Result { + if batch_size == 0 { + return Ok(0); + } + let cutoff = to_i64(cutoff_unix_secs, "usage counter cleanup cutoff")?; + let limit = i64::try_from(batch_size).map_err(|_| { + DataLayerError::InvalidInput(format!( + "usage counter cleanup batch size is out of range: {batch_size}" + )) + })?; + let deleted = sqlx::query( + r#" +DELETE FROM usage_counter_deltas +WHERE id IN ( + SELECT id FROM ( + SELECT id + FROM usage_counter_deltas + WHERE processed_at IS NOT NULL AND processed_at < ? + ORDER BY processed_at ASC, created_at ASC, id ASC + LIMIT ? + ) AS doomed +) +"#, + ) + .bind(cutoff) + .bind(limit) + .execute(pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(usize::try_from(deleted).unwrap_or(usize::MAX)) +} + +pub(super) async fn read_health( + pool: &SqlitePool, +) -> Result { + let row = sqlx::query( + r#" +SELECT + (SELECT COUNT(*) FROM usage_counter_deltas WHERE processed_at IS NULL) + AS pending_rows, + (SELECT COUNT(*) FROM usage_counter_deltas WHERE processed_at IS NOT NULL) + AS processed_rows, + (SELECT MIN(created_at) FROM usage_counter_deltas WHERE processed_at IS NULL) + AS oldest_pending_created_at_unix_secs, + (SELECT MAX(processed_at) FROM usage_counter_deltas WHERE processed_at IS NOT NULL) + AS latest_processed_at_unix_secs +"#, + ) + .fetch_one(pool) + .await + .map_sql_err()?; + let mut snapshot = UsageCounterHealthSnapshot { + pending_rows: nonnegative_u64(row.try_get("pending_rows").map_sql_err()?), + processed_rows: nonnegative_u64(row.try_get("processed_rows").map_sql_err()?), + oldest_pending_created_at_unix_secs: optional_nonnegative_u64( + row.try_get("oldest_pending_created_at_unix_secs") + .map_sql_err()?, + ), + latest_processed_at_unix_secs: optional_nonnegative_u64( + row.try_get("latest_processed_at_unix_secs").map_sql_err()?, + ), + pending_by_kind: BTreeMap::new(), + }; + for row in pending_health_rows(pool).await? { + snapshot.pending_by_kind.insert(row.0, row.1); + } + Ok(snapshot) +} + +pub(super) async fn read_pending_health( + pool: &SqlitePool, +) -> Result { + let mut snapshot = UsageCounterPendingHealthSnapshot::default(); + for (kind, pending_rows, oldest) in pending_health_rows(pool).await? { + snapshot.pending_rows = snapshot.pending_rows.saturating_add(pending_rows); + if let Some(oldest) = oldest { + snapshot.oldest_pending_created_at_unix_secs = Some( + snapshot + .oldest_pending_created_at_unix_secs + .map_or(oldest, |current| current.min(oldest)), + ); + } + snapshot.pending_by_kind.insert(kind, pending_rows); + } + Ok(snapshot) +} + +async fn pending_health_rows( + pool: &SqlitePool, +) -> Result)>, DataLayerError> { + let rows = sqlx::query( + r#" +SELECT + kind, + COUNT(*) AS pending_rows, + MIN(created_at) AS oldest_pending_created_at_unix_secs +FROM usage_counter_deltas +WHERE processed_at IS NULL +GROUP BY kind +ORDER BY kind ASC +"#, + ) + .fetch_all(pool) + .await + .map_sql_err()?; + rows.iter() + .map(|row| { + Ok(( + row.try_get("kind").map_sql_err()?, + nonnegative_u64(row.try_get("pending_rows").map_sql_err()?), + optional_nonnegative_u64( + row.try_get("oldest_pending_created_at_unix_secs") + .map_sql_err()?, + ), + )) + }) + .collect() +} + +pub(super) async fn enqueue_usage_transition( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, + before: Option<&StoredRequestUsageAudit>, + after: &StoredRequestUsageAudit, +) -> Result<(), DataLayerError> { + let before_api_key = before.and_then(api_key_usage_contribution); + let after_api_key = api_key_usage_contribution(after); + match (before_api_key.as_ref(), after_api_key.as_ref()) { + (Some(before), Some(after)) if before.api_key_id == after.api_key_id => { + enqueue_api_key_delta( + tx, + request_id, + &before.api_key_id, + &ApiKeyUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_api_key.as_ref() { + enqueue_api_key_delta( + tx, + request_id, + &before.api_key_id, + &ApiKeyUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_api_key.as_ref() { + enqueue_api_key_delta( + tx, + request_id, + &after.api_key_id, + &ApiKeyUsageDelta::addition(after), + ) + .await?; + } + } + } + + let before_model = before.and_then(model_usage_contribution); + let after_model = model_usage_contribution(after); + match (before_model.as_ref(), after_model.as_ref()) { + (Some(before), Some(after)) if before.model == after.model => { + enqueue_model_delta( + tx, + request_id, + &before.model, + &ModelUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_model.as_ref() { + enqueue_model_delta( + tx, + request_id, + &before.model, + &ModelUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_model.as_ref() { + enqueue_model_delta( + tx, + request_id, + &after.model, + &ModelUsageDelta::addition(after), + ) + .await?; + } + } + } + + let before_provider = before.and_then(provider_api_key_usage_contribution); + let after_provider = provider_api_key_usage_contribution(after); + match (before_provider.as_ref(), after_provider.as_ref()) { + (Some(before), Some(after)) if before.key_id == after.key_id => { + enqueue_provider_api_key_delta( + tx, + request_id, + &before.key_id, + &ProviderApiKeyUsageDelta::between(before, after), + ) + .await?; + } + _ => { + if let Some(before) = before_provider.as_ref() { + enqueue_provider_api_key_delta( + tx, + request_id, + &before.key_id, + &ProviderApiKeyUsageDelta::removal(before), + ) + .await?; + } + if let Some(after) = after_provider.as_ref() { + enqueue_provider_api_key_delta( + tx, + request_id, + &after.key_id, + &ProviderApiKeyUsageDelta::addition(after), + ) + .await?; + } + } + } + Ok(()) +} + +pub(super) async fn enqueue_usage_transition_for_request( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, + before: Option<&StoredRequestUsageAudit>, +) -> Result<(), DataLayerError> { + let row = sqlx::query(&format!( + "{} WHERE \"usage\".request_id = ? LIMIT 1", + super::USAGE_COLUMNS + )) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()? + .ok_or_else(|| { + DataLayerError::UnexpectedValue(format!( + "usage row missing while preparing counter delta: {request_id}" + )) + })?; + let after = super::map_usage_row(&row, false)?; + enqueue_usage_transition(tx, request_id, before, &after).await +} + +pub(super) async fn lock_and_load_usage( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, +) -> Result, DataLayerError> { + // A write statement upgrades the deferred transaction before reading the old contribution. + // SQLite then serializes concurrent upserts for every request ID until this transaction ends. + sqlx::query("UPDATE \"usage\" SET request_id = request_id WHERE request_id = ?") + .bind(request_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + let row = sqlx::query(&format!( + "{} WHERE \"usage\".request_id = ? LIMIT 1", + super::USAGE_COLUMNS + )) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + row.as_ref() + .map(|row| super::map_usage_row(row, false)) + .transpose() +} + +async fn enqueue_api_key_delta( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, + target_id: &str, + delta: &ApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_API_KEY, + target_id, + total_requests_delta: delta.total_requests, + total_tokens_delta: delta.total_tokens, + total_cost_usd_delta: finite_or_zero(delta.total_cost_usd), + candidate_last_used_at_unix_secs: delta.candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs: delta.removed_last_used_at_unix_secs, + ..DeltaInsert::default() + }, + ) + .await +} + +async fn enqueue_model_delta( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, + target_id: &str, + delta: &ModelUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_MODEL, + target_id, + request_count_delta: delta.request_count, + ..DeltaInsert::default() + }, + ) + .await +} + +async fn enqueue_provider_api_key_delta( + tx: &mut sqlx::Transaction<'_, Sqlite>, + request_id: &str, + target_id: &str, + delta: &ProviderApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + insert_delta( + tx, + DeltaInsert { + request_id, + kind: KIND_PROVIDER_API_KEY, + target_id, + request_count_delta: delta.request_count, + success_count_delta: delta.success_count, + error_count_delta: delta.error_count, + total_tokens_delta: delta.total_tokens, + total_cost_usd_delta: finite_or_zero(delta.total_cost_usd), + total_response_time_ms_delta: delta.total_response_time_ms, + candidate_last_used_at_unix_secs: delta.candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs: delta.removed_last_used_at_unix_secs, + usage_created_at_unix_secs: delta.usage_created_at_unix_secs, + ..DeltaInsert::default() + }, + ) + .await +} + +#[derive(Default)] +struct DeltaInsert<'a> { + request_id: &'a str, + kind: &'a str, + target_id: &'a str, + request_count_delta: i64, + total_requests_delta: i64, + success_count_delta: i64, + error_count_delta: i64, + dns_failures_delta: i64, + stream_errors_delta: i64, + total_tokens_delta: i64, + total_cost_usd_delta: f64, + total_response_time_ms_delta: i64, + last_used_at_unix_secs: Option, + last_used_ip: Option<&'a str>, + candidate_last_used_at_unix_secs: Option, + removed_last_used_at_unix_secs: Option, + usage_created_at_unix_secs: Option, +} + +async fn insert_delta( + tx: &mut sqlx::Transaction<'_, Sqlite>, + input: DeltaInsert<'_>, +) -> Result<(), DataLayerError> { + let request_id = input.request_id.trim(); + let target_id = input.target_id.trim(); + if request_id.is_empty() || target_id.is_empty() { + return Ok(()); + } + sqlx::query( + r#" +INSERT INTO usage_counter_deltas ( + id, request_id, kind, target_id, request_count_delta, total_requests_delta, + success_count_delta, error_count_delta, dns_failures_delta, stream_errors_delta, + total_tokens_delta, total_cost_usd_delta, total_response_time_ms_delta, + last_used_at_unix_secs, last_used_ip, candidate_last_used_at_unix_secs, + removed_last_used_at_unix_secs, usage_created_at_unix_secs, created_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +"#, + ) + .bind(uuid::Uuid::new_v4().to_string()) + .bind(request_id) + .bind(input.kind) + .bind(target_id) + .bind(input.request_count_delta) + .bind(input.total_requests_delta) + .bind(input.success_count_delta) + .bind(input.error_count_delta) + .bind(input.dns_failures_delta) + .bind(input.stream_errors_delta) + .bind(input.total_tokens_delta) + .bind(finite_or_zero(input.total_cost_usd_delta)) + .bind(input.total_response_time_ms_delta) + .bind(optional_to_i64( + input.last_used_at_unix_secs, + "usage counter last_used_at_unix_secs", + )?) + .bind( + input + .last_used_ip + .map(str::trim) + .filter(|value| !value.is_empty()), + ) + .bind(optional_to_i64( + input.candidate_last_used_at_unix_secs, + "usage counter candidate_last_used_at_unix_secs", + )?) + .bind(optional_to_i64( + input.removed_last_used_at_unix_secs, + "usage counter removed_last_used_at_unix_secs", + )?) + .bind(optional_to_i64( + input.usage_created_at_unix_secs, + "usage counter usage_created_at_unix_secs", + )?) + .bind(current_unix_secs()) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +fn map_row(row: &sqlx::sqlite::SqliteRow) -> Result { + Ok(DeltaRow { + id: row.try_get("id").map_sql_err()?, + kind: row.try_get("kind").map_sql_err()?, + target_id: row.try_get("target_id").map_sql_err()?, + request_count_delta: row.try_get("request_count_delta").map_sql_err()?, + total_requests_delta: row.try_get("total_requests_delta").map_sql_err()?, + success_count_delta: row.try_get("success_count_delta").map_sql_err()?, + error_count_delta: row.try_get("error_count_delta").map_sql_err()?, + dns_failures_delta: row.try_get("dns_failures_delta").map_sql_err()?, + stream_errors_delta: row.try_get("stream_errors_delta").map_sql_err()?, + total_tokens_delta: row.try_get("total_tokens_delta").map_sql_err()?, + total_cost_usd_delta: sqlite_real(row, "total_cost_usd_delta")?, + total_response_time_ms_delta: row.try_get("total_response_time_ms_delta").map_sql_err()?, + last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.last_used_at_unix_secs", + row.try_get("last_used_at_unix_secs").map_sql_err()?, + )?, + last_used_ip: row.try_get("last_used_ip").map_sql_err()?, + candidate_last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.candidate_last_used_at_unix_secs", + row.try_get("candidate_last_used_at_unix_secs") + .map_sql_err()?, + )?, + removed_last_used_at_unix_secs: optional_u64( + "usage_counter_deltas.removed_last_used_at_unix_secs", + row.try_get("removed_last_used_at_unix_secs") + .map_sql_err()?, + )?, + usage_created_at_unix_secs: optional_u64( + "usage_counter_deltas.usage_created_at_unix_secs", + row.try_get("usage_created_at_unix_secs").map_sql_err()?, + )?, + }) +} + +async fn apply_api_key( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let candidate = optional_to_i64( + delta.candidate_last_used_at_unix_secs, + "api key candidate last used at", + )?; + let removed = optional_to_i64( + delta.removed_last_used_at_unix_secs, + "api key removed last used at", + )?; + sqlx::query( + r#" +UPDATE api_keys +SET total_requests = MAX(COALESCE(total_requests, 0) + ?, 0), + total_tokens = MAX(COALESCE(total_tokens, 0) + ?, 0), + total_cost_usd = MAX(CAST(COALESCE(total_cost_usd, 0) AS REAL) + ?, 0), + last_used_at = CASE + WHEN ? IS NOT NULL THEN MAX(COALESCE(last_used_at, 0), ?) + WHEN ? IS NOT NULL AND last_used_at = ? THEN ( + SELECT MAX(created_at_unix_ms) + FROM "usage" + WHERE api_key_id = ? AND status NOT IN ('pending', 'streaming') + ) + ELSE last_used_at + END +WHERE id = ? +"#, + ) + .bind(delta.total_requests) + .bind(delta.total_tokens) + .bind(finite_or_zero(delta.total_cost_usd)) + .bind(candidate) + .bind(candidate) + .bind(removed) + .bind(removed) + .bind(target_id) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_model( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ModelUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + "UPDATE global_models SET usage_count = MAX(COALESCE(usage_count, 0) + ?, 0), updated_at = ? WHERE name = ?", + ) + .bind(delta.request_count) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_provider_api_key( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ProviderApiKeyUsageDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let candidate = optional_to_i64( + delta.candidate_last_used_at_unix_secs, + "provider api key candidate last used at", + )?; + let removed = optional_to_i64( + delta.removed_last_used_at_unix_secs, + "provider api key removed last used at", + )?; + sqlx::query( + r#" +UPDATE provider_api_keys +SET request_count = MAX(COALESCE(request_count, 0) + ?, 0), + success_count = MAX(COALESCE(success_count, 0) + ?, 0), + error_count = MAX(COALESCE(error_count, 0) + ?, 0), + total_tokens = MAX(COALESCE(total_tokens, 0) + ?, 0), + total_cost_usd = MAX(CAST(COALESCE(total_cost_usd, 0) AS REAL) + ?, 0), + total_response_time_ms = MAX(COALESCE(total_response_time_ms, 0) + ?, 0), + last_used_at = CASE + WHEN ? IS NOT NULL THEN MAX(COALESCE(last_used_at, 0), ?) + WHEN ? IS NOT NULL AND last_used_at = ? THEN ( + SELECT MAX(created_at_unix_ms) + FROM "usage" + WHERE provider_api_key_id = ? AND status NOT IN ('pending', 'streaming') + ) + ELSE last_used_at + END +WHERE id = ? +"#, + ) + .bind(delta.request_count) + .bind(delta.success_count) + .bind(delta.error_count) + .bind(delta.total_tokens) + .bind(finite_or_zero(delta.total_cost_usd)) + .bind(delta.total_response_time_ms) + .bind(candidate) + .bind(candidate) + .bind(removed) + .bind(removed) + .bind(target_id) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_provider_monthly( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: f64, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta == 0.0 { + return Ok(()); + } + if !delta.is_finite() { + return Err(DataLayerError::UnexpectedValue(format!( + "providers.monthly_used_usd delta is not finite for {target_id}" + ))); + } + sqlx::query( + "UPDATE providers SET monthly_used_usd = COALESCE(monthly_used_usd, 0) + ?, updated_at = ? WHERE id = ?", + ) + .bind(delta) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_proxy_node( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ProxyNodeCounterDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + r#" +UPDATE proxy_nodes +SET total_requests = total_requests + MAX(?, 0), + failed_requests = failed_requests + MAX(?, 0), + dns_failures = dns_failures + MAX(?, 0), + stream_errors = stream_errors + MAX(?, 0), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(delta.total_requests_delta) + .bind(delta.failed_requests_delta) + .bind(delta.dns_failures_delta) + .bind(delta.stream_errors_delta) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_management_token( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ManagementTokenCounterDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + let last_used_at = optional_to_i64( + delta.last_used_at_unix_secs, + "management token last used at", + )?; + sqlx::query( + r#" +UPDATE management_tokens +SET usage_count = COALESCE(usage_count, 0) + MAX(?, 0), + last_used_at = CASE + WHEN ? IS NULL THEN last_used_at + ELSE MAX(COALESCE(last_used_at, 0), ?) + END, + last_used_ip = COALESCE(?, last_used_ip), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(delta.usage_count_delta) + .bind(last_used_at) + .bind(last_used_at) + .bind(delta.last_used_ip.as_deref()) + .bind(current_unix_secs()) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn apply_api_key_last_used( + tx: &mut sqlx::Transaction<'_, Sqlite>, + target_id: &str, + delta: &ApiKeyLastUsedDelta, +) -> Result<(), DataLayerError> { + if target_id.trim().is_empty() || delta.is_noop() { + return Ok(()); + } + sqlx::query( + "UPDATE api_keys SET last_used_at = MAX(COALESCE(last_used_at, 0), ?) WHERE id = ?", + ) + .bind(to_i64( + delta.last_used_at_unix_secs, + "api key last used at", + )?) + .bind(target_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +fn merge_optional_max(target: &mut Option, value: Option) { + if let Some(value) = value { + if target.is_none_or(|current| value > current) { + *target = Some(value); + } + } +} + +fn merge_latest_timestamp_with_value( + target_timestamp: &mut Option, + target_value: &mut Option, + timestamp: Option, + value: Option, +) { + let Some(timestamp) = timestamp else { + return; + }; + if target_timestamp.is_none_or(|current| timestamp >= current) { + *target_timestamp = Some(timestamp); + if value + .as_deref() + .map(str::trim) + .is_some_and(|v| !v.is_empty()) + { + *target_value = value; + } + } +} + +fn finite_or_zero(value: f64) -> f64 { + if value.is_finite() { + value + } else { + 0.0 + } +} + +fn current_unix_secs() -> i64 { + chrono::Utc::now().timestamp().max(0) +} + +fn to_i64(value: u64, field: &str) -> Result { + i64::try_from(value) + .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds i64: {value}"))) +} + +fn optional_to_i64(value: Option, field: &str) -> Result, DataLayerError> { + value.map(|value| to_i64(value, field)).transpose() +} + +fn optional_u64(field: &str, value: Option) -> Result, DataLayerError> { + value + .map(|value| { + u64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!("{field} is negative: {value}")) + }) + }) + .transpose() +} + +fn nonnegative_u64(value: i64) -> u64 { + value.max(0) as u64 +} + +fn optional_nonnegative_u64(value: Option) -> Option { + value.map(nonnegative_u64) +} + +#[cfg(test)] +mod tests { + use super::{ + cleanup_processed, enqueue_api_key_last_used, enqueue_management_token, enqueue_proxy_node, + flush, read_health, read_pending_health, + }; + use aether_data_contracts::repository::usage::{ + ApiKeyLastUsedDelta, ManagementTokenCounterDelta, ProxyNodeCounterDelta, + }; + + #[tokio::test] + async fn auxiliary_counters_flush_report_health_and_cleanup() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + crate::run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + sqlx::query( + r#" +INSERT INTO users (id, auth_source, created_at, updated_at) +VALUES ('counter-user', 'local', 1, 1); +INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) +VALUES ('counter-api-key', 'counter-user', 'counter-hash', 1, 1); +INSERT INTO management_tokens ( + id, user_id, name, token_hash, created_at, updated_at +) VALUES ( + 'counter-token', 'counter-user', 'counter token', 'counter-token-hash', 1, 1 +); +INSERT INTO proxy_nodes (id, name, ip, port, created_at, updated_at) +VALUES ('counter-node', 'counter node', '127.0.0.1', 8080, 1, 1); +"#, + ) + .execute(&pool) + .await + .expect("counter targets should seed"); + + assert!(enqueue_proxy_node( + &pool, + ProxyNodeCounterDelta { + node_id: "counter-node".to_string(), + total_requests_delta: 3, + failed_requests_delta: 1, + dns_failures_delta: 2, + stream_errors_delta: 1, + }, + ) + .await + .expect("proxy counter should enqueue")); + assert!(enqueue_management_token( + &pool, + ManagementTokenCounterDelta { + token_id: "counter-token".to_string(), + usage_count_delta: 2, + last_used_at_unix_secs: Some(100), + last_used_ip: Some("127.0.0.2".to_string()), + }, + ) + .await + .expect("management token counter should enqueue")); + assert!(enqueue_api_key_last_used( + &pool, + ApiKeyLastUsedDelta { + api_key_id: "counter-api-key".to_string(), + last_used_at_unix_secs: 110, + }, + ) + .await + .expect("api key last-used counter should enqueue")); + + let pending = read_pending_health(&pool) + .await + .expect("pending health should load"); + assert_eq!(pending.pending_rows, 3); + assert_eq!(pending.pending_by_kind.get("proxy_node"), Some(&1)); + assert_eq!(pending.pending_by_kind.get("management_token"), Some(&1)); + assert_eq!(pending.pending_by_kind.get("api_key_last_used"), Some(&1)); + + let summary = flush(&pool, 100).await.expect("counters should flush"); + assert_eq!(summary.rows_claimed, 3); + assert_eq!(summary.proxy_node_targets, 1); + assert_eq!(summary.management_token_targets, 1); + assert_eq!(summary.api_key_last_used_targets, 1); + + let proxy = sqlx::query_as::<_, (i64, i64, i64, i64)>( + "SELECT total_requests, failed_requests, dns_failures, stream_errors FROM proxy_nodes WHERE id = 'counter-node'", + ) + .fetch_one(&pool) + .await + .expect("proxy counters should load"); + assert_eq!(proxy, (3, 1, 2, 1)); + let token = sqlx::query_as::<_, (i64, Option, Option)>( + "SELECT usage_count, last_used_at, last_used_ip FROM management_tokens WHERE id = 'counter-token'", + ) + .fetch_one(&pool) + .await + .expect("management token counters should load"); + assert_eq!(token, (2, Some(100), Some("127.0.0.2".to_string()))); + let api_key_last_used: Option = + sqlx::query_scalar("SELECT last_used_at FROM api_keys WHERE id = 'counter-api-key'") + .fetch_one(&pool) + .await + .expect("api key last-used should load"); + assert_eq!(api_key_last_used, Some(110)); + + let health = read_health(&pool).await.expect("full health should load"); + assert_eq!(health.pending_rows, 0); + assert_eq!(health.processed_rows, 3); + assert!(health.latest_processed_at_unix_secs.is_some()); + + let deleted = + cleanup_processed(&pool, chrono::Utc::now().timestamp().max(0) as u64 + 1, 100) + .await + .expect("processed counters should clean up"); + assert_eq!(deleted, 3); + assert_eq!( + read_health(&pool) + .await + .expect("health should load after cleanup") + .processed_rows, + 0 + ); + } +} diff --git a/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs b/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs new file mode 100644 index 000000000..9e66588e7 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs @@ -0,0 +1,816 @@ +use std::io::Write; + +use aether_data_contracts::repository::usage::{ + parse_usage_body_ref, usage_body_ref, StoredRequestUsageAudit, UpsertUsageRecord, + UsageBodyCaptureState, UsageBodyField, +}; +use flate2::{write::GzEncoder, Compression}; +use serde_json::{Map, Value}; +use sqlx::{sqlite::SqliteRow, Row}; + +use crate::error::SqlResultExt; +use crate::{DataLayerError, SqlitePool}; + +#[derive(Debug)] +pub(crate) struct PreparedUsageHttpCapture { + request_headers: Option, + provider_request_headers: Option, + response_headers: Option, + client_response_headers: Option, + request_body: PreparedBody, + provider_request_body: PreparedBody, + response_body: PreparedBody, + client_response_body: PreparedBody, + refs: HttpAuditRefs, + states: HttpAuditStates, + capture_mode: &'static str, +} + +#[derive(Debug)] +struct PreparedBody { + field: UsageBodyField, + payload_gzip: Option>, + clear_existing: bool, +} + +#[derive(Debug, Default)] +struct HttpAuditRefs { + request_body_ref: Option, + provider_request_body_ref: Option, + response_body_ref: Option, + client_response_body_ref: Option, +} + +impl HttpAuditRefs { + fn any_present(&self) -> bool { + self.request_body_ref.is_some() + || self.provider_request_body_ref.is_some() + || self.response_body_ref.is_some() + || self.client_response_body_ref.is_some() + } +} + +#[derive(Debug, Default)] +struct HttpAuditStates { + request_body_state: Option, + provider_request_body_state: Option, + response_body_state: Option, + client_response_body_state: Option, +} + +impl HttpAuditStates { + fn any_present(&self) -> bool { + self.request_body_state.is_some() + || self.provider_request_body_state.is_some() + || self.response_body_state.is_some() + || self.client_response_body_state.is_some() + } +} + +pub(crate) fn capture_update_allowed( + previous: Option<&StoredRequestUsageAudit>, + incoming_status: &str, +) -> bool { + let Some(previous) = previous else { + return true; + }; + if previous.billing_status != "pending" { + return false; + } + if matches!( + previous.status.as_str(), + "completed" | "failed" | "cancelled" + ) && matches!(incoming_status, "pending" | "streaming") + { + return false; + } + !(previous.status == "streaming" && incoming_status == "pending") +} + +pub(crate) fn apply_previous_metadata_tombstones( + usage: &mut UpsertUsageRecord, + previous: Option<&StoredRequestUsageAudit>, +) { + if usage.request_metadata.is_some() { + return; + } + let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None); + let clear_provider_request = + usage.provider_request_body_state == Some(UsageBodyCaptureState::None); + if !clear_request && !clear_provider_request { + return; + } + let mut metadata = previous + .and_then(|previous| previous.request_metadata.as_ref()) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + if clear_request { + metadata.remove("requested_reasoning_effort"); + metadata.remove("request_body_ref"); + } + if clear_provider_request { + metadata.remove("provider_reasoning_effort"); + metadata.remove("provider_service_tier"); + metadata.remove("provider_cache_ttl_minutes"); + metadata.remove("provider_request_body_ref"); + } + usage.request_metadata = Some(Value::Object(metadata)); +} + +pub(crate) fn prepare_usage_http_capture( + usage: &mut UpsertUsageRecord, +) -> Result { + let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None); + let clear_provider_request = + usage.provider_request_body_state == Some(UsageBodyCaptureState::None); + let clear_response = usage.response_body_state == Some(UsageBodyCaptureState::None); + let clear_client_response = + usage.client_response_body_state == Some(UsageBodyCaptureState::None); + + let request_body_value = (!clear_request) + .then_some(usage.request_body.as_ref()) + .flatten(); + let provider_request_body_value = (!clear_provider_request) + .then_some(usage.provider_request_body.as_ref()) + .flatten(); + let response_body_value = (!clear_response) + .then_some(usage.response_body.as_ref()) + .flatten(); + let client_response_body_value = (!clear_client_response) + .then_some(usage.client_response_body.as_ref()) + .flatten(); + + let request_body = prepare_body( + UsageBodyField::RequestBody, + request_body_value, + clear_request, + )?; + let provider_request_body = prepare_body( + UsageBodyField::ProviderRequestBody, + provider_request_body_value, + clear_provider_request, + )?; + let response_body = prepare_body( + UsageBodyField::ResponseBody, + response_body_value, + clear_response, + )?; + let client_response_body = prepare_body( + UsageBodyField::ClientResponseBody, + client_response_body_value, + clear_client_response, + )?; + + let refs = HttpAuditRefs { + request_body_ref: resolved_write_ref( + (!clear_request) + .then_some(usage.request_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::RequestBody, + request_body.payload_gzip.is_some(), + ), + provider_request_body_ref: resolved_write_ref( + (!clear_provider_request) + .then_some(usage.provider_request_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ProviderRequestBody, + provider_request_body.payload_gzip.is_some(), + ), + response_body_ref: resolved_write_ref( + (!clear_response) + .then_some(usage.response_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ResponseBody, + response_body.payload_gzip.is_some(), + ), + client_response_body_ref: resolved_write_ref( + (!clear_client_response) + .then_some(usage.client_response_body_ref.as_deref()) + .flatten(), + &usage.request_id, + UsageBodyField::ClientResponseBody, + client_response_body.payload_gzip.is_some(), + ), + }; + let states = HttpAuditStates { + request_body_state: state_for_storage( + usage.request_body_state, + &request_body, + refs.request_body_ref.as_deref(), + ), + provider_request_body_state: state_for_storage( + usage.provider_request_body_state, + &provider_request_body, + refs.provider_request_body_ref.as_deref(), + ), + response_body_state: state_for_storage( + usage.response_body_state, + &response_body, + refs.response_body_ref.as_deref(), + ), + client_response_body_state: state_for_storage( + usage.client_response_body_state, + &client_response_body, + refs.client_response_body_ref.as_deref(), + ), + }; + + usage.request_metadata = prepare_metadata_for_body_storage( + usage.request_metadata.take(), + [ + ( + UsageBodyField::RequestBody, + request_body_value.is_some(), + usage.request_body_ref.as_deref(), + ), + ( + UsageBodyField::ProviderRequestBody, + provider_request_body_value.is_some(), + usage.provider_request_body_ref.as_deref(), + ), + ( + UsageBodyField::ResponseBody, + response_body_value.is_some(), + usage.response_body_ref.as_deref(), + ), + ( + UsageBodyField::ClientResponseBody, + client_response_body_value.is_some(), + usage.client_response_body_ref.as_deref(), + ), + ], + ); + + let capture_mode = if refs.any_present() { + "ref_backed" + } else if [ + request_body_value, + provider_request_body_value, + response_body_value, + client_response_body_value, + ] + .iter() + .any(Option::is_some) + { + "inline_legacy" + } else { + "none" + }; + + Ok(PreparedUsageHttpCapture { + request_headers: json_text(usage.request_headers.as_ref())?, + provider_request_headers: json_text(usage.provider_request_headers.as_ref())?, + response_headers: json_text(usage.response_headers.as_ref())?, + client_response_headers: json_text(usage.client_response_headers.as_ref())?, + request_body, + provider_request_body, + response_body, + client_response_body, + refs, + states, + capture_mode, + }) +} + +fn prepare_body( + field: UsageBodyField, + value: Option<&Value>, + clear_existing: bool, +) -> Result { + let payload_gzip = value.map(compress_json).transpose()?; + Ok(PreparedBody { + field, + payload_gzip, + clear_existing, + }) +} + +fn compress_json(value: &Value) -> Result, DataLayerError> { + let bytes = serde_json::to_vec(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage body: {err}")) + })?; + let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6)); + encoder.write_all(&bytes).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to gzip usage body: {err}")) + })?; + encoder.finish().map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to finish usage body gzip: {err}")) + }) +} + +fn resolved_write_ref( + explicit_ref: Option<&str>, + request_id: &str, + field: UsageBodyField, + has_blob: bool, +) -> Option { + explicit_ref + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .or_else(|| has_blob.then(|| usage_body_ref(request_id, field))) +} + +fn state_for_storage( + incoming: Option, + body: &PreparedBody, + body_ref: Option<&str>, +) -> Option { + if matches!( + incoming, + Some( + UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable + | UsageBodyCaptureState::None + ) + ) { + return incoming; + } + if body.payload_gzip.is_some() || body_ref.is_some() { + return Some(UsageBodyCaptureState::Reference); + } + incoming +} + +fn json_text(value: Option<&Value>) -> Result, DataLayerError> { + value + .map(|value| { + serde_json::to_string(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!("failed to serialize usage JSON: {err}")) + }) + }) + .transpose() +} + +fn prepare_metadata_for_body_storage( + metadata: Option, + body_fields: [(UsageBodyField, bool, Option<&str>); 4], +) -> Option { + let mut object = match metadata { + Some(Value::Object(object)) => object, + Some(value) => { + let mut object = Map::new(); + object.insert("request_metadata".to_string(), value); + object + } + None => Map::new(), + }; + let should_replace = !object.is_empty() + || body_fields + .iter() + .any(|(_, has_value, explicit_ref)| *has_value || explicit_ref.is_some()); + if !should_replace { + return None; + } + for (field, has_value, explicit_ref) in body_fields { + if has_value || explicit_ref.is_some() { + object.remove(field.as_ref_key()); + } + } + Some(Value::Object(object)) +} + +pub(crate) async fn sync_usage_http_capture( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, + prepared: &PreparedUsageHttpCapture, +) -> Result<(), DataLayerError> { + for body in [ + &prepared.request_body, + &prepared.provider_request_body, + &prepared.response_body, + &prepared.client_response_body, + ] { + sync_body(tx, request_id, body).await?; + } + + let headers_present = prepared.request_headers.is_some() + || prepared.provider_request_headers.is_some() + || prepared.response_headers.is_some() + || prepared.client_response_headers.is_some(); + if !headers_present + && !prepared.refs.any_present() + && !prepared.states.any_present() + && prepared.capture_mode == "none" + { + return Ok(()); + } + + sqlx::query( + r#" +INSERT INTO usage_http_audits ( + request_id, + request_headers, + provider_request_headers, + response_headers, + client_response_headers, + request_body_ref, + provider_request_body_ref, + response_body_ref, + client_response_body_ref, + request_body_state, + provider_request_body_state, + response_body_state, + client_response_body_state, + body_capture_mode +) +VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(request_id) DO UPDATE SET + request_headers = COALESCE(excluded.request_headers, usage_http_audits.request_headers), + provider_request_headers = COALESCE( + excluded.provider_request_headers, + usage_http_audits.provider_request_headers + ), + response_headers = COALESCE(excluded.response_headers, usage_http_audits.response_headers), + client_response_headers = COALESCE( + excluded.client_response_headers, + usage_http_audits.client_response_headers + ), + request_body_ref = CASE + WHEN excluded.request_body_state = 'none' THEN NULL + ELSE COALESCE(excluded.request_body_ref, usage_http_audits.request_body_ref) + END, + provider_request_body_ref = CASE + WHEN excluded.provider_request_body_state = 'none' THEN NULL + ELSE COALESCE( + excluded.provider_request_body_ref, + usage_http_audits.provider_request_body_ref + ) + END, + response_body_ref = CASE + WHEN excluded.response_body_state = 'none' THEN NULL + ELSE COALESCE(excluded.response_body_ref, usage_http_audits.response_body_ref) + END, + client_response_body_ref = CASE + WHEN excluded.client_response_body_state = 'none' THEN NULL + ELSE COALESCE( + excluded.client_response_body_ref, + usage_http_audits.client_response_body_ref + ) + END, + request_body_state = COALESCE( + excluded.request_body_state, + usage_http_audits.request_body_state + ), + provider_request_body_state = COALESCE( + excluded.provider_request_body_state, + usage_http_audits.provider_request_body_state + ), + response_body_state = COALESCE( + excluded.response_body_state, + usage_http_audits.response_body_state + ), + client_response_body_state = COALESCE( + excluded.client_response_body_state, + usage_http_audits.client_response_body_state + ), + body_capture_mode = CASE + WHEN excluded.body_capture_mode = 'none' + AND ( + excluded.request_body_state = 'none' + OR excluded.provider_request_body_state = 'none' + OR excluded.response_body_state = 'none' + OR excluded.client_response_body_state = 'none' + ) + THEN 'none' + ELSE COALESCE( + NULLIF(excluded.body_capture_mode, 'none'), + usage_http_audits.body_capture_mode, + 'none' + ) + END, + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +"#, + ) + .bind(request_id) + .bind(&prepared.request_headers) + .bind(&prepared.provider_request_headers) + .bind(&prepared.response_headers) + .bind(&prepared.client_response_headers) + .bind(prepared.refs.request_body_ref.as_deref()) + .bind(prepared.refs.provider_request_body_ref.as_deref()) + .bind(prepared.refs.response_body_ref.as_deref()) + .bind(prepared.refs.client_response_body_ref.as_deref()) + .bind( + prepared + .states + .request_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .provider_request_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .response_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind( + prepared + .states + .client_response_body_state + .map(UsageBodyCaptureState::as_str), + ) + .bind(prepared.capture_mode) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn sync_body( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, + body: &PreparedBody, +) -> Result<(), DataLayerError> { + let body_ref = usage_body_ref(request_id, body.field); + if body.clear_existing || body.payload_gzip.is_some() { + sqlx::query(clear_legacy_body_sql(body.field)) + .bind(request_id) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + if body.clear_existing { + sqlx::query("DELETE FROM usage_body_blobs WHERE body_ref = ?") + .bind(body_ref) + .execute(&mut **tx) + .await + .map_sql_err()?; + return Ok(()); + } + if let Some(payload_gzip) = body.payload_gzip.as_deref() { + sqlx::query( + r#" +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES (?, ?, ?, ?) +ON CONFLICT(body_ref) DO UPDATE SET + request_id = excluded.request_id, + body_field = excluded.body_field, + payload_gzip = excluded.payload_gzip, + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +"#, + ) + .bind(body_ref) + .bind(request_id) + .bind(body.field.as_storage_field()) + .bind(payload_gzip) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + Ok(()) +} + +fn clear_legacy_body_sql(field: UsageBodyField) -> &'static str { + match field { + UsageBodyField::RequestBody => { + "UPDATE \"usage\" SET request_body = NULL, request_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ProviderRequestBody => { + "UPDATE \"usage\" SET provider_request_body = NULL, provider_request_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ResponseBody => { + "UPDATE \"usage\" SET response_body = NULL, response_body_compressed = NULL WHERE request_id = ?" + } + UsageBodyField::ClientResponseBody => { + "UPDATE \"usage\" SET client_response_body = NULL, client_response_body_compressed = NULL WHERE request_id = ?" + } + } +} + +pub(crate) fn hydrate_usage_row( + row: &SqliteRow, + usage: &mut StoredRequestUsageAudit, + resolve_legacy_compressed: bool, +) -> Result<(), DataLayerError> { + usage.request_headers = optional_json(row, "request_headers")?; + usage.provider_request_headers = optional_json(row, "provider_request_headers")?; + usage.response_headers = optional_json(row, "response_headers")?; + usage.client_response_headers = optional_json(row, "client_response_headers")?; + + let request_body = legacy_body_column( + row, + "request_body", + "request_body_compressed", + resolve_legacy_compressed, + )?; + let provider_request_body = legacy_body_column( + row, + "provider_request_body", + "provider_request_body_compressed", + resolve_legacy_compressed, + )?; + let response_body = legacy_body_column( + row, + "response_body", + "response_body_compressed", + resolve_legacy_compressed, + )?; + let client_response_body = legacy_body_column( + row, + "client_response_body", + "client_response_body_compressed", + resolve_legacy_compressed, + )?; + usage.request_body = request_body.0; + usage.provider_request_body = provider_request_body.0; + usage.response_body = response_body.0; + usage.client_response_body = client_response_body.0; + + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + usage.request_body_ref = resolved_read_ref( + row.try_get("http_request_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::RequestBody, + request_body.1, + ); + usage.provider_request_body_ref = resolved_read_ref( + row.try_get("http_provider_request_body_ref") + .map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ProviderRequestBody, + provider_request_body.1, + ); + usage.response_body_ref = resolved_read_ref( + row.try_get("http_response_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ResponseBody, + response_body.1, + ); + usage.client_response_body_ref = resolved_read_ref( + row.try_get("http_client_response_body_ref").map_sql_err()?, + metadata, + &usage.request_id, + UsageBodyField::ClientResponseBody, + client_response_body.1, + ); + usage.request_body_state = optional_state(row, "http_request_body_state")?; + usage.provider_request_body_state = optional_state(row, "http_provider_request_body_state")?; + usage.response_body_state = optional_state(row, "http_response_body_state")?; + usage.client_response_body_state = optional_state(row, "http_client_response_body_state")?; + Ok(()) +} + +fn optional_json(row: &SqliteRow, column: &str) -> Result, DataLayerError> { + row.try_get::, _>(column) + .map_sql_err()? + .map(|raw| super::parse_usage_json_text(&raw)) + .transpose() +} + +fn legacy_body_column( + row: &SqliteRow, + inline_column: &str, + compressed_column: &str, + resolve_compressed: bool, +) -> Result<(Option, bool), DataLayerError> { + let inline = optional_json(row, inline_column)?; + if inline.is_some() { + return Ok((inline, false)); + } + let compressed = row + .try_get::>, _>(compressed_column) + .map_sql_err()?; + let has_compressed = compressed.is_some(); + let value = if resolve_compressed { + compressed + .map(|bytes| super::inflate_usage_json_value(&bytes)) + .transpose()? + } else { + None + }; + Ok((value, has_compressed)) +} + +fn resolved_read_ref( + audit_ref: Option, + metadata: Option<&Map>, + request_id: &str, + field: UsageBodyField, + has_compressed: bool, +) -> Option { + audit_ref + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .or_else(|| has_compressed.then(|| usage_body_ref(request_id, field))) + .or_else(|| metadata_body_ref(metadata, request_id, field)) +} + +fn metadata_body_ref( + metadata: Option<&Map>, + request_id: &str, + field: UsageBodyField, +) -> Option { + metadata + .and_then(|metadata| metadata.get(field.as_ref_key())) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(parse_usage_body_ref) + .filter(|(parsed_request_id, parsed_field)| { + parsed_request_id == request_id && *parsed_field == field + }) + .map(|(parsed_request_id, parsed_field)| usage_body_ref(&parsed_request_id, parsed_field)) +} + +fn optional_state( + row: &SqliteRow, + column: &str, +) -> Result, DataLayerError> { + Ok(row + .try_get::, _>(column) + .map_sql_err()? + .as_deref() + .and_then(parse_state)) +} + +fn parse_state(value: &str) -> Option { + match value.trim() { + "none" => Some(UsageBodyCaptureState::None), + "inline" => Some(UsageBodyCaptureState::Inline), + "reference" => Some(UsageBodyCaptureState::Reference), + "truncated" => Some(UsageBodyCaptureState::Truncated), + "disabled" => Some(UsageBodyCaptureState::Disabled), + "unavailable" => Some(UsageBodyCaptureState::Unavailable), + _ => None, + } +} + +pub(crate) async fn hydrate_usage_body_refs( + pool: &SqlitePool, + mut usage: StoredRequestUsageAudit, +) -> Result { + for field in [ + UsageBodyField::RequestBody, + UsageBodyField::ProviderRequestBody, + UsageBodyField::ResponseBody, + UsageBodyField::ClientResponseBody, + ] { + if usage.body_value(field).is_some() { + continue; + } + let Some(body_ref) = usage.body_ref(field) else { + continue; + }; + let value = resolve_body_ref(pool, body_ref).await?; + match field { + UsageBodyField::RequestBody => usage.request_body = value, + UsageBodyField::ProviderRequestBody => usage.provider_request_body = value, + UsageBodyField::ResponseBody => usage.response_body = value, + UsageBodyField::ClientResponseBody => usage.client_response_body = value, + } + } + Ok(usage) +} + +pub(crate) async fn resolve_body_ref( + pool: &SqlitePool, + body_ref: &str, +) -> Result, DataLayerError> { + if let Some(payload_gzip) = sqlx::query_scalar::<_, Vec>( + "SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? LIMIT 1", + ) + .bind(body_ref) + .fetch_optional(pool) + .await + .map_sql_err()? + { + return super::inflate_usage_json_value(&payload_gzip).map(Some); + } + let Some((request_id, field)) = parse_usage_body_ref(body_ref) else { + return Ok(None); + }; + let (inline_column, compressed_column) = super::sqlite_usage_body_sql_columns(field); + let row = sqlx::query(&format!( + "SELECT {inline_column} AS inline_body, {compressed_column} AS compressed_body FROM \"usage\" WHERE request_id = ? LIMIT 1" + )) + .bind(request_id) + .fetch_optional(pool) + .await + .map_sql_err()?; + let Some(row) = row.as_ref() else { + return Ok(None); + }; + if let Some(raw) = row + .try_get::, _>("inline_body") + .map_sql_err()? + { + return super::parse_usage_json_text(&raw).map(Some); + } + row.try_get::>, _>("compressed_body") + .map_sql_err()? + .map(|bytes| super::inflate_usage_json_value(&bytes)) + .transpose() +} diff --git a/crates/aether-data/adapters/sqlite/src/usage/snapshots.rs b/crates/aether-data/adapters/sqlite/src/usage/snapshots.rs new file mode 100644 index 000000000..212e9e960 --- /dev/null +++ b/crates/aether-data/adapters/sqlite/src/usage/snapshots.rs @@ -0,0 +1,1047 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord}; +use aether_data_contracts::DataLayerError; +use serde_json::{Map, Value}; +use sqlx::{QueryBuilder, Row, Sqlite, Transaction}; + +use crate::error::SqlResultExt; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct UsageRoutingSnapshot { + candidate_id: Option, + candidate_index: Option, + key_name: Option, + planner_kind: Option, + route_family: Option, + route_kind: Option, + execution_path: Option, + local_execution_runtime_miss_reason: Option, + selected_provider_id: Option, + selected_endpoint_id: Option, + selected_provider_api_key_id: Option, + has_format_conversion: Option, +} + +impl UsageRoutingSnapshot { + fn has_metadata_fields(&self) -> bool { + self.candidate_id.is_some() + || self.candidate_index.is_some() + || self.key_name.is_some() + || self.planner_kind.is_some() + || self.route_family.is_some() + || self.route_kind.is_some() + || self.execution_path.is_some() + || self.local_execution_runtime_miss_reason.is_some() + } + + fn any_present(&self) -> bool { + self.has_metadata_fields() + || self.selected_provider_id.is_some() + || self.selected_endpoint_id.is_some() + || self.selected_provider_api_key_id.is_some() + || self.has_format_conversion.is_some() + } +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub(crate) struct UsageSettlementPricingSnapshot { + billing_status: Option, + billing_snapshot_schema_version: Option, + billing_snapshot_status: Option, + settlement_snapshot_schema_version: Option, + settlement_snapshot: Option, + billing_dimensions: Option, + billing_input_tokens: Option, + billing_effective_input_tokens: Option, + billing_output_tokens: Option, + billing_cache_creation_tokens: Option, + billing_cache_creation_5m_tokens: Option, + billing_cache_creation_1h_tokens: Option, + billing_cache_read_tokens: Option, + billing_total_input_context: Option, + billing_cache_creation_cost_usd: Option, + billing_cache_read_cost_usd: Option, + billing_total_cost_usd: Option, + billing_actual_total_cost_usd: Option, + billing_pricing_source: Option, + billing_rule_id: Option, + billing_rule_version: Option, + rate_multiplier: Option, + is_free_tier: Option, + input_price_per_1m: Option, + output_price_per_1m: Option, + cache_creation_price_per_1m: Option, + cache_read_price_per_1m: Option, + price_per_request: Option, +} + +impl UsageSettlementPricingSnapshot { + fn any_present(&self) -> bool { + self.billing_status.is_some() + || self.billing_snapshot_schema_version.is_some() + || self.billing_snapshot_status.is_some() + || self.settlement_snapshot_schema_version.is_some() + || self.settlement_snapshot.is_some() + || self.billing_dimensions.is_some() + || self.billing_input_tokens.is_some() + || self.billing_effective_input_tokens.is_some() + || self.billing_output_tokens.is_some() + || self.billing_cache_creation_tokens.is_some() + || self.billing_cache_creation_5m_tokens.is_some() + || self.billing_cache_creation_1h_tokens.is_some() + || self.billing_cache_read_tokens.is_some() + || self.billing_total_input_context.is_some() + || self.billing_cache_creation_cost_usd.is_some() + || self.billing_cache_read_cost_usd.is_some() + || self.billing_total_cost_usd.is_some() + || self.billing_actual_total_cost_usd.is_some() + || self.billing_pricing_source.is_some() + || self.billing_rule_id.is_some() + || self.billing_rule_version.is_some() + || self.rate_multiplier.is_some() + || self.is_free_tier.is_some() + || self.input_price_per_1m.is_some() + || self.output_price_per_1m.is_some() + || self.cache_creation_price_per_1m.is_some() + || self.cache_read_price_per_1m.is_some() + || self.price_per_request.is_some() + } +} + +pub(crate) fn from_usage( + usage: &UpsertUsageRecord, +) -> Result<(UsageRoutingSnapshot, UsageSettlementPricingSnapshot), DataLayerError> { + Ok(( + routing_snapshot_from_usage(usage), + settlement_snapshot_from_usage(usage)?, + )) +} + +fn routing_snapshot_from_usage(usage: &UpsertUsageRecord) -> UsageRoutingSnapshot { + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + let mut snapshot = UsageRoutingSnapshot { + candidate_id: usage + .candidate_id + .clone() + .or_else(|| metadata_string(metadata, "candidate_id")), + candidate_index: usage + .candidate_index + .or_else(|| metadata_u64(metadata, "candidate_index")), + key_name: usage + .key_name + .clone() + .or_else(|| metadata_string(metadata, "key_name")), + planner_kind: usage + .planner_kind + .clone() + .or_else(|| metadata_string(metadata, "planner_kind")), + route_family: usage + .route_family + .clone() + .or_else(|| metadata_string(metadata, "route_family")), + route_kind: usage + .route_kind + .clone() + .or_else(|| metadata_string(metadata, "route_kind")), + execution_path: usage + .execution_path + .clone() + .or_else(|| metadata_string(metadata, "execution_path")), + local_execution_runtime_miss_reason: usage + .local_execution_runtime_miss_reason + .clone() + .or_else(|| metadata_string(metadata, "local_execution_runtime_miss_reason")), + selected_provider_id: None, + selected_endpoint_id: None, + selected_provider_api_key_id: None, + has_format_conversion: None, + }; + if snapshot.has_metadata_fields() { + snapshot.selected_provider_id = usage.provider_id.clone(); + snapshot.selected_endpoint_id = usage.provider_endpoint_id.clone(); + snapshot.selected_provider_api_key_id = usage.provider_api_key_id.clone(); + snapshot.has_format_conversion = usage.has_format_conversion; + } + snapshot +} + +fn settlement_snapshot_from_usage( + usage: &UpsertUsageRecord, +) -> Result { + let metadata = usage.request_metadata.as_ref().and_then(Value::as_object); + let billing_dimensions = metadata_or_snapshot_dimensions(metadata); + let has_billing_dimensions = billing_dimensions.is_some(); + let usage_input_tokens = optional_i64(usage.input_tokens, "input_tokens")?; + let usage_output_tokens = optional_i64(usage.output_tokens, "output_tokens")?; + let usage_cache_creation_uncategorized_tokens = optional_i64( + usage.cache_creation_input_tokens, + "cache_creation_input_tokens", + )?; + let usage_cache_creation_5m_tokens = optional_i64( + usage.cache_creation_ephemeral_5m_input_tokens, + "cache_creation_ephemeral_5m_input_tokens", + )?; + let usage_cache_creation_1h_tokens = optional_i64( + usage.cache_creation_ephemeral_1h_input_tokens, + "cache_creation_ephemeral_1h_input_tokens", + )?; + let usage_cache_read_tokens = + optional_i64(usage.cache_read_input_tokens, "cache_read_input_tokens")?; + let usage_cache_creation_tokens = cache_creation_tokens_from_parts( + usage_cache_creation_uncategorized_tokens, + usage_cache_creation_5m_tokens, + usage_cache_creation_1h_tokens, + ); + let billing_cache_creation_tokens = billing_dimension_i64(metadata, "cache_creation_tokens") + .or_else(|| { + cache_creation_tokens_from_parts( + billing_dimension_i64(metadata, "cache_creation_uncategorized_tokens"), + billing_dimension_i64(metadata, "cache_creation_ephemeral_5m_tokens"), + billing_dimension_i64(metadata, "cache_creation_ephemeral_1h_tokens"), + ) + }) + .or(usage_cache_creation_tokens); + let billing_cache_creation_5m_tokens = + billing_dimension_i64(metadata, "cache_creation_ephemeral_5m_tokens") + .or(usage_cache_creation_5m_tokens); + let billing_cache_creation_1h_tokens = + billing_dimension_i64(metadata, "cache_creation_ephemeral_1h_tokens") + .or(usage_cache_creation_1h_tokens); + let billing_input_tokens = + billing_dimension_i64(metadata, "input_tokens").or(usage_input_tokens); + let billing_output_tokens = + billing_dimension_i64(metadata, "output_tokens").or(usage_output_tokens); + let billing_cache_read_tokens = + billing_dimension_i64(metadata, "cache_read_tokens").or(usage_cache_read_tokens); + let api_family = normalized_api_family(usage); + let billing_effective_input_tokens = billing_dimension_i64(metadata, "effective_input_tokens") + .or_else(|| { + has_billing_dimensions + .then(|| billing_dimension_i64(metadata, "input_tokens")) + .flatten() + }) + .or_else(|| { + effective_input_tokens( + billing_input_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens, + &api_family, + ) + }); + let billing_total_input_context = billing_dimension_i64(metadata, "total_input_context") + .or_else(|| { + total_input_context( + billing_input_tokens, + billing_effective_input_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens, + &api_family, + ) + }); + + Ok(UsageSettlementPricingSnapshot { + billing_status: Some(usage.billing_status.clone()), + billing_snapshot_schema_version: metadata_string( + metadata, + "billing_snapshot_schema_version", + ) + .or_else(|| billing_snapshot_string(metadata, "schema_version")), + billing_snapshot_status: metadata_string(metadata, "billing_snapshot_status") + .or_else(|| billing_snapshot_string(metadata, "status")), + settlement_snapshot_schema_version: settlement_snapshot_schema_version(metadata), + settlement_snapshot: settlement_snapshot_value(metadata), + billing_dimensions, + billing_input_tokens, + billing_effective_input_tokens, + billing_output_tokens, + billing_cache_creation_tokens, + billing_cache_creation_5m_tokens, + billing_cache_creation_1h_tokens, + billing_cache_read_tokens, + billing_total_input_context, + billing_cache_creation_cost_usd: settlement_cache_creation_cost(metadata) + .or(usage.cache_creation_cost_usd), + billing_cache_read_cost_usd: settlement_cost_breakdown_number(metadata, "cache_read_cost") + .or(usage.cache_read_cost_usd), + billing_total_cost_usd: settlement_snapshot_number(metadata, "total_cost") + .or_else(|| billing_snapshot_number(metadata, "total_cost")) + .or(usage.total_cost_usd), + billing_actual_total_cost_usd: settlement_snapshot_number(metadata, "actual_total_cost") + .or(usage.actual_total_cost_usd), + billing_pricing_source: settlement_nested_string( + metadata, + "pricing_snapshot", + "pricing_source", + ), + billing_rule_id: settlement_nested_string(metadata, "billing_plan_snapshot", "rule_id") + .or_else(|| billing_snapshot_string_field(metadata, "rule_id")), + billing_rule_version: settlement_nested_string( + metadata, + "billing_plan_snapshot", + "rule_version", + ), + rate_multiplier: metadata_number(metadata, "rate_multiplier"), + is_free_tier: metadata_bool(metadata, "is_free_tier"), + input_price_per_1m: metadata_number(metadata, "input_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "input_price_per_1m")), + output_price_per_1m: metadata_number(metadata, "output_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "output_price_per_1m")) + .or(usage.output_price_per_1m), + cache_creation_price_per_1m: metadata_number(metadata, "cache_creation_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "cache_creation_price_per_1m")), + cache_read_price_per_1m: metadata_number(metadata, "cache_read_price_per_1m") + .or_else(|| billing_snapshot_resolved_number(metadata, "cache_read_price_per_1m")), + price_per_request: metadata_number(metadata, "price_per_request") + .or_else(|| billing_snapshot_resolved_number(metadata, "price_per_request")), + }) +} + +pub(crate) async fn sync( + tx: &mut Transaction<'_, Sqlite>, + request_id: &str, + routing: &UsageRoutingSnapshot, + settlement: &UsageSettlementPricingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + sync_routing(tx, request_id, routing, replace_existing).await?; + sync_settlement(tx, request_id, settlement, replace_existing).await +} + +async fn sync_routing( + tx: &mut Transaction<'_, Sqlite>, + request_id: &str, + snapshot: &UsageRoutingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + if !snapshot.any_present() && !replace_existing { + return Ok(()); + } + let now = unix_now()?; + let mut query = QueryBuilder::::new( + "INSERT INTO usage_routing_snapshots (request_id, candidate_id, candidate_index, \ + key_name, planner_kind, route_family, route_kind, execution_path, \ + local_execution_runtime_miss_reason, selected_provider_id, selected_endpoint_id, \ + selected_provider_api_key_id, has_format_conversion, created_at, updated_at) VALUES (", + ); + { + let mut values = query.separated(", "); + values + .push_bind(request_id) + .push_bind(snapshot.candidate_id.as_deref()) + .push_bind(optional_i64(snapshot.candidate_index, "candidate_index")?) + .push_bind(snapshot.key_name.as_deref()) + .push_bind(snapshot.planner_kind.as_deref()) + .push_bind(snapshot.route_family.as_deref()) + .push_bind(snapshot.route_kind.as_deref()) + .push_bind(snapshot.execution_path.as_deref()) + .push_bind(snapshot.local_execution_runtime_miss_reason.as_deref()) + .push_bind(snapshot.selected_provider_id.as_deref()) + .push_bind(snapshot.selected_endpoint_id.as_deref()) + .push_bind(snapshot.selected_provider_api_key_id.as_deref()) + .push_bind(snapshot.has_format_conversion) + .push_bind(now) + .push_bind(now); + } + query.push(") ON CONFLICT (request_id) DO UPDATE SET "); + push_sqlite_updates( + &mut query, + &[ + "candidate_id", + "candidate_index", + "key_name", + "planner_kind", + "route_family", + "route_kind", + "execution_path", + "local_execution_runtime_miss_reason", + "selected_provider_id", + "selected_endpoint_id", + "selected_provider_api_key_id", + "has_format_conversion", + ], + "usage_routing_snapshots", + replace_existing, + ); + query.push(", updated_at = excluded.updated_at"); + query.build().execute(&mut **tx).await.map_sql_err()?; + Ok(()) +} + +async fn sync_settlement( + tx: &mut Transaction<'_, Sqlite>, + request_id: &str, + snapshot: &UsageSettlementPricingSnapshot, + replace_existing: bool, +) -> Result<(), DataLayerError> { + if !snapshot.any_present() && !replace_existing { + return Ok(()); + } + let now = unix_now()?; + let settlement_json = json_text(snapshot.settlement_snapshot.as_ref())?; + let dimensions_json = json_text(snapshot.billing_dimensions.as_ref())?; + let mut query = QueryBuilder::::new( + "INSERT INTO usage_settlement_snapshots (request_id, billing_status, \ + billing_snapshot_schema_version, billing_snapshot_status, \ + settlement_snapshot_schema_version, settlement_snapshot, billing_dimensions, \ + billing_input_tokens, billing_effective_input_tokens, billing_output_tokens, \ + billing_cache_creation_tokens, billing_cache_creation_5m_tokens, \ + billing_cache_creation_1h_tokens, billing_cache_read_tokens, \ + billing_total_input_context, billing_cache_creation_cost_usd, \ + billing_cache_read_cost_usd, billing_total_cost_usd, \ + billing_actual_total_cost_usd, billing_pricing_source, billing_rule_id, \ + billing_rule_version, rate_multiplier, is_free_tier, input_price_per_1m, \ + output_price_per_1m, cache_creation_price_per_1m, cache_read_price_per_1m, \ + price_per_request, created_at, updated_at) VALUES (", + ); + { + let mut values = query.separated(", "); + values + .push_bind(request_id) + .push_bind(snapshot.billing_status.as_deref().unwrap_or("pending")) + .push_bind(snapshot.billing_snapshot_schema_version.as_deref()) + .push_bind(snapshot.billing_snapshot_status.as_deref()) + .push_bind(snapshot.settlement_snapshot_schema_version.as_deref()) + .push_bind(settlement_json.as_deref()) + .push_bind(dimensions_json.as_deref()) + .push_bind(snapshot.billing_input_tokens) + .push_bind(snapshot.billing_effective_input_tokens) + .push_bind(snapshot.billing_output_tokens) + .push_bind(snapshot.billing_cache_creation_tokens) + .push_bind(snapshot.billing_cache_creation_5m_tokens) + .push_bind(snapshot.billing_cache_creation_1h_tokens) + .push_bind(snapshot.billing_cache_read_tokens) + .push_bind(snapshot.billing_total_input_context) + .push_bind(snapshot.billing_cache_creation_cost_usd) + .push_bind(snapshot.billing_cache_read_cost_usd) + .push_bind(snapshot.billing_total_cost_usd) + .push_bind(snapshot.billing_actual_total_cost_usd) + .push_bind(snapshot.billing_pricing_source.as_deref()) + .push_bind(snapshot.billing_rule_id.as_deref()) + .push_bind(snapshot.billing_rule_version.as_deref()) + .push_bind(snapshot.rate_multiplier) + .push_bind(snapshot.is_free_tier) + .push_bind(snapshot.input_price_per_1m) + .push_bind(snapshot.output_price_per_1m) + .push_bind(snapshot.cache_creation_price_per_1m) + .push_bind(snapshot.cache_read_price_per_1m) + .push_bind(snapshot.price_per_request) + .push_bind(now) + .push_bind(now); + } + query.push(") ON CONFLICT (request_id) DO UPDATE SET "); + if replace_existing { + query.push("billing_status = excluded.billing_status, "); + } + push_sqlite_updates( + &mut query, + &[ + "billing_snapshot_schema_version", + "billing_snapshot_status", + "settlement_snapshot_schema_version", + "settlement_snapshot", + "billing_dimensions", + "billing_input_tokens", + "billing_effective_input_tokens", + "billing_output_tokens", + "billing_cache_creation_tokens", + "billing_cache_creation_5m_tokens", + "billing_cache_creation_1h_tokens", + "billing_cache_read_tokens", + "billing_total_input_context", + "billing_cache_creation_cost_usd", + "billing_cache_read_cost_usd", + "billing_total_cost_usd", + "billing_actual_total_cost_usd", + "billing_pricing_source", + "billing_rule_id", + "billing_rule_version", + "rate_multiplier", + "is_free_tier", + "input_price_per_1m", + "output_price_per_1m", + "cache_creation_price_per_1m", + "cache_read_price_per_1m", + "price_per_request", + ], + "usage_settlement_snapshots", + replace_existing, + ); + query.push(", updated_at = excluded.updated_at"); + query.build().execute(&mut **tx).await.map_sql_err()?; + Ok(()) +} + +fn push_sqlite_updates( + query: &mut QueryBuilder<'_, Sqlite>, + fields: &[&str], + table: &str, + replace_existing: bool, +) { + for (index, field) in fields.iter().enumerate() { + if index > 0 { + query.push(", "); + } + query.push(*field).push(" = "); + if replace_existing { + query.push("excluded.").push(*field); + } else { + query + .push("COALESCE(excluded.") + .push(*field) + .push(", ") + .push(table) + .push(".") + .push(*field) + .push(")"); + } + } +} + +pub(crate) fn hydrate_row( + row: &sqlx::sqlite::SqliteRow, + audit: &mut StoredRequestUsageAudit, +) -> Result<(), DataLayerError> { + audit.candidate_id = row.try_get("routing_candidate_id").map_sql_err()?; + audit.candidate_index = row + .try_get::, _>("routing_candidate_index") + .map_sql_err()? + .map(|value| { + u64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!( + "usage routing candidate_index is negative: {value}" + )) + }) + }) + .transpose()?; + audit.key_name = row.try_get("routing_key_name").map_sql_err()?; + audit.planner_kind = row.try_get("routing_planner_kind").map_sql_err()?; + audit.route_family = row.try_get("routing_route_family").map_sql_err()?; + audit.route_kind = row.try_get("routing_route_kind").map_sql_err()?; + audit.execution_path = row.try_get("routing_execution_path").map_sql_err()?; + audit.local_execution_runtime_miss_reason = row + .try_get("routing_local_execution_runtime_miss_reason") + .map_sql_err()?; + + let snapshot = settlement_snapshot_from_row(row)?; + if let Some(effective) = nonnegative_u64(snapshot.billing_effective_input_tokens) { + audit.total_tokens = effective + .saturating_add(audit.output_tokens) + .saturating_add(audit.cache_creation_input_tokens) + .saturating_add(audit.cache_read_input_tokens); + } else if let Some(context) = nonnegative_u64(snapshot.billing_total_input_context) { + audit.total_tokens = context.saturating_add(audit.output_tokens); + } + audit.request_metadata = attach_settlement_metadata(audit.request_metadata.take(), &snapshot); + Ok(()) +} + +fn settlement_snapshot_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> Result { + Ok(UsageSettlementPricingSnapshot { + billing_status: None, + billing_snapshot_schema_version: row + .try_get("settlement_billing_snapshot_schema_version") + .map_sql_err()?, + billing_snapshot_status: row + .try_get("settlement_billing_snapshot_status") + .map_sql_err()?, + settlement_snapshot_schema_version: row + .try_get("settlement_snapshot_schema_version") + .map_sql_err()?, + settlement_snapshot: json_value_from_row(row, "settlement_snapshot")?, + billing_dimensions: json_value_from_row(row, "settlement_billing_dimensions")?, + billing_input_tokens: row + .try_get("settlement_billing_input_tokens") + .map_sql_err()?, + billing_effective_input_tokens: row + .try_get("settlement_billing_effective_input_tokens") + .map_sql_err()?, + billing_output_tokens: row + .try_get("settlement_billing_output_tokens") + .map_sql_err()?, + billing_cache_creation_tokens: row + .try_get("settlement_billing_cache_creation_tokens") + .map_sql_err()?, + billing_cache_creation_5m_tokens: row + .try_get("settlement_billing_cache_creation_5m_tokens") + .map_sql_err()?, + billing_cache_creation_1h_tokens: row + .try_get("settlement_billing_cache_creation_1h_tokens") + .map_sql_err()?, + billing_cache_read_tokens: row + .try_get("settlement_billing_cache_read_tokens") + .map_sql_err()?, + billing_total_input_context: row + .try_get("settlement_billing_total_input_context") + .map_sql_err()?, + billing_cache_creation_cost_usd: row + .try_get("settlement_billing_cache_creation_cost_usd") + .map_sql_err()?, + billing_cache_read_cost_usd: row + .try_get("settlement_billing_cache_read_cost_usd") + .map_sql_err()?, + billing_total_cost_usd: row + .try_get("settlement_billing_total_cost_usd") + .map_sql_err()?, + billing_actual_total_cost_usd: row + .try_get("settlement_billing_actual_total_cost_usd") + .map_sql_err()?, + billing_pricing_source: row + .try_get("settlement_billing_pricing_source") + .map_sql_err()?, + billing_rule_id: row.try_get("settlement_billing_rule_id").map_sql_err()?, + billing_rule_version: row + .try_get("settlement_billing_rule_version") + .map_sql_err()?, + rate_multiplier: row.try_get("settlement_rate_multiplier").map_sql_err()?, + is_free_tier: row + .try_get::, _>("settlement_is_free_tier") + .map_sql_err()? + .map(|value| value != 0), + input_price_per_1m: row.try_get("settlement_input_price_per_1m").map_sql_err()?, + output_price_per_1m: row + .try_get("settlement_output_price_per_1m") + .map_sql_err()?, + cache_creation_price_per_1m: row + .try_get("settlement_cache_creation_price_per_1m") + .map_sql_err()?, + cache_read_price_per_1m: row + .try_get("settlement_cache_read_price_per_1m") + .map_sql_err()?, + price_per_request: row.try_get("settlement_price_per_request").map_sql_err()?, + }) +} + +fn attach_settlement_metadata( + metadata: Option, + snapshot: &UsageSettlementPricingSnapshot, +) -> Option { + if !snapshot.any_present() { + return metadata; + } + let mut metadata = match metadata { + Some(Value::Object(object)) => object, + Some(value) => return Some(value), + None => Map::new(), + }; + insert_string( + &mut metadata, + "billing_snapshot_schema_version", + snapshot.billing_snapshot_schema_version.as_deref(), + ); + insert_string( + &mut metadata, + "billing_snapshot_status", + snapshot.billing_snapshot_status.as_deref(), + ); + insert_string( + &mut metadata, + "settlement_snapshot_schema_version", + snapshot.settlement_snapshot_schema_version.as_deref(), + ); + insert_value( + &mut metadata, + "settlement_snapshot", + snapshot.settlement_snapshot.as_ref(), + ); + insert_value( + &mut metadata, + "billing_dimensions", + snapshot.billing_dimensions.as_ref(), + ); + insert_number(&mut metadata, "rate_multiplier", snapshot.rate_multiplier); + insert_bool(&mut metadata, "is_free_tier", snapshot.is_free_tier); + insert_number( + &mut metadata, + "input_price_per_1m", + snapshot.input_price_per_1m, + ); + insert_number( + &mut metadata, + "output_price_per_1m", + snapshot.output_price_per_1m, + ); + insert_number( + &mut metadata, + "cache_creation_price_per_1m", + snapshot.cache_creation_price_per_1m, + ); + insert_number( + &mut metadata, + "cache_read_price_per_1m", + snapshot.cache_read_price_per_1m, + ); + insert_number( + &mut metadata, + "price_per_request", + snapshot.price_per_request, + ); + (!metadata.is_empty()).then_some(Value::Object(metadata)) +} + +fn insert_string(metadata: &mut Map, key: &str, value: Option<&str>) { + if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) { + metadata.insert(key.to_string(), Value::String(value.to_string())); + } +} + +fn insert_number(metadata: &mut Map, key: &str, value: Option) { + if let Some(number) = value + .filter(|value| value.is_finite()) + .and_then(serde_json::Number::from_f64) + { + metadata.insert(key.to_string(), Value::Number(number)); + } +} + +fn insert_bool(metadata: &mut Map, key: &str, value: Option) { + if let Some(value) = value { + metadata.insert(key.to_string(), Value::Bool(value)); + } +} + +fn insert_value(metadata: &mut Map, key: &str, value: Option<&Value>) { + if let Some(value) = value { + metadata.insert(key.to_string(), value.clone()); + } +} + +fn metadata_string(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn metadata_number(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn metadata_u64(metadata: Option<&Map>, key: &str) -> Option { + metadata.and_then(|object| { + object.get(key).and_then(|value| { + value + .as_u64() + .or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok())) + }) + }) +} + +fn metadata_bool(metadata: Option<&Map>, key: &str) -> Option { + metadata + .and_then(|object| object.get(key)) + .and_then(Value::as_bool) +} + +fn billing_snapshot_object(metadata: Option<&Map>) -> Option<&Map> { + metadata + .and_then(|object| object.get("billing_snapshot")) + .and_then(Value::as_object) +} + +fn billing_snapshot_string(metadata: Option<&Map>, key: &str) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn billing_snapshot_resolved_number( + metadata: Option<&Map>, + key: &str, +) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("resolved_variables")) + .and_then(Value::as_object) + .and_then(|variables| variables.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_snapshot_object( + metadata: Option<&Map>, +) -> Option<&Map> { + metadata + .and_then(|object| object.get("settlement_snapshot")) + .and_then(Value::as_object) +} + +fn settlement_snapshot_schema_version(metadata: Option<&Map>) -> Option { + metadata_string(metadata, "settlement_snapshot_schema_version").or_else(|| { + settlement_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("schema_version")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + }) +} + +fn settlement_snapshot_value(metadata: Option<&Map>) -> Option { + metadata + .and_then(|object| object.get("settlement_snapshot")) + .cloned() +} + +fn settlement_child_value<'a>( + metadata: Option<&'a Map>, + child: &str, +) -> Option<&'a Value> { + settlement_snapshot_object(metadata).and_then(|snapshot| snapshot.get(child)) +} + +fn settlement_child_object<'a>( + metadata: Option<&'a Map>, + child: &str, +) -> Option<&'a Map> { + settlement_child_value(metadata, child).and_then(Value::as_object) +} + +fn metadata_or_snapshot_dimensions(metadata: Option<&Map>) -> Option { + metadata + .and_then(|object| object.get("billing_dimensions")) + .cloned() + .or_else(|| settlement_child_value(metadata, "resolved_dimensions").cloned()) + .or_else(|| { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("resolved_dimensions")) + .cloned() + }) +} + +fn billing_dimension_i64(metadata: Option<&Map>, key: &str) -> Option { + metadata_or_snapshot_dimensions(metadata) + .and_then(|dimensions| dimensions.get(key).and_then(json_i64)) + .filter(|value| *value >= 0) +} + +fn json_i64(value: &Value) -> Option { + value + .as_i64() + .or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok())) +} + +fn settlement_snapshot_number(metadata: Option<&Map>, key: &str) -> Option { + settlement_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn billing_snapshot_number(metadata: Option<&Map>, key: &str) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_cost_breakdown_number( + metadata: Option<&Map>, + key: &str, +) -> Option { + settlement_child_object(metadata, "cost_breakdown") + .or_else(|| { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get("cost_breakdown")) + .and_then(Value::as_object) + }) + .and_then(|breakdown| breakdown.get(key)) + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn settlement_cache_creation_cost(metadata: Option<&Map>) -> Option { + let mut found = false; + let total = [ + "cache_creation_uncategorized_cost", + "cache_creation_ephemeral_5m_cost", + "cache_creation_ephemeral_1h_cost", + "cache_creation_cost", + ] + .into_iter() + .fold(0.0, |sum, key| { + if let Some(value) = settlement_cost_breakdown_number(metadata, key) { + found = true; + sum + value + } else { + sum + } + }); + found.then_some(total) +} + +fn settlement_nested_string( + metadata: Option<&Map>, + child: &str, + key: &str, +) -> Option { + settlement_child_object(metadata, child) + .and_then(|object| object.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn billing_snapshot_string_field( + metadata: Option<&Map>, + key: &str, +) -> Option { + billing_snapshot_object(metadata) + .and_then(|snapshot| snapshot.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn optional_i64(value: Option, field: &str) -> Result, DataLayerError> { + value + .map(|value| { + i64::try_from(value).map_err(|_| { + DataLayerError::UnexpectedValue(format!("usage {field} exceeds bigint: {value}")) + }) + }) + .transpose() +} + +fn cache_creation_tokens_from_parts( + uncategorized: Option, + ephemeral_5m: Option, + ephemeral_1h: Option, +) -> Option { + let categorized = ephemeral_5m + .unwrap_or_default() + .saturating_add(ephemeral_1h.unwrap_or_default()); + match uncategorized { + Some(0) if categorized > 0 => Some(categorized), + Some(value) => Some(value), + None if categorized > 0 => Some(categorized), + None => None, + } +} + +fn normalized_api_family(usage: &UpsertUsageRecord) -> String { + usage + .endpoint_api_format + .as_deref() + .or(usage.api_format.as_deref()) + .unwrap_or_default() + .split(':') + .next() + .unwrap_or_default() + .trim() + .to_ascii_lowercase() +} + +fn effective_input_tokens( + input_tokens: Option, + cache_creation_tokens: Option, + cache_read_tokens: Option, + api_family: &str, +) -> Option { + let input_tokens = input_tokens?; + let cache_creation_tokens = cache_creation_tokens.unwrap_or_default(); + let cache_read_tokens = cache_read_tokens.unwrap_or_default(); + if input_tokens > 0 { + if api_family == "openai" && (cache_creation_tokens > 0 || cache_read_tokens > 0) { + return Some( + input_tokens + .saturating_sub(cache_creation_tokens) + .saturating_sub(cache_read_tokens), + ); + } + if matches!(api_family, "gemini" | "google") && cache_read_tokens > 0 { + return Some(input_tokens.saturating_sub(cache_read_tokens)); + } + } + Some(input_tokens) +} + +fn total_input_context( + input_tokens: Option, + effective_input_tokens: Option, + cache_creation_tokens: Option, + cache_read_tokens: Option, + api_family: &str, +) -> Option { + if input_tokens.is_none() + && effective_input_tokens.is_none() + && cache_creation_tokens.is_none() + && cache_read_tokens.is_none() + { + return None; + } + let input_tokens = input_tokens.unwrap_or_default(); + let effective_input_tokens = effective_input_tokens.unwrap_or(input_tokens); + let cache_creation_tokens = cache_creation_tokens.unwrap_or_default(); + let cache_read_tokens = cache_read_tokens.unwrap_or_default(); + match api_family { + "claude" | "anthropic" => Some( + input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + "openai" => Some( + effective_input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + "gemini" | "google" => Some(effective_input_tokens.saturating_add(cache_read_tokens)), + _ => Some( + input_tokens + .saturating_add(cache_creation_tokens) + .saturating_add(cache_read_tokens), + ), + } +} + +fn json_text(value: Option<&Value>) -> Result, DataLayerError> { + value + .map(|value| { + serde_json::to_string(value).map_err(|error| { + DataLayerError::UnexpectedValue(format!( + "failed to serialize usage settlement snapshot: {error}" + )) + }) + }) + .transpose() +} + +fn json_value_from_row( + row: &sqlx::sqlite::SqliteRow, + column: &str, +) -> Result, DataLayerError> { + row.try_get::, _>(column) + .map_sql_err()? + .map(|value| { + serde_json::from_str(&value).map_err(|error| { + DataLayerError::UnexpectedValue(format!( + "invalid usage settlement JSON in {column}: {error}" + )) + }) + }) + .transpose() +} + +fn nonnegative_u64(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + +fn unix_now() -> Result { + let seconds = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))? + .as_secs(); + i64::try_from(seconds) + .map_err(|_| DataLayerError::UnexpectedValue("unix timestamp overflow".to_string())) +} diff --git a/crates/aether-data/adapters/sqlite/src/usage/tests.rs b/crates/aether-data/adapters/sqlite/src/usage/tests.rs index aebfb0c2d..b0782b8b6 100644 --- a/crates/aether-data/adapters/sqlite/src/usage/tests.rs +++ b/crates/aether-data/adapters/sqlite/src/usage/tests.rs @@ -1,10 +1,14 @@ use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository}; use crate::run_migrations; use aether_data_contracts::repository::usage::{ - UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery, + ProviderApiKeyWindowUsageRequest, UpsertUsageRecord, UsageAuditAggregationGroupBy, + UsageAuditAggregationQuery, UsageAuditListQuery, UsageAuditSummaryQuery, UsageBodyCaptureState, + UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCleanupExecutionMode, + UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageProviderPerformanceQuery, UsageReadRepository, UsageTimeSeriesGranularity, UsageWriteRepository, }; +use chrono::{DateTime, Utc}; #[test] fn sqlite_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() { @@ -106,7 +110,7 @@ async fn sqlite_provider_performance_can_skip_timeline() { } #[tokio::test] -async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() { +async fn sqlite_usage_write_repository_upserts_and_flushes_counter_deltas() { let pool = sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") @@ -125,7 +129,7 @@ async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() { assert_eq!(record.request_id, "request-1"); assert_eq!(record.api_key_id.as_deref(), Some("api-key-1")); - assert_eq!(record.total_tokens, 7); + assert_eq!(record.total_tokens, 5); assert_eq!(record.cache_read_input_tokens, 2); assert_eq!( record.request_metadata.as_ref().unwrap()["trace_id"], @@ -143,6 +147,11 @@ async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() { .expect("usage stream mode should load"); assert_eq!(upstream_is_stream, Some(1)); + repository + .flush_usage_counter_deltas(100) + .await + .expect("usage counter deltas should flush"); + let loaded = repository .find_by_request_id("request-1") .await @@ -159,7 +168,7 @@ async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() { .fetch_one(&pool) .await .expect("api key stats should load"); - assert_eq!(stats, (1, 7, 0.5, Some(1_000))); + assert_eq!(stats, (1, 5, 0.5, Some(1_000))); let provider_stats = sqlx::query_as::<_, (i64, i64, i64, i64, f64, i64, Option)>( "SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, total_response_time_ms, last_used_at FROM provider_api_keys WHERE id = 'provider-key-1'", @@ -167,7 +176,1011 @@ async fn sqlite_usage_write_repository_upserts_and_rebuilds_stats() { .fetch_one(&pool) .await .expect("provider key stats should load"); - assert_eq!(provider_stats, (1, 1, 0, 7, 0.5, 42, Some(1_000))); + assert_eq!(provider_stats, (1, 1, 0, 5, 0.5, 42, Some(1_000))); + let model_usage_count: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = 'model-1'") + .fetch_one(&pool) + .await + .expect("global model usage count should load"); + assert_eq!(model_usage_count, 1); + + repository + .upsert(sample_usage("request-1", "completed", "pending", 1_000)) + .await + .expect("identical terminal usage should remain idempotent"); + repository + .flush_usage_counter_deltas(100) + .await + .expect("idempotent counter flush should succeed"); + let repeated_stats = sqlx::query_as::<_, (i64, i64, f64)>( + "SELECT total_requests, total_tokens, total_cost_usd FROM api_keys WHERE id = 'api-key-1'", + ) + .fetch_one(&pool) + .await + .expect("repeated api key stats should load"); + assert_eq!(repeated_stats, (1, 5, 0.5)); + let repeated_provider_requests: i64 = sqlx::query_scalar( + "SELECT request_count FROM provider_api_keys WHERE id = 'provider-key-1'", + ) + .fetch_one(&pool) + .await + .expect("repeated provider stats should load"); + assert_eq!(repeated_provider_requests, 1); + let repeated_model_usage_count: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = 'model-1'") + .fetch_one(&pool) + .await + .expect("repeated global model usage count should load"); + assert_eq!(repeated_model_usage_count, 1); +} + +#[tokio::test] +async fn sqlite_usage_stats_rebuild_uses_canonical_terminal_totals() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let repository = SqliteUsageWriteRepository::new(pool.clone()); + + repository + .upsert(sample_usage( + "rebuild-completed", + "completed", + "pending", + 1_000, + )) + .await + .expect("completed usage should upsert"); + repository + .upsert(sample_usage("rebuild-pending", "pending", "pending", 2_000)) + .await + .expect("pending usage should upsert"); + sqlx::query( + r#" +UPDATE "usage" +SET total_tokens = 99 +WHERE request_id = 'rebuild-completed'; +UPDATE usage_settlement_snapshots +SET billing_effective_input_tokens = 11, + billing_output_tokens = 13, + billing_cache_creation_tokens = 2, + billing_cache_read_tokens = 3, + billing_total_input_context = NULL +WHERE request_id = 'rebuild-completed'; +"#, + ) + .execute(&pool) + .await + .expect("conflicting raw and settlement token totals should seed"); + + let rebuilt = repository + .rebuild_api_key_usage_stats() + .await + .expect("api key stats should rebuild"); + let stats: (i64, i64, f64, Option) = sqlx::query_as( + "SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = 'api-key-1'", + ) + .fetch_one(&pool) + .await + .expect("rebuilt api key stats should load"); + + assert_eq!(rebuilt, 1); + assert_eq!(stats, (1, 29, 0.5, Some(1_000))); + + let provider_rebuilt = repository + .rebuild_provider_api_key_usage_stats() + .await + .expect("provider api key stats should rebuild"); + let provider_stats: (i64, i64, i64, i64, f64, Option) = sqlx::query_as( + "SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, last_used_at FROM provider_api_keys WHERE id = 'provider-key-1'", + ) + .fetch_one(&pool) + .await + .expect("rebuilt provider api key stats should load"); + + assert_eq!(provider_rebuilt, 1); + assert_eq!(provider_stats, (2, 1, 0, 29, 0.5, Some(2_000))); + + let reader = SqliteUsageReadRepository::new(pool); + let user_totals = reader + .summarize_usage_totals_by_user_ids(&["user-1".to_string()]) + .await + .expect("user totals should load"); + assert_eq!(user_totals[0].request_count, 1); + assert_eq!(user_totals[0].total_tokens, 29); + + let api_key_totals = reader + .summarize_total_tokens_by_api_key_ids(&["api-key-1".to_string()]) + .await + .expect("api key totals should load"); + assert_eq!(api_key_totals["api-key-1"], 34); + + let provider_key_totals = reader + .summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()]) + .await + .expect("provider key totals should load"); + assert_eq!(provider_key_totals["provider-key-1"].total_tokens, 34); + + let provider_window = reader + .summarize_usage_by_provider_api_key_windows(&[ProviderApiKeyWindowUsageRequest { + provider_api_key_id: "provider-key-1".to_string(), + window_code: "test".to_string(), + start_unix_secs: 0, + end_unix_secs: 3_000, + }]) + .await + .expect("provider key window should load"); + assert_eq!(provider_window[0].total_tokens, 34); + + let audit_summary = reader + .summarize_usage_audits(&UsageAuditSummaryQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 3_000, + ..UsageAuditSummaryQuery::default() + }) + .await + .expect("usage audit summary should load"); + assert_eq!(audit_summary.recorded_total_tokens, 34); + + let aggregation = reader + .aggregate_usage_audits(&UsageAuditAggregationQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 3_000, + group_by: UsageAuditAggregationGroupBy::Model, + limit: 10, + exclude_reserved_provider_labels: false, + }) + .await + .expect("usage audit aggregation should load"); + assert_eq!(aggregation[0].total_tokens, 29); + + let breakdown = reader + .summarize_usage_breakdown(&UsageBreakdownSummaryQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 3_000, + group_by: UsageBreakdownGroupBy::Model, + ..UsageBreakdownSummaryQuery::default() + }) + .await + .expect("usage breakdown should load"); + assert_eq!(breakdown[0].total_tokens, 29); + + let daily = reader + .list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 3_000, + tz_offset_minutes: 0, + user_id: Some("user-1".to_string()), + }) + .await + .expect("dashboard daily breakdown should load"); + assert_eq!(daily[0].total_tokens, 29); +} + +#[tokio::test] +async fn sqlite_usage_http_capture_round_trips_and_preserves_sparse_updates() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let writer = SqliteUsageWriteRepository::new(pool.clone()); + + let mut rich = sample_usage("canonical-capture", "pending", "pending", 1_000); + rich.request_headers = Some(serde_json::json!({"x-client": "one"})); + rich.provider_request_headers = Some(serde_json::json!({"x-provider": "two"})); + rich.response_headers = Some(serde_json::json!({"x-upstream": "three"})); + rich.client_response_headers = Some(serde_json::json!({"x-response": "four"})); + rich.request_body = Some(serde_json::json!({"request": 1})); + rich.provider_request_body = Some(serde_json::json!({"provider_request": 2})); + rich.response_body = Some(serde_json::json!({"response": 3})); + rich.client_response_body = Some(serde_json::json!({"client_response": 4})); + rich.request_body_state = Some(UsageBodyCaptureState::Inline); + rich.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + rich.response_body_state = Some(UsageBodyCaptureState::Inline); + rich.client_response_body_state = Some(UsageBodyCaptureState::Inline); + rich.request_metadata = Some(serde_json::json!({ + "trace_id": "canonical-trace", + "request_body_ref": "usage://request/stale/request_body" + })); + + let stored = writer + .upsert(rich) + .await + .expect("canonical capture should upsert"); + assert_eq!( + stored.request_headers, + Some(serde_json::json!({"x-client": "one"})) + ); + assert_eq!(stored.request_body, Some(serde_json::json!({"request": 1}))); + assert_eq!( + stored.provider_request_body, + Some(serde_json::json!({"provider_request": 2})) + ); + assert_eq!( + stored.response_body, + Some(serde_json::json!({"response": 3})) + ); + assert_eq!( + stored.client_response_body, + Some(serde_json::json!({"client_response": 4})) + ); + assert_eq!( + stored.request_body_state, + Some(UsageBodyCaptureState::Reference) + ); + assert_eq!( + stored.request_body_ref.as_deref(), + Some("usage://request/canonical-capture/request_body") + ); + assert_eq!( + stored.request_metadata.as_ref().unwrap()["trace_id"], + "canonical-trace" + ); + assert!(stored + .request_metadata + .as_ref() + .unwrap() + .get("request_body_ref") + .is_none()); + + let legacy_columns: (Option, Option, Option>) = sqlx::query_as( + "SELECT request_headers, request_body, request_body_compressed FROM \"usage\" WHERE request_id = 'canonical-capture'", + ) + .fetch_one(&pool) + .await + .expect("legacy columns should load"); + assert_eq!(legacy_columns, (None, None, None)); + let audit: (String, String, String) = sqlx::query_as( + "SELECT request_headers, request_body_ref, request_body_state FROM usage_http_audits WHERE request_id = 'canonical-capture'", + ) + .fetch_one(&pool) + .await + .expect("canonical audit should load"); + assert_eq!( + serde_json::from_str::(&audit.0).expect("header JSON should decode"), + serde_json::json!({"x-client": "one"}) + ); + assert_eq!(audit.1, "usage://request/canonical-capture/request_body"); + assert_eq!(audit.2, "reference"); + let blob_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'canonical-capture'", + ) + .fetch_one(&pool) + .await + .expect("canonical blobs should count"); + assert_eq!(blob_count, 4); + + let sparse = sample_usage("canonical-capture", "streaming", "pending", 1_001); + let sparse_stored = writer + .upsert(sparse) + .await + .expect("sparse lifecycle update should upsert"); + assert_eq!(sparse_stored.request_headers, stored.request_headers); + assert_eq!(sparse_stored.request_body, stored.request_body); + assert_eq!(sparse_stored.response_body, stored.response_body); + let sparse_blob_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'canonical-capture'", + ) + .fetch_one(&pool) + .await + .expect("preserved blobs should count"); + assert_eq!(sparse_blob_count, 4); + + let mut clear = sample_usage("canonical-capture", "streaming", "pending", 1_002); + clear.request_body = Some(serde_json::json!({"residual": true})); + clear.request_body_ref = Some("usage://request/canonical-capture/request_body".to_string()); + clear.request_body_state = Some(UsageBodyCaptureState::None); + let cleared = writer + .upsert(clear) + .await + .expect("explicit none capture should clear"); + assert!(cleared.request_body.is_none()); + assert!(cleared.request_body_ref.is_none()); + assert_eq!( + cleared.request_body_state, + Some(UsageBodyCaptureState::None) + ); + assert_eq!(cleared.provider_request_body, stored.provider_request_body); + let cleared_blob_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'canonical-capture'", + ) + .fetch_one(&pool) + .await + .expect("remaining blobs should count"); + assert_eq!(cleared_blob_count, 3); + + let reader = SqliteUsageReadRepository::new(pool.clone()); + let resolved = reader + .resolve_body_ref("usage://request/canonical-capture/provider_request_body") + .await + .expect("body ref should resolve"); + assert_eq!(resolved, stored.provider_request_body); + let loaded = reader + .find_by_request_id("canonical-capture") + .await + .expect("canonical usage should load") + .expect("canonical usage should exist"); + assert_eq!( + loaded.provider_request_headers, + stored.provider_request_headers + ); + assert_eq!(loaded.provider_request_body, stored.provider_request_body); +} + +#[tokio::test] +async fn sqlite_usage_http_read_falls_back_to_legacy_inline_and_compressed_columns() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let writer = SqliteUsageWriteRepository::new(pool.clone()); + let mut captured = sample_usage("legacy-capture", "pending", "pending", 2_000); + captured.request_body = Some(serde_json::json!({"compressed": true})); + writer + .upsert(captured) + .await + .expect("temporary canonical body should upsert"); + let payload: Vec = sqlx::query_scalar( + "SELECT payload_gzip FROM usage_body_blobs WHERE request_id = 'legacy-capture' AND body_field = 'request_body'", + ) + .fetch_one(&pool) + .await + .expect("temporary gzip should load"); + sqlx::query( + r#" +DELETE FROM usage_http_audits WHERE request_id = 'legacy-capture'; +DELETE FROM usage_body_blobs WHERE request_id = 'legacy-capture'; +UPDATE "usage" +SET request_headers = '{"legacy":true}', + request_body_compressed = ?, + response_body = '{"inline":true}', + request_metadata = '{"request_body_ref":"usage://request/legacy-capture/request_body"}' +WHERE request_id = 'legacy-capture'; +"#, + ) + .bind(payload) + .execute(&pool) + .await + .expect("legacy capture should seed"); + + let reader = SqliteUsageReadRepository::new(pool.clone()); + let loaded = reader + .find_by_request_id("legacy-capture") + .await + .expect("legacy usage should load") + .expect("legacy usage should exist"); + assert_eq!( + loaded.request_headers, + Some(serde_json::json!({"legacy": true})) + ); + assert_eq!( + loaded.request_body, + Some(serde_json::json!({"compressed": true})) + ); + assert_eq!( + loaded.response_body, + Some(serde_json::json!({"inline": true})) + ); + assert_eq!( + loaded.request_body_ref.as_deref(), + Some("usage://request/legacy-capture/request_body") + ); + assert!(loaded.request_body_state.is_none()); + + let mut clear = sample_usage("legacy-capture", "streaming", "pending", 2_001); + clear.request_metadata = None; + clear.request_body_state = Some(UsageBodyCaptureState::None); + let cleared = writer + .upsert(clear) + .await + .expect("explicit none should clear legacy fallback storage"); + assert!(cleared.request_body.is_none()); + assert!(cleared.request_body_ref.is_none()); + assert_eq!( + cleared.request_body_state, + Some(UsageBodyCaptureState::None) + ); + assert!(cleared + .request_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .is_some_and(|metadata| !metadata.contains_key("request_body_ref"))); + let compressed_after_clear: Option> = sqlx::query_scalar( + "SELECT request_body_compressed FROM \"usage\" WHERE request_id = 'legacy-capture'", + ) + .fetch_one(&pool) + .await + .expect("legacy compressed body should load after clear"); + assert!(compressed_after_clear.is_none()); +} + +#[tokio::test] +async fn sqlite_usage_canonical_snapshots_round_trip_preserve_sparse_and_clear_terminal() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let writer = SqliteUsageWriteRepository::new(pool.clone()); + let reader = SqliteUsageReadRepository::new(pool.clone()); + + let mut rich = sample_usage("canonical-snapshots", "streaming", "pending", 1_000); + rich.candidate_id = None; + rich.candidate_index = None; + rich.key_name = None; + rich.planner_kind = None; + rich.route_family = None; + rich.route_kind = None; + rich.execution_path = None; + rich.input_tokens = Some(1_000); + rich.output_tokens = Some(100); + rich.cache_creation_input_tokens = Some(0); + rich.cache_creation_ephemeral_5m_input_tokens = Some(100); + rich.cache_creation_ephemeral_1h_input_tokens = Some(50); + rich.cache_read_input_tokens = Some(200); + rich.request_metadata = Some(serde_json::json!({ + "trace_id": "rich", + "candidate_id": "candidate-canonical", + "candidate_index": 4, + "key_name": "key-canonical", + "planner_kind": "fallback", + "route_family": "chat", + "route_kind": "remote", + "execution_path": "converted", + "local_execution_runtime_miss_reason": "runtime_busy", + "billing_snapshot_schema_version": "v3", + "billing_snapshot_status": "resolved", + "rate_multiplier": 0.75, + "is_free_tier": false, + "input_price_per_1m": 1.1, + "output_price_per_1m": 2.2, + "cache_creation_price_per_1m": 3.3, + "cache_read_price_per_1m": 4.4, + "price_per_request": 0.05, + "billing_dimensions": { + "input_tokens": 1000, + "effective_input_tokens": 650, + "output_tokens": 100, + "cache_creation_tokens": 150, + "cache_creation_ephemeral_5m_tokens": 100, + "cache_creation_ephemeral_1h_tokens": 50, + "cache_read_tokens": 200, + "total_input_context": 1000 + }, + "settlement_snapshot": { + "schema_version": "v3", + "total_cost": 1.25, + "actual_total_cost": 1.0, + "cost_breakdown": { + "cache_creation_ephemeral_5m_cost": 0.02, + "cache_creation_ephemeral_1h_cost": 0.03, + "cache_read_cost": 0.04 + }, + "pricing_snapshot": {"pricing_source": "catalog"}, + "billing_plan_snapshot": {"rule_id": "rule-1", "rule_version": "7"} + } + })); + let stored = writer + .upsert(rich) + .await + .expect("canonical snapshots should upsert"); + assert_eq!(stored.routing_candidate_id(), Some("candidate-canonical")); + assert_eq!(stored.routing_candidate_index(), Some(4)); + assert_eq!(stored.provider_id.as_deref(), Some("provider-1")); + assert_eq!(stored.output_tokens, 100); + assert_eq!(stored.cache_creation_input_tokens, 150); + assert_eq!(stored.cache_read_input_tokens, 200); + assert_eq!(stored.total_tokens, 1_100); + assert_eq!(stored.total_cost_usd, 1.25); + assert_eq!(stored.actual_total_cost_usd, 1.0); + assert_eq!(stored.cache_creation_cost_usd, 0.05); + assert_eq!(stored.cache_read_cost_usd, 0.04); + assert_eq!(stored.settlement_rate_multiplier(), Some(0.75)); + assert_eq!(stored.settlement_input_price_per_1m(), Some(1.1)); + assert_eq!(stored.settlement_output_price_per_1m(), Some(2.2)); + assert_eq!(stored.settlement_price_per_request(), Some(0.05)); + + sqlx::query( + r#" +UPDATE usage_settlement_snapshots +SET wallet_id = 'wallet-sentinel', + wallet_balance_before = 10, + wallet_balance_after = 9, + provider_monthly_used_usd = 8, + finalized_at = 2000 +WHERE request_id = 'canonical-snapshots' +"#, + ) + .execute(&pool) + .await + .expect("wallet settlement facts should seed"); + + let mut sparse = sample_usage("canonical-snapshots", "streaming", "pending", 1_001); + sparse.provider_id = None; + sparse.provider_endpoint_id = None; + sparse.provider_api_key_id = None; + sparse.has_format_conversion = None; + sparse.candidate_id = None; + sparse.candidate_index = None; + sparse.key_name = None; + sparse.planner_kind = None; + sparse.route_family = None; + sparse.route_kind = None; + sparse.execution_path = None; + sparse.local_execution_runtime_miss_reason = None; + sparse.input_tokens = None; + sparse.output_tokens = None; + sparse.total_tokens = None; + sparse.cache_creation_input_tokens = None; + sparse.cache_creation_ephemeral_5m_input_tokens = None; + sparse.cache_creation_ephemeral_1h_input_tokens = None; + sparse.cache_read_input_tokens = None; + sparse.cache_creation_cost_usd = None; + sparse.cache_read_cost_usd = None; + sparse.output_price_per_1m = None; + sparse.total_cost_usd = None; + sparse.actual_total_cost_usd = None; + sparse.request_metadata = Some(serde_json::json!({"trace_id": "sparse"})); + let sparse_stored = writer + .upsert(sparse) + .await + .expect("sparse snapshots should merge"); + assert_eq!( + sparse_stored.routing_candidate_id(), + Some("candidate-canonical") + ); + assert_eq!(sparse_stored.provider_id.as_deref(), Some("provider-1")); + assert_eq!(sparse_stored.output_tokens, 100); + assert_eq!(sparse_stored.total_tokens, 1_100); + assert_eq!(sparse_stored.total_cost_usd, 1.25); + assert_eq!(sparse_stored.trace_id(), Some("sparse")); + assert_eq!(sparse_stored.settlement_rate_multiplier(), Some(0.75)); + + sqlx::query( + r#" +UPDATE "usage" +SET candidate_id = 'legacy-candidate', + route_family = 'legacy-route', + total_cost_usd = 99, + output_price_per_1m = 99, + request_metadata = '{"trace_id":"legacy","rate_multiplier":9}' +WHERE request_id = 'canonical-snapshots' +"#, + ) + .execute(&pool) + .await + .expect("legacy mirrors should be corruptible for precedence test"); + let canonical = reader + .find_by_request_id("canonical-snapshots") + .await + .expect("canonical usage should load") + .expect("canonical usage should exist"); + assert_eq!( + canonical.routing_candidate_id(), + Some("candidate-canonical") + ); + assert_eq!(canonical.routing_route_family(), Some("chat")); + assert_eq!(canonical.total_cost_usd, 1.25); + assert_eq!(canonical.settlement_output_price_per_1m(), Some(2.2)); + assert_eq!(canonical.settlement_rate_multiplier(), Some(0.75)); + + let mut terminal = sample_usage("canonical-snapshots", "completed", "settled", 1_002); + terminal.provider_id = None; + terminal.provider_endpoint_id = None; + terminal.provider_api_key_id = None; + terminal.has_format_conversion = None; + terminal.candidate_id = None; + terminal.candidate_index = None; + terminal.key_name = None; + terminal.planner_kind = None; + terminal.route_family = None; + terminal.route_kind = None; + terminal.execution_path = None; + terminal.local_execution_runtime_miss_reason = None; + terminal.input_tokens = None; + terminal.output_tokens = None; + terminal.total_tokens = None; + terminal.cache_creation_input_tokens = None; + terminal.cache_creation_ephemeral_5m_input_tokens = None; + terminal.cache_creation_ephemeral_1h_input_tokens = None; + terminal.cache_read_input_tokens = None; + terminal.cache_creation_cost_usd = None; + terminal.cache_read_cost_usd = None; + terminal.output_price_per_1m = None; + terminal.total_cost_usd = None; + terminal.actual_total_cost_usd = None; + terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal"})); + let terminal_stored = writer + .upsert(terminal) + .await + .expect("terminal snapshots should replace"); + assert_eq!(terminal_stored.status, "completed"); + assert_eq!(terminal_stored.billing_status, "settled"); + assert!(terminal_stored.candidate_id.is_none()); + assert!(terminal_stored.provider_id.is_none()); + assert_eq!(terminal_stored.total_tokens, 0); + assert_eq!(terminal_stored.total_cost_usd, 0.0); + assert_eq!(terminal_stored.settlement_rate_multiplier(), None); + assert_eq!(terminal_stored.trace_id(), Some("terminal")); + + let terminal_row: (Option, Option, String, Option) = sqlx::query_as( + r#" +SELECT routing.candidate_id, settlement.settlement_snapshot, + settlement.billing_status, settlement.billing_total_cost_usd +FROM usage_routing_snapshots routing +JOIN usage_settlement_snapshots settlement USING (request_id) +WHERE routing.request_id = 'canonical-snapshots' +"#, + ) + .fetch_one(&pool) + .await + .expect("terminal canonical rows should load"); + assert_eq!(terminal_row, (None, None, "settled".to_string(), None)); + let wallet_row = sqlx::query( + r#" +SELECT wallet_id, wallet_balance_before, wallet_balance_after, + provider_monthly_used_usd, finalized_at +FROM usage_settlement_snapshots +WHERE request_id = 'canonical-snapshots' +"#, + ) + .fetch_one(&pool) + .await + .expect("wallet settlement facts should load"); + assert_eq!( + sqlx::Row::try_get::, _>(&wallet_row, "wallet_id") + .expect("wallet id should decode"), + Some("wallet-sentinel".to_string()) + ); + assert_eq!( + sqlx::Row::try_get::, _>(&wallet_row, "wallet_balance_before") + .expect("wallet balance before should decode"), + Some(10.0) + ); + assert_eq!( + sqlx::Row::try_get::, _>(&wallet_row, "wallet_balance_after") + .expect("wallet balance after should decode"), + Some(9.0) + ); + assert_eq!( + sqlx::Row::try_get::, _>(&wallet_row, "provider_monthly_used_usd") + .expect("provider monthly usage should decode"), + Some(8.0) + ); + assert_eq!( + sqlx::Row::try_get::, _>(&wallet_row, "finalized_at") + .expect("settlement finalized_at should decode"), + Some(2_000) + ); + assert_eq!(terminal_stored.finalized_at_unix_secs, Some(2_000)); + + let mut late = sample_usage("canonical-snapshots", "pending", "pending", 1_003); + late.candidate_id = Some("late-candidate".to_string()); + late.request_metadata = Some(serde_json::json!({ + "trace_id": "late", + "rate_multiplier": 9, + "settlement_snapshot": {"schema_version": "late", "total_cost": 99} + })); + let after_late = writer + .upsert(late) + .await + .expect("late pending usage should return terminal record"); + assert_eq!(after_late.status, "completed"); + assert_eq!(after_late.billing_status, "settled"); + assert!(after_late.candidate_id.is_none()); + assert_eq!(after_late.total_cost_usd, 0.0); + assert_eq!(after_late.trace_id(), Some("terminal")); +} + +#[tokio::test] +async fn sqlite_usage_cleanup_matches_policy_windows_and_targets() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let repository = SqliteUsageWriteRepository::new(pool.clone()); + + for (request_id, created_at) in [ + ("cleanup-log", 10), + ("cleanup-stale-body", 30), + ("cleanup-header", 50), + ("cleanup-detail", 70), + ("cleanup-legacy", 75), + ("cleanup-new", 90), + ] { + repository + .upsert(sample_usage(request_id, "completed", "settled", created_at)) + .await + .expect("usage should seed"); + } + + sqlx::query( + r#" +UPDATE "usage" +SET request_headers = '{"old":true}', request_body = '{"delete":true}' +WHERE request_id = 'cleanup-log'; +UPDATE "usage" +SET request_headers = '{"stale":true}', request_body_compressed = X'1F8B' +WHERE request_id = 'cleanup-stale-body'; +UPDATE "usage" +SET response_headers = '{"header":true}' +WHERE request_id = 'cleanup-header'; +UPDATE "usage" +SET request_body = '{"detail":true}' +WHERE request_id = 'cleanup-detail'; +UPDATE "usage" +SET request_metadata = '{"trace":"kept","request_body_ref":"usage://request/cleanup-legacy/request_body"}' +WHERE request_id = 'cleanup-legacy'; +UPDATE "usage" +SET request_headers = '{"new":true}', request_body = '{"new":true}' +WHERE request_id = 'cleanup-new'; + +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES ('usage-body://cleanup-stale-body/request_body', 'cleanup-stale-body', 'request_body', X'1F8B'); +INSERT INTO usage_http_audits ( + request_id, request_headers, request_body_ref, body_capture_mode +) +VALUES ( + 'cleanup-stale-body', '{"audit":true}', + 'usage-body://cleanup-stale-body/request_body', 'ref_backed' +); + +INSERT INTO api_keys ( + id, user_id, key_hash, is_active, auto_delete_on_expiry, + expires_at, created_at, updated_at +) +VALUES + ('cleanup-disable-key', 'user-1', 'cleanup-disable-hash', 1, 0, 1, 1, 1), + ('cleanup-delete-key', 'user-1', 'cleanup-delete-hash', 1, 1, 1, 1, 1); +INSERT INTO wallets ( + id, api_key_id, balance, gift_balance, limit_mode, currency, status, + created_at, updated_at +) +VALUES ( + 'cleanup-delete-wallet', 'cleanup-delete-key', 0, 0, 'finite', 'USD', 'active', 1, 1 +); +"#, + ) + .execute(&pool) + .await + .expect("cleanup fixtures should seed"); + + let window = cleanup_window(80, 60, 70, 20); + let preview = repository + .preview_usage_cleanup( + &window, + UsageCleanupTargets::all_policy_targets(), + UsageCleanupExecutionMode::Policy, + ) + .await + .expect("cleanup preview should load"); + assert_eq!(preview.detail, 2); + assert_eq!(preview.compressed, 1); + assert_eq!(preview.header, 2); + assert_eq!(preview.log, 1); + + let summary = repository + .cleanup_usage( + &window, + 1, + true, + UsageCleanupTargets::all_policy_targets(), + UsageCleanupExecutionMode::Policy, + ) + .await + .expect("usage cleanup should succeed"); + assert_eq!(summary.records_deleted, 1); + assert_eq!(summary.header_cleaned, 2); + assert_eq!(summary.body_cleaned, 1); + assert_eq!(summary.legacy_body_refs_migrated, 1); + assert_eq!(summary.body_externalized, 1); + assert_eq!(summary.keys_cleaned, 2); + + let deleted_log: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM \"usage\" WHERE request_id = 'cleanup-log'") + .fetch_one(&pool) + .await + .expect("deleted usage should count"); + assert_eq!(deleted_log, 0); + let stale_fields: (Option, Option>) = sqlx::query_as( + "SELECT request_headers, request_body_compressed FROM \"usage\" WHERE request_id = 'cleanup-stale-body'", + ) + .fetch_one(&pool) + .await + .expect("stale usage should load"); + assert_eq!(stale_fields, (None, None)); + let stale_blobs: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'cleanup-stale-body'", + ) + .fetch_one(&pool) + .await + .expect("stale blobs should count"); + assert_eq!(stale_blobs, 0); + + let detail_blob: Vec = sqlx::query_scalar( + "SELECT payload_gzip FROM usage_body_blobs WHERE request_id = 'cleanup-detail'", + ) + .fetch_one(&pool) + .await + .expect("externalized body should load"); + assert_eq!( + super::inflate_usage_json_value(&detail_blob).expect("body gzip should decode"), + serde_json::json!({"detail": true}) + ); + let detail_inline: Option = sqlx::query_scalar( + "SELECT request_body FROM \"usage\" WHERE request_id = 'cleanup-detail'", + ) + .fetch_one(&pool) + .await + .expect("detail inline body should load"); + assert!(detail_inline.is_none()); + + let legacy_metadata: String = sqlx::query_scalar( + "SELECT request_metadata FROM \"usage\" WHERE request_id = 'cleanup-legacy'", + ) + .fetch_one(&pool) + .await + .expect("legacy metadata should load"); + assert_eq!( + serde_json::from_str::(&legacy_metadata).expect("valid metadata"), + serde_json::json!({"trace": "kept"}) + ); + let legacy_ref: String = sqlx::query_scalar( + "SELECT request_body_ref FROM usage_http_audits WHERE request_id = 'cleanup-legacy'", + ) + .fetch_one(&pool) + .await + .expect("legacy ref should migrate"); + assert_eq!( + legacy_ref, + "usage://request/cleanup-legacy/request_body".to_string() + ); + + let disabled_key: i64 = + sqlx::query_scalar("SELECT is_active FROM api_keys WHERE id = 'cleanup-disable-key'") + .fetch_one(&pool) + .await + .expect("disabled key should remain"); + assert_eq!(disabled_key, 0); + let deleted_key: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM api_keys WHERE id = 'cleanup-delete-key'") + .fetch_one(&pool) + .await + .expect("deleted key should count"); + assert_eq!(deleted_key, 0); + let wallet_status: String = + sqlx::query_scalar("SELECT status FROM wallets WHERE id = 'cleanup-delete-wallet'") + .fetch_one(&pool) + .await + .expect("expired key wallet should load"); + assert_eq!(wallet_status, "disabled"); + + let new_fields: (Option, Option) = sqlx::query_as( + "SELECT request_headers, request_body FROM \"usage\" WHERE request_id = 'cleanup-new'", + ) + .fetch_one(&pool) + .await + .expect("new usage should load"); + assert!(new_fields.0.is_some()); + assert!(new_fields.1.is_some()); +} + +#[tokio::test] +async fn sqlite_usage_cleanup_before_now_only_clears_selected_body_fields() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let repository = SqliteUsageWriteRepository::new(pool.clone()); + repository + .upsert(sample_usage( + "cleanup-before-now", + "completed", + "settled", + 10, + )) + .await + .expect("usage should seed"); + sqlx::query( + r#" +UPDATE "usage" +SET request_headers = '{"keep":true}', + request_body = '{"raw":true}', + request_body_compressed = X'1F8B' +WHERE request_id = 'cleanup-before-now'; +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) +VALUES ( + 'usage-body://cleanup-before-now/request_body', + 'cleanup-before-now', 'request_body', X'1F8B' +); +INSERT INTO usage_http_audits ( + request_id, request_headers, request_body_ref, body_capture_mode +) +VALUES ( + 'cleanup-before-now', '{"keep":true}', + 'usage-body://cleanup-before-now/request_body', 'ref_backed' +); +"#, + ) + .execute(&pool) + .await + .expect("body fixtures should seed"); + + let window = cleanup_window(20, 20, 20, 20); + let preview = repository + .preview_usage_cleanup( + &window, + UsageCleanupTargets::all_policy_targets(), + UsageCleanupExecutionMode::BeforeNowBodyFields, + ) + .await + .expect("cleanup preview should load"); + assert_eq!(preview.detail, 1); + assert_eq!(preview.compressed, 1); + assert_eq!(preview.header, 0); + assert_eq!(preview.log, 0); + + let summary = repository + .cleanup_usage( + &window, + 1, + false, + UsageCleanupTargets::all_policy_targets(), + UsageCleanupExecutionMode::BeforeNowBodyFields, + ) + .await + .expect("before-now cleanup should succeed"); + assert_eq!(summary.body_externalized, 1); + assert_eq!(summary.body_cleaned, 1); + assert_eq!(summary.header_cleaned, 0); + assert_eq!(summary.records_deleted, 0); + assert_eq!(summary.keys_cleaned, 0); + + let fields: (Option, Option>, Option) = sqlx::query_as( + "SELECT request_body, request_body_compressed, request_headers FROM \"usage\" WHERE request_id = 'cleanup-before-now'", + ) + .fetch_one(&pool) + .await + .expect("cleaned usage should load"); + assert!(fields.0.is_none()); + assert!(fields.1.is_none()); + assert!(fields.2.is_some()); + let audit_headers: Option = sqlx::query_scalar( + "SELECT request_headers FROM usage_http_audits WHERE request_id = 'cleanup-before-now'", + ) + .fetch_one(&pool) + .await + .expect("audit headers should remain"); + assert!(audit_headers.is_some()); } #[tokio::test] @@ -299,7 +1312,7 @@ async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late assert_eq!(current.status, "completed"); assert_eq!(current.billing_status, "pending"); - assert_eq!(current.total_tokens, 7); + assert_eq!(current.total_tokens, 5); assert_eq!(current.cache_read_input_tokens, 2); assert_eq!(current.total_cost_usd, 0.5); assert_eq!(current.actual_total_cost_usd, 0.4); @@ -686,6 +1699,16 @@ async fn sqlite_usage_read_repository_reads_usage_contract_views() { .upsert(sample_usage("request-2", "failed", "void", 1_010)) .await .expect("usage should upsert"); + sqlx::query( + r#" +UPDATE "usage" +SET username = 'legacy-user', api_key_name = 'legacy-key' +WHERE request_id = 'request-1' +"#, + ) + .execute(&pool) + .await + .expect("legacy display fields should seed"); let reader = SqliteUsageReadRepository::new(pool); let loaded = reader @@ -693,8 +1716,10 @@ async fn sqlite_usage_read_repository_reads_usage_contract_views() { .await .expect("usage should load") .expect("usage should exist"); - assert_eq!(loaded.total_tokens, 7); + assert_eq!(loaded.total_tokens, 5); assert_eq!(loaded.billing_status, "settled"); + assert_eq!(loaded.username.as_deref(), Some("legacy-user")); + assert_eq!(loaded.api_key_name.as_deref(), Some("legacy-key")); let listed = reader .list_usage_audits(&UsageAuditListQuery { @@ -883,6 +1908,331 @@ INSERT INTO stats_daily ( assert_eq!(rows[0].total_tokens, 37); } +#[tokio::test] +async fn sqlite_first_byte_fast_path_preserves_lifecycle_state_and_counters() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let repository = SqliteUsageWriteRepository::new(pool.clone()); + + assert!(repository.supports_first_byte_usage_fast_path()); + assert!(repository.supports_first_byte_usage_batch()); + + for request_id in ["first-byte-duplicate", "first-byte-unique"] { + let mut pending = sample_usage(request_id, "pending", "pending", 1_000); + pending.status_code = None; + pending.response_time_ms = None; + pending.first_byte_time_ms = None; + pending.finalized_at_unix_secs = None; + pending.is_stream = Some(true); + pending.request_body = Some(serde_json::json!({"prompt": request_id})); + pending.request_body_state = Some(UsageBodyCaptureState::Inline); + pending.request_metadata = Some(serde_json::json!({ + "trace_id": format!("pending-{request_id}"), + "upstream_is_stream": false + })); + repository + .upsert(pending) + .await + .expect("pending usage should seed"); + } + + let mut terminal = sample_usage("first-byte-terminal", "completed", "pending", 1_000); + terminal.first_byte_time_ms = Some(44); + terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal"})); + repository + .upsert(terminal) + .await + .expect("terminal usage should seed"); + + let mut first = sample_usage("first-byte-duplicate", "streaming", "pending", 1_001); + first.first_byte_time_ms = Some(30); + first.response_time_ms = Some(31); + first.finalized_at_unix_secs = None; + first.request_metadata = Some(serde_json::json!({"trace_id": "incoming-first"})); + let mut replay = first.clone(); + replay.first_byte_time_ms = Some(7); + replay.response_time_ms = Some(99); + + let mut unique = sample_usage("first-byte-unique", "streaming", "pending", 1_001); + unique.first_byte_time_ms = Some(18); + unique.finalized_at_unix_secs = None; + unique.request_metadata = Some(serde_json::json!({"trace_id": "incoming-unique"})); + + let mut late_terminal = sample_usage("first-byte-terminal", "streaming", "pending", 1_002); + late_terminal.first_byte_time_ms = Some(3); + late_terminal.finalized_at_unix_secs = None; + late_terminal.request_metadata = Some(serde_json::json!({"trace_id": "late"})); + + let mut missing = sample_usage("first-byte-missing", "streaming", "pending", 1_001); + missing.first_byte_time_ms = Some(12); + missing.finalized_at_unix_secs = None; + missing.request_metadata = Some(serde_json::json!({ + "trace_id": "missing", + "upstream_is_stream": false + })); + + repository + .upsert_first_byte_many(vec![first, unique, late_terminal, replay, missing]) + .await + .expect("first-byte batch should persist"); + + let duplicate = repository + .find_by_request_id("first-byte-duplicate") + .await + .expect("duplicate request should load") + .expect("duplicate request should exist"); + assert_eq!(duplicate.status, "streaming"); + assert_eq!(duplicate.first_byte_time_ms, Some(30)); + assert_eq!(duplicate.response_time_ms, Some(99)); + assert_eq!( + duplicate.request_metadata.as_ref().unwrap()["trace_id"], + "pending-first-byte-duplicate" + ); + assert_eq!( + duplicate.request_body, + Some(serde_json::json!({"prompt": "first-byte-duplicate"})) + ); + + let unique = repository + .find_by_request_id("first-byte-unique") + .await + .expect("unique request should load") + .expect("unique request should exist"); + assert_eq!(unique.status, "streaming"); + assert_eq!(unique.first_byte_time_ms, Some(18)); + assert_eq!( + unique.request_metadata.as_ref().unwrap()["trace_id"], + "pending-first-byte-unique" + ); + + let terminal = repository + .find_by_request_id("first-byte-terminal") + .await + .expect("terminal request should load") + .expect("terminal request should exist"); + assert_eq!(terminal.status, "completed"); + assert_eq!(terminal.first_byte_time_ms, Some(44)); + assert_eq!( + terminal.request_metadata.as_ref().unwrap()["trace_id"], + "terminal" + ); + + let missing = repository + .find_by_request_id("first-byte-missing") + .await + .expect("missing request should load") + .expect("missing request should have been inserted"); + assert_eq!(missing.status, "streaming"); + assert_eq!(missing.billing_status, "pending"); + assert_eq!(missing.first_byte_time_ms, Some(12)); + assert_eq!( + missing.request_metadata.as_ref().unwrap()["upstream_is_stream"], + false + ); + + let missing_counter_delta: i64 = sqlx::query_scalar( + r#" +SELECT COALESCE(SUM(request_count_delta), 0) +FROM usage_counter_deltas +WHERE request_id = 'first-byte-missing' + AND kind = 'provider_api_key' + AND target_id = 'provider-key-1' +"#, + ) + .fetch_one(&pool) + .await + .expect("missing first-byte counter delta should load"); + assert_eq!(missing_counter_delta, 1); +} + +#[tokio::test] +async fn sqlite_pending_batch_is_atomic_and_persists_auxiliary_state() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + let repository = SqliteUsageWriteRepository::new(pool.clone()); + assert!(repository.supports_pending_usage_batch()); + + let mut first = sample_usage("pending-batch-first", "pending", "pending", 1_000); + first.finalized_at_unix_secs = None; + first.request_headers = Some(serde_json::json!({"x-request": "first"})); + first.request_body = Some(serde_json::json!({"prompt": "first"})); + first.request_body_state = Some(UsageBodyCaptureState::Inline); + let mut second = sample_usage("pending-batch-second", "pending", "pending", 1_001); + second.finalized_at_unix_secs = None; + second.request_headers = Some(serde_json::json!({"x-request": "second"})); + + sqlx::query( + r#" +CREATE TRIGGER reject_second_pending_audit +BEFORE INSERT ON usage_http_audits +WHEN NEW.request_id = 'pending-batch-second' +BEGIN + SELECT RAISE(ABORT, 'reject pending batch test row'); +END +"#, + ) + .execute(&pool) + .await + .expect("rollback trigger should install"); + repository + .upsert_pending_many(vec![first.clone(), second.clone()]) + .await + .expect_err("auxiliary write failure should roll back the pending batch"); + let rolled_back_usage: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM \"usage\" WHERE request_id LIKE 'pending-batch-%'", + ) + .fetch_one(&pool) + .await + .expect("rolled back usage should count"); + let rolled_back_deltas: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id LIKE 'pending-batch-%'", + ) + .fetch_one(&pool) + .await + .expect("rolled back deltas should count"); + assert_eq!(rolled_back_usage, 0); + assert_eq!(rolled_back_deltas, 0); + + sqlx::query("DROP TRIGGER reject_second_pending_audit") + .execute(&pool) + .await + .expect("rollback trigger should drop"); + repository + .upsert_pending_many(vec![first, second]) + .await + .expect("pending batch should commit"); + + let committed: (i64, i64, i64, i64, i64) = sqlx::query_as( + r#" +SELECT + (SELECT COUNT(*) FROM "usage" WHERE request_id LIKE 'pending-batch-%'), + (SELECT COUNT(*) FROM usage_http_audits WHERE request_id LIKE 'pending-batch-%'), + (SELECT COUNT(*) FROM usage_body_blobs WHERE request_id LIKE 'pending-batch-%'), + (SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id LIKE 'pending-batch-%'), + (SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id LIKE 'pending-batch-%') +"#, + ) + .fetch_one(&pool) + .await + .expect("pending batch auxiliary rows should count"); + assert_eq!(committed, (2, 2, 1, 2, 2)); + + let provider_deltas: i64 = sqlx::query_scalar( + r#" +SELECT COALESCE(SUM(request_count_delta), 0) +FROM usage_counter_deltas +WHERE request_id LIKE 'pending-batch-%' AND kind = 'provider_api_key' +"#, + ) + .fetch_one(&pool) + .await + .expect("pending batch provider deltas should load"); + assert_eq!(provider_deltas, 2); +} + +#[tokio::test] +async fn sqlite_concurrent_same_request_upserts_enqueue_counters_once() { + let database_path = std::env::temp_dir().join(format!( + "aether-usage-counter-concurrency-{}.sqlite", + uuid::Uuid::new_v4() + )); + let options = sqlx::sqlite::SqliteConnectOptions::new() + .filename(&database_path) + .create_if_missing(true) + .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) + .busy_timeout(std::time::Duration::from_secs(30)); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(4) + .connect_with(options) + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let repository = SqliteUsageWriteRepository::new(pool.clone()); + let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let repository = repository.clone(); + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + let usage = sample_usage("concurrent-counter-request", "completed", "pending", 1_000); + barrier.wait().await; + repository.upsert(usage).await + })); + } + for task in tasks { + task.await + .expect("concurrent usage writer should join") + .expect("concurrent usage should persist"); + } + + repository + .flush_usage_counter_deltas(100) + .await + .expect("usage counter deltas should flush"); + let api_key_requests: i64 = + sqlx::query_scalar("SELECT total_requests FROM api_keys WHERE id = 'api-key-1'") + .fetch_one(&pool) + .await + .expect("api key counter should load"); + let provider_key_requests: i64 = sqlx::query_scalar( + "SELECT request_count FROM provider_api_keys WHERE id = 'provider-key-1'", + ) + .fetch_one(&pool) + .await + .expect("provider key counter should load"); + let model_requests: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = 'model-1'") + .fetch_one(&pool) + .await + .expect("model counter should load"); + let outbox_rows: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas") + .fetch_one(&pool) + .await + .expect("usage counter outbox should load"); + let routing_rows: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id = 'concurrent-counter-request'", + ) + .fetch_one(&pool) + .await + .expect("routing snapshot should count"); + let settlement_rows: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = 'concurrent-counter-request'", + ) + .fetch_one(&pool) + .await + .expect("settlement snapshot should count"); + assert_eq!(api_key_requests, 1); + assert_eq!(provider_key_requests, 1); + assert_eq!(model_requests, 1); + assert_eq!(outbox_rows, 3); + assert_eq!(routing_rows, 1); + assert_eq!(settlement_rows, 1); + + drop(repository); + pool.close().await; + let _ = std::fs::remove_file(&database_path); + let _ = std::fs::remove_file(format!("{}-wal", database_path.display())); + let _ = std::fs::remove_file(format!("{}-shm", database_path.display())); +} + async fn seed_stats_targets(pool: &sqlx::SqlitePool) { sqlx::query( r#" @@ -894,6 +2244,8 @@ INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES ('provider-1', 'Provider One', 'openai', 1, 1); INSERT INTO provider_api_keys (id, provider_id, name, created_at, updated_at) VALUES ('provider-key-1', 'provider-1', 'Provider Key One', 1, 1); +INSERT INTO global_models (id, name, created_at, updated_at) +VALUES ('global-model-1', 'model-1', 1, 1); "#, ) .execute(pool) @@ -980,3 +2332,21 @@ fn sample_usage( updated_at_unix_secs: updated_at, } } + +fn cleanup_window( + detail_cutoff: i64, + compressed_cutoff: i64, + header_cutoff: i64, + log_cutoff: i64, +) -> UsageCleanupWindow { + fn timestamp(value: i64) -> DateTime { + DateTime::from_timestamp(value, 0).expect("test timestamp should be valid") + } + + UsageCleanupWindow { + detail_cutoff: timestamp(detail_cutoff), + compressed_cutoff: timestamp(compressed_cutoff), + header_cutoff: timestamp(header_cutoff), + log_cutoff: timestamp(log_cutoff), + } +} diff --git a/crates/aether-data/adapters/sqlite/src/video_tasks.rs b/crates/aether-data/adapters/sqlite/src/video_tasks.rs index 5789a6802..a4095fc54 100644 --- a/crates/aether-data/adapters/sqlite/src/video_tasks.rs +++ b/crates/aether-data/adapters/sqlite/src/video_tasks.rs @@ -99,32 +99,6 @@ impl SqliteVideoTaskRepository { .map_sql_err()?; row.as_ref().map(map_video_task_row).transpose() } - - async fn reload_ids(&self, ids: &[String]) -> Result, DataLayerError> { - if ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(VIDEO_TASK_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for id in ids { - separated.push_bind(id); - } - } - builder.push(")"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut tasks = rows - .iter() - .map(map_video_task_row) - .collect::, _>>()?; - tasks.sort_by(|left, right| { - left.next_poll_at_unix_secs - .cmp(&right.next_poll_at_unix_secs) - .then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs)) - }); - Ok(tasks) - } } #[async_trait] @@ -320,20 +294,70 @@ impl VideoTaskWriteRepository for SqliteVideoTaskRepository { if limit == 0 { return Ok(Vec::new()); } - let due = self.list_due(now_unix_secs, limit).await?; - let ids = due.iter().map(|task| task.id.clone()).collect::>(); - for id in &ids { - sqlx::query( - "UPDATE video_tasks SET next_poll_at = ?, updated_at = MAX(updated_at, ?) WHERE id = ?", - ) - .bind(u64_to_i64(claim_until_unix_secs, "video task claim_until")?) - .bind(u64_to_i64(now_unix_secs, "video task now")?) - .bind(id) - .execute(&self.pool) - .await - .map_sql_err()?; + let now = u64_to_i64(now_unix_secs, "video task now")?; + let claim_until = u64_to_i64(claim_until_unix_secs, "video task claim_until")?; + let limit = limit_i64(limit, "due video task claim limit")?; + let mut tx = self.pool.begin().await.map_sql_err()?; + let claimed_rows = sqlx::query( + r#" +UPDATE video_tasks +SET next_poll_at = ?, + updated_at = MAX(updated_at, ?) +WHERE id IN ( + SELECT id + FROM video_tasks + WHERE status IN ('submitted', 'queued', 'processing') + AND next_poll_at IS NOT NULL + AND next_poll_at <= ? + AND poll_count < max_poll_count + ORDER BY next_poll_at ASC, updated_at ASC + LIMIT ? +) + AND status IN ('submitted', 'queued', 'processing') + AND next_poll_at IS NOT NULL + AND next_poll_at <= ? + AND poll_count < max_poll_count +RETURNING id +"#, + ) + .bind(claim_until) + .bind(now) + .bind(now) + .bind(limit) + .bind(now) + .fetch_all(&mut *tx) + .await + .map_sql_err()?; + let ids = claimed_rows + .iter() + .map(|row| row.try_get::("id").map_sql_err()) + .collect::, _>>()?; + if ids.is_empty() { + tx.commit().await.map_sql_err()?; + return Ok(Vec::new()); } - self.reload_ids(&ids).await + + let mut reload = QueryBuilder::::new(VIDEO_TASK_COLUMNS); + reload.push(" WHERE id IN ("); + { + let mut separated = reload.separated(", "); + for id in &ids { + separated.push_bind(id); + } + } + reload.push(")"); + let rows = reload.build().fetch_all(&mut *tx).await.map_sql_err()?; + let mut tasks = rows + .iter() + .map(map_video_task_row) + .collect::, _>>()?; + tasks.sort_by(|left, right| { + left.next_poll_at_unix_secs + .cmp(&right.next_poll_at_unix_secs) + .then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs)) + }); + tx.commit().await.map_sql_err()?; + Ok(tasks) } } @@ -681,6 +705,8 @@ mod tests { UpsertVideoTask, VideoTaskLookupKey, VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, }; + use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; + use std::{sync::Arc, time::Duration}; #[tokio::test] async fn sqlite_repository_writes_and_reads_video_tasks() { @@ -777,6 +803,69 @@ mod tests { assert_eq!(updated.progress_percent, 50); } + #[tokio::test] + async fn sqlite_claim_due_does_not_return_one_task_to_multiple_workers() { + const WORKERS: usize = 8; + + let database_path = std::env::temp_dir().join(format!( + "aether-sqlite-video-claim-{}.db", + uuid::Uuid::new_v4() + )); + let options = SqliteConnectOptions::new() + .filename(&database_path) + .create_if_missing(true) + .foreign_keys(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)); + let pool = SqlitePoolOptions::new() + .max_connections(WORKERS as u32) + .connect_with(options) + .await + .expect("concurrent sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + let task_id = "task-concurrent-claim"; + let repository = SqliteVideoTaskRepository::new(pool.clone()); + repository + .upsert(sample_task(task_id, VideoTaskStatus::Submitted, 100)) + .await + .expect("claimable sqlite task should insert"); + + let barrier = Arc::new(tokio::sync::Barrier::new(WORKERS + 1)); + let mut workers = tokio::task::JoinSet::new(); + for _ in 0..WORKERS { + let worker_repository = repository.clone(); + let worker_barrier = barrier.clone(); + workers.spawn(async move { + worker_barrier.wait().await; + worker_repository.claim_due(100, 130, 1).await + }); + } + barrier.wait().await; + + let mut results = Vec::with_capacity(WORKERS); + while let Some(result) = workers.join_next().await { + results.push(result); + } + let followup_result = repository.claim_due(100, 130, 1).await; + pool.close().await; + let _ = std::fs::remove_file(&database_path); + + let mut claimed_ids = Vec::new(); + for result in results { + let claimed = result + .expect("sqlite claim worker should join") + .expect("sqlite claim should execute"); + claimed_ids.extend(claimed.into_iter().map(|task| task.id)); + } + assert_eq!(claimed_ids, vec![task_id.to_string()]); + assert!(followup_result + .expect("follow-up sqlite claim should execute") + .is_empty()); + } + fn sample_task( id: &str, status: VideoTaskStatus, diff --git a/crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql b/crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql new file mode 100644 index 000000000..5a81e59c2 --- /dev/null +++ b/crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql @@ -0,0 +1,36 @@ +UPDATE api_keys AS target +LEFT JOIN ( + SELECT + `usage`.api_key_id, + COUNT(*) AS total_requests, + COALESCE( + SUM( + GREATEST( + COALESCE( + `usage`.total_tokens, + COALESCE(`usage`.input_tokens, 0) + COALESCE(`usage`.output_tokens, 0) + ), + 0 + ) + ), + 0 + ) AS total_tokens, + COALESCE(SUM(COALESCE(`usage`.total_cost_usd, 0)), 0) AS total_cost_usd, + MAX( + COALESCE( + `usage`.created_at, + `usage`.created_at_unix_ms, + `usage`.updated_at_unix_secs + ) + ) AS last_used_at + FROM `usage` + WHERE `usage`.api_key_id IS NOT NULL + AND TRIM(`usage`.api_key_id) <> '' + GROUP BY `usage`.api_key_id +) AS aggregated + ON aggregated.api_key_id = target.id +SET + target.total_requests = COALESCE(aggregated.total_requests, 0), + target.total_tokens = COALESCE(aggregated.total_tokens, 0), + target.total_cost_usd = COALESCE(aggregated.total_cost_usd, 0), + target.last_used_at = aggregated.last_used_at; diff --git a/crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql b/crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql new file mode 100644 index 000000000..c0e41fef5 --- /dev/null +++ b/crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql @@ -0,0 +1,15 @@ +UPDATE global_models AS target +LEFT JOIN ( + SELECT + `usage`.model, + COUNT(*) AS usage_count + FROM `usage` + WHERE `usage`.model IS NOT NULL + AND TRIM(`usage`.model) <> '' + AND `usage`.status NOT IN ('pending', 'streaming') + GROUP BY `usage`.model +) AS aggregated + ON aggregated.model = target.name +SET + target.usage_count = COALESCE(aggregated.usage_count, 0), + target.updated_at = UNIX_TIMESTAMP(); diff --git a/crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql b/crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql new file mode 100644 index 000000000..88626dfd7 --- /dev/null +++ b/crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql @@ -0,0 +1,11 @@ +UPDATE providers +SET enabled = is_active +WHERE enabled <> is_active; + +UPDATE provider_endpoints +SET enabled = is_active +WHERE enabled <> is_active; + +UPDATE models +SET enabled = is_active +WHERE enabled <> is_active; diff --git a/crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql b/crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql new file mode 100644 index 000000000..bdb4d9190 --- /dev/null +++ b/crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql @@ -0,0 +1,229 @@ +UPDATE api_keys AS target +LEFT JOIN ( + SELECT + source.api_key_id, + COALESCE(SUM(source.canonical_total_tokens), 0) AS total_tokens + FROM ( + SELECT + `usage`.api_key_id, + GREATEST( + COALESCE( + CASE + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + GREATEST(settlement.billing_effective_input_tokens, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + + GREATEST( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE(`usage`.cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + GREATEST( + COALESCE( + settlement.billing_cache_read_tokens, + `usage`.cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + GREATEST(settlement.billing_total_input_context, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + END, + NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0), + CASE + WHEN SUBSTRING_INDEX( + LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')), + ':', + 1 + ) IN ('openai', 'gemini', 'google') + THEN GREATEST(COALESCE(`usage`.input_tokens, 0), 0) + + GREATEST(COALESCE(`usage`.output_tokens, 0), 0) + ELSE GREATEST(COALESCE(`usage`.input_tokens, 0), 0) + + GREATEST(COALESCE(`usage`.output_tokens, 0), 0) + + GREATEST( + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE(`usage`.cache_creation_input_tokens, 0) + END, + 0 + ) + + GREATEST(COALESCE(`usage`.cache_read_input_tokens, 0), 0) + END, + 0 + ), + 0 + ) AS canonical_total_tokens + FROM `usage` + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id + WHERE `usage`.status NOT IN ('pending', 'streaming') + ) AS source + WHERE source.api_key_id IS NOT NULL + AND TRIM(source.api_key_id) <> '' + GROUP BY source.api_key_id +) AS aggregated + ON aggregated.api_key_id = target.id +SET target.total_tokens = COALESCE(aggregated.total_tokens, 0); + +UPDATE provider_api_keys AS target +LEFT JOIN ( + SELECT + source.provider_api_key_id, + COALESCE(SUM(source.canonical_total_tokens), 0) AS total_tokens + FROM ( + SELECT + `usage`.provider_api_key_id, + GREATEST( + COALESCE( + CASE + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + GREATEST(settlement.billing_effective_input_tokens, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + + GREATEST( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE(`usage`.cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + GREATEST( + COALESCE( + settlement.billing_cache_read_tokens, + `usage`.cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + GREATEST(settlement.billing_total_input_context, 0) + + GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0) + END, + NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0), + CASE + WHEN SUBSTRING_INDEX( + LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')), + ':', + 1 + ) IN ('openai', 'gemini', 'google') + THEN GREATEST(COALESCE(`usage`.input_tokens, 0), 0) + + GREATEST(COALESCE(`usage`.output_tokens, 0), 0) + ELSE GREATEST(COALESCE(`usage`.input_tokens, 0), 0) + + GREATEST(COALESCE(`usage`.output_tokens, 0), 0) + + GREATEST( + CASE + WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0 + AND ( + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_5m, 0), + COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0) + ) + + GREATEST( + COALESCE(`usage`.cache_creation_input_tokens_1h, 0), + COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE(`usage`.cache_creation_input_tokens, 0) + END, + 0 + ) + + GREATEST(COALESCE(`usage`.cache_read_input_tokens, 0), 0) + END, + 0 + ), + 0 + ) AS canonical_total_tokens + FROM `usage` + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = `usage`.request_id + WHERE `usage`.status NOT IN ('pending', 'streaming') + ) AS source + WHERE source.provider_api_key_id IS NOT NULL + AND TRIM(source.provider_api_key_id) <> '' + GROUP BY source.provider_api_key_id +) AS aggregated + ON aggregated.provider_api_key_id = target.id +SET target.total_tokens = COALESCE(aggregated.total_tokens, 0); diff --git a/crates/aether-data/runtime/backfills/mysql/README.md b/crates/aether-data/runtime/backfills/mysql/README.md index 7f06bb116..733ea7188 100644 --- a/crates/aether-data/runtime/backfills/mysql/README.md +++ b/crates/aether-data/runtime/backfills/mysql/README.md @@ -1 +1,2 @@ -MySQL-specific backfills live here when they are needed. +MySQL-specific lifecycle backfills live here. Versions intentionally match the +equivalent PostgreSQL backfills when they repair the same logical data. diff --git a/crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql b/crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql new file mode 100644 index 000000000..b73cb7f8d --- /dev/null +++ b/crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql @@ -0,0 +1,28 @@ +UPDATE api_keys AS target +SET + total_requests = ( + SELECT COUNT(*) + FROM "usage" + WHERE "usage".api_key_id = target.id + ), + total_tokens = COALESCE(( + SELECT SUM(MAX(COALESCE("usage".total_tokens, 0), 0)) + FROM "usage" + WHERE "usage".api_key_id = target.id + ), 0), + total_cost_usd = COALESCE(( + SELECT SUM(COALESCE("usage".total_cost_usd, 0)) + FROM "usage" + WHERE "usage".api_key_id = target.id + ), 0), + last_used_at = ( + SELECT MAX( + COALESCE( + "usage".created_at, + "usage".created_at_unix_ms, + "usage".updated_at_unix_secs + ) + ) + FROM "usage" + WHERE "usage".api_key_id = target.id + ); diff --git a/crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql b/crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql new file mode 100644 index 000000000..f67d6cfaf --- /dev/null +++ b/crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql @@ -0,0 +1,9 @@ +UPDATE global_models AS target +SET + usage_count = ( + SELECT COUNT(*) + FROM "usage" + WHERE "usage".model = target.name + AND "usage".status NOT IN ('pending', 'streaming') + ), + updated_at = CAST(strftime('%s', 'now') AS INTEGER); diff --git a/crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql b/crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql new file mode 100644 index 000000000..88626dfd7 --- /dev/null +++ b/crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql @@ -0,0 +1,11 @@ +UPDATE providers +SET enabled = is_active +WHERE enabled <> is_active; + +UPDATE provider_endpoints +SET enabled = is_active +WHERE enabled <> is_active; + +UPDATE models +SET enabled = is_active +WHERE enabled <> is_active; diff --git a/crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql b/crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql new file mode 100644 index 000000000..75caf9057 --- /dev/null +++ b/crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql @@ -0,0 +1,225 @@ +WITH canonical_usage AS ( + SELECT + "usage".api_key_id, + MAX( + COALESCE( + CASE + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + MAX(settlement.billing_effective_input_tokens, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + + MAX( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE("usage".cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + MAX( + COALESCE( + settlement.billing_cache_read_tokens, + "usage".cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + MAX(settlement.billing_total_input_context, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + END, + NULLIF(MAX(COALESCE("usage".total_tokens, 0), 0), 0), + CASE + WHEN LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + IN ('openai', 'gemini', 'google') + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'openai:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'gemini:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'google:%' + THEN MAX(COALESCE("usage".input_tokens, 0), 0) + + MAX(COALESCE("usage".output_tokens, 0), 0) + ELSE MAX(COALESCE("usage".input_tokens, 0), 0) + + MAX(COALESCE("usage".output_tokens, 0), 0) + + MAX( + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE("usage".cache_creation_input_tokens, 0) + END, + 0 + ) + + MAX(COALESCE("usage".cache_read_input_tokens, 0), 0) + END, + 0 + ), + 0 + ) AS canonical_total_tokens + FROM "usage" + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id + WHERE "usage".status NOT IN ('pending', 'streaming') +) +UPDATE api_keys AS target +SET total_tokens = COALESCE(( + SELECT SUM(canonical_usage.canonical_total_tokens) + FROM canonical_usage + WHERE canonical_usage.api_key_id = target.id +), 0); + +WITH canonical_usage AS ( + SELECT + "usage".provider_api_key_id, + MAX( + COALESCE( + CASE + WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN + MAX(settlement.billing_effective_input_tokens, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + + MAX( + COALESCE( + settlement.billing_cache_creation_tokens, + CASE + WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL + OR settlement.billing_cache_creation_1h_tokens IS NOT NULL + THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0) + + COALESCE(settlement.billing_cache_creation_1h_tokens, 0) + END, + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE("usage".cache_creation_input_tokens, 0) + END, + 0 + ), + 0 + ) + + MAX( + COALESCE( + settlement.billing_cache_read_tokens, + "usage".cache_read_input_tokens, + 0 + ), + 0 + ) + WHEN settlement.billing_total_input_context IS NOT NULL THEN + MAX(settlement.billing_total_input_context, 0) + + MAX(COALESCE(settlement.billing_output_tokens, "usage".output_tokens, 0), 0) + END, + NULLIF(MAX(COALESCE("usage".total_tokens, 0), 0), 0), + CASE + WHEN LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + IN ('openai', 'gemini', 'google') + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'openai:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'gemini:%' + OR LOWER(COALESCE("usage".endpoint_api_format, "usage".api_format, '')) + LIKE 'google:%' + THEN MAX(COALESCE("usage".input_tokens, 0), 0) + + MAX(COALESCE("usage".output_tokens, 0), 0) + ELSE MAX(COALESCE("usage".input_tokens, 0), 0) + + MAX(COALESCE("usage".output_tokens, 0), 0) + + MAX( + CASE + WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0 + AND ( + MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ) > 0 + THEN MAX( + COALESCE("usage".cache_creation_input_tokens_5m, 0), + COALESCE("usage".cache_creation_ephemeral_5m_input_tokens, 0) + ) + + MAX( + COALESCE("usage".cache_creation_input_tokens_1h, 0), + COALESCE("usage".cache_creation_ephemeral_1h_input_tokens, 0) + ) + ELSE COALESCE("usage".cache_creation_input_tokens, 0) + END, + 0 + ) + + MAX(COALESCE("usage".cache_read_input_tokens, 0), 0) + END, + 0 + ), + 0 + ) AS canonical_total_tokens + FROM "usage" + LEFT JOIN usage_settlement_snapshots AS settlement + ON settlement.request_id = "usage".request_id + WHERE "usage".status NOT IN ('pending', 'streaming') +) +UPDATE provider_api_keys AS target +SET total_tokens = COALESCE(( + SELECT SUM(canonical_usage.canonical_total_tokens) + FROM canonical_usage + WHERE canonical_usage.provider_api_key_id = target.id +), 0); diff --git a/crates/aether-data/runtime/backfills/sqlite/README.md b/crates/aether-data/runtime/backfills/sqlite/README.md index 659fb7c71..bd3728e7e 100644 --- a/crates/aether-data/runtime/backfills/sqlite/README.md +++ b/crates/aether-data/runtime/backfills/sqlite/README.md @@ -1 +1,2 @@ -SQLite-specific backfills live here when they are needed. +SQLite-specific lifecycle backfills live here. Versions intentionally match +the equivalent PostgreSQL backfills when they repair the same logical data. diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql index 4a5330ca1..1a2b1dbb6 100644 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql +++ b/crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql @@ -111,6 +111,63 @@ CREATE TABLE IF NOT EXISTS `usage` ( KEY usage_wallet_id_idx (`wallet_id`) ); +CREATE TABLE IF NOT EXISTS usage_body_blobs ( + `body_ref` VARCHAR(160) NOT NULL, + `request_id` VARCHAR(128) NOT NULL, + `body_field` VARCHAR(50) NOT NULL, + `payload_gzip` LONGBLOB NOT NULL, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`body_ref`), + UNIQUE KEY usage_body_blobs_request_id_field_key (`request_id`, `body_field`), + KEY ix_usage_body_blobs_request_id (`request_id`), + CONSTRAINT usage_body_blobs_request_id_fkey FOREIGN KEY (`request_id`) REFERENCES usage (`request_id`) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS usage_http_audits ( + `request_id` VARCHAR(128) NOT NULL, + `request_headers` JSON, + `provider_request_headers` JSON, + `response_headers` JSON, + `client_response_headers` JSON, + `request_body_ref` VARCHAR(160), + `provider_request_body_ref` VARCHAR(160), + `response_body_ref` VARCHAR(160), + `client_response_body_ref` VARCHAR(160), + `request_body_state` VARCHAR(32), + `provider_request_body_state` VARCHAR(32), + `response_body_state` VARCHAR(32), + `client_response_body_state` VARCHAR(32), + `body_capture_mode` VARCHAR(32) NOT NULL DEFAULT 'none', + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`request_id`), + KEY ix_usage_http_audits_updated_at (`updated_at`), + CONSTRAINT usage_http_audits_request_id_fkey FOREIGN KEY (`request_id`) REFERENCES usage (`request_id`) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS usage_routing_snapshots ( + `request_id` VARCHAR(128) NOT NULL, + `candidate_id` VARCHAR(160), + `candidate_index` BIGINT, + `key_name` VARCHAR(255), + `planner_kind` VARCHAR(120), + `route_family` VARCHAR(80), + `route_kind` VARCHAR(80), + `execution_path` VARCHAR(80), + `local_execution_runtime_miss_reason` VARCHAR(255), + `selected_provider_id` VARCHAR(100), + `selected_endpoint_id` VARCHAR(100), + `selected_provider_api_key_id` VARCHAR(100), + `has_format_conversion` TINYINT(1), + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`request_id`), + KEY ix_usage_routing_snapshots_route_family_kind (`route_family`, `route_kind`), + KEY ix_usage_routing_snapshots_candidate_id (`candidate_id`), + CONSTRAINT usage_routing_snapshots_request_id_fkey FOREIGN KEY (`request_id`) REFERENCES usage (`request_id`) ON DELETE CASCADE +); + CREATE TABLE IF NOT EXISTS usage_counter_deltas ( `id` VARCHAR(36) NOT NULL, `request_id` VARCHAR(128) NOT NULL, @@ -149,11 +206,40 @@ CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( `wallet_gift_balance_before` DOUBLE, `wallet_gift_balance_after` DOUBLE, `provider_monthly_used_usd` DOUBLE, + `billing_snapshot_schema_version` VARCHAR(20), + `billing_snapshot_status` VARCHAR(20), + `rate_multiplier` DOUBLE, + `is_free_tier` TINYINT(1), + `input_price_per_1m` DOUBLE, + `output_price_per_1m` DOUBLE, + `cache_creation_price_per_1m` DOUBLE, + `cache_read_price_per_1m` DOUBLE, + `price_per_request` DOUBLE, + `settlement_snapshot_schema_version` VARCHAR(20), + `settlement_snapshot` JSON, + `billing_dimensions` JSON, + `billing_input_tokens` BIGINT, + `billing_effective_input_tokens` BIGINT, + `billing_output_tokens` BIGINT, + `billing_cache_creation_tokens` BIGINT, + `billing_cache_creation_5m_tokens` BIGINT, + `billing_cache_creation_1h_tokens` BIGINT, + `billing_cache_read_tokens` BIGINT, + `billing_total_input_context` BIGINT, + `billing_cache_creation_cost_usd` DOUBLE, + `billing_cache_read_cost_usd` DOUBLE, + `billing_total_cost_usd` DOUBLE, + `billing_actual_total_cost_usd` DOUBLE, + `billing_pricing_source` VARCHAR(50), + `billing_rule_id` VARCHAR(100), + `billing_rule_version` VARCHAR(50), `finalized_at` BIGINT, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, PRIMARY KEY (`request_id`), KEY usage_settlement_snapshots_billing_status_idx (`billing_status`), - KEY usage_settlement_snapshots_wallet_id_idx (`wallet_id`) + KEY usage_settlement_snapshots_wallet_id_idx (`wallet_id`), + KEY ix_usage_settlement_snapshots_schema_version (`settlement_snapshot_schema_version`), + KEY ix_usage_settlement_snapshots_pricing_source (`billing_pricing_source`) ); diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql index 22e58f0c1..ae5d2aee2 100644 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql +++ b/crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql @@ -18,6 +18,26 @@ CREATE TABLE IF NOT EXISTS stats_hourly ( `aggregated_at` BIGINT, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `cache_hit_total_requests` BIGINT NOT NULL DEFAULT 0, + `cache_hit_requests` BIGINT NOT NULL DEFAULT 0, + `completed_total_requests` BIGINT NOT NULL DEFAULT 0, + `completed_cache_hit_requests` BIGINT NOT NULL DEFAULT 0, + `completed_input_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_total_input_context` BIGINT NOT NULL DEFAULT 0, + `completed_cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `completed_cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_requests` BIGINT NOT NULL DEFAULT 0, + `settled_input_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_output_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_first_finalized_at_unix_secs` BIGINT, + `settled_last_finalized_at_unix_secs` BIGINT, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_hourly_hour (`hour_utc`) ); @@ -55,6 +75,19 @@ CREATE TABLE IF NOT EXISTS stats_hourly_user ( `total_cost` DOUBLE NOT NULL DEFAULT 0, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `actual_total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `settled_total_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_requests` BIGINT NOT NULL DEFAULT 0, + `settled_input_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_output_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_first_finalized_at_unix_secs` BIGINT, + `settled_last_finalized_at_unix_secs` BIGINT, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_hourly_user (`hour_utc`, `user_id`) ); @@ -70,6 +103,8 @@ CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( `total_cost` DOUBLE NOT NULL DEFAULT 0, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_hourly_user_model (`hour_utc`, `user_id`, `model`) ); @@ -98,6 +133,8 @@ CREATE TABLE IF NOT EXISTS stats_hourly_model ( `avg_response_time_ms` DOUBLE NOT NULL DEFAULT 0, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_hourly_model (`hour_utc`, `model`) ); @@ -146,6 +183,30 @@ CREATE TABLE IF NOT EXISTS stats_daily ( `p50_first_byte_time_ms` BIGINT, `p90_first_byte_time_ms` BIGINT, `p99_first_byte_time_ms` BIGINT, + `effective_input_tokens` BIGINT NOT NULL DEFAULT 0, + `total_input_context` BIGINT NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_hit_total_requests` BIGINT NOT NULL DEFAULT 0, + `cache_hit_requests` BIGINT NOT NULL DEFAULT 0, + `completed_total_requests` BIGINT NOT NULL DEFAULT 0, + `completed_cache_hit_requests` BIGINT NOT NULL DEFAULT 0, + `completed_input_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `completed_total_input_context` BIGINT NOT NULL DEFAULT 0, + `completed_cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `completed_cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_requests` BIGINT NOT NULL DEFAULT 0, + `settled_input_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_output_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_first_finalized_at_unix_secs` BIGINT, + `settled_last_finalized_at_unix_secs` BIGINT, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_daily_date (`date`) ); @@ -163,6 +224,10 @@ CREATE TABLE IF NOT EXISTS stats_daily_model ( `avg_response_time_ms` DOUBLE NOT NULL DEFAULT 0, `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_daily_model (`date`, `model`) ); @@ -230,7 +295,305 @@ CREATE TABLE IF NOT EXISTS stats_user_daily ( `username` VARCHAR(255), `created_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL, + `actual_total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `effective_input_tokens` BIGINT NOT NULL DEFAULT 0, + `total_input_context` BIGINT NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_total_cost` DOUBLE NOT NULL DEFAULT 0, + `settled_total_requests` BIGINT NOT NULL DEFAULT 0, + `settled_input_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_output_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `settled_first_finalized_at_unix_secs` BIGINT, + `settled_last_finalized_at_unix_secs` BIGINT, PRIMARY KEY (`id`), UNIQUE KEY uq_stats_user_daily (`date`, `user_id`) ); +CREATE TABLE IF NOT EXISTS stats_user_summary ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `cutoff_date` BIGINT NOT NULL, + `all_time_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_success_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_error_requests` BIGINT NOT NULL DEFAULT 0, + `all_time_input_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_output_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `all_time_cost` DOUBLE NOT NULL DEFAULT 0, + `all_time_actual_cost` DOUBLE NOT NULL DEFAULT 0, + `active_days` BIGINT NOT NULL DEFAULT 0, + `first_active_date` BIGINT, + `last_active_date` BIGINT, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_summary_user_id (`user_id`), + KEY idx_stats_user_summary_cutoff_date (`cutoff_date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `total_requests` BIGINT NOT NULL DEFAULT 0, + `success_requests` BIGINT NOT NULL DEFAULT 0, + `input_tokens` BIGINT NOT NULL DEFAULT 0, + `effective_input_tokens` BIGINT NOT NULL DEFAULT 0, + `output_tokens` BIGINT NOT NULL DEFAULT 0, + `total_tokens` BIGINT NOT NULL DEFAULT 0, + `total_input_context` BIGINT NOT NULL DEFAULT 0, + `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `total_cost` DOUBLE NOT NULL DEFAULT 0, + `actual_total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `successful_response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `successful_response_time_samples` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_model (`user_id`, `date`, `model`), + KEY idx_stats_user_daily_model_date (`date`), + KEY idx_stats_user_daily_model_user_id (`user_id`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_provider ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `total_requests` BIGINT NOT NULL DEFAULT 0, + `success_requests` BIGINT NOT NULL DEFAULT 0, + `input_tokens` BIGINT NOT NULL DEFAULT 0, + `effective_input_tokens` BIGINT NOT NULL DEFAULT 0, + `output_tokens` BIGINT NOT NULL DEFAULT 0, + `total_tokens` BIGINT NOT NULL DEFAULT 0, + `total_input_context` BIGINT NOT NULL DEFAULT 0, + `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `total_cost` DOUBLE NOT NULL DEFAULT 0, + `actual_total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `successful_response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `successful_response_time_samples` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_provider (`user_id`, `date`, `provider_name`), + KEY idx_stats_user_daily_provider_date (`date`), + KEY idx_stats_user_daily_provider_user_id (`user_id`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_api_format ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `api_format` VARCHAR(128) NOT NULL, + `total_requests` BIGINT NOT NULL DEFAULT 0, + `success_requests` BIGINT NOT NULL DEFAULT 0, + `input_tokens` BIGINT NOT NULL DEFAULT 0, + `effective_input_tokens` BIGINT NOT NULL DEFAULT 0, + `output_tokens` BIGINT NOT NULL DEFAULT 0, + `total_tokens` BIGINT NOT NULL DEFAULT 0, + `total_input_context` BIGINT NOT NULL DEFAULT 0, + `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_5m_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_creation_ephemeral_1h_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `total_cost` DOUBLE NOT NULL DEFAULT 0, + `actual_total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `successful_response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `successful_response_time_samples` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_api_format (`user_id`, `date`, `api_format`), + KEY idx_stats_user_daily_api_format_date (`date`), + KEY idx_stats_user_daily_api_format_user_id (`user_id`) +); + +CREATE TABLE IF NOT EXISTS stats_daily_model_provider ( + `id` VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `total_requests` BIGINT NOT NULL DEFAULT 0, + `total_tokens` BIGINT NOT NULL DEFAULT 0, + `total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_daily_model_provider (`date`, `model`, `provider_name`), + KEY idx_stats_daily_model_provider_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model_provider ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `total_requests` BIGINT NOT NULL DEFAULT 0, + `total_tokens` BIGINT NOT NULL DEFAULT 0, + `total_cost` DOUBLE NOT NULL DEFAULT 0, + `response_time_sum_ms` DOUBLE NOT NULL DEFAULT 0, + `response_time_samples` BIGINT NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_model_provider (`user_id`, `date`, `model`, `provider_name`), + KEY idx_stats_user_daily_model_provider_date (`date`), + KEY idx_stats_user_daily_model_provider_user_date (`user_id`, `date`) +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings ( + `id` VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_daily_cost_savings_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_provider ( + `id` VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_daily_cost_savings_provider (`date`, `provider_name`), + KEY idx_stats_daily_cost_savings_provider_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model ( + `id` VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_daily_cost_savings_model (`date`, `model`), + KEY idx_stats_daily_cost_savings_model_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model_provider ( + `id` VARCHAR(64) NOT NULL, + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_daily_cost_savings_model_provider (`date`, `model`, `provider_name`), + KEY idx_stats_daily_cost_savings_model_provider_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_cost_savings (`user_id`, `date`), + KEY idx_stats_user_daily_cost_savings_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_provider ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_cost_savings_provider (`user_id`, `date`, `provider_name`), + KEY idx_stats_user_daily_cost_savings_provider_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_cost_savings_model (`user_id`, `date`, `model`), + KEY idx_stats_user_daily_cost_savings_model_date (`date`) +); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model_provider ( + `id` VARCHAR(64) NOT NULL, + `user_id` VARCHAR(64) NOT NULL, + `username` VARCHAR(255), + `date` BIGINT NOT NULL, + `model` VARCHAR(255) NOT NULL, + `provider_name` VARCHAR(255) NOT NULL, + `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, + `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, + `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, + `estimated_full_cost` DOUBLE NOT NULL DEFAULT 0, + `created_at` BIGINT NOT NULL, + `updated_at` BIGINT NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY uq_stats_user_daily_cost_savings_model_provider (`user_id`, `date`, `model`, `provider_name`), + KEY idx_stats_user_daily_cost_savings_model_provider_date (`date`) +); + diff --git a/crates/aether-data/runtime/schema/generated/postgres/baseline/006_usage.sql b/crates/aether-data/runtime/schema/generated/postgres/baseline/006_usage.sql index b0321ae63..3b3dbf2ac 100644 --- a/crates/aether-data/runtime/schema/generated/postgres/baseline/006_usage.sql +++ b/crates/aether-data/runtime/schema/generated/postgres/baseline/006_usage.sql @@ -112,6 +112,66 @@ CREATE INDEX IF NOT EXISTS usage_request_id_idx ON public.usage USING btree (req CREATE INDEX IF NOT EXISTS usage_user_id_idx ON public.usage USING btree (user_id); CREATE INDEX IF NOT EXISTS usage_wallet_id_idx ON public.usage USING btree (wallet_id); +CREATE TABLE IF NOT EXISTS public.usage_body_blobs ( + body_ref character varying(160) NOT NULL, + request_id character varying(128) NOT NULL, + body_field character varying(50) NOT NULL, + payload_gzip bytea NOT NULL, + created_at timestamp with time zone NOT NULL, + updated_at timestamp with time zone NOT NULL +); + +ALTER TABLE ONLY public.usage_body_blobs ADD CONSTRAINT usage_body_blobs_pkey PRIMARY KEY (body_ref); +ALTER TABLE ONLY public.usage_body_blobs ADD CONSTRAINT usage_body_blobs_request_id_field_key UNIQUE (request_id, body_field); +CREATE INDEX IF NOT EXISTS ix_usage_body_blobs_request_id ON public.usage_body_blobs USING btree (request_id); +ALTER TABLE ONLY public.usage_body_blobs ADD CONSTRAINT usage_body_blobs_request_id_fkey FOREIGN KEY (request_id) REFERENCES public.usage(request_id) ON DELETE CASCADE; + +CREATE TABLE IF NOT EXISTS public.usage_http_audits ( + request_id character varying(128) NOT NULL, + request_headers jsonb, + provider_request_headers jsonb, + response_headers jsonb, + client_response_headers jsonb, + request_body_ref character varying(160), + provider_request_body_ref character varying(160), + response_body_ref character varying(160), + client_response_body_ref character varying(160), + request_body_state character varying(32), + provider_request_body_state character varying(32), + response_body_state character varying(32), + client_response_body_state character varying(32), + body_capture_mode character varying(32) DEFAULT 'none' NOT NULL, + created_at timestamp with time zone NOT NULL, + updated_at timestamp with time zone NOT NULL +); + +ALTER TABLE ONLY public.usage_http_audits ADD CONSTRAINT usage_http_audits_pkey PRIMARY KEY (request_id); +CREATE INDEX IF NOT EXISTS ix_usage_http_audits_updated_at ON public.usage_http_audits USING btree (updated_at); +ALTER TABLE ONLY public.usage_http_audits ADD CONSTRAINT usage_http_audits_request_id_fkey FOREIGN KEY (request_id) REFERENCES public.usage(request_id) ON DELETE CASCADE; + +CREATE TABLE IF NOT EXISTS public.usage_routing_snapshots ( + request_id character varying(128) NOT NULL, + candidate_id character varying(160), + candidate_index bigint, + key_name character varying(255), + planner_kind character varying(120), + route_family character varying(80), + route_kind character varying(80), + execution_path character varying(80), + local_execution_runtime_miss_reason character varying(255), + selected_provider_id character varying(100), + selected_endpoint_id character varying(100), + selected_provider_api_key_id character varying(100), + has_format_conversion boolean, + created_at timestamp with time zone NOT NULL, + updated_at timestamp with time zone NOT NULL +); + +ALTER TABLE ONLY public.usage_routing_snapshots ADD CONSTRAINT usage_routing_snapshots_pkey PRIMARY KEY (request_id); +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_route_family_kind ON public.usage_routing_snapshots USING btree (route_family, route_kind); +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_candidate_id ON public.usage_routing_snapshots USING btree (candidate_id); +ALTER TABLE ONLY public.usage_routing_snapshots ADD CONSTRAINT usage_routing_snapshots_request_id_fkey FOREIGN KEY (request_id) REFERENCES public.usage(request_id) ON DELETE CASCADE; + CREATE TABLE IF NOT EXISTS public.usage_counter_deltas ( id character varying(36) NOT NULL, request_id character varying(128) NOT NULL, @@ -151,6 +211,33 @@ CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots ( wallet_gift_balance_before double precision, wallet_gift_balance_after double precision, provider_monthly_used_usd double precision, + billing_snapshot_schema_version character varying(20), + billing_snapshot_status character varying(20), + rate_multiplier double precision, + is_free_tier boolean, + input_price_per_1m double precision, + output_price_per_1m double precision, + cache_creation_price_per_1m double precision, + cache_read_price_per_1m double precision, + price_per_request double precision, + settlement_snapshot_schema_version character varying(20), + settlement_snapshot jsonb, + billing_dimensions jsonb, + billing_input_tokens bigint, + billing_effective_input_tokens bigint, + billing_output_tokens bigint, + billing_cache_creation_tokens bigint, + billing_cache_creation_5m_tokens bigint, + billing_cache_creation_1h_tokens bigint, + billing_cache_read_tokens bigint, + billing_total_input_context bigint, + billing_cache_creation_cost_usd double precision, + billing_cache_read_cost_usd double precision, + billing_total_cost_usd double precision, + billing_actual_total_cost_usd double precision, + billing_pricing_source character varying(50), + billing_rule_id character varying(100), + billing_rule_version character varying(50), finalized_at bigint, created_at bigint NOT NULL, updated_at bigint NOT NULL @@ -159,4 +246,6 @@ CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots ( ALTER TABLE ONLY public.usage_settlement_snapshots ADD CONSTRAINT usage_settlement_snapshots_pkey PRIMARY KEY (request_id); CREATE INDEX IF NOT EXISTS usage_settlement_snapshots_billing_status_idx ON public.usage_settlement_snapshots USING btree (billing_status); CREATE INDEX IF NOT EXISTS usage_settlement_snapshots_wallet_id_idx ON public.usage_settlement_snapshots USING btree (wallet_id); +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_schema_version ON public.usage_settlement_snapshots USING btree (settlement_snapshot_schema_version); +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_pricing_source ON public.usage_settlement_snapshots USING btree (billing_pricing_source); diff --git a/crates/aether-data/runtime/schema/generated/postgres/baseline/007_stats.sql b/crates/aether-data/runtime/schema/generated/postgres/baseline/007_stats.sql index 6ab7880a2..2319906e5 100644 --- a/crates/aether-data/runtime/schema/generated/postgres/baseline/007_stats.sql +++ b/crates/aether-data/runtime/schema/generated/postgres/baseline/007_stats.sql @@ -17,7 +17,27 @@ CREATE TABLE IF NOT EXISTS public.stats_hourly ( is_complete boolean DEFAULT false NOT NULL, aggregated_at bigint, created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + cache_hit_total_requests bigint DEFAULT 0 NOT NULL, + cache_hit_requests bigint DEFAULT 0 NOT NULL, + completed_total_requests bigint DEFAULT 0 NOT NULL, + completed_cache_hit_requests bigint DEFAULT 0 NOT NULL, + completed_input_tokens bigint DEFAULT 0 NOT NULL, + completed_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + completed_cache_read_tokens bigint DEFAULT 0 NOT NULL, + completed_total_input_context bigint DEFAULT 0 NOT NULL, + completed_cache_creation_cost double precision DEFAULT 0 NOT NULL, + completed_cache_read_cost double precision DEFAULT 0 NOT NULL, + settled_total_cost double precision DEFAULT 0 NOT NULL, + settled_total_requests bigint DEFAULT 0 NOT NULL, + settled_input_tokens bigint DEFAULT 0 NOT NULL, + settled_output_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_read_tokens bigint DEFAULT 0 NOT NULL, + settled_first_finalized_at_unix_secs bigint, + settled_last_finalized_at_unix_secs bigint ); ALTER TABLE ONLY public.stats_hourly ADD CONSTRAINT stats_hourly_pkey PRIMARY KEY (id); @@ -56,7 +76,20 @@ CREATE TABLE IF NOT EXISTS public.stats_hourly_user ( output_tokens bigint DEFAULT 0 NOT NULL, total_cost double precision DEFAULT 0 NOT NULL, created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + cache_creation_tokens bigint DEFAULT 0 NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + actual_total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + settled_total_cost double precision DEFAULT 0 NOT NULL, + settled_total_requests bigint DEFAULT 0 NOT NULL, + settled_input_tokens bigint DEFAULT 0 NOT NULL, + settled_output_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_read_tokens bigint DEFAULT 0 NOT NULL, + settled_first_finalized_at_unix_secs bigint, + settled_last_finalized_at_unix_secs bigint ); ALTER TABLE ONLY public.stats_hourly_user ADD CONSTRAINT stats_hourly_user_pkey PRIMARY KEY (id); @@ -72,7 +105,9 @@ CREATE TABLE IF NOT EXISTS public.stats_hourly_user_model ( output_tokens bigint DEFAULT 0 NOT NULL, total_cost double precision DEFAULT 0 NOT NULL, created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL ); ALTER TABLE ONLY public.stats_hourly_user_model ADD CONSTRAINT stats_hourly_user_model_pkey PRIMARY KEY (id); @@ -102,7 +137,9 @@ CREATE TABLE IF NOT EXISTS public.stats_hourly_model ( total_cost double precision DEFAULT 0 NOT NULL, avg_response_time_ms double precision DEFAULT 0 NOT NULL, created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL ); ALTER TABLE ONLY public.stats_hourly_model ADD CONSTRAINT stats_hourly_model_pkey PRIMARY KEY (id); @@ -152,7 +189,31 @@ CREATE TABLE IF NOT EXISTS public.stats_daily ( p99_response_time_ms bigint, p50_first_byte_time_ms bigint, p90_first_byte_time_ms bigint, - p99_first_byte_time_ms bigint + p99_first_byte_time_ms bigint, + effective_input_tokens bigint DEFAULT 0 NOT NULL, + total_input_context bigint DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL, + cache_hit_total_requests bigint DEFAULT 0 NOT NULL, + cache_hit_requests bigint DEFAULT 0 NOT NULL, + completed_total_requests bigint DEFAULT 0 NOT NULL, + completed_cache_hit_requests bigint DEFAULT 0 NOT NULL, + completed_input_tokens bigint DEFAULT 0 NOT NULL, + completed_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + completed_cache_read_tokens bigint DEFAULT 0 NOT NULL, + completed_total_input_context bigint DEFAULT 0 NOT NULL, + completed_cache_creation_cost double precision DEFAULT 0 NOT NULL, + completed_cache_read_cost double precision DEFAULT 0 NOT NULL, + settled_total_cost double precision DEFAULT 0 NOT NULL, + settled_total_requests bigint DEFAULT 0 NOT NULL, + settled_input_tokens bigint DEFAULT 0 NOT NULL, + settled_output_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_read_tokens bigint DEFAULT 0 NOT NULL, + settled_first_finalized_at_unix_secs bigint, + settled_last_finalized_at_unix_secs bigint ); ALTER TABLE ONLY public.stats_daily ADD CONSTRAINT stats_daily_pkey PRIMARY KEY (id); @@ -170,7 +231,11 @@ CREATE TABLE IF NOT EXISTS public.stats_daily_model ( total_cost double precision DEFAULT 0 NOT NULL, avg_response_time_ms double precision DEFAULT 0 NOT NULL, created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL ); ALTER TABLE ONLY public.stats_daily_model ADD CONSTRAINT stats_daily_model_pkey PRIMARY KEY (id); @@ -241,9 +306,321 @@ CREATE TABLE IF NOT EXISTS public.stats_user_daily ( total_cost double precision DEFAULT 0 NOT NULL, username character varying(255), created_at bigint NOT NULL, - updated_at bigint NOT NULL + updated_at bigint NOT NULL, + actual_total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + effective_input_tokens bigint DEFAULT 0 NOT NULL, + total_input_context bigint DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL, + settled_total_cost double precision DEFAULT 0 NOT NULL, + settled_total_requests bigint DEFAULT 0 NOT NULL, + settled_input_tokens bigint DEFAULT 0 NOT NULL, + settled_output_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + settled_cache_read_tokens bigint DEFAULT 0 NOT NULL, + settled_first_finalized_at_unix_secs bigint, + settled_last_finalized_at_unix_secs bigint ); ALTER TABLE ONLY public.stats_user_daily ADD CONSTRAINT stats_user_daily_pkey PRIMARY KEY (id); ALTER TABLE ONLY public.stats_user_daily ADD CONSTRAINT uq_stats_user_daily UNIQUE (date, user_id); +CREATE TABLE IF NOT EXISTS public.stats_user_summary ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + cutoff_date bigint NOT NULL, + all_time_requests bigint DEFAULT 0 NOT NULL, + all_time_success_requests bigint DEFAULT 0 NOT NULL, + all_time_error_requests bigint DEFAULT 0 NOT NULL, + all_time_input_tokens bigint DEFAULT 0 NOT NULL, + all_time_output_tokens bigint DEFAULT 0 NOT NULL, + all_time_cache_creation_tokens bigint DEFAULT 0 NOT NULL, + all_time_cache_read_tokens bigint DEFAULT 0 NOT NULL, + all_time_cost double precision DEFAULT 0 NOT NULL, + all_time_actual_cost double precision DEFAULT 0 NOT NULL, + active_days bigint DEFAULT 0 NOT NULL, + first_active_date bigint, + last_active_date bigint, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_summary ADD CONSTRAINT stats_user_summary_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_summary ADD CONSTRAINT uq_stats_user_summary_user_id UNIQUE (user_id); +CREATE INDEX IF NOT EXISTS idx_stats_user_summary_cutoff_date ON public.stats_user_summary USING btree (cutoff_date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_model ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + model character varying(255) NOT NULL, + total_requests bigint DEFAULT 0 NOT NULL, + success_requests bigint DEFAULT 0 NOT NULL, + input_tokens bigint DEFAULT 0 NOT NULL, + effective_input_tokens bigint DEFAULT 0 NOT NULL, + output_tokens bigint DEFAULT 0 NOT NULL, + total_tokens bigint DEFAULT 0 NOT NULL, + total_input_context bigint DEFAULT 0 NOT NULL, + cache_creation_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + total_cost double precision DEFAULT 0 NOT NULL, + actual_total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + successful_response_time_sum_ms double precision DEFAULT 0 NOT NULL, + successful_response_time_samples bigint DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_model ADD CONSTRAINT stats_user_daily_model_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_model ADD CONSTRAINT uq_stats_user_daily_model UNIQUE (user_id, date, model); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_date ON public.stats_user_daily_model USING btree (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_user_id ON public.stats_user_daily_model USING btree (user_id); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_provider ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + provider_name character varying(255) NOT NULL, + total_requests bigint DEFAULT 0 NOT NULL, + success_requests bigint DEFAULT 0 NOT NULL, + input_tokens bigint DEFAULT 0 NOT NULL, + effective_input_tokens bigint DEFAULT 0 NOT NULL, + output_tokens bigint DEFAULT 0 NOT NULL, + total_tokens bigint DEFAULT 0 NOT NULL, + total_input_context bigint DEFAULT 0 NOT NULL, + cache_creation_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + total_cost double precision DEFAULT 0 NOT NULL, + actual_total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + successful_response_time_sum_ms double precision DEFAULT 0 NOT NULL, + successful_response_time_samples bigint DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_provider ADD CONSTRAINT stats_user_daily_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_provider ADD CONSTRAINT uq_stats_user_daily_provider UNIQUE (user_id, date, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_provider_date ON public.stats_user_daily_provider USING btree (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_provider_user_id ON public.stats_user_daily_provider USING btree (user_id); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_api_format ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + api_format character varying(128) NOT NULL, + total_requests bigint DEFAULT 0 NOT NULL, + success_requests bigint DEFAULT 0 NOT NULL, + input_tokens bigint DEFAULT 0 NOT NULL, + effective_input_tokens bigint DEFAULT 0 NOT NULL, + output_tokens bigint DEFAULT 0 NOT NULL, + total_tokens bigint DEFAULT 0 NOT NULL, + total_input_context bigint DEFAULT 0 NOT NULL, + cache_creation_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_5m_tokens bigint DEFAULT 0 NOT NULL, + cache_creation_ephemeral_1h_tokens bigint DEFAULT 0 NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + total_cost double precision DEFAULT 0 NOT NULL, + actual_total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + successful_response_time_sum_ms double precision DEFAULT 0 NOT NULL, + successful_response_time_samples bigint DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_api_format ADD CONSTRAINT stats_user_daily_api_format_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_api_format ADD CONSTRAINT uq_stats_user_daily_api_format UNIQUE (user_id, date, api_format); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_api_format_date ON public.stats_user_daily_api_format USING btree (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_api_format_user_id ON public.stats_user_daily_api_format USING btree (user_id); + +CREATE TABLE IF NOT EXISTS public.stats_daily_model_provider ( + id character varying(64) NOT NULL, + date bigint NOT NULL, + model character varying(255) NOT NULL, + provider_name character varying(255) NOT NULL, + total_requests bigint DEFAULT 0 NOT NULL, + total_tokens bigint DEFAULT 0 NOT NULL, + total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_daily_model_provider ADD CONSTRAINT stats_daily_model_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_daily_model_provider ADD CONSTRAINT uq_stats_daily_model_provider UNIQUE (date, model, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_daily_model_provider_date ON public.stats_daily_model_provider USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_model_provider ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + model character varying(255) NOT NULL, + provider_name character varying(255) NOT NULL, + total_requests bigint DEFAULT 0 NOT NULL, + total_tokens bigint DEFAULT 0 NOT NULL, + total_cost double precision DEFAULT 0 NOT NULL, + response_time_sum_ms double precision DEFAULT 0 NOT NULL, + response_time_samples bigint DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_model_provider ADD CONSTRAINT stats_user_daily_model_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_model_provider ADD CONSTRAINT uq_stats_user_daily_model_provider UNIQUE (user_id, date, model, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_provider_date ON public.stats_user_daily_model_provider USING btree (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_provider_user_date ON public.stats_user_daily_model_provider USING btree (user_id, date); + +CREATE TABLE IF NOT EXISTS public.stats_daily_cost_savings ( + id character varying(64) NOT NULL, + date bigint NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_daily_cost_savings ADD CONSTRAINT stats_daily_cost_savings_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_daily_cost_savings ADD CONSTRAINT uq_stats_daily_cost_savings_date UNIQUE (date); + +CREATE TABLE IF NOT EXISTS public.stats_daily_cost_savings_provider ( + id character varying(64) NOT NULL, + date bigint NOT NULL, + provider_name character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_daily_cost_savings_provider ADD CONSTRAINT stats_daily_cost_savings_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_daily_cost_savings_provider ADD CONSTRAINT uq_stats_daily_cost_savings_provider UNIQUE (date, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_provider_date ON public.stats_daily_cost_savings_provider USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_daily_cost_savings_model ( + id character varying(64) NOT NULL, + date bigint NOT NULL, + model character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_daily_cost_savings_model ADD CONSTRAINT stats_daily_cost_savings_model_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_daily_cost_savings_model ADD CONSTRAINT uq_stats_daily_cost_savings_model UNIQUE (date, model); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_model_date ON public.stats_daily_cost_savings_model USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_daily_cost_savings_model_provider ( + id character varying(64) NOT NULL, + date bigint NOT NULL, + model character varying(255) NOT NULL, + provider_name character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_daily_cost_savings_model_provider ADD CONSTRAINT stats_daily_cost_savings_model_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_daily_cost_savings_model_provider ADD CONSTRAINT uq_stats_daily_cost_savings_model_provider UNIQUE (date, model, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_model_provider_date ON public.stats_daily_cost_savings_model_provider USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_cost_savings ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_cost_savings ADD CONSTRAINT stats_user_daily_cost_savings_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_cost_savings ADD CONSTRAINT uq_stats_user_daily_cost_savings UNIQUE (user_id, date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_date ON public.stats_user_daily_cost_savings USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_cost_savings_provider ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + provider_name character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_cost_savings_provider ADD CONSTRAINT stats_user_daily_cost_savings_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_cost_savings_provider ADD CONSTRAINT uq_stats_user_daily_cost_savings_provider UNIQUE (user_id, date, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_provider_date ON public.stats_user_daily_cost_savings_provider USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_cost_savings_model ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + model character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_cost_savings_model ADD CONSTRAINT stats_user_daily_cost_savings_model_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_cost_savings_model ADD CONSTRAINT uq_stats_user_daily_cost_savings_model UNIQUE (user_id, date, model); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_model_date ON public.stats_user_daily_cost_savings_model USING btree (date); + +CREATE TABLE IF NOT EXISTS public.stats_user_daily_cost_savings_model_provider ( + id character varying(64) NOT NULL, + user_id character varying(64) NOT NULL, + username character varying(255), + date bigint NOT NULL, + model character varying(255) NOT NULL, + provider_name character varying(255) NOT NULL, + cache_read_tokens bigint DEFAULT 0 NOT NULL, + cache_read_cost double precision DEFAULT 0 NOT NULL, + cache_creation_cost double precision DEFAULT 0 NOT NULL, + estimated_full_cost double precision DEFAULT 0 NOT NULL, + created_at bigint NOT NULL, + updated_at bigint NOT NULL +); + +ALTER TABLE ONLY public.stats_user_daily_cost_savings_model_provider ADD CONSTRAINT stats_user_daily_cost_savings_model_provider_pkey PRIMARY KEY (id); +ALTER TABLE ONLY public.stats_user_daily_cost_savings_model_provider ADD CONSTRAINT uq_stats_user_daily_cost_savings_model_provider UNIQUE (user_id, date, model, provider_name); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_model_provider_date ON public.stats_user_daily_cost_savings_model_provider USING btree (date); + diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql index 8de8633ff..f4495c760 100644 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql +++ b/crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql @@ -110,6 +110,60 @@ CREATE INDEX IF NOT EXISTS usage_request_id_idx ON "usage" (request_id); CREATE INDEX IF NOT EXISTS usage_user_id_idx ON "usage" (user_id); CREATE INDEX IF NOT EXISTS usage_wallet_id_idx ON "usage" (wallet_id); +CREATE TABLE IF NOT EXISTS usage_body_blobs ( + body_ref TEXT PRIMARY KEY NOT NULL, + request_id TEXT NOT NULL, + body_field TEXT NOT NULL, + payload_gzip BLOB NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (request_id, body_field), + CONSTRAINT usage_body_blobs_request_id_fkey FOREIGN KEY (request_id) REFERENCES usage (request_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS ix_usage_body_blobs_request_id ON usage_body_blobs (request_id); + +CREATE TABLE IF NOT EXISTS usage_http_audits ( + request_id TEXT PRIMARY KEY NOT NULL, + request_headers TEXT, + provider_request_headers TEXT, + response_headers TEXT, + client_response_headers TEXT, + request_body_ref TEXT, + provider_request_body_ref TEXT, + response_body_ref TEXT, + client_response_body_ref TEXT, + request_body_state TEXT, + provider_request_body_state TEXT, + response_body_state TEXT, + client_response_body_state TEXT, + body_capture_mode TEXT NOT NULL DEFAULT 'none', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + CONSTRAINT usage_http_audits_request_id_fkey FOREIGN KEY (request_id) REFERENCES usage (request_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS ix_usage_http_audits_updated_at ON usage_http_audits (updated_at); + +CREATE TABLE IF NOT EXISTS usage_routing_snapshots ( + request_id TEXT PRIMARY KEY NOT NULL, + candidate_id TEXT, + candidate_index INTEGER, + key_name TEXT, + planner_kind TEXT, + route_family TEXT, + route_kind TEXT, + execution_path TEXT, + local_execution_runtime_miss_reason TEXT, + selected_provider_id TEXT, + selected_endpoint_id TEXT, + selected_provider_api_key_id TEXT, + has_format_conversion INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + CONSTRAINT usage_routing_snapshots_request_id_fkey FOREIGN KEY (request_id) REFERENCES usage (request_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_route_family_kind ON usage_routing_snapshots (route_family, route_kind); +CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_candidate_id ON usage_routing_snapshots (candidate_id); + CREATE TABLE IF NOT EXISTS usage_counter_deltas ( id TEXT PRIMARY KEY NOT NULL, request_id TEXT NOT NULL, @@ -147,10 +201,39 @@ CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( wallet_gift_balance_before REAL, wallet_gift_balance_after REAL, provider_monthly_used_usd REAL, + billing_snapshot_schema_version TEXT, + billing_snapshot_status TEXT, + rate_multiplier REAL, + is_free_tier INTEGER, + input_price_per_1m REAL, + output_price_per_1m REAL, + cache_creation_price_per_1m REAL, + cache_read_price_per_1m REAL, + price_per_request REAL, + settlement_snapshot_schema_version TEXT, + settlement_snapshot TEXT, + billing_dimensions TEXT, + billing_input_tokens INTEGER, + billing_effective_input_tokens INTEGER, + billing_output_tokens INTEGER, + billing_cache_creation_tokens INTEGER, + billing_cache_creation_5m_tokens INTEGER, + billing_cache_creation_1h_tokens INTEGER, + billing_cache_read_tokens INTEGER, + billing_total_input_context INTEGER, + billing_cache_creation_cost_usd REAL, + billing_cache_read_cost_usd REAL, + billing_total_cost_usd REAL, + billing_actual_total_cost_usd REAL, + billing_pricing_source TEXT, + billing_rule_id TEXT, + billing_rule_version TEXT, finalized_at INTEGER, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS usage_settlement_snapshots_billing_status_idx ON usage_settlement_snapshots (billing_status); CREATE INDEX IF NOT EXISTS usage_settlement_snapshots_wallet_id_idx ON usage_settlement_snapshots (wallet_id); +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_schema_version ON usage_settlement_snapshots (settlement_snapshot_schema_version); +CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_pricing_source ON usage_settlement_snapshots (billing_pricing_source); diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql index 6f5ac00cb..60af13fe3 100644 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql +++ b/crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql @@ -18,6 +18,26 @@ CREATE TABLE IF NOT EXISTS stats_hourly ( aggregated_at INTEGER, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + cache_hit_total_requests INTEGER NOT NULL DEFAULT 0, + cache_hit_requests INTEGER NOT NULL DEFAULT 0, + completed_total_requests INTEGER NOT NULL DEFAULT 0, + completed_cache_hit_requests INTEGER NOT NULL DEFAULT 0, + completed_input_tokens INTEGER NOT NULL DEFAULT 0, + completed_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + completed_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + completed_total_input_context INTEGER NOT NULL DEFAULT 0, + completed_cache_creation_cost REAL NOT NULL DEFAULT 0, + completed_cache_read_cost REAL NOT NULL DEFAULT 0, + settled_total_cost REAL NOT NULL DEFAULT 0, + settled_total_requests INTEGER NOT NULL DEFAULT 0, + settled_input_tokens INTEGER NOT NULL DEFAULT 0, + settled_output_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + settled_first_finalized_at_unix_secs INTEGER, + settled_last_finalized_at_unix_secs INTEGER, UNIQUE (hour_utc) ); @@ -53,6 +73,19 @@ CREATE TABLE IF NOT EXISTS stats_hourly_user ( total_cost REAL NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + settled_total_cost REAL NOT NULL DEFAULT 0, + settled_total_requests INTEGER NOT NULL DEFAULT 0, + settled_input_tokens INTEGER NOT NULL DEFAULT 0, + settled_output_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + settled_first_finalized_at_unix_secs INTEGER, + settled_last_finalized_at_unix_secs INTEGER, UNIQUE (hour_utc, user_id) ); @@ -67,6 +100,8 @@ CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( total_cost REAL NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, UNIQUE (hour_utc, user_id, model) ); @@ -93,6 +128,8 @@ CREATE TABLE IF NOT EXISTS stats_hourly_model ( avg_response_time_ms REAL NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, UNIQUE (hour_utc, model) ); @@ -139,6 +176,30 @@ CREATE TABLE IF NOT EXISTS stats_daily ( p50_first_byte_time_ms INTEGER, p90_first_byte_time_ms INTEGER, p99_first_byte_time_ms INTEGER, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_hit_total_requests INTEGER NOT NULL DEFAULT 0, + cache_hit_requests INTEGER NOT NULL DEFAULT 0, + completed_total_requests INTEGER NOT NULL DEFAULT 0, + completed_cache_hit_requests INTEGER NOT NULL DEFAULT 0, + completed_input_tokens INTEGER NOT NULL DEFAULT 0, + completed_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + completed_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + completed_total_input_context INTEGER NOT NULL DEFAULT 0, + completed_cache_creation_cost REAL NOT NULL DEFAULT 0, + completed_cache_read_cost REAL NOT NULL DEFAULT 0, + settled_total_cost REAL NOT NULL DEFAULT 0, + settled_total_requests INTEGER NOT NULL DEFAULT 0, + settled_input_tokens INTEGER NOT NULL DEFAULT 0, + settled_output_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + settled_first_finalized_at_unix_secs INTEGER, + settled_last_finalized_at_unix_secs INTEGER, UNIQUE (date) ); @@ -155,6 +216,10 @@ CREATE TABLE IF NOT EXISTS stats_daily_model ( avg_response_time_ms REAL NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, UNIQUE (date, model) ); @@ -218,6 +283,290 @@ CREATE TABLE IF NOT EXISTS stats_user_daily ( username TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + settled_total_cost REAL NOT NULL DEFAULT 0, + settled_total_requests INTEGER NOT NULL DEFAULT 0, + settled_input_tokens INTEGER NOT NULL DEFAULT 0, + settled_output_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + settled_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + settled_first_finalized_at_unix_secs INTEGER, + settled_last_finalized_at_unix_secs INTEGER, UNIQUE (date, user_id) ); +CREATE TABLE IF NOT EXISTS stats_user_summary ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + cutoff_date INTEGER NOT NULL, + all_time_requests INTEGER NOT NULL DEFAULT 0, + all_time_success_requests INTEGER NOT NULL DEFAULT 0, + all_time_error_requests INTEGER NOT NULL DEFAULT 0, + all_time_input_tokens INTEGER NOT NULL DEFAULT 0, + all_time_output_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cache_read_tokens INTEGER NOT NULL DEFAULT 0, + all_time_cost REAL NOT NULL DEFAULT 0, + all_time_actual_cost REAL NOT NULL DEFAULT 0, + active_days INTEGER NOT NULL DEFAULT 0, + first_active_date INTEGER, + last_active_date INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_summary_cutoff_date ON stats_user_summary (cutoff_date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + model TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, model) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_date ON stats_user_daily_model (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_user_id ON stats_user_daily_model (user_id); + +CREATE TABLE IF NOT EXISTS stats_user_daily_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_provider_date ON stats_user_daily_provider (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_provider_user_id ON stats_user_daily_provider (user_id); + +CREATE TABLE IF NOT EXISTS stats_user_daily_api_format ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + api_format TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + success_requests INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + effective_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_input_context INTEGER NOT NULL DEFAULT 0, + cache_creation_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_tokens INTEGER NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + actual_total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + successful_response_time_sum_ms REAL NOT NULL DEFAULT 0, + successful_response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, api_format) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_api_format_date ON stats_user_daily_api_format (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_api_format_user_id ON stats_user_daily_api_format (user_id); + +CREATE TABLE IF NOT EXISTS stats_daily_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + date INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (date, model, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_daily_model_provider_date ON stats_daily_model_provider (date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + total_requests INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + total_cost REAL NOT NULL DEFAULT 0, + response_time_sum_ms REAL NOT NULL DEFAULT 0, + response_time_samples INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, model, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_provider_date ON stats_user_daily_model_provider (date); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_model_provider_user_date ON stats_user_daily_model_provider (user_id, date); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings ( + id TEXT PRIMARY KEY NOT NULL, + date INTEGER NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (date) +); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_provider ( + id TEXT PRIMARY KEY NOT NULL, + date INTEGER NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (date, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_provider_date ON stats_daily_cost_savings_provider (date); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model ( + id TEXT PRIMARY KEY NOT NULL, + date INTEGER NOT NULL, + model TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (date, model) +); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_model_date ON stats_daily_cost_savings_model (date); + +CREATE TABLE IF NOT EXISTS stats_daily_cost_savings_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + date INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (date, model, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_daily_cost_savings_model_provider_date ON stats_daily_cost_savings_model_provider (date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_date ON stats_user_daily_cost_savings (date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_provider_date ON stats_user_daily_cost_savings_provider (date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + model TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, model) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_model_date ON stats_user_daily_cost_savings_model (date); + +CREATE TABLE IF NOT EXISTS stats_user_daily_cost_savings_model_provider ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + username TEXT, + date INTEGER NOT NULL, + model TEXT NOT NULL, + provider_name TEXT NOT NULL, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_cost REAL NOT NULL DEFAULT 0, + cache_creation_cost REAL NOT NULL DEFAULT 0, + estimated_full_cost REAL NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, date, model, provider_name) +); +CREATE INDEX IF NOT EXISTS idx_stats_user_daily_cost_savings_model_provider_date ON stats_user_daily_cost_savings_model_provider (date); + diff --git a/crates/aether-data/runtime/schema/logical/006_usage.toml b/crates/aether-data/runtime/schema/logical/006_usage.toml index 018c367cc..f137bae1d 100644 --- a/crates/aether-data/runtime/schema/logical/006_usage.toml +++ b/crates/aether-data/runtime/schema/logical/006_usage.toml @@ -554,6 +554,259 @@ columns = ["user_id"] name = "usage_wallet_id_idx" columns = ["wallet_id"] +[table.usage_body_blobs] +domain = "usage" +order = 11 +primary_key = ["body_ref"] + +[[table.usage_body_blobs.columns]] +name = "body_ref" +type = "text" +length = 160 + +[[table.usage_body_blobs.columns]] +name = "request_id" +type = "text" +length = 128 + +[[table.usage_body_blobs.columns]] +name = "body_field" +type = "text" +length = 50 + +[[table.usage_body_blobs.columns]] +name = "payload_gzip" +type = "bytes" + +[[table.usage_body_blobs.columns]] +name = "created_at" +type = "timestamp" + +[[table.usage_body_blobs.columns]] +name = "updated_at" +type = "timestamp" + +[[table.usage_body_blobs.uniques]] +name = "usage_body_blobs_request_id_field_key" +columns = ["request_id", "body_field"] + +[[table.usage_body_blobs.indexes]] +name = "ix_usage_body_blobs_request_id" +columns = ["request_id"] + +[[table.usage_body_blobs.foreign_keys]] +name = "usage_body_blobs_request_id_fkey" +columns = ["request_id"] +references_table = "usage" +references_columns = ["request_id"] +on_delete = "cascade" + +[table.usage_http_audits] +domain = "usage" +order = 12 +primary_key = ["request_id"] + +[[table.usage_http_audits.columns]] +name = "request_id" +type = "text" +length = 128 + +[[table.usage_http_audits.columns]] +name = "request_headers" +type = "json" +nullable = true + +[[table.usage_http_audits.columns]] +name = "provider_request_headers" +type = "json" +nullable = true + +[[table.usage_http_audits.columns]] +name = "response_headers" +type = "json" +nullable = true + +[[table.usage_http_audits.columns]] +name = "client_response_headers" +type = "json" +nullable = true + +[[table.usage_http_audits.columns]] +name = "request_body_ref" +type = "text" +length = 160 +nullable = true + +[[table.usage_http_audits.columns]] +name = "provider_request_body_ref" +type = "text" +length = 160 +nullable = true + +[[table.usage_http_audits.columns]] +name = "response_body_ref" +type = "text" +length = 160 +nullable = true + +[[table.usage_http_audits.columns]] +name = "client_response_body_ref" +type = "text" +length = 160 +nullable = true + +[[table.usage_http_audits.columns]] +name = "request_body_state" +type = "text" +length = 32 +nullable = true + +[[table.usage_http_audits.columns]] +name = "provider_request_body_state" +type = "text" +length = 32 +nullable = true + +[[table.usage_http_audits.columns]] +name = "response_body_state" +type = "text" +length = 32 +nullable = true + +[[table.usage_http_audits.columns]] +name = "client_response_body_state" +type = "text" +length = 32 +nullable = true + +[[table.usage_http_audits.columns]] +name = "body_capture_mode" +type = "text" +length = 32 +default = "none" + +[[table.usage_http_audits.columns]] +name = "created_at" +type = "timestamp" + +[[table.usage_http_audits.columns]] +name = "updated_at" +type = "timestamp" + +[[table.usage_http_audits.indexes]] +name = "ix_usage_http_audits_updated_at" +columns = ["updated_at"] + +[[table.usage_http_audits.foreign_keys]] +name = "usage_http_audits_request_id_fkey" +columns = ["request_id"] +references_table = "usage" +references_columns = ["request_id"] +on_delete = "cascade" + +[table.usage_routing_snapshots] +domain = "usage" +order = 13 +primary_key = ["request_id"] + +[[table.usage_routing_snapshots.columns]] +name = "request_id" +type = "text" +length = 128 + +[[table.usage_routing_snapshots.columns]] +name = "candidate_id" +type = "text" +length = 160 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "candidate_index" +type = "int64" +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "key_name" +type = "text" +length = 255 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "planner_kind" +type = "text" +length = 120 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "route_family" +type = "text" +length = 80 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "route_kind" +type = "text" +length = 80 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "execution_path" +type = "text" +length = 80 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "local_execution_runtime_miss_reason" +type = "text" +length = 255 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "selected_provider_id" +type = "text_id" +length = 100 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "selected_endpoint_id" +type = "text_id" +length = 100 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "selected_provider_api_key_id" +type = "text_id" +length = 100 +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "has_format_conversion" +type = "bool" +nullable = true + +[[table.usage_routing_snapshots.columns]] +name = "created_at" +type = "timestamp" + +[[table.usage_routing_snapshots.columns]] +name = "updated_at" +type = "timestamp" + +[[table.usage_routing_snapshots.indexes]] +name = "ix_usage_routing_snapshots_route_family_kind" +columns = ["route_family", "route_kind"] + +[[table.usage_routing_snapshots.indexes]] +name = "ix_usage_routing_snapshots_candidate_id" +columns = ["candidate_id"] + +[[table.usage_routing_snapshots.foreign_keys]] +name = "usage_routing_snapshots_request_id_fkey" +columns = ["request_id"] +references_table = "usage" +references_columns = ["request_id"] +on_delete = "cascade" + [table.usage_counter_deltas] domain = "usage" order = 15 @@ -725,6 +978,147 @@ name = "provider_monthly_used_usd" type = "float64" nullable = true +[[table.usage_settlement_snapshots.columns]] +name = "billing_snapshot_schema_version" +type = "text" +length = 20 +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_snapshot_status" +type = "text" +length = 20 +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "rate_multiplier" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "is_free_tier" +type = "bool" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "input_price_per_1m" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "output_price_per_1m" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "cache_creation_price_per_1m" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "cache_read_price_per_1m" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "price_per_request" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "settlement_snapshot_schema_version" +type = "text" +length = 20 +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "settlement_snapshot" +type = "json" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_dimensions" +type = "json" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_input_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_effective_input_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_output_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_creation_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_creation_5m_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_creation_1h_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_read_tokens" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_total_input_context" +type = "int64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_creation_cost_usd" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_cache_read_cost_usd" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_total_cost_usd" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_actual_total_cost_usd" +type = "float64" +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_pricing_source" +type = "text" +length = 50 +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_rule_id" +type = "text" +length = 100 +nullable = true + +[[table.usage_settlement_snapshots.columns]] +name = "billing_rule_version" +type = "text" +length = 50 +nullable = true + [[table.usage_settlement_snapshots.columns]] name = "finalized_at" type = "unix_seconds" @@ -745,3 +1139,11 @@ columns = ["billing_status"] [[table.usage_settlement_snapshots.indexes]] name = "usage_settlement_snapshots_wallet_id_idx" columns = ["wallet_id"] + +[[table.usage_settlement_snapshots.indexes]] +name = "ix_usage_settlement_snapshots_schema_version" +columns = ["settlement_snapshot_schema_version"] + +[[table.usage_settlement_snapshots.indexes]] +name = "ix_usage_settlement_snapshots_pricing_source" +columns = ["billing_pricing_source"] diff --git a/crates/aether-data/runtime/schema/logical/007_stats.toml b/crates/aether-data/runtime/schema/logical/007_stats.toml index 9190e6e0a..70b7ccfe6 100644 --- a/crates/aether-data/runtime/schema/logical/007_stats.toml +++ b/crates/aether-data/runtime/schema/logical/007_stats.toml @@ -919,3 +919,1545 @@ type = "unix_seconds" [[table.stats_user_daily.uniques]] name = "uq_stats_user_daily" columns = ["date", "user_id"] + +# Advanced rollup columns are deployed incrementally by the portable drivers. +[[table.stats_user_daily.columns]] +name = "actual_total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "effective_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "total_input_context" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_total_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_output_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily.columns]] +name = "settled_first_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_user_daily.columns]] +name = "settled_last_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_hourly_user.columns]] +name = "cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "actual_total_cost" +type = "float64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_total_cost" +type = "float64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_total_requests" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_input_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_output_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_user.columns]] +name = "settled_first_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_hourly_user.columns]] +name = "settled_last_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_daily.columns]] +name = "effective_input_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "total_input_context" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_daily.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "cache_hit_total_requests" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "cache_hit_requests" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_total_requests" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_cache_hit_requests" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_input_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_total_input_context" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_daily.columns]] +name = "completed_cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_total_cost" +type = "float64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_total_requests" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_input_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_output_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily.columns]] +name = "settled_first_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_daily.columns]] +name = "settled_last_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_hourly.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_hourly.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "cache_hit_total_requests" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "cache_hit_requests" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_total_requests" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_cache_hit_requests" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_input_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_total_input_context" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_hourly.columns]] +name = "completed_cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_total_cost" +type = "float64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_total_requests" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_input_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_output_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly.columns]] +name = "settled_first_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_hourly.columns]] +name = "settled_last_finalized_at_unix_secs" +type = "unix_seconds" +nullable = true + +[[table.stats_daily_model.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_daily_model.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_daily_model.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_model.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_hourly_model.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_hourly_model.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_hourly_user_model.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_hourly_user_model.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[table.stats_user_summary] +domain = "stats" +order = 120 +primary_key = ["id"] + +[[table.stats_user_summary.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_summary.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_summary.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_summary.columns]] +name = "cutoff_date" +type = "unix_seconds" + +[[table.stats_user_summary.columns]] +name = "all_time_requests" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_success_requests" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_error_requests" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_output_tokens" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_cost" +type = "float64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "all_time_actual_cost" +type = "float64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "active_days" +type = "int64" +default = 0 + +[[table.stats_user_summary.columns]] +name = "first_active_date" +type = "unix_seconds" +nullable = true + +[[table.stats_user_summary.columns]] +name = "last_active_date" +type = "unix_seconds" +nullable = true + +[[table.stats_user_summary.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_summary.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_summary.uniques]] +name = "uq_stats_user_summary_user_id" +columns = ["user_id"] + +[[table.stats_user_summary.indexes]] +name = "idx_stats_user_summary_cutoff_date" +columns = ["cutoff_date"] + +[table.stats_user_daily_model] +domain = "stats" +order = 130 +primary_key = ["id"] + +[[table.stats_user_daily_model.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_model.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_model.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_model.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_model.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_user_daily_model.columns]] +name = "total_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "success_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "effective_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "output_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "total_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "total_input_context" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "actual_total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "successful_response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "successful_response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_model.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_model.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_model.uniques]] +name = "uq_stats_user_daily_model" +columns = ["user_id", "date", "model"] + +[[table.stats_user_daily_model.indexes]] +name = "idx_stats_user_daily_model_date" +columns = ["date"] + +[[table.stats_user_daily_model.indexes]] +name = "idx_stats_user_daily_model_user_id" +columns = ["user_id"] + +[table.stats_user_daily_provider] +domain = "stats" +order = 140 +primary_key = ["id"] + +[[table.stats_user_daily_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_provider.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_provider.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_user_daily_provider.columns]] +name = "total_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "success_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "effective_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "output_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "total_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "total_input_context" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "actual_total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "successful_response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "successful_response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_provider.uniques]] +name = "uq_stats_user_daily_provider" +columns = ["user_id", "date", "provider_name"] + +[[table.stats_user_daily_provider.indexes]] +name = "idx_stats_user_daily_provider_date" +columns = ["date"] + +[[table.stats_user_daily_provider.indexes]] +name = "idx_stats_user_daily_provider_user_id" +columns = ["user_id"] + +[table.stats_user_daily_api_format] +domain = "stats" +order = 150 +primary_key = ["id"] + +[[table.stats_user_daily_api_format.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_api_format.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_api_format.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_api_format.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_api_format.columns]] +name = "api_format" +type = "text" +length = 128 + +[[table.stats_user_daily_api_format.columns]] +name = "total_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "success_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "effective_input_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "output_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "total_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "total_input_context" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "cache_creation_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "cache_creation_ephemeral_5m_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "cache_creation_ephemeral_1h_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "actual_total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "successful_response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "successful_response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_api_format.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_api_format.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_api_format.uniques]] +name = "uq_stats_user_daily_api_format" +columns = ["user_id", "date", "api_format"] + +[[table.stats_user_daily_api_format.indexes]] +name = "idx_stats_user_daily_api_format_date" +columns = ["date"] + +[[table.stats_user_daily_api_format.indexes]] +name = "idx_stats_user_daily_api_format_user_id" +columns = ["user_id"] + +[table.stats_daily_model_provider] +domain = "stats" +order = 160 +primary_key = ["id"] + +[[table.stats_daily_model_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_daily_model_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_daily_model_provider.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_daily_model_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_daily_model_provider.columns]] +name = "total_requests" +type = "int64" +default = 0 + +[[table.stats_daily_model_provider.columns]] +name = "total_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_model_provider.columns]] +name = "total_cost" +type = "float64" +default = 0 + +[[table.stats_daily_model_provider.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_daily_model_provider.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_daily_model_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_daily_model_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_daily_model_provider.uniques]] +name = "uq_stats_daily_model_provider" +columns = ["date", "model", "provider_name"] + +[[table.stats_daily_model_provider.indexes]] +name = "idx_stats_daily_model_provider_date" +columns = ["date"] + +[table.stats_user_daily_model_provider] +domain = "stats" +order = 170 +primary_key = ["id"] + +[[table.stats_user_daily_model_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_model_provider.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_model_provider.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_model_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_model_provider.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_user_daily_model_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_user_daily_model_provider.columns]] +name = "total_requests" +type = "int64" +default = 0 + +[[table.stats_user_daily_model_provider.columns]] +name = "total_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_model_provider.columns]] +name = "total_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_model_provider.columns]] +name = "response_time_sum_ms" +type = "float64" +default = 0 + +[[table.stats_user_daily_model_provider.columns]] +name = "response_time_samples" +type = "int64" +default = 0 + +[[table.stats_user_daily_model_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_model_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_model_provider.uniques]] +name = "uq_stats_user_daily_model_provider" +columns = ["user_id", "date", "model", "provider_name"] + +[[table.stats_user_daily_model_provider.indexes]] +name = "idx_stats_user_daily_model_provider_date" +columns = ["date"] + +[[table.stats_user_daily_model_provider.indexes]] +name = "idx_stats_user_daily_model_provider_user_date" +columns = ["user_id", "date"] + +[table.stats_daily_cost_savings] +domain = "stats" +order = 180 +primary_key = ["id"] + +[[table.stats_daily_cost_savings.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_daily_cost_savings.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_daily_cost_savings.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_cost_savings.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings.uniques]] +name = "uq_stats_daily_cost_savings_date" +columns = ["date"] + +[table.stats_daily_cost_savings_provider] +domain = "stats" +order = 190 +primary_key = ["id"] + +[[table.stats_daily_cost_savings_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_provider.uniques]] +name = "uq_stats_daily_cost_savings_provider" +columns = ["date", "provider_name"] + +[[table.stats_daily_cost_savings_provider.indexes]] +name = "idx_stats_daily_cost_savings_provider_date" +columns = ["date"] + +[table.stats_daily_cost_savings_model] +domain = "stats" +order = 200 +primary_key = ["id"] + +[[table.stats_daily_cost_savings_model.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_daily_cost_savings_model.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_daily_cost_savings_model.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_cost_savings_model.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model.uniques]] +name = "uq_stats_daily_cost_savings_model" +columns = ["date", "model"] + +[[table.stats_daily_cost_savings_model.indexes]] +name = "idx_stats_daily_cost_savings_model_date" +columns = ["date"] + +[table.stats_daily_cost_savings_model_provider] +domain = "stats" +order = 210 +primary_key = ["id"] + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_daily_cost_savings_model_provider.uniques]] +name = "uq_stats_daily_cost_savings_model_provider" +columns = ["date", "model", "provider_name"] + +[[table.stats_daily_cost_savings_model_provider.indexes]] +name = "idx_stats_daily_cost_savings_model_provider_date" +columns = ["date"] + +[table.stats_user_daily_cost_savings] +domain = "stats" +order = 220 +primary_key = ["id"] + +[[table.stats_user_daily_cost_savings.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_cost_savings.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_cost_savings.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings.uniques]] +name = "uq_stats_user_daily_cost_savings" +columns = ["user_id", "date"] + +[[table.stats_user_daily_cost_savings.indexes]] +name = "idx_stats_user_daily_cost_savings_date" +columns = ["date"] + +[table.stats_user_daily_cost_savings_provider] +domain = "stats" +order = 230 +primary_key = ["id"] + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_provider.uniques]] +name = "uq_stats_user_daily_cost_savings_provider" +columns = ["user_id", "date", "provider_name"] + +[[table.stats_user_daily_cost_savings_provider.indexes]] +name = "idx_stats_user_daily_cost_savings_provider_date" +columns = ["date"] + +[table.stats_user_daily_cost_savings_model] +domain = "stats" +order = 240 +primary_key = ["id"] + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model.uniques]] +name = "uq_stats_user_daily_cost_savings_model" +columns = ["user_id", "date", "model"] + +[[table.stats_user_daily_cost_savings_model.indexes]] +name = "idx_stats_user_daily_cost_savings_model_date" +columns = ["date"] + +[table.stats_user_daily_cost_savings_model_provider] +domain = "stats" +order = 250 +primary_key = ["id"] + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "user_id" +type = "text_id" +length = 64 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "username" +type = "text" +length = 255 +nullable = true + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "date" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "model" +type = "text" +length = 255 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "provider_name" +type = "text" +length = 255 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "cache_read_tokens" +type = "int64" +default = 0 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "cache_read_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "cache_creation_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "estimated_full_cost" +type = "float64" +default = 0 + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "created_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model_provider.columns]] +name = "updated_at" +type = "unix_seconds" + +[[table.stats_user_daily_cost_savings_model_provider.uniques]] +name = "uq_stats_user_daily_cost_savings_model_provider" +columns = ["user_id", "date", "model", "provider_name"] + +[[table.stats_user_daily_cost_savings_model_provider.indexes]] +name = "idx_stats_user_daily_cost_savings_model_provider_date" +columns = ["date"] diff --git a/crates/aether-data/runtime/src/backend/maintenance/mysql.rs b/crates/aether-data/runtime/src/backend/maintenance/mysql.rs index bf14f7736..87ecbeeb0 100644 --- a/crates/aether-data/runtime/src/backend/maintenance/mysql.rs +++ b/crates/aether-data/runtime/src/backend/maintenance/mysql.rs @@ -14,11 +14,11 @@ impl MysqlBackend { let table_name = maintenance_identifier(table_name)?; summary.attempted += 1; let statement = format!("ANALYZE TABLE `{table_name}`"); - if sqlx::raw_sql(&statement) - .execute(self.pool()) + if sqlx::query_as::<_, (String, String, String, String)>(&statement) + .fetch_all(self.pool()) .await .map_sql_err() - .is_ok() + .is_ok_and(|rows| mysql_analyze_succeeded(&rows)) { summary.succeeded += 1; } @@ -26,3 +26,42 @@ impl MysqlBackend { Ok(summary) } } + +fn mysql_analyze_succeeded(rows: &[(String, String, String, String)]) -> bool { + !rows.is_empty() + && rows.iter().any(|(_, _, message_type, message)| { + message_type.eq_ignore_ascii_case("status") && message.eq_ignore_ascii_case("ok") + }) + && rows + .iter() + .all(|(_, _, message_type, _)| !message_type.eq_ignore_ascii_case("error")) +} + +#[cfg(test)] +mod tests { + use super::mysql_analyze_succeeded; + + fn row(message_type: &str, message: &str) -> (String, String, String, String) { + ( + "aether.usage".to_string(), + "analyze".to_string(), + message_type.to_string(), + message.to_string(), + ) + } + + #[test] + fn analyze_requires_an_explicit_ok_status() { + assert!(mysql_analyze_succeeded(&[row("status", "OK")])); + assert!(!mysql_analyze_succeeded(&[])); + assert!(!mysql_analyze_succeeded(&[row("note", "skipped")])); + } + + #[test] + fn analyze_rejects_error_rows_even_when_an_ok_row_is_present() { + assert!(!mysql_analyze_succeeded(&[ + row("Error", "Table does not exist"), + row("status", "OK"), + ])); + } +} diff --git a/crates/aether-data/runtime/src/backend/mysql.rs b/crates/aether-data/runtime/src/backend/mysql.rs index fb881c4be..9ee96bc02 100644 --- a/crates/aether-data/runtime/src/backend/mysql.rs +++ b/crates/aether-data/runtime/src/backend/mysql.rs @@ -538,6 +538,20 @@ WHERE wallet_id = ? for sql in [ "DELETE FROM stats_daily WHERE `date` = 0", "DELETE FROM stats_hourly WHERE hour_utc = 3600", + "DELETE FROM stats_user_summary WHERE user_id LIKE 'user-%'", + "DELETE FROM stats_user_daily_model WHERE `date` = 0", + "DELETE FROM stats_user_daily_provider WHERE `date` = 0", + "DELETE FROM stats_user_daily_api_format WHERE `date` = 0", + "DELETE FROM stats_daily_model_provider WHERE `date` = 0", + "DELETE FROM stats_user_daily_model_provider WHERE `date` = 0", + "DELETE FROM stats_daily_cost_savings WHERE `date` = 0", + "DELETE FROM stats_daily_cost_savings_provider WHERE `date` = 0", + "DELETE FROM stats_daily_cost_savings_model WHERE `date` = 0", + "DELETE FROM stats_daily_cost_savings_model_provider WHERE `date` = 0", + "DELETE FROM stats_user_daily_cost_savings WHERE `date` = 0", + "DELETE FROM stats_user_daily_cost_savings_provider WHERE `date` = 0", + "DELETE FROM stats_user_daily_cost_savings_model WHERE `date` = 0", + "DELETE FROM stats_user_daily_cost_savings_model_provider WHERE `date` = 0", "DELETE FROM usage_settlement_snapshots WHERE request_id LIKE 'request-daily-%' OR request_id LIKE 'stats-%'", "DELETE FROM `usage` WHERE request_id LIKE 'request-%' OR request_id LIKE 'export-request-%' OR request_id LIKE 'stats-%'", ] { @@ -550,19 +564,21 @@ WHERE wallet_id = ? sqlx::query( r#" INSERT INTO `usage` ( - request_id, user_id, api_key_id, provider_name, model, status, billing_status, + request_id, user_id, api_key_id, provider_name, model, api_format, status, billing_status, status_code, error_category, input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens, total_cost_usd, - actual_total_cost_usd, response_time_ms, created_at_unix_ms, updated_at_unix_secs + actual_total_cost_usd, cache_creation_cost_usd, cache_read_cost_usd, + input_price_per_1m, response_time_ms, first_byte_time_ms, + created_at_unix_ms, updated_at_unix_secs ) VALUES - ('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'completed', 'settled', - 200, NULL, 10, 20, 1, 2, 0.30, 0.25, 100, 3600000, 3600), - ('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'failed', 'void', - 500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 300, 3610000, 3610), - ('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'pending', 'pending', - NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3620000, 3620), - ('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'completed', 'settled', - 200, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3630000, 3630) + ('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'openai', 'completed', 'settled', + 200, NULL, 10, 20, 1, 2, 0.30, 0.25, 0.01, 0.02, 10.0, 100, 50, 3600, 3600), + ('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'claude', 'failed', 'void', + 500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 0.00, 0.01, 20.0, 300, 200, 3610, 3610), + ('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'openai', 'pending', 'pending', + NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 0.00, 0.00, 0.0, 50, 25, 3620, 3620), + ('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'openai', 'completed', 'settled', + 200, NULL, 100, 100, 0, 0, 9.99, 9.99, 0.00, 0.00, 0.0, 50, 25, 3630, 3630) "#, ) .execute(backend.pool()) @@ -627,7 +643,7 @@ WHERE hour_utc = 3600 assert_eq!(daily.total_requests, 2); assert_eq!(daily.model_rows, 2); assert_eq!(daily.provider_rows, 2); - assert_eq!(daily.api_key_rows, 2); + assert_eq!(daily.api_key_rows, 4); assert_eq!(daily.error_rows, 1); assert_eq!(daily.user_rows, 2); @@ -642,5 +658,46 @@ WHERE `date` = 0 .await .expect("daily stats row should load"); assert_eq!(daily_row, (2, 1, 1, 2)); + + let enriched_daily = sqlx::query_as::<_, (i64, i64, i64, i64, i64, Option)>( + r#" +SELECT effective_input_tokens, total_input_context, cache_hit_total_requests, + completed_total_requests, settled_total_requests, p50_response_time_ms +FROM stats_daily +WHERE `date` = 0 +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("mysql enriched daily stats row should load"); + assert_eq!(enriched_daily, (13, 17, 4, 2, 2, None)); + + for table in [ + "stats_user_summary", + "stats_user_daily_model", + "stats_user_daily_provider", + "stats_user_daily_api_format", + "stats_daily_model_provider", + "stats_user_daily_model_provider", + "stats_daily_cost_savings", + "stats_daily_cost_savings_provider", + "stats_daily_cost_savings_model", + "stats_daily_cost_savings_model_provider", + "stats_user_daily_cost_savings", + "stats_user_daily_cost_savings_provider", + "stats_user_daily_cost_savings_model", + "stats_user_daily_cost_savings_model_provider", + ] { + let sql = if table == "stats_user_summary" { + format!("SELECT COUNT(*) FROM {table} WHERE user_id IN ('user-1', 'user-2')") + } else { + format!("SELECT COUNT(*) FROM {table} WHERE `date` = 0") + }; + let count: i64 = sqlx::query_scalar(&sql) + .fetch_one(backend.pool()) + .await + .expect("mysql advanced stats count should load"); + assert!(count > 0, "{table} should be populated"); + } } } diff --git a/crates/aether-data/runtime/src/backend/sqlite.rs b/crates/aether-data/runtime/src/backend/sqlite.rs index 3f5e2f0f1..e92e83dc8 100644 --- a/crates/aether-data/runtime/src/backend/sqlite.rs +++ b/crates/aether-data/runtime/src/backend/sqlite.rs @@ -586,22 +586,6 @@ VALUES ('target-key-1', 'target-user-1', 'hash-target-key', 'target key', 1, 1) .await .expect("sqlite migrations should run"); - for (column, ty) in [ - ("request_body", "TEXT"), - ("response_body", "TEXT"), - ("provider_request_body", "TEXT"), - ("client_response_body", "TEXT"), - ("request_body_compressed", "BLOB"), - ("response_body_compressed", "BLOB"), - ("provider_request_body_compressed", "BLOB"), - ("client_response_body_compressed", "BLOB"), - ] { - sqlx::query(&format!(r#"ALTER TABLE "usage" ADD COLUMN {column} {ty}"#)) - .execute(backend.pool()) - .await - .expect("legacy body column should be added"); - } - sqlx::query( r#" INSERT INTO "usage" ( @@ -842,25 +826,40 @@ WHERE billing_date = '2026-05-03' sqlx::query( r#" INSERT INTO "usage" ( - request_id, user_id, api_key_id, provider_name, model, status, billing_status, + request_id, user_id, api_key_id, provider_name, model, api_format, status, billing_status, status_code, error_category, input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens, total_cost_usd, - actual_total_cost_usd, response_time_ms, created_at_unix_ms, updated_at_unix_secs + actual_total_cost_usd, cache_creation_cost_usd, cache_read_cost_usd, + input_price_per_1m, response_time_ms, first_byte_time_ms, + created_at_unix_ms, updated_at_unix_secs ) VALUES - ('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'completed', 'settled', - 200, NULL, 10, 20, 1, 2, 0.30, 0.25, 100, 3600000, 3600), - ('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'failed', 'void', - 500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 300, 3610000, 3610), - ('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'pending', 'pending', - NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3620000, 3620), - ('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'completed', 'settled', - 200, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3630000, 3630) + ('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'openai', 'completed', 'settled', + 200, NULL, 10, 20, 1, 2, 0.30, 0.25, 0.01, 0.02, 10.0, 100, 50, 3600, 3600), + ('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'claude', 'failed', 'void', + 500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 0.00, 0.01, 20.0, 300, 200, 3610, 3610), + ('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'openai', 'pending', 'pending', + NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 0.00, 0.00, 0.0, 50, 25, 3620, 3620), + ('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'openai', 'completed', 'settled', + 200, NULL, 100, 100, 0, 0, 9.99, 9.99, 0.00, 0.00, 0.0, 50, 25, 3630, 3630) "#, ) .execute(backend.pool()) .await .expect("usage stats rows should seed"); + sqlx::query( + r#" +INSERT INTO request_candidates ( + id, request_id, candidate_index, retry_index, status, created_at +) VALUES + ('stats-candidate-1', 'stats-fallback', 0, 0, 'failed', 3600000), + ('stats-candidate-2', 'stats-fallback', 1, 0, 'success', 3610000) +"#, + ) + .execute(backend.pool()) + .await + .expect("fallback candidates should seed"); + let target_hour = chrono::DateTime::::from_timestamp(3600, 0) .expect("target hour should be valid"); let aggregated_at = chrono::DateTime::::from_timestamp(7200, 0) @@ -896,6 +895,66 @@ WHERE hour_utc = 3600 assert_eq!(hourly_row.3, 15); assert!((hourly_row.4 - 0.50).abs() < f64::EPSILON); + let enriched_hourly: (f64, i64, i64, i64, i64, i64) = sqlx::query_as( + r#" +SELECT response_time_sum_ms, response_time_samples, cache_hit_total_requests, + cache_hit_requests, completed_total_requests, settled_total_requests +FROM stats_hourly +WHERE hour_utc = 3600 +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("enriched hourly stats row should load"); + assert!((enriched_hourly.0 - 400.0).abs() < f64::EPSILON); + assert_eq!(enriched_hourly.1, 2); + assert_eq!(enriched_hourly.2, 4); + assert_eq!(enriched_hourly.3, 2); + assert_eq!(enriched_hourly.4, 2); + assert_eq!(enriched_hourly.5, 2); + + assert_eq!(sqlite_count(backend.pool(), "stats_hourly_user").await, 2); + assert_eq!( + sqlite_count(backend.pool(), "stats_hourly_user_model").await, + 2 + ); + assert_eq!(sqlite_count(backend.pool(), "stats_hourly_model").await, 2); + assert_eq!( + sqlite_count(backend.pool(), "stats_hourly_provider").await, + 2 + ); + let hourly_user = sqlx::query_as::<_, (i64, i64, i64, i64, i64, f64)>( + r#" +SELECT total_requests, success_requests, error_requests, input_tokens, output_tokens, total_cost +FROM stats_hourly_user +WHERE hour_utc = 3600 AND user_id = 'user-2' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("hourly user stats row should load"); + assert_eq!(hourly_user.0, 1); + assert_eq!(hourly_user.1, 0); + assert_eq!(hourly_user.2, 1); + assert_eq!(hourly_user.3, 5); + assert_eq!(hourly_user.4, 7); + assert!((hourly_user.5 - 0.20).abs() < f64::EPSILON); + let hourly_model = sqlx::query_as::<_, (i64, i64, i64, f64, f64)>( + r#" +SELECT total_requests, input_tokens, output_tokens, total_cost, avg_response_time_ms +FROM stats_hourly_model +WHERE hour_utc = 3600 AND model = 'model-a' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("hourly model stats row should load"); + assert_eq!(hourly_model.0, 1); + assert_eq!(hourly_model.1, 10); + assert_eq!(hourly_model.2, 20); + assert!((hourly_model.3 - 0.30).abs() < f64::EPSILON); + assert!((hourly_model.4 - 100.0).abs() < f64::EPSILON); + let second_hourly = backend .aggregate_stats_hourly(&StatsHourlyAggregationInput { target_hour_utc: target_hour, @@ -919,13 +978,13 @@ WHERE hour_utc = 3600 assert_eq!(daily.total_requests, 2); assert_eq!(daily.model_rows, 2); assert_eq!(daily.provider_rows, 2); - assert_eq!(daily.api_key_rows, 2); + assert_eq!(daily.api_key_rows, 4); assert_eq!(daily.error_rows, 1); assert_eq!(daily.user_rows, 2); - let daily_row = sqlx::query_as::<_, (i64, i64, i64, i64)>( + let daily_row = sqlx::query_as::<_, (i64, i64, i64, i64, i64)>( r#" -SELECT total_requests, success_requests, error_requests, unique_models +SELECT total_requests, success_requests, error_requests, unique_models, fallback_count FROM stats_daily WHERE "date" = 0 "#, @@ -933,6 +992,167 @@ WHERE "date" = 0 .fetch_one(backend.pool()) .await .expect("daily stats row should load"); - assert_eq!(daily_row, (2, 1, 1, 2)); + assert_eq!(daily_row, (2, 1, 1, 2, 1)); + + assert_eq!(sqlite_count(backend.pool(), "stats_daily_model").await, 2); + assert_eq!( + sqlite_count(backend.pool(), "stats_daily_provider").await, + 2 + ); + assert_eq!(sqlite_count(backend.pool(), "stats_daily_api_key").await, 4); + assert_eq!(sqlite_count(backend.pool(), "stats_daily_error").await, 1); + assert_eq!(sqlite_count(backend.pool(), "stats_user_daily").await, 2); + let daily_model = sqlx::query_as::<_, (i64, i64, i64, i64, i64, f64, f64)>( + r#" +SELECT total_requests, input_tokens, output_tokens, cache_creation_tokens, + cache_read_tokens, total_cost, avg_response_time_ms +FROM stats_daily_model +WHERE "date" = 0 AND model = 'model-a' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("daily model stats row should load"); + assert_eq!(daily_model.0, 1); + assert_eq!(daily_model.1, 10); + assert_eq!(daily_model.2, 20); + assert_eq!(daily_model.3, 1); + assert_eq!(daily_model.4, 2); + assert!((daily_model.5 - 0.30).abs() < f64::EPSILON); + assert!((daily_model.6 - 100.0).abs() < f64::EPSILON); + let daily_error = sqlx::query_as::<_, (String, Option, Option, i64)>( + r#" +SELECT error_category, provider_name, model, count +FROM stats_daily_error +WHERE "date" = 0 +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("daily error stats row should load"); + assert_eq!( + daily_error, + ( + "upstream_error".to_string(), + Some("provider-b".to_string()), + Some("model-b".to_string()), + 1, + ) + ); + let daily_user = sqlx::query_as::<_, (i64, i64, i64, i64, i64, f64)>( + r#" +SELECT total_requests, success_requests, error_requests, input_tokens, output_tokens, total_cost +FROM stats_user_daily +WHERE "date" = 0 AND user_id = 'user-2' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("daily user stats row should load"); + assert_eq!(daily_user.0, 1); + assert_eq!(daily_user.1, 0); + assert_eq!(daily_user.2, 1); + assert_eq!(daily_user.3, 5); + assert_eq!(daily_user.4, 7); + assert!((daily_user.5 - 0.20).abs() < f64::EPSILON); + + let enriched_daily = + sqlx::query_as::<_, (i64, i64, f64, i64, i64, i64, i64, i64, i64, Option)>( + r#" +SELECT effective_input_tokens, total_input_context, response_time_sum_ms, + response_time_samples, cache_hit_total_requests, cache_hit_requests, + completed_total_requests, completed_cache_hit_requests, + settled_total_requests, p50_response_time_ms +FROM stats_daily +WHERE "date" = 0 +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("enriched daily stats row should load"); + assert_eq!(enriched_daily.0, 13); + assert_eq!(enriched_daily.1, 17); + assert!((enriched_daily.2 - 400.0).abs() < f64::EPSILON); + assert_eq!(enriched_daily.3, 2); + assert_eq!(enriched_daily.4, 4); + assert_eq!(enriched_daily.5, 2); + assert_eq!(enriched_daily.6, 2); + assert_eq!(enriched_daily.7, 1); + assert_eq!(enriched_daily.8, 2); + assert_eq!(enriched_daily.9, None); + + for (table, expected) in [ + ("stats_user_summary", 2), + ("stats_user_daily_model", 2), + ("stats_user_daily_provider", 2), + ("stats_user_daily_api_format", 2), + ("stats_daily_model_provider", 2), + ("stats_user_daily_model_provider", 2), + ("stats_daily_cost_savings", 1), + ("stats_daily_cost_savings_provider", 3), + ("stats_daily_cost_savings_model", 2), + ("stats_daily_cost_savings_model_provider", 3), + ("stats_user_daily_cost_savings", 4), + ("stats_user_daily_cost_savings_provider", 4), + ("stats_user_daily_cost_savings_model", 4), + ("stats_user_daily_cost_savings_model_provider", 4), + ] { + assert_eq!( + sqlite_count(backend.pool(), table).await, + expected, + "{table}" + ); + } + + let model_rollup: (i64, i64, i64, f64, i64) = sqlx::query_as( + r#" +SELECT total_requests, effective_input_tokens, total_tokens, + response_time_sum_ms, successful_response_time_samples +FROM stats_user_daily_model +WHERE user_id = 'user-1' AND "date" = 0 AND model = 'model-a' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("advanced user model row should load"); + assert_eq!(model_rollup.0, 1); + assert_eq!(model_rollup.1, 8); + assert_eq!(model_rollup.2, 31); + assert!((model_rollup.3 - 100.0).abs() < f64::EPSILON); + assert_eq!(model_rollup.4, 1); + + let savings: (i64, f64, f64, f64) = sqlx::query_as( + r#" +SELECT cache_read_tokens, cache_read_cost, cache_creation_cost, estimated_full_cost +FROM stats_daily_cost_savings +WHERE "date" = 0 +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("daily cost savings row should load"); + assert_eq!(savings.0, 3); + assert!((savings.1 - 0.03).abs() < 1e-12); + assert!((savings.2 - 0.01).abs() < 1e-12); + assert!((savings.3 - 0.00004).abs() < 1e-12); + + let summary: (i64, i64, i64) = sqlx::query_as( + r#" +SELECT all_time_requests, all_time_input_tokens, active_days +FROM stats_user_summary +WHERE user_id = 'user-1' +"#, + ) + .fetch_one(backend.pool()) + .await + .expect("user summary row should load"); + assert_eq!(summary, (1, 10, 1)); + + let global_summary: (i64, i64) = + sqlx::query_as("SELECT all_time_requests, all_time_input_tokens FROM stats_summary") + .fetch_one(backend.pool()) + .await + .expect("global stats summary should load"); + assert_eq!(global_summary, (2, 15)); } } diff --git a/crates/aether-data/runtime/src/backend/stats/mysql.rs b/crates/aether-data/runtime/src/backend/stats/mysql.rs index 8b9cb4238..3292ba968 100644 --- a/crates/aether-data/runtime/src/backend/stats/mysql.rs +++ b/crates/aether-data/runtime/src/backend/stats/mysql.rs @@ -1,7 +1,7 @@ use chrono::{DateTime, Utc}; use sqlx::Row; -use crate::backend::stats_common::{stats_id, unix_ms, unix_secs, utc_from_unix_secs}; +use crate::backend::stats_common::{stats_id, unix_secs, utc_from_unix_secs}; use crate::backend::MysqlBackend; use crate::driver::mysql::MysqlPool; use crate::error::SqlResultExt; @@ -10,6 +10,8 @@ use crate::{ StatsHourlyAggregationInput, StatsHourlyAggregationSummary, }; +mod advanced; + impl MysqlBackend { pub async fn aggregate_stats_hourly( &self, @@ -56,7 +58,7 @@ async fn next_mysql_stats_hourly_bucket( } let next_bucket: Option = sqlx::query_scalar( r#" -SELECT CAST(MIN(FLOOR(created_at_unix_ms / 3600000) * 3600) AS SIGNED) +SELECT CAST(MIN(FLOOR(created_at_unix_ms / 3600) * 3600) AS SIGNED) FROM `usage` WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? @@ -64,8 +66,8 @@ WHERE created_at_unix_ms >= ? AND provider_name NOT IN ('unknown', 'pending') "#, ) - .bind(unix_ms(search_from)?) - .bind(unix_ms(search_until)?) + .bind(search_from) + .bind(search_until) .fetch_one(pool) .await .map_sql_err()?; @@ -88,7 +90,7 @@ async fn next_mysql_stats_daily_bucket( } let next_bucket: Option = sqlx::query_scalar( r#" -SELECT CAST(MIN(FLOOR(created_at_unix_ms / 86400000) * 86400) AS SIGNED) +SELECT CAST(MIN(FLOOR(created_at_unix_ms / 86400) * 86400) AS SIGNED) FROM `usage` WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? @@ -96,8 +98,8 @@ WHERE created_at_unix_ms >= ? AND provider_name NOT IN ('unknown', 'pending') "#, ) - .bind(unix_ms(search_from)?) - .bind(unix_ms(search_until)?) + .bind(search_from) + .bind(search_until) .fetch_one(pool) .await .map_sql_err()?; @@ -110,7 +112,7 @@ SELECT CAST(COALESCE(SUM(CASE WHEN status = 'failed' OR status_code >= 400 - OR (error_category IS NOT NULL AND error_category <> '') + OR error_message IS NOT NULL THEN 1 ELSE 0 END), 0) AS SIGNED) AS error_requests, CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens, @@ -131,13 +133,13 @@ async fn perform_mysql_stats_hourly_aggregation( hour_utc_unix_secs: i64, aggregated_at: DateTime, ) -> Result { - let start_ms = unix_ms(hour_utc_unix_secs)?; - let end_ms = unix_ms(hour_utc_unix_secs + 3600)?; + let start_unix_secs = hour_utc_unix_secs; + let end_unix_secs = hour_utc_unix_secs + 3600; let aggregated_at_unix_secs = unix_secs(aggregated_at); let mut tx = pool.begin().await.map_sql_err()?; let row = sqlx::query(MYSQL_STATS_AGGREGATE_SQL) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut *tx) .await .map_sql_err()?; @@ -193,10 +195,39 @@ ON DUPLICATE KEY UPDATE .await .map_sql_err()?; - let user_rows = mysql_group_count(&mut tx, "user_id", start_ms, end_ms).await?; - let user_model_rows = mysql_group_count(&mut tx, "user_id, model", start_ms, end_ms).await?; - let model_rows = mysql_group_count(&mut tx, "model", start_ms, end_ms).await?; - let provider_rows = mysql_group_count(&mut tx, "provider_name", start_ms, end_ms).await?; + let user_rows = upsert_mysql_stats_hourly_user_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let user_model_rows = upsert_mysql_stats_hourly_user_model_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let model_rows = upsert_mysql_stats_hourly_model_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let provider_rows = upsert_mysql_stats_hourly_provider_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + advanced::refresh_hourly(&mut tx, hour_utc_unix_secs, start_unix_secs, end_unix_secs).await?; tx.commit().await.map_sql_err()?; Ok(StatsHourlyAggregationSummary { @@ -214,21 +245,24 @@ async fn perform_mysql_stats_daily_aggregation( day_start_unix_secs: i64, aggregated_at: DateTime, ) -> Result { - let start_ms = unix_ms(day_start_unix_secs)?; - let end_ms = unix_ms(day_start_unix_secs + 86_400)?; + let start_unix_secs = day_start_unix_secs; + let end_unix_secs = day_start_unix_secs + 86_400; let aggregated_at_unix_secs = unix_secs(aggregated_at); let mut tx = pool.begin().await.map_sql_err()?; let row = sqlx::query(MYSQL_STATS_AGGREGATE_SQL) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut *tx) .await .map_sql_err()?; let total_requests: i64 = row.try_get("total_requests").map_sql_err()?; let error_requests: i64 = row.try_get("error_requests").map_sql_err()?; - let unique_models = mysql_group_count(&mut tx, "model", start_ms, end_ms).await? as i64; + let unique_models = + mysql_group_count(&mut tx, "model", start_unix_secs, end_unix_secs).await? as i64; let unique_providers = - mysql_group_count(&mut tx, "provider_name", start_ms, end_ms).await? as i64; + mysql_group_count(&mut tx, "provider_name", start_unix_secs, end_unix_secs).await? as i64; + let fallback_count = + mysql_daily_fallback_count(&mut tx, start_unix_secs, end_unix_secs).await?; sqlx::query( r#" @@ -237,7 +271,7 @@ INSERT INTO stats_daily ( input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, total_cost, actual_total_cost, avg_response_time_ms, fallback_count, unique_models, unique_providers, is_complete, aggregated_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, TRUE, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, TRUE, ?, ?, ?) ON DUPLICATE KEY UPDATE total_requests = VALUES(total_requests), success_requests = VALUES(success_requests), @@ -275,6 +309,7 @@ ON DUPLICATE KEY UPDATE row.try_get::("avg_response_time_ms") .map_sql_err()?, ) + .bind(fallback_count) .bind(unique_models) .bind(unique_providers) .bind(aggregated_at_unix_secs) @@ -284,11 +319,54 @@ ON DUPLICATE KEY UPDATE .await .map_sql_err()?; - let model_rows = usize::try_from(unique_models).unwrap_or(usize::MAX); - let provider_rows = usize::try_from(unique_providers).unwrap_or(usize::MAX); - let api_key_rows = mysql_group_count(&mut tx, "api_key_id", start_ms, end_ms).await?; - let error_rows = mysql_error_group_count(&mut tx, start_ms, end_ms).await?; - let user_rows = mysql_group_count(&mut tx, "user_id", start_ms, end_ms).await?; + let model_rows = upsert_mysql_stats_daily_model_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let provider_rows = upsert_mysql_stats_daily_provider_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let api_key_rows = upsert_mysql_stats_daily_api_key_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let error_rows = refresh_mysql_stats_daily_error_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let user_rows = upsert_mysql_stats_user_daily_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + advanced::refresh_daily( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; tx.commit().await.map_sql_err()?; Ok(StatsDailyAggregationSummary { @@ -302,11 +380,425 @@ ON DUPLICATE KEY UPDATE }) } +async fn upsert_mysql_stats_hourly_user_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_user ( + id, hour_utc, user_id, total_requests, success_requests, error_requests, + input_tokens, output_tokens, total_cost, created_at, updated_at +) +SELECT + SHA2(UUID(), 256), ?, user_id, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN status = 'failed' OR status_code >= 400 OR error_message IS NOT NULL + THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN status = 'failed' OR status_code >= 400 OR error_message IS NOT NULL + THEN 1 ELSE 0 END), 0), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(total_cost_usd), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND user_id IS NOT NULL AND user_id <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY user_id +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), success_requests = VALUES(success_requests), + error_requests = VALUES(error_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), total_cost = VALUES(total_cost), + updated_at = VALUES(updated_at) +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_hourly_user_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_user_model ( + id, hour_utc, user_id, model, total_requests, input_tokens, output_tokens, + total_cost, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, user_id, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(total_cost_usd), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND user_id IS NOT NULL AND user_id <> '' AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY user_id, model +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), total_cost = VALUES(total_cost), + updated_at = VALUES(updated_at) +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_hourly_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_model ( + id, hour_utc, model, total_requests, input_tokens, output_tokens, total_cost, + avg_response_time_ms, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(total_cost_usd), 0), COALESCE(AVG(response_time_ms), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY model +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), total_cost = VALUES(total_cost), + avg_response_time_ms = VALUES(avg_response_time_ms), updated_at = VALUES(updated_at) +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_hourly_provider_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_provider ( + id, hour_utc, provider_name, total_requests, input_tokens, output_tokens, + total_cost, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, provider_name, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(total_cost_usd), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY provider_name +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), total_cost = VALUES(total_cost), + updated_at = VALUES(updated_at) +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_daily_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_model ( + id, `date`, model, total_requests, input_tokens, output_tokens, + cache_creation_tokens, cache_read_tokens, total_cost, avg_response_time_ms, + created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_creation_input_tokens), 0), + COALESCE(SUM(cache_read_input_tokens), 0), COALESCE(SUM(total_cost_usd), 0), + COALESCE(AVG(response_time_ms), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY model +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), cache_creation_tokens = VALUES(cache_creation_tokens), + cache_read_tokens = VALUES(cache_read_tokens), total_cost = VALUES(total_cost), + avg_response_time_ms = VALUES(avg_response_time_ms), updated_at = VALUES(updated_at) +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_daily_provider_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_provider ( + id, `date`, provider_name, total_requests, input_tokens, output_tokens, + cache_creation_tokens, cache_read_tokens, total_cost, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, provider_name, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_creation_input_tokens), 0), + COALESCE(SUM(cache_read_input_tokens), 0), COALESCE(SUM(total_cost_usd), 0), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY provider_name +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), cache_creation_tokens = VALUES(cache_creation_tokens), + cache_read_tokens = VALUES(cache_read_tokens), total_cost = VALUES(total_cost), + updated_at = VALUES(updated_at) +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_daily_api_key_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_api_key ( + id, api_key_id, `date`, total_requests, success_requests, error_requests, + input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, + total_cost, api_key_name, created_at, updated_at +) +SELECT SHA2(UUID(), 256), usage.api_key_id, ?, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(usage.input_tokens), 0), COALESCE(SUM(usage.output_tokens), 0), + COALESCE(SUM(usage.cache_creation_input_tokens), 0), + COALESCE(SUM(usage.cache_read_input_tokens), 0), + COALESCE(SUM(usage.total_cost_usd), 0), MAX(api_keys.name), ?, ? +FROM `usage` AS usage +LEFT JOIN api_keys ON api_keys.id = usage.api_key_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.api_key_id IS NOT NULL AND usage.api_key_id <> '' +GROUP BY usage.api_key_id +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), success_requests = VALUES(success_requests), + error_requests = VALUES(error_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), cache_creation_tokens = VALUES(cache_creation_tokens), + cache_read_tokens = VALUES(cache_read_tokens), total_cost = VALUES(total_cost), + api_key_name = COALESCE(VALUES(api_key_name), stats_daily_api_key.api_key_name), + updated_at = VALUES(updated_at) +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn refresh_mysql_stats_daily_error_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + sqlx::query("DELETE FROM stats_daily_error WHERE `date` = ?") + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + let result = sqlx::query( + r#" +INSERT INTO stats_daily_error ( + id, `date`, error_category, provider_name, model, count, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, error_category, provider_name, model, COUNT(*), ?, ? +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND error_category IS NOT NULL AND error_category <> '' +GROUP BY error_category, provider_name, model +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_mysql_stats_user_daily_rows( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_user_daily ( + id, user_id, `date`, total_requests, success_requests, error_requests, + input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, + total_cost, username, created_at, updated_at +) +SELECT SHA2(UUID(), 256), usage.user_id, ?, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(usage.input_tokens), 0), COALESCE(SUM(usage.output_tokens), 0), + COALESCE(SUM(usage.cache_creation_input_tokens), 0), + COALESCE(SUM(usage.cache_read_input_tokens), 0), + COALESCE(SUM(usage.total_cost_usd), 0), MAX(users.username), ?, ? +FROM `usage` AS usage +LEFT JOIN users ON users.id = usage.user_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND usage.status NOT IN ('pending', 'streaming') + AND usage.provider_name NOT IN ('unknown', 'pending') +GROUP BY usage.user_id +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), success_requests = VALUES(success_requests), + error_requests = VALUES(error_requests), input_tokens = VALUES(input_tokens), + output_tokens = VALUES(output_tokens), cache_creation_tokens = VALUES(cache_creation_tokens), + cache_read_tokens = VALUES(cache_read_tokens), total_cost = VALUES(total_cost), + username = COALESCE(VALUES(username), stats_user_daily.username), + updated_at = VALUES(updated_at) +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn mysql_daily_fallback_count( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result { + let start_unix_ms = start_unix_secs.checked_mul(1000).ok_or_else(|| { + DataLayerError::InvalidInput("stats fallback window start overflows milliseconds".into()) + })?; + let end_unix_ms = end_unix_secs.checked_mul(1000).ok_or_else(|| { + DataLayerError::InvalidInput("stats fallback window end overflows milliseconds".into()) + })?; + sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM ( + SELECT request_id + FROM request_candidates + WHERE created_at >= ? AND created_at < ? + AND status IN ('success', 'failed') + GROUP BY request_id + HAVING COUNT(id) > 1 +) AS fallback_requests +"#, + ) + .bind(start_unix_ms) + .bind(end_unix_ms) + .fetch_one(&mut **tx) + .await + .map_sql_err() +} + async fn mysql_group_count( tx: &mut sqlx::Transaction<'_, sqlx::MySql>, group_columns: &str, - start_ms: i64, - end_ms: i64, + start_unix_secs: i64, + end_unix_secs: i64, ) -> Result { let not_empty = group_columns .split(',') @@ -330,42 +822,10 @@ FROM ( "# ); let count: i64 = sqlx::query_scalar(&sql) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut **tx) .await .map_sql_err()?; Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX)) } - -async fn mysql_error_group_count( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - start_ms: i64, - end_ms: i64, -) -> Result { - let count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM ( - SELECT 1 - FROM `usage` - WHERE created_at_unix_ms >= ? - AND created_at_unix_ms < ? - AND status NOT IN ('pending', 'streaming') - AND provider_name NOT IN ('unknown', 'pending') - AND ( - status = 'failed' - OR status_code >= 400 - OR (error_category IS NOT NULL AND error_category <> '') - ) - GROUP BY COALESCE(NULLIF(error_category, ''), 'unknown_error'), provider_name, model -) AS grouped -"#, - ) - .bind(start_ms) - .bind(end_ms) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX)) -} diff --git a/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs b/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs new file mode 100644 index 000000000..df9fa97b5 --- /dev/null +++ b/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs @@ -0,0 +1,973 @@ +use sqlx::MySql; + +use crate::backend::stats_common::stats_id; +use crate::error::SqlResultExt; +use crate::DataLayerError; + +const CACHE_5M: &str = r#"GREATEST( + COALESCE(usage.cache_creation_input_tokens_5m, 0), + COALESCE(usage.cache_creation_ephemeral_5m_input_tokens, 0) +)"#; +const CACHE_1H: &str = r#"GREATEST( + COALESCE(usage.cache_creation_input_tokens_1h, 0), + COALESCE(usage.cache_creation_ephemeral_1h_input_tokens, 0) +)"#; +const CACHE_CREATION: &str = r#"CASE + WHEN COALESCE(usage.cache_creation_input_tokens, 0) = 0 + AND ({cache_5m} + {cache_1h}) > 0 + THEN {cache_5m} + {cache_1h} + ELSE GREATEST(COALESCE(usage.cache_creation_input_tokens, 0), 0) +END"#; +const EFFECTIVE_INPUT: &str = r#"CASE + WHEN SUBSTRING_INDEX( + LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, '')), ':', 1 + ) IN ('openai', 'gemini', 'google') + AND COALESCE(usage.input_tokens, 0) > 0 + AND COALESCE(usage.cache_read_input_tokens, 0) > 0 + THEN GREATEST(COALESCE(usage.input_tokens, 0) - COALESCE(usage.cache_read_input_tokens, 0), 0) + ELSE GREATEST(COALESCE(usage.input_tokens, 0), 0) +END"#; +const SUCCESS: &str = r#"CASE + WHEN usage.status <> 'failed' + AND (usage.status_code IS NULL OR usage.status_code < 400) + AND usage.error_message IS NULL + THEN 1 ELSE 0 +END"#; +const AGGREGATABLE: &str = r#"usage.status NOT IN ('pending', 'streaming') + AND usage.provider_name NOT IN ('unknown', 'pending')"#; +const SETTLED: &str = r#"COALESCE(settlement.billing_status, usage.billing_status) = 'settled' + AND COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) > 0"#; + +fn cache_creation_expr() -> String { + CACHE_CREATION + .replace("{cache_5m}", CACHE_5M) + .replace("{cache_1h}", CACHE_1H) +} + +fn total_input_context_expr() -> String { + format!( + "({EFFECTIVE_INPUT}) + ({}) + GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0)", + cache_creation_expr() + ) +} + +fn total_tokens_expr() -> String { + format!( + r#"COALESCE( + NULLIF(GREATEST(COALESCE(usage.total_tokens, 0), 0), 0), + ({EFFECTIVE_INPUT}) + + GREATEST(COALESCE(usage.output_tokens, 0), 0) + + ({}) + + GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0), + 0 +)"#, + cache_creation_expr() + ) +} + +fn percentile_cont(sorted: &[i64], percentile: f64) -> Option { + if sorted.is_empty() { + return None; + } + let position = percentile * (sorted.len().saturating_sub(1) as f64); + let lower = position.floor() as usize; + let upper = position.ceil() as usize; + let fraction = position - lower as f64; + let value = sorted[lower] as f64 + (sorted[upper] - sorted[lower]) as f64 * fraction; + Some(value.round() as i64) +} + +async fn load_percentiles( + tx: &mut sqlx::Transaction<'_, MySql>, + column: &str, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(Option, Option, Option), DataLayerError> { + let sql = format!( + r#" +SELECT {column} +FROM `usage` +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status = 'completed' + AND provider_name NOT IN ('unknown', 'pending') + AND {column} IS NOT NULL +ORDER BY {column} +"# + ); + let values: Vec = sqlx::query_scalar(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .fetch_all(&mut **tx) + .await + .map_sql_err()?; + if values.len() < 10 { + return Ok((None, None, None)); + } + Ok(( + percentile_cont(&values, 0.50), + percentile_cont(&values, 0.90), + percentile_cont(&values, 0.99), + )) +} + +pub(super) async fn refresh_hourly( + tx: &mut sqlx::Transaction<'_, MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let sql = format!( + r#" +UPDATE stats_hourly AS target +JOIN ( + SELECT + COUNT(*) AS cache_hit_total_requests, + COALESCE(SUM(CASE WHEN COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_total_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' AND COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS completed_cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS completed_input_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({cache_creation}) ELSE 0 END), 0) AS completed_cache_creation_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS completed_cache_read_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({total_context}) ELSE 0 END), 0) AS completed_total_input_context, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_creation_cost, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_read_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples + FROM `usage` AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? +) AS aggregated +SET + target.cache_hit_total_requests = aggregated.cache_hit_total_requests, + target.cache_hit_requests = aggregated.cache_hit_requests, + target.completed_total_requests = aggregated.completed_total_requests, + target.completed_cache_hit_requests = aggregated.completed_cache_hit_requests, + target.completed_input_tokens = aggregated.completed_input_tokens, + target.completed_cache_creation_tokens = aggregated.completed_cache_creation_tokens, + target.completed_cache_read_tokens = aggregated.completed_cache_read_tokens, + target.completed_total_input_context = aggregated.completed_total_input_context, + target.completed_cache_creation_cost = aggregated.completed_cache_creation_cost, + target.completed_cache_read_cost = aggregated.completed_cache_read_cost, + target.settled_total_cost = aggregated.settled_total_cost, + target.settled_total_requests = aggregated.settled_total_requests, + target.settled_input_tokens = aggregated.settled_input_tokens, + target.settled_output_tokens = aggregated.settled_output_tokens, + target.settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + target.settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + target.settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + target.settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs, + target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples +WHERE target.hour_utc = ? +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + refresh_hourly_user(tx, hour_utc, start_unix_secs, end_unix_secs).await?; + refresh_hourly_response_dimensions(tx, hour_utc, start_unix_secs, end_unix_secs).await +} + +async fn refresh_hourly_user( + tx: &mut sqlx::Transaction<'_, MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let sql = format!( + r#" +UPDATE stats_hourly_user AS target +JOIN ( + SELECT usage.user_id, + COALESCE(SUM({cache_creation}), 0) AS cache_creation_tokens, + COALESCE(SUM(GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0) AS actual_total_cost, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM `usage` AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} + GROUP BY usage.user_id +) AS aggregated ON target.user_id = aggregated.user_id +SET target.cache_creation_tokens = aggregated.cache_creation_tokens, + target.cache_read_tokens = aggregated.cache_read_tokens, + target.actual_total_cost = aggregated.actual_total_cost, + target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples, + target.settled_total_cost = aggregated.settled_total_cost, + target.settled_total_requests = aggregated.settled_total_requests, + target.settled_input_tokens = aggregated.settled_input_tokens, + target.settled_output_tokens = aggregated.settled_output_tokens, + target.settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + target.settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + target.settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + target.settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs +WHERE target.hour_utc = ? +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_hourly_response_dimensions( + tx: &mut sqlx::Transaction<'_, MySql>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + for (table, select_dimensions, group_by, join) in [ + ( + "stats_hourly_model", + "usage.model AS model", + "usage.model", + "target.model = aggregated.model", + ), + ( + "stats_hourly_user_model", + "usage.user_id AS user_id, usage.model AS model", + "usage.user_id, usage.model", + "target.user_id = aggregated.user_id AND target.model = aggregated.model", + ), + ] { + let sql = format!( + r#" +UPDATE {table} AS target +JOIN ( + SELECT {select_dimensions}, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples + FROM `usage` AS usage + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? AND {AGGREGATABLE} + GROUP BY {group_by} +) AS aggregated ON {join} +SET target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples +WHERE target.hour_utc = ? +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + Ok(()) +} + +pub(super) async fn refresh_daily( + tx: &mut sqlx::Transaction<'_, MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let response = load_percentiles(tx, "response_time_ms", start_unix_secs, end_unix_secs).await?; + let first_byte = + load_percentiles(tx, "first_byte_time_ms", start_unix_secs, end_unix_secs).await?; + refresh_daily_root( + tx, + day_start, + start_unix_secs, + end_unix_secs, + response, + first_byte, + ) + .await?; + refresh_daily_existing_dimensions(tx, day_start, start_unix_secs, end_unix_secs).await?; + upsert_user_dimension( + tx, + "stats_user_daily_model", + "model", + "usage.model", + "usage.model IS NOT NULL AND usage.model <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_user_dimension( + tx, + "stats_user_daily_provider", + "provider_name", + "usage.provider_name", + "usage.provider_name IS NOT NULL AND usage.provider_name <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_user_dimension( + tx, + "stats_user_daily_api_format", + "api_format", + "LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, ''))", + "COALESCE(usage.endpoint_api_format, usage.api_format, '') <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_model_provider_rows(tx, day_start, start_unix_secs, end_unix_secs, now_unix_secs) + .await?; + upsert_cost_savings_rows(tx, day_start, start_unix_secs, end_unix_secs, now_unix_secs).await?; + refresh_user_summary(tx, end_unix_secs, now_unix_secs).await +} + +async fn refresh_daily_root( + tx: &mut sqlx::Transaction<'_, MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + response: (Option, Option, Option), + first_byte: (Option, Option, Option), +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let sql = format!( + r#" +UPDATE stats_daily AS target +JOIN ( + SELECT + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN ({EFFECTIVE_INPUT}) ELSE 0 END), 0) AS effective_input_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN ({total_context}) ELSE 0 END), 0) AS total_input_context, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN {CACHE_5M} ELSE 0 END), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN {CACHE_1H} ELSE 0 END), 0) AS cache_creation_ephemeral_1h_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.input_cost_usd, 0) ELSE 0 END), 0) AS input_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.output_cost_usd, 0) ELSE 0 END), 0) AS output_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS cache_creation_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS cache_read_cost, + COUNT(*) AS cache_hit_total_requests, + COALESCE(SUM(CASE WHEN COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_total_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' AND COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS completed_cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS completed_input_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({cache_creation}) ELSE 0 END), 0) AS completed_cache_creation_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS completed_cache_read_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({total_context}) ELSE 0 END), 0) AS completed_total_input_context, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_creation_cost, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_read_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM `usage` AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? +) AS aggregated +SET target.effective_input_tokens = aggregated.effective_input_tokens, + target.total_input_context = aggregated.total_input_context, + target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples, + target.cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + target.cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens, + target.input_cost = aggregated.input_cost, + target.output_cost = aggregated.output_cost, + target.cache_creation_cost = aggregated.cache_creation_cost, + target.cache_read_cost = aggregated.cache_read_cost, + target.cache_hit_total_requests = aggregated.cache_hit_total_requests, + target.cache_hit_requests = aggregated.cache_hit_requests, + target.completed_total_requests = aggregated.completed_total_requests, + target.completed_cache_hit_requests = aggregated.completed_cache_hit_requests, + target.completed_input_tokens = aggregated.completed_input_tokens, + target.completed_cache_creation_tokens = aggregated.completed_cache_creation_tokens, + target.completed_cache_read_tokens = aggregated.completed_cache_read_tokens, + target.completed_total_input_context = aggregated.completed_total_input_context, + target.completed_cache_creation_cost = aggregated.completed_cache_creation_cost, + target.completed_cache_read_cost = aggregated.completed_cache_read_cost, + target.settled_total_cost = aggregated.settled_total_cost, + target.settled_total_requests = aggregated.settled_total_requests, + target.settled_input_tokens = aggregated.settled_input_tokens, + target.settled_output_tokens = aggregated.settled_output_tokens, + target.settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + target.settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + target.settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + target.settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs, + target.p50_response_time_ms = ?, target.p90_response_time_ms = ?, target.p99_response_time_ms = ?, + target.p50_first_byte_time_ms = ?, target.p90_first_byte_time_ms = ?, target.p99_first_byte_time_ms = ? +WHERE target.`date` = ? +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(response.0) + .bind(response.1) + .bind(response.2) + .bind(first_byte.0) + .bind(first_byte.1) + .bind(first_byte.2) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_daily_existing_dimensions( + tx: &mut sqlx::Transaction<'_, MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let model_sql = format!( + r#" +UPDATE stats_daily_model AS target +JOIN ( + SELECT usage.model, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens + FROM `usage` AS usage + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND {AGGREGATABLE} AND usage.model IS NOT NULL AND usage.model <> '' + GROUP BY usage.model +) AS aggregated ON target.model = aggregated.model +SET target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples, + target.cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + target.cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens +WHERE target.`date` = ? +"# + ); + sqlx::query(&model_sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + + let user_sql = format!( + r#" +UPDATE stats_user_daily AS target +JOIN ( + SELECT usage.user_id, + COALESCE(SUM({EFFECTIVE_INPUT}), 0) AS effective_input_tokens, + COALESCE(SUM({total_context}), 0) AS total_input_context, + COALESCE(SUM(COALESCE(usage.cache_creation_cost_usd, 0)), 0) AS cache_creation_cost, + COALESCE(SUM(COALESCE(usage.cache_read_cost_usd, 0)), 0) AS cache_read_cost, + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0) AS actual_total_cost, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM `usage` AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} + GROUP BY usage.user_id +) AS aggregated ON target.user_id = aggregated.user_id +SET target.effective_input_tokens = aggregated.effective_input_tokens, + target.total_input_context = aggregated.total_input_context, + target.cache_creation_cost = aggregated.cache_creation_cost, + target.cache_read_cost = aggregated.cache_read_cost, + target.actual_total_cost = aggregated.actual_total_cost, + target.response_time_sum_ms = aggregated.response_time_sum_ms, + target.response_time_samples = aggregated.response_time_samples, + target.cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + target.cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens, + target.settled_total_cost = aggregated.settled_total_cost, + target.settled_total_requests = aggregated.settled_total_requests, + target.settled_input_tokens = aggregated.settled_input_tokens, + target.settled_output_tokens = aggregated.settled_output_tokens, + target.settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + target.settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + target.settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + target.settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs +WHERE target.`date` = ? +"# + ); + sqlx::query(&user_sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn upsert_user_dimension( + tx: &mut sqlx::Transaction<'_, MySql>, + table: &str, + dimension_column: &str, + dimension_expr: &str, + dimension_filter: &str, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let total_tokens = total_tokens_expr(); + let sql = format!( + r#" +INSERT INTO {table} ( + id, user_id, username, `date`, {dimension_column}, total_requests, success_requests, + input_tokens, effective_input_tokens, output_tokens, total_tokens, total_input_context, + cache_creation_tokens, cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens, cache_read_tokens, total_cost, actual_total_cost, + response_time_sum_ms, response_time_samples, successful_response_time_sum_ms, + successful_response_time_samples, created_at, updated_at +) +SELECT SHA2(UUID(), 256), usage.user_id, + MAX(COALESCE(usage.username, users.username)), ?, {dimension_expr}, COUNT(*), + COALESCE(SUM({SUCCESS}), 0), + COALESCE(SUM(GREATEST(COALESCE(usage.input_tokens, 0), 0)), 0), + COALESCE(SUM({EFFECTIVE_INPUT}), 0), + COALESCE(SUM(GREATEST(COALESCE(usage.output_tokens, 0), 0)), 0), + COALESCE(SUM({total_tokens}), 0), COALESCE(SUM({total_context}), 0), + COALESCE(SUM({cache_creation}), 0), COALESCE(SUM({CACHE_5M}), 0), + COALESCE(SUM({CACHE_1H}), 0), + COALESCE(SUM(GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), + ?, ? +FROM `usage` AS usage +LEFT JOIN users ON users.id = usage.user_id +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND {dimension_filter} AND {AGGREGATABLE} +GROUP BY usage.user_id, {dimension_expr} +ON DUPLICATE KEY UPDATE + username = COALESCE(VALUES(username), {table}.username), + total_requests = VALUES(total_requests), success_requests = VALUES(success_requests), + input_tokens = VALUES(input_tokens), effective_input_tokens = VALUES(effective_input_tokens), + output_tokens = VALUES(output_tokens), total_tokens = VALUES(total_tokens), + total_input_context = VALUES(total_input_context), + cache_creation_tokens = VALUES(cache_creation_tokens), + cache_creation_ephemeral_5m_tokens = VALUES(cache_creation_ephemeral_5m_tokens), + cache_creation_ephemeral_1h_tokens = VALUES(cache_creation_ephemeral_1h_tokens), + cache_read_tokens = VALUES(cache_read_tokens), total_cost = VALUES(total_cost), + actual_total_cost = VALUES(actual_total_cost), + response_time_sum_ms = VALUES(response_time_sum_ms), + response_time_samples = VALUES(response_time_samples), + successful_response_time_sum_ms = VALUES(successful_response_time_sum_ms), + successful_response_time_samples = VALUES(successful_response_time_samples), + updated_at = VALUES(updated_at) +"# + ); + sqlx::query(&sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn upsert_model_provider_rows( + tx: &mut sqlx::Transaction<'_, MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let total_tokens = total_tokens_expr(); + let model_provider_sql = format!( + r#" +INSERT INTO stats_daily_model_provider ( + id, `date`, model, provider_name, total_requests, total_tokens, total_cost, + response_time_sum_ms, response_time_samples, created_at, updated_at +) +SELECT SHA2(UUID(), 256), ?, usage.model, usage.provider_name, COUNT(*), + COALESCE(SUM({total_tokens}), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? +FROM `usage` AS usage +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} +GROUP BY usage.model, usage.provider_name +ON DUPLICATE KEY UPDATE + total_requests = VALUES(total_requests), total_tokens = VALUES(total_tokens), + total_cost = VALUES(total_cost), response_time_sum_ms = VALUES(response_time_sum_ms), + response_time_samples = VALUES(response_time_samples), updated_at = VALUES(updated_at) +"# + ); + sqlx::query(&model_provider_sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + + let user_model_provider_sql = format!( + r#" +INSERT INTO stats_user_daily_model_provider ( + id, user_id, username, `date`, model, provider_name, total_requests, total_tokens, + total_cost, response_time_sum_ms, response_time_samples, created_at, updated_at +) +SELECT SHA2(UUID(), 256), usage.user_id, MAX(COALESCE(usage.username, users.username)), + ?, usage.model, usage.provider_name, COUNT(*), COALESCE(SUM({total_tokens}), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN GREATEST(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? +FROM `usage` AS usage +LEFT JOIN users ON users.id = usage.user_id +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} +GROUP BY usage.user_id, usage.model, usage.provider_name +ON DUPLICATE KEY UPDATE + username = COALESCE(VALUES(username), stats_user_daily_model_provider.username), + total_requests = VALUES(total_requests), total_tokens = VALUES(total_tokens), + total_cost = VALUES(total_cost), response_time_sum_ms = VALUES(response_time_sum_ms), + response_time_samples = VALUES(response_time_samples), updated_at = VALUES(updated_at) +"# + ); + sqlx::query(&user_model_provider_sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn upsert_cost_savings_rows( + tx: &mut sqlx::Transaction<'_, MySql>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + for (table, dimensions) in [ + ("stats_daily_cost_savings", Vec::new()), + ( + "stats_daily_cost_savings_provider", + vec![("provider_name", "COALESCE(usage.provider_name, '')")], + ), + ( + "stats_daily_cost_savings_model", + vec![("model", "COALESCE(usage.model, '')")], + ), + ( + "stats_daily_cost_savings_model_provider", + vec![ + ("model", "COALESCE(usage.model, '')"), + ("provider_name", "COALESCE(usage.provider_name, '')"), + ], + ), + ] { + upsert_cost_savings_dimension( + tx, + table, + false, + &dimensions, + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + } + for (table, dimensions) in [ + ("stats_user_daily_cost_savings", Vec::new()), + ( + "stats_user_daily_cost_savings_provider", + vec![("provider_name", "COALESCE(usage.provider_name, '')")], + ), + ( + "stats_user_daily_cost_savings_model", + vec![("model", "COALESCE(usage.model, '')")], + ), + ( + "stats_user_daily_cost_savings_model_provider", + vec![ + ("model", "COALESCE(usage.model, '')"), + ("provider_name", "COALESCE(usage.provider_name, '')"), + ], + ), + ] { + upsert_cost_savings_dimension( + tx, + table, + true, + &dimensions, + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn upsert_cost_savings_dimension( + tx: &mut sqlx::Transaction<'_, MySql>, + table: &str, + per_user: bool, + dimensions: &[(&str, &str)], + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let dimension_columns = dimensions + .iter() + .map(|(column, _)| *column) + .collect::>(); + let dimension_exprs = dimensions + .iter() + .map(|(_, expression)| *expression) + .collect::>(); + let user_columns = if per_user { "user_id, username, " } else { "" }; + let user_select = if per_user { + "usage.user_id, MAX(COALESCE(usage.username, users.username)), " + } else { + "" + }; + let user_join = if per_user { + "LEFT JOIN users ON users.id = usage.user_id" + } else { + "" + }; + let user_filter = if per_user { + "AND usage.user_id IS NOT NULL AND usage.user_id <> ''" + } else { + "" + }; + let mut group_by = Vec::new(); + if per_user { + group_by.push("usage.user_id"); + } + group_by.extend(dimension_exprs.iter().copied()); + let dimension_columns_sql = if dimension_columns.is_empty() { + String::new() + } else { + format!("{}, ", dimension_columns.join(", ")) + }; + let dimension_select_sql = if dimension_exprs.is_empty() { + String::new() + } else { + format!("{}, ", dimension_exprs.join(", ")) + }; + let group_by_sql = if group_by.is_empty() { + String::new() + } else { + format!("GROUP BY {}", group_by.join(", ")) + }; + let sql = format!( + r#" +INSERT INTO {table} ( + id, {user_columns}`date`, {dimension_columns_sql}cache_read_tokens, + cache_read_cost, cache_creation_cost, estimated_full_cost, created_at, updated_at +) +SELECT SHA2(UUID(), 256), {user_select}?, {dimension_select_sql} + COALESCE(SUM(GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0), + COALESCE(SUM(COALESCE(usage.cache_read_cost_usd, 0)), 0), + COALESCE(SUM(COALESCE(usage.cache_creation_cost_usd, 0)), 0), + COALESCE(SUM( + COALESCE(settlement.input_price_per_1m, usage.input_price_per_1m, 0) + * GREATEST(COALESCE(usage.cache_read_input_tokens, 0), 0) / 1000000.0 + ), 0), ?, ? +FROM `usage` AS usage +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +{user_join} +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? {user_filter} +{group_by_sql} +ON DUPLICATE KEY UPDATE + {}cache_read_tokens = VALUES(cache_read_tokens), + cache_read_cost = VALUES(cache_read_cost), + cache_creation_cost = VALUES(cache_creation_cost), + estimated_full_cost = VALUES(estimated_full_cost), updated_at = VALUES(updated_at) +"#, + if per_user { + format!("username = COALESCE(VALUES(username), {table}.username), ") + } else { + String::new() + } + ); + sqlx::query(&sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_user_summary( + tx: &mut sqlx::Transaction<'_, MySql>, + cutoff_date: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO stats_user_summary ( + id, user_id, username, cutoff_date, all_time_requests, all_time_success_requests, + all_time_error_requests, all_time_input_tokens, all_time_output_tokens, + all_time_cache_creation_tokens, all_time_cache_read_tokens, all_time_cost, + all_time_actual_cost, active_days, first_active_date, last_active_date, + created_at, updated_at +) +SELECT SHA2(UUID(), 256), user_id, MAX(username), ?, + COALESCE(SUM(total_requests), 0), COALESCE(SUM(success_requests), 0), + COALESCE(SUM(error_requests), 0), COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), COALESCE(SUM(cache_creation_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(total_cost), 0), + COALESCE(SUM(actual_total_cost), 0), + COALESCE(SUM(CASE WHEN total_requests > 0 THEN 1 ELSE 0 END), 0), + MIN(CASE WHEN total_requests > 0 THEN `date` END), + MAX(CASE WHEN total_requests > 0 THEN `date` END), ?, ? +FROM stats_user_daily +WHERE `date` < ? +GROUP BY user_id +ON DUPLICATE KEY UPDATE + username = COALESCE(VALUES(username), stats_user_summary.username), + cutoff_date = VALUES(cutoff_date), all_time_requests = VALUES(all_time_requests), + all_time_success_requests = VALUES(all_time_success_requests), + all_time_error_requests = VALUES(all_time_error_requests), + all_time_input_tokens = VALUES(all_time_input_tokens), + all_time_output_tokens = VALUES(all_time_output_tokens), + all_time_cache_creation_tokens = VALUES(all_time_cache_creation_tokens), + all_time_cache_read_tokens = VALUES(all_time_cache_read_tokens), + all_time_cost = VALUES(all_time_cost), all_time_actual_cost = VALUES(all_time_actual_cost), + active_days = VALUES(active_days), first_active_date = VALUES(first_active_date), + last_active_date = VALUES(last_active_date), updated_at = VALUES(updated_at) +"#, + ) + .bind(cutoff_date) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(cutoff_date) + .execute(&mut **tx) + .await + .map_sql_err()?; + refresh_global_summary(tx, cutoff_date, now_unix_secs).await?; + Ok(()) +} + +async fn refresh_global_summary( + tx: &mut sqlx::Transaction<'_, MySql>, + cutoff_date: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let existing_id: Option = + sqlx::query_scalar("SELECT id FROM stats_summary ORDER BY created_at, id LIMIT 1") + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + let summary_id = existing_id.unwrap_or_else(|| stats_id("stats-summary")); + sqlx::query( + r#" +INSERT INTO stats_summary ( + id, cutoff_date, all_time_requests, all_time_success_requests, + all_time_error_requests, all_time_input_tokens, all_time_output_tokens, + all_time_cache_creation_tokens, all_time_cache_read_tokens, all_time_cost, + all_time_actual_cost, total_users, active_users, total_api_keys, + active_api_keys, created_at, updated_at +) +SELECT ?, ?, COALESCE(SUM(total_requests), 0), COALESCE(SUM(success_requests), 0), + COALESCE(SUM(error_requests), 0), COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), COALESCE(SUM(cache_creation_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(total_cost), 0), + COALESCE(SUM(actual_total_cost), 0), + (SELECT COUNT(*) FROM users), + (SELECT COUNT(*) FROM users WHERE is_active <> 0), + (SELECT COUNT(*) FROM api_keys), + (SELECT COUNT(*) FROM api_keys WHERE is_active <> 0), ?, ? +FROM stats_daily +WHERE `date` < ? +ON DUPLICATE KEY UPDATE + cutoff_date = VALUES(cutoff_date), + all_time_requests = VALUES(all_time_requests), + all_time_success_requests = VALUES(all_time_success_requests), + all_time_error_requests = VALUES(all_time_error_requests), + all_time_input_tokens = VALUES(all_time_input_tokens), + all_time_output_tokens = VALUES(all_time_output_tokens), + all_time_cache_creation_tokens = VALUES(all_time_cache_creation_tokens), + all_time_cache_read_tokens = VALUES(all_time_cache_read_tokens), + all_time_cost = VALUES(all_time_cost), + all_time_actual_cost = VALUES(all_time_actual_cost), + total_users = VALUES(total_users), active_users = VALUES(active_users), + total_api_keys = VALUES(total_api_keys), active_api_keys = VALUES(active_api_keys), + updated_at = VALUES(updated_at) +"#, + ) + .bind(summary_id) + .bind(cutoff_date) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(cutoff_date) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} diff --git a/crates/aether-data/runtime/src/backend/stats/sqlite.rs b/crates/aether-data/runtime/src/backend/stats/sqlite.rs index 04f9dfb18..571f8c82d 100644 --- a/crates/aether-data/runtime/src/backend/stats/sqlite.rs +++ b/crates/aether-data/runtime/src/backend/stats/sqlite.rs @@ -1,7 +1,7 @@ use chrono::{DateTime, Utc}; use sqlx::Row; -use crate::backend::stats_common::{stats_id, unix_ms, unix_secs, utc_from_unix_secs}; +use crate::backend::stats_common::{stats_id, unix_secs, utc_from_unix_secs}; use crate::backend::SqliteBackend; use crate::driver::sqlite::{sqlite_real, SqlitePool}; use crate::error::SqlResultExt; @@ -10,6 +10,8 @@ use crate::{ StatsHourlyAggregationInput, StatsHourlyAggregationSummary, }; +mod advanced; + impl SqliteBackend { pub async fn aggregate_stats_hourly( &self, @@ -64,7 +66,7 @@ async fn next_sqlite_stats_hourly_bucket( } let next_bucket: Option = sqlx::query_scalar( r#" -SELECT MIN(CAST(created_at_unix_ms / 3600000 AS INTEGER) * 3600) +SELECT MIN(CAST(created_at_unix_ms / 3600 AS INTEGER) * 3600) FROM "usage" WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? @@ -72,8 +74,8 @@ WHERE created_at_unix_ms >= ? AND provider_name NOT IN ('unknown', 'pending') "#, ) - .bind(unix_ms(search_from)?) - .bind(unix_ms(search_until)?) + .bind(search_from) + .bind(search_until) .fetch_one(pool) .await .map_sql_err()?; @@ -96,7 +98,7 @@ async fn next_sqlite_stats_daily_bucket( } let next_bucket: Option = sqlx::query_scalar( r#" -SELECT MIN(CAST(created_at_unix_ms / 86400000 AS INTEGER) * 86400) +SELECT MIN(CAST(created_at_unix_ms / 86400 AS INTEGER) * 86400) FROM "usage" WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? @@ -104,8 +106,8 @@ WHERE created_at_unix_ms >= ? AND provider_name NOT IN ('unknown', 'pending') "#, ) - .bind(unix_ms(search_from)?) - .bind(unix_ms(search_until)?) + .bind(search_from) + .bind(search_until) .fetch_one(pool) .await .map_sql_err()?; @@ -118,7 +120,7 @@ SELECT COALESCE(SUM(CASE WHEN status = 'failed' OR status_code >= 400 - OR (error_category IS NOT NULL AND error_category <> '') + OR error_message IS NOT NULL THEN 1 ELSE 0 END), 0) AS error_requests, COALESCE(SUM(input_tokens), 0) AS input_tokens, COALESCE(SUM(output_tokens), 0) AS output_tokens, @@ -139,13 +141,13 @@ async fn perform_sqlite_stats_hourly_aggregation( hour_utc_unix_secs: i64, aggregated_at: DateTime, ) -> Result { - let start_ms = unix_ms(hour_utc_unix_secs)?; - let end_ms = unix_ms(hour_utc_unix_secs + 3600)?; + let start_unix_secs = hour_utc_unix_secs; + let end_unix_secs = hour_utc_unix_secs + 3600; let aggregated_at_unix_secs = unix_secs(aggregated_at); let mut tx = pool.begin().await.map_sql_err()?; let row = sqlx::query(SQLITE_STATS_AGGREGATE_SQL) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut *tx) .await .map_sql_err()?; @@ -198,10 +200,39 @@ ON CONFLICT (hour_utc) DO UPDATE SET .await .map_sql_err()?; - let user_rows = sqlite_group_count(&mut tx, "user_id", start_ms, end_ms).await?; - let user_model_rows = sqlite_group_count(&mut tx, "user_id, model", start_ms, end_ms).await?; - let model_rows = sqlite_group_count(&mut tx, "model", start_ms, end_ms).await?; - let provider_rows = sqlite_group_count(&mut tx, "provider_name", start_ms, end_ms).await?; + let user_rows = upsert_sqlite_stats_hourly_user_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let user_model_rows = upsert_sqlite_stats_hourly_user_model_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let model_rows = upsert_sqlite_stats_hourly_model_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let provider_rows = upsert_sqlite_stats_hourly_provider_rows( + &mut tx, + hour_utc_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + advanced::refresh_hourly(&mut tx, hour_utc_unix_secs, start_unix_secs, end_unix_secs).await?; tx.commit().await.map_sql_err()?; Ok(StatsHourlyAggregationSummary { @@ -219,21 +250,24 @@ async fn perform_sqlite_stats_daily_aggregation( day_start_unix_secs: i64, aggregated_at: DateTime, ) -> Result { - let start_ms = unix_ms(day_start_unix_secs)?; - let end_ms = unix_ms(day_start_unix_secs + 86_400)?; + let start_unix_secs = day_start_unix_secs; + let end_unix_secs = day_start_unix_secs + 86_400; let aggregated_at_unix_secs = unix_secs(aggregated_at); let mut tx = pool.begin().await.map_sql_err()?; let row = sqlx::query(SQLITE_STATS_AGGREGATE_SQL) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut *tx) .await .map_sql_err()?; let total_requests: i64 = row.try_get("total_requests").map_sql_err()?; let error_requests: i64 = row.try_get("error_requests").map_sql_err()?; - let unique_models = sqlite_group_count(&mut tx, "model", start_ms, end_ms).await? as i64; + let unique_models = + sqlite_group_count(&mut tx, "model", start_unix_secs, end_unix_secs).await? as i64; let unique_providers = - sqlite_group_count(&mut tx, "provider_name", start_ms, end_ms).await? as i64; + sqlite_group_count(&mut tx, "provider_name", start_unix_secs, end_unix_secs).await? as i64; + let fallback_count = + sqlite_daily_fallback_count(&mut tx, start_unix_secs, end_unix_secs).await?; sqlx::query( r#" @@ -242,7 +276,7 @@ INSERT INTO stats_daily ( input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, total_cost, actual_total_cost, avg_response_time_ms, fallback_count, unique_models, unique_providers, is_complete, aggregated_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, 1, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?) ON CONFLICT ("date") DO UPDATE SET total_requests = excluded.total_requests, success_requests = excluded.success_requests, @@ -277,6 +311,7 @@ ON CONFLICT ("date") DO UPDATE SET .bind(sqlite_real(&row, "total_cost")?) .bind(sqlite_real(&row, "actual_total_cost")?) .bind(sqlite_real(&row, "avg_response_time_ms")?) + .bind(fallback_count) .bind(unique_models) .bind(unique_providers) .bind(aggregated_at_unix_secs) @@ -286,11 +321,54 @@ ON CONFLICT ("date") DO UPDATE SET .await .map_sql_err()?; - let model_rows = usize::try_from(unique_models).unwrap_or(usize::MAX); - let provider_rows = usize::try_from(unique_providers).unwrap_or(usize::MAX); - let api_key_rows = sqlite_group_count(&mut tx, "api_key_id", start_ms, end_ms).await?; - let error_rows = sqlite_error_group_count(&mut tx, start_ms, end_ms).await?; - let user_rows = sqlite_group_count(&mut tx, "user_id", start_ms, end_ms).await?; + let model_rows = upsert_sqlite_stats_daily_model_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let provider_rows = upsert_sqlite_stats_daily_provider_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let api_key_rows = upsert_sqlite_stats_daily_api_key_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let error_rows = refresh_sqlite_stats_daily_error_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + let user_rows = upsert_sqlite_stats_user_daily_rows( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; + advanced::refresh_daily( + &mut tx, + day_start_unix_secs, + start_unix_secs, + end_unix_secs, + aggregated_at_unix_secs, + ) + .await?; tx.commit().await.map_sql_err()?; Ok(StatsDailyAggregationSummary { @@ -304,11 +382,463 @@ ON CONFLICT ("date") DO UPDATE SET }) } +async fn upsert_sqlite_stats_hourly_user_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_user ( + id, hour_utc, user_id, total_requests, success_requests, error_requests, + input_tokens, output_tokens, total_cost, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, user_id, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN status = 'failed' OR status_code >= 400 OR error_message IS NOT NULL + THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN status = 'failed' OR status_code >= 400 OR error_message IS NOT NULL + THEN 1 ELSE 0 END), 0), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND user_id IS NOT NULL AND user_id <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY user_id +ON CONFLICT (hour_utc, user_id) DO UPDATE SET + total_requests = excluded.total_requests, + success_requests = excluded.success_requests, + error_requests = excluded.error_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + total_cost = excluded.total_cost, + updated_at = excluded.updated_at +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_hourly_user_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_user_model ( + id, hour_utc, user_id, model, total_requests, input_tokens, output_tokens, + total_cost, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, user_id, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND user_id IS NOT NULL AND user_id <> '' + AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY user_id, model +ON CONFLICT (hour_utc, user_id, model) DO UPDATE SET + total_requests = excluded.total_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + total_cost = excluded.total_cost, + updated_at = excluded.updated_at +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_hourly_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_model ( + id, hour_utc, model, total_requests, input_tokens, output_tokens, total_cost, + avg_response_time_ms, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), + CAST(COALESCE(AVG(response_time_ms), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY model +ON CONFLICT (hour_utc, model) DO UPDATE SET + total_requests = excluded.total_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + total_cost = excluded.total_cost, + avg_response_time_ms = excluded.avg_response_time_ms, + updated_at = excluded.updated_at +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_hourly_provider_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_hourly_provider ( + id, hour_utc, provider_name, total_requests, input_tokens, output_tokens, + total_cost, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, provider_name, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY provider_name +ON CONFLICT (hour_utc, provider_name) DO UPDATE SET + total_requests = excluded.total_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + total_cost = excluded.total_cost, + updated_at = excluded.updated_at +"#, + ) + .bind(hour_utc) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_daily_model_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_model ( + id, "date", model, total_requests, input_tokens, output_tokens, + cache_creation_tokens, cache_read_tokens, total_cost, avg_response_time_ms, + created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, model, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_creation_input_tokens), 0), + COALESCE(SUM(cache_read_input_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), + CAST(COALESCE(AVG(response_time_ms), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND model IS NOT NULL AND model <> '' + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY model +ON CONFLICT ("date", model) DO UPDATE SET + total_requests = excluded.total_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + cache_read_tokens = excluded.cache_read_tokens, + total_cost = excluded.total_cost, + avg_response_time_ms = excluded.avg_response_time_ms, + updated_at = excluded.updated_at +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_daily_provider_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_provider ( + id, "date", provider_name, total_requests, input_tokens, output_tokens, + cache_creation_tokens, cache_read_tokens, total_cost, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, provider_name, COUNT(*), + COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_creation_input_tokens), 0), + COALESCE(SUM(cache_read_input_tokens), 0), + CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status NOT IN ('pending', 'streaming') + AND provider_name NOT IN ('unknown', 'pending') +GROUP BY provider_name +ON CONFLICT ("date", provider_name) DO UPDATE SET + total_requests = excluded.total_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + cache_read_tokens = excluded.cache_read_tokens, + total_cost = excluded.total_cost, + updated_at = excluded.updated_at +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_daily_api_key_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_daily_api_key ( + id, api_key_id, "date", total_requests, success_requests, error_requests, + input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, + total_cost, api_key_name, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), usage.api_key_id, ?, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(usage.input_tokens), 0), COALESCE(SUM(usage.output_tokens), 0), + COALESCE(SUM(usage.cache_creation_input_tokens), 0), + COALESCE(SUM(usage.cache_read_input_tokens), 0), + CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS REAL), MAX(api_keys.name), ?, ? +FROM "usage" AS usage +LEFT JOIN api_keys ON api_keys.id = usage.api_key_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.api_key_id IS NOT NULL AND usage.api_key_id <> '' +GROUP BY usage.api_key_id +ON CONFLICT ("date", api_key_id) DO UPDATE SET + total_requests = excluded.total_requests, + success_requests = excluded.success_requests, + error_requests = excluded.error_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + cache_read_tokens = excluded.cache_read_tokens, + total_cost = excluded.total_cost, + api_key_name = COALESCE(excluded.api_key_name, stats_daily_api_key.api_key_name), + updated_at = excluded.updated_at +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn refresh_sqlite_stats_daily_error_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + sqlx::query(r#"DELETE FROM stats_daily_error WHERE "date" = ?"#) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + let result = sqlx::query( + r#" +INSERT INTO stats_daily_error ( + id, "date", error_category, provider_name, model, count, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), ?, error_category, provider_name, model, + COUNT(*), ?, ? +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND error_category IS NOT NULL AND error_category <> '' +GROUP BY error_category, provider_name, model +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn upsert_sqlite_stats_user_daily_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result { + let result = sqlx::query( + r#" +INSERT INTO stats_user_daily ( + id, user_id, "date", total_requests, success_requests, error_requests, + input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, + total_cost, username, created_at, updated_at +) +SELECT + lower(hex(randomblob(32))), usage.user_id, ?, COUNT(*), + COUNT(*) - COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE + WHEN usage.status = 'failed' OR usage.status_code >= 400 + OR usage.error_message IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(usage.input_tokens), 0), COALESCE(SUM(usage.output_tokens), 0), + COALESCE(SUM(usage.cache_creation_input_tokens), 0), + COALESCE(SUM(usage.cache_read_input_tokens), 0), + CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS REAL), MAX(users.username), ?, ? +FROM "usage" AS usage +LEFT JOIN users ON users.id = usage.user_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND usage.status NOT IN ('pending', 'streaming') + AND usage.provider_name NOT IN ('unknown', 'pending') +GROUP BY usage.user_id +ON CONFLICT ("date", user_id) DO UPDATE SET + total_requests = excluded.total_requests, + success_requests = excluded.success_requests, + error_requests = excluded.error_requests, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + cache_read_tokens = excluded.cache_read_tokens, + total_cost = excluded.total_cost, + username = COALESCE(excluded.username, stats_user_daily.username), + updated_at = excluded.updated_at +"#, + ) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(usize::MAX)) +} + +async fn sqlite_daily_fallback_count( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result { + let start_unix_ms = start_unix_secs.checked_mul(1000).ok_or_else(|| { + DataLayerError::InvalidInput("stats fallback window start overflows milliseconds".into()) + })?; + let end_unix_ms = end_unix_secs.checked_mul(1000).ok_or_else(|| { + DataLayerError::InvalidInput("stats fallback window end overflows milliseconds".into()) + })?; + sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM ( + SELECT request_id + FROM request_candidates + WHERE created_at >= ? AND created_at < ? + AND status IN ('success', 'failed') + GROUP BY request_id + HAVING COUNT(id) > 1 +) +"#, + ) + .bind(start_unix_ms) + .bind(end_unix_ms) + .fetch_one(&mut **tx) + .await + .map_sql_err() +} + async fn sqlite_group_count( tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, group_columns: &str, - start_ms: i64, - end_ms: i64, + start_unix_secs: i64, + end_unix_secs: i64, ) -> Result { let not_empty = group_columns .split(',') @@ -332,42 +862,10 @@ FROM ( "# ); let count: i64 = sqlx::query_scalar(&sql) - .bind(start_ms) - .bind(end_ms) + .bind(start_unix_secs) + .bind(end_unix_secs) .fetch_one(&mut **tx) .await .map_sql_err()?; Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX)) } - -async fn sqlite_error_group_count( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - start_ms: i64, - end_ms: i64, -) -> Result { - let count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM ( - SELECT 1 - FROM "usage" - WHERE created_at_unix_ms >= ? - AND created_at_unix_ms < ? - AND status NOT IN ('pending', 'streaming') - AND provider_name NOT IN ('unknown', 'pending') - AND ( - status = 'failed' - OR status_code >= 400 - OR (error_category IS NOT NULL AND error_category <> '') - ) - GROUP BY COALESCE(NULLIF(error_category, ''), 'unknown_error'), provider_name, model -) -"#, - ) - .bind(start_ms) - .bind(end_ms) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX)) -} diff --git a/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs b/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs new file mode 100644 index 000000000..75b5839c0 --- /dev/null +++ b/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs @@ -0,0 +1,986 @@ +use sqlx::Sqlite; + +use crate::backend::stats_common::stats_id; +use crate::error::SqlResultExt; +use crate::DataLayerError; + +const CACHE_5M: &str = r#"MAX( + COALESCE(usage.cache_creation_input_tokens_5m, 0), + COALESCE(usage.cache_creation_ephemeral_5m_input_tokens, 0) +)"#; +const CACHE_1H: &str = r#"MAX( + COALESCE(usage.cache_creation_input_tokens_1h, 0), + COALESCE(usage.cache_creation_ephemeral_1h_input_tokens, 0) +)"#; +const CACHE_CREATION: &str = r#"CASE + WHEN COALESCE(usage.cache_creation_input_tokens, 0) = 0 + AND ({cache_5m} + {cache_1h}) > 0 + THEN {cache_5m} + {cache_1h} + ELSE MAX(COALESCE(usage.cache_creation_input_tokens, 0), 0) +END"#; +const EFFECTIVE_INPUT: &str = r#"CASE + WHEN ( + LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, '')) IN ('openai', 'gemini', 'google') + OR LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, '')) LIKE 'openai:%' + OR LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, '')) LIKE 'gemini:%' + OR LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, '')) LIKE 'google:%' + ) + AND COALESCE(usage.input_tokens, 0) > 0 + AND COALESCE(usage.cache_read_input_tokens, 0) > 0 + THEN MAX(COALESCE(usage.input_tokens, 0) - COALESCE(usage.cache_read_input_tokens, 0), 0) + ELSE MAX(COALESCE(usage.input_tokens, 0), 0) +END"#; +const SUCCESS: &str = r#"CASE + WHEN usage.status <> 'failed' + AND (usage.status_code IS NULL OR usage.status_code < 400) + AND usage.error_message IS NULL + THEN 1 ELSE 0 +END"#; +const AGGREGATABLE: &str = r#"usage.status NOT IN ('pending', 'streaming') + AND usage.provider_name NOT IN ('unknown', 'pending')"#; +const SETTLED: &str = r#"COALESCE(settlement.billing_status, usage.billing_status) = 'settled' + AND COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) > 0"#; + +fn cache_creation_expr() -> String { + CACHE_CREATION + .replace("{cache_5m}", CACHE_5M) + .replace("{cache_1h}", CACHE_1H) +} + +fn total_input_context_expr() -> String { + format!( + "({EFFECTIVE_INPUT}) + ({}) + MAX(COALESCE(usage.cache_read_input_tokens, 0), 0)", + cache_creation_expr() + ) +} + +fn total_tokens_expr() -> String { + format!( + r#"COALESCE( + NULLIF(MAX(COALESCE(usage.total_tokens, 0), 0), 0), + ({EFFECTIVE_INPUT}) + + MAX(COALESCE(usage.output_tokens, 0), 0) + + ({}) + + MAX(COALESCE(usage.cache_read_input_tokens, 0), 0), + 0 +)"#, + cache_creation_expr() + ) +} + +fn percentile_cont(sorted: &[i64], percentile: f64) -> Option { + if sorted.is_empty() { + return None; + } + let position = percentile * (sorted.len().saturating_sub(1) as f64); + let lower = position.floor() as usize; + let upper = position.ceil() as usize; + let fraction = position - lower as f64; + let value = sorted[lower] as f64 + (sorted[upper] - sorted[lower]) as f64 * fraction; + Some(value.round() as i64) +} + +async fn load_percentiles( + tx: &mut sqlx::Transaction<'_, Sqlite>, + column: &str, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(Option, Option, Option), DataLayerError> { + let sql = format!( + r#" +SELECT {column} +FROM "usage" +WHERE created_at_unix_ms >= ? AND created_at_unix_ms < ? + AND status = 'completed' + AND provider_name NOT IN ('unknown', 'pending') + AND {column} IS NOT NULL +ORDER BY {column} +"# + ); + let values: Vec = sqlx::query_scalar(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .fetch_all(&mut **tx) + .await + .map_sql_err()?; + if values.len() < 10 { + return Ok((None, None, None)); + } + Ok(( + percentile_cont(&values, 0.50), + percentile_cont(&values, 0.90), + percentile_cont(&values, 0.99), + )) +} + +pub(super) async fn refresh_hourly( + tx: &mut sqlx::Transaction<'_, Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let sql = format!( + r#" +UPDATE stats_hourly AS target +SET + cache_hit_total_requests = aggregated.cache_hit_total_requests, + cache_hit_requests = aggregated.cache_hit_requests, + completed_total_requests = aggregated.completed_total_requests, + completed_cache_hit_requests = aggregated.completed_cache_hit_requests, + completed_input_tokens = aggregated.completed_input_tokens, + completed_cache_creation_tokens = aggregated.completed_cache_creation_tokens, + completed_cache_read_tokens = aggregated.completed_cache_read_tokens, + completed_total_input_context = aggregated.completed_total_input_context, + completed_cache_creation_cost = aggregated.completed_cache_creation_cost, + completed_cache_read_cost = aggregated.completed_cache_read_cost, + settled_total_cost = aggregated.settled_total_cost, + settled_total_requests = aggregated.settled_total_requests, + settled_input_tokens = aggregated.settled_input_tokens, + settled_output_tokens = aggregated.settled_output_tokens, + settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs, + response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples +FROM ( + SELECT + COUNT(*) AS cache_hit_total_requests, + COALESCE(SUM(CASE WHEN COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_total_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' AND COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS completed_cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS completed_input_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({cache_creation}) ELSE 0 END), 0) AS completed_cache_creation_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS completed_cache_read_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({total_context}) ELSE 0 END), 0) AS completed_total_input_context, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_creation_cost, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_read_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples + FROM "usage" AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? +) AS aggregated +WHERE target.hour_utc = ? +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + + refresh_hourly_user(tx, hour_utc, start_unix_secs, end_unix_secs).await?; + refresh_hourly_response_dimensions(tx, hour_utc, start_unix_secs, end_unix_secs).await +} + +async fn refresh_hourly_user( + tx: &mut sqlx::Transaction<'_, Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let sql = format!( + r#" +UPDATE stats_hourly_user AS target +SET + cache_creation_tokens = aggregated.cache_creation_tokens, + cache_read_tokens = aggregated.cache_read_tokens, + actual_total_cost = aggregated.actual_total_cost, + response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples, + settled_total_cost = aggregated.settled_total_cost, + settled_total_requests = aggregated.settled_total_requests, + settled_input_tokens = aggregated.settled_input_tokens, + settled_output_tokens = aggregated.settled_output_tokens, + settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs +FROM ( + SELECT + usage.user_id, + COALESCE(SUM({cache_creation}), 0) AS cache_creation_tokens, + COALESCE(SUM(MAX(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0) AS actual_total_cost, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM "usage" AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} + GROUP BY usage.user_id +) AS aggregated +WHERE target.hour_utc = ? AND target.user_id = aggregated.user_id +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_hourly_response_dimensions( + tx: &mut sqlx::Transaction<'_, Sqlite>, + hour_utc: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + for (table, dimensions, group_by, target_match) in [ + ( + "stats_hourly_model", + "usage.model AS model", + "usage.model", + "target.model = aggregated.model", + ), + ( + "stats_hourly_user_model", + "usage.user_id AS user_id, usage.model AS model", + "usage.user_id, usage.model", + "target.user_id = aggregated.user_id AND target.model = aggregated.model", + ), + ] { + let sql = format!( + r#" +UPDATE {table} AS target +SET response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples +FROM ( + SELECT {dimensions}, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples + FROM "usage" AS usage + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND {AGGREGATABLE} + GROUP BY {group_by} +) AS aggregated +WHERE target.hour_utc = ? AND {target_match} +"# + ); + sqlx::query(&sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(hour_utc) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + Ok(()) +} + +pub(super) async fn refresh_daily( + tx: &mut sqlx::Transaction<'_, Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let response = load_percentiles(tx, "response_time_ms", start_unix_secs, end_unix_secs).await?; + let first_byte = + load_percentiles(tx, "first_byte_time_ms", start_unix_secs, end_unix_secs).await?; + refresh_daily_root( + tx, + day_start, + start_unix_secs, + end_unix_secs, + response, + first_byte, + ) + .await?; + refresh_daily_existing_dimensions(tx, day_start, start_unix_secs, end_unix_secs).await?; + upsert_user_dimension( + tx, + "stats_user_daily_model", + "model", + "usage.model", + "usage.model IS NOT NULL AND usage.model <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_user_dimension( + tx, + "stats_user_daily_provider", + "provider_name", + "usage.provider_name", + "usage.provider_name IS NOT NULL AND usage.provider_name <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_user_dimension( + tx, + "stats_user_daily_api_format", + "api_format", + "LOWER(COALESCE(usage.endpoint_api_format, usage.api_format, ''))", + "COALESCE(usage.endpoint_api_format, usage.api_format, '') <> ''", + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + upsert_model_provider_rows(tx, day_start, start_unix_secs, end_unix_secs, now_unix_secs) + .await?; + upsert_cost_savings_rows(tx, day_start, start_unix_secs, end_unix_secs, now_unix_secs).await?; + refresh_user_summary(tx, end_unix_secs, now_unix_secs).await +} + +async fn refresh_daily_root( + tx: &mut sqlx::Transaction<'_, Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + response: (Option, Option, Option), + first_byte: (Option, Option, Option), +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let sql = format!( + r#" +UPDATE stats_daily AS target +SET + effective_input_tokens = aggregated.effective_input_tokens, + total_input_context = aggregated.total_input_context, + response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples, + cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens, + input_cost = aggregated.input_cost, + output_cost = aggregated.output_cost, + cache_creation_cost = aggregated.cache_creation_cost, + cache_read_cost = aggregated.cache_read_cost, + cache_hit_total_requests = aggregated.cache_hit_total_requests, + cache_hit_requests = aggregated.cache_hit_requests, + completed_total_requests = aggregated.completed_total_requests, + completed_cache_hit_requests = aggregated.completed_cache_hit_requests, + completed_input_tokens = aggregated.completed_input_tokens, + completed_cache_creation_tokens = aggregated.completed_cache_creation_tokens, + completed_cache_read_tokens = aggregated.completed_cache_read_tokens, + completed_total_input_context = aggregated.completed_total_input_context, + completed_cache_creation_cost = aggregated.completed_cache_creation_cost, + completed_cache_read_cost = aggregated.completed_cache_read_cost, + settled_total_cost = aggregated.settled_total_cost, + settled_total_requests = aggregated.settled_total_requests, + settled_input_tokens = aggregated.settled_input_tokens, + settled_output_tokens = aggregated.settled_output_tokens, + settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs, + p50_response_time_ms = ?, p90_response_time_ms = ?, p99_response_time_ms = ?, + p50_first_byte_time_ms = ?, p90_first_byte_time_ms = ?, p99_first_byte_time_ms = ? +FROM ( + SELECT + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN ({EFFECTIVE_INPUT}) ELSE 0 END), 0) AS effective_input_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN ({total_context}) ELSE 0 END), 0) AS total_input_context, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL AND {AGGREGATABLE} THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN {CACHE_5M} ELSE 0 END), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN {CACHE_1H} ELSE 0 END), 0) AS cache_creation_ephemeral_1h_tokens, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.input_cost_usd, 0) ELSE 0 END), 0) AS input_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.output_cost_usd, 0) ELSE 0 END), 0) AS output_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS cache_creation_cost, + COALESCE(SUM(CASE WHEN {AGGREGATABLE} THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS cache_read_cost, + COUNT(*) AS cache_hit_total_requests, + COALESCE(SUM(CASE WHEN COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_total_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' AND COALESCE(usage.cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) AS completed_cache_hit_requests, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS completed_input_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({cache_creation}) ELSE 0 END), 0) AS completed_cache_creation_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS completed_cache_read_tokens, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN ({total_context}) ELSE 0 END), 0) AS completed_total_input_context, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_creation_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_creation_cost, + COALESCE(SUM(CASE WHEN usage.status = 'completed' THEN COALESCE(usage.cache_read_cost_usd, 0) ELSE 0 END), 0) AS completed_cache_read_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM "usage" AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? +) AS aggregated +WHERE target."date" = ? +"# + ); + sqlx::query(&sql) + .bind(response.0) + .bind(response.1) + .bind(response.2) + .bind(first_byte.0) + .bind(first_byte.1) + .bind(first_byte.2) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_daily_existing_dimensions( + tx: &mut sqlx::Transaction<'_, Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let daily_model_sql = format!( + r#" +UPDATE stats_daily_model AS target +SET response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples, + cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens +FROM ( + SELECT usage.model, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens + FROM "usage" AS usage + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND {AGGREGATABLE} AND usage.model IS NOT NULL AND usage.model <> '' + GROUP BY usage.model +) AS aggregated +WHERE target."date" = ? AND target.model = aggregated.model +"# + ); + sqlx::query(&daily_model_sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + + let user_sql = format!( + r#" +UPDATE stats_user_daily AS target +SET effective_input_tokens = aggregated.effective_input_tokens, + total_input_context = aggregated.total_input_context, + cache_creation_cost = aggregated.cache_creation_cost, + cache_read_cost = aggregated.cache_read_cost, + actual_total_cost = aggregated.actual_total_cost, + response_time_sum_ms = aggregated.response_time_sum_ms, + response_time_samples = aggregated.response_time_samples, + cache_creation_ephemeral_5m_tokens = aggregated.cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens = aggregated.cache_creation_ephemeral_1h_tokens, + settled_total_cost = aggregated.settled_total_cost, + settled_total_requests = aggregated.settled_total_requests, + settled_input_tokens = aggregated.settled_input_tokens, + settled_output_tokens = aggregated.settled_output_tokens, + settled_cache_creation_tokens = aggregated.settled_cache_creation_tokens, + settled_cache_read_tokens = aggregated.settled_cache_read_tokens, + settled_first_finalized_at_unix_secs = aggregated.settled_first_finalized_at_unix_secs, + settled_last_finalized_at_unix_secs = aggregated.settled_last_finalized_at_unix_secs +FROM ( + SELECT usage.user_id, + COALESCE(SUM({EFFECTIVE_INPUT}), 0) AS effective_input_tokens, + COALESCE(SUM({total_context}), 0) AS total_input_context, + COALESCE(SUM(COALESCE(usage.cache_creation_cost_usd, 0)), 0) AS cache_creation_cost, + COALESCE(SUM(COALESCE(usage.cache_read_cost_usd, 0)), 0) AS cache_read_cost, + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0) AS actual_total_cost, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0) AS response_time_sum_ms, + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) AS response_time_samples, + COALESCE(SUM({CACHE_5M}), 0) AS cache_creation_ephemeral_5m_tokens, + COALESCE(SUM({CACHE_1H}), 0) AS cache_creation_ephemeral_1h_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0) ELSE 0 END), 0) AS settled_total_cost, + COALESCE(SUM(CASE WHEN {SETTLED} THEN 1 ELSE 0 END), 0) AS settled_total_requests, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.input_tokens, 0), 0) ELSE 0 END), 0) AS settled_input_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.output_tokens, 0), 0) ELSE 0 END), 0) AS settled_output_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN ({cache_creation}) ELSE 0 END), 0) AS settled_cache_creation_tokens, + COALESCE(SUM(CASE WHEN {SETTLED} THEN MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) ELSE 0 END), 0) AS settled_cache_read_tokens, + MIN(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_first_finalized_at_unix_secs, + MAX(CASE WHEN {SETTLED} THEN COALESCE(settlement.finalized_at, usage.finalized_at) END) AS settled_last_finalized_at_unix_secs + FROM "usage" AS usage + LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id + WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' AND {AGGREGATABLE} + GROUP BY usage.user_id +) AS aggregated +WHERE target."date" = ? AND target.user_id = aggregated.user_id +"# + ); + sqlx::query(&user_sql) + .bind(start_unix_secs) + .bind(end_unix_secs) + .bind(day_start) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn upsert_user_dimension( + tx: &mut sqlx::Transaction<'_, Sqlite>, + table: &str, + dimension_column: &str, + dimension_expr: &str, + dimension_filter: &str, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let cache_creation = cache_creation_expr(); + let total_context = total_input_context_expr(); + let total_tokens = total_tokens_expr(); + let sql = format!( + r#" +INSERT INTO {table} ( + id, user_id, username, "date", {dimension_column}, total_requests, success_requests, + input_tokens, effective_input_tokens, output_tokens, total_tokens, total_input_context, + cache_creation_tokens, cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens, cache_read_tokens, total_cost, actual_total_cost, + response_time_sum_ms, response_time_samples, successful_response_time_sum_ms, + successful_response_time_samples, created_at, updated_at +) +SELECT lower(hex(randomblob(32))), usage.user_id, + MAX(COALESCE(usage.username, users.username)), ?, {dimension_expr}, COUNT(*), + COALESCE(SUM({SUCCESS}), 0), + COALESCE(SUM(MAX(COALESCE(usage.input_tokens, 0), 0)), 0), + COALESCE(SUM({EFFECTIVE_INPUT}), 0), + COALESCE(SUM(MAX(COALESCE(usage.output_tokens, 0), 0)), 0), + COALESCE(SUM({total_tokens}), 0), COALESCE(SUM({total_context}), 0), + COALESCE(SUM({cache_creation}), 0), COALESCE(SUM({CACHE_5M}), 0), + COALESCE(SUM({CACHE_1H}), 0), + COALESCE(SUM(MAX(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(COALESCE(settlement.billing_actual_total_cost_usd, usage.actual_total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN ({SUCCESS}) = 1 AND usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), + ?, ? +FROM "usage" AS usage +LEFT JOIN users ON users.id = usage.user_id +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND {dimension_filter} AND {AGGREGATABLE} +GROUP BY usage.user_id, {dimension_expr} +ON CONFLICT (user_id, "date", {dimension_column}) DO UPDATE SET + username = COALESCE(excluded.username, {table}.username), + total_requests = excluded.total_requests, success_requests = excluded.success_requests, + input_tokens = excluded.input_tokens, effective_input_tokens = excluded.effective_input_tokens, + output_tokens = excluded.output_tokens, total_tokens = excluded.total_tokens, + total_input_context = excluded.total_input_context, + cache_creation_tokens = excluded.cache_creation_tokens, + cache_creation_ephemeral_5m_tokens = excluded.cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h_tokens = excluded.cache_creation_ephemeral_1h_tokens, + cache_read_tokens = excluded.cache_read_tokens, total_cost = excluded.total_cost, + actual_total_cost = excluded.actual_total_cost, + response_time_sum_ms = excluded.response_time_sum_ms, + response_time_samples = excluded.response_time_samples, + successful_response_time_sum_ms = excluded.successful_response_time_sum_ms, + successful_response_time_samples = excluded.successful_response_time_samples, + updated_at = excluded.updated_at +"# + ); + sqlx::query(&sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn upsert_model_provider_rows( + tx: &mut sqlx::Transaction<'_, Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let total_tokens = total_tokens_expr(); + let model_provider_sql = format!( + r#" +INSERT INTO stats_daily_model_provider ( + id, "date", model, provider_name, total_requests, total_tokens, total_cost, + response_time_sum_ms, response_time_samples, created_at, updated_at +) +SELECT lower(hex(randomblob(32))), ?, usage.model, usage.provider_name, COUNT(*), + COALESCE(SUM({total_tokens}), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? +FROM "usage" AS usage +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} +GROUP BY usage.model, usage.provider_name +ON CONFLICT ("date", model, provider_name) DO UPDATE SET + total_requests = excluded.total_requests, total_tokens = excluded.total_tokens, + total_cost = excluded.total_cost, response_time_sum_ms = excluded.response_time_sum_ms, + response_time_samples = excluded.response_time_samples, updated_at = excluded.updated_at +"# + ); + sqlx::query(&model_provider_sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + + let user_model_provider_sql = format!( + r#" +INSERT INTO stats_user_daily_model_provider ( + id, user_id, username, "date", model, provider_name, total_requests, total_tokens, + total_cost, response_time_sum_ms, response_time_samples, created_at, updated_at +) +SELECT lower(hex(randomblob(32))), usage.user_id, MAX(COALESCE(usage.username, users.username)), + ?, usage.model, usage.provider_name, COUNT(*), COALESCE(SUM({total_tokens}), 0), + COALESCE(SUM(COALESCE(settlement.billing_total_cost_usd, usage.total_cost_usd, 0)), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN MAX(usage.response_time_ms, 0) ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN usage.response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0), ?, ? +FROM "usage" AS usage +LEFT JOIN users ON users.id = usage.user_id +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? + AND usage.user_id IS NOT NULL AND usage.user_id <> '' + AND usage.model IS NOT NULL AND usage.model <> '' AND {AGGREGATABLE} +GROUP BY usage.user_id, usage.model, usage.provider_name +ON CONFLICT (user_id, "date", model, provider_name) DO UPDATE SET + username = COALESCE(excluded.username, stats_user_daily_model_provider.username), + total_requests = excluded.total_requests, total_tokens = excluded.total_tokens, + total_cost = excluded.total_cost, response_time_sum_ms = excluded.response_time_sum_ms, + response_time_samples = excluded.response_time_samples, updated_at = excluded.updated_at +"# + ); + sqlx::query(&user_model_provider_sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn upsert_cost_savings_rows( + tx: &mut sqlx::Transaction<'_, Sqlite>, + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + for (table, dimensions) in [ + ("stats_daily_cost_savings", Vec::new()), + ( + "stats_daily_cost_savings_provider", + vec![("provider_name", "COALESCE(usage.provider_name, '')")], + ), + ( + "stats_daily_cost_savings_model", + vec![("model", "COALESCE(usage.model, '')")], + ), + ( + "stats_daily_cost_savings_model_provider", + vec![ + ("model", "COALESCE(usage.model, '')"), + ("provider_name", "COALESCE(usage.provider_name, '')"), + ], + ), + ] { + upsert_cost_savings_dimension( + tx, + table, + false, + &dimensions, + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + } + for (table, dimensions) in [ + ("stats_user_daily_cost_savings", Vec::new()), + ( + "stats_user_daily_cost_savings_provider", + vec![("provider_name", "COALESCE(usage.provider_name, '')")], + ), + ( + "stats_user_daily_cost_savings_model", + vec![("model", "COALESCE(usage.model, '')")], + ), + ( + "stats_user_daily_cost_savings_model_provider", + vec![ + ("model", "COALESCE(usage.model, '')"), + ("provider_name", "COALESCE(usage.provider_name, '')"), + ], + ), + ] { + upsert_cost_savings_dimension( + tx, + table, + true, + &dimensions, + day_start, + start_unix_secs, + end_unix_secs, + now_unix_secs, + ) + .await?; + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn upsert_cost_savings_dimension( + tx: &mut sqlx::Transaction<'_, Sqlite>, + table: &str, + per_user: bool, + dimensions: &[(&str, &str)], + day_start: i64, + start_unix_secs: i64, + end_unix_secs: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let dimension_columns = dimensions + .iter() + .map(|(column, _)| *column) + .collect::>(); + let dimension_exprs = dimensions + .iter() + .map(|(_, expression)| *expression) + .collect::>(); + let user_columns = if per_user { "user_id, username, " } else { "" }; + let user_select = if per_user { + "usage.user_id, MAX(COALESCE(usage.username, users.username)), " + } else { + "" + }; + let user_join = if per_user { + "LEFT JOIN users ON users.id = usage.user_id" + } else { + "" + }; + let user_filter = if per_user { + "AND usage.user_id IS NOT NULL AND usage.user_id <> ''" + } else { + "" + }; + let mut conflict_columns = vec!["\"date\""]; + let mut group_by = Vec::new(); + if per_user { + conflict_columns.insert(0, "user_id"); + group_by.push("usage.user_id"); + } + conflict_columns.extend(dimension_columns.iter().copied()); + group_by.extend(dimension_exprs.iter().copied()); + let dimension_columns_sql = if dimension_columns.is_empty() { + String::new() + } else { + format!("{}, ", dimension_columns.join(", ")) + }; + let dimension_select_sql = if dimension_exprs.is_empty() { + String::new() + } else { + format!("{}, ", dimension_exprs.join(", ")) + }; + let group_by_sql = if group_by.is_empty() { + String::new() + } else { + format!("GROUP BY {}", group_by.join(", ")) + }; + let sql = format!( + r#" +INSERT INTO {table} ( + id, {user_columns}"date", {dimension_columns_sql}cache_read_tokens, + cache_read_cost, cache_creation_cost, estimated_full_cost, created_at, updated_at +) +SELECT lower(hex(randomblob(32))), {user_select}?, {dimension_select_sql} + COALESCE(SUM(MAX(COALESCE(usage.cache_read_input_tokens, 0), 0)), 0), + COALESCE(SUM(COALESCE(usage.cache_read_cost_usd, 0)), 0), + COALESCE(SUM(COALESCE(usage.cache_creation_cost_usd, 0)), 0), + COALESCE(SUM( + COALESCE(settlement.input_price_per_1m, usage.input_price_per_1m, 0) + * MAX(COALESCE(usage.cache_read_input_tokens, 0), 0) / 1000000.0 + ), 0), ?, ? +FROM "usage" AS usage +LEFT JOIN usage_settlement_snapshots AS settlement ON settlement.request_id = usage.request_id +{user_join} +WHERE usage.created_at_unix_ms >= ? AND usage.created_at_unix_ms < ? {user_filter} +{group_by_sql} +ON CONFLICT ({}) DO UPDATE SET + {}cache_read_tokens = excluded.cache_read_tokens, + cache_read_cost = excluded.cache_read_cost, + cache_creation_cost = excluded.cache_creation_cost, + estimated_full_cost = excluded.estimated_full_cost, + updated_at = excluded.updated_at +"#, + conflict_columns.join(", "), + if per_user { + format!("username = COALESCE(excluded.username, {table}.username), ") + } else { + String::new() + } + ); + sqlx::query(&sql) + .bind(day_start) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(start_unix_secs) + .bind(end_unix_secs) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} + +async fn refresh_user_summary( + tx: &mut sqlx::Transaction<'_, Sqlite>, + cutoff_date: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + sqlx::query( + r#" +INSERT INTO stats_user_summary ( + id, user_id, username, cutoff_date, all_time_requests, all_time_success_requests, + all_time_error_requests, all_time_input_tokens, all_time_output_tokens, + all_time_cache_creation_tokens, all_time_cache_read_tokens, all_time_cost, + all_time_actual_cost, active_days, first_active_date, last_active_date, + created_at, updated_at +) +SELECT lower(hex(randomblob(32))), user_id, MAX(username), ?, + COALESCE(SUM(total_requests), 0), COALESCE(SUM(success_requests), 0), + COALESCE(SUM(error_requests), 0), COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), COALESCE(SUM(cache_creation_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(total_cost), 0), + COALESCE(SUM(actual_total_cost), 0), + COALESCE(SUM(CASE WHEN total_requests > 0 THEN 1 ELSE 0 END), 0), + MIN(CASE WHEN total_requests > 0 THEN "date" END), + MAX(CASE WHEN total_requests > 0 THEN "date" END), ?, ? +FROM stats_user_daily +WHERE "date" < ? +GROUP BY user_id +ON CONFLICT (user_id) DO UPDATE SET + username = COALESCE(excluded.username, stats_user_summary.username), + cutoff_date = excluded.cutoff_date, all_time_requests = excluded.all_time_requests, + all_time_success_requests = excluded.all_time_success_requests, + all_time_error_requests = excluded.all_time_error_requests, + all_time_input_tokens = excluded.all_time_input_tokens, + all_time_output_tokens = excluded.all_time_output_tokens, + all_time_cache_creation_tokens = excluded.all_time_cache_creation_tokens, + all_time_cache_read_tokens = excluded.all_time_cache_read_tokens, + all_time_cost = excluded.all_time_cost, all_time_actual_cost = excluded.all_time_actual_cost, + active_days = excluded.active_days, first_active_date = excluded.first_active_date, + last_active_date = excluded.last_active_date, updated_at = excluded.updated_at +"#, + ) + .bind(cutoff_date) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(cutoff_date) + .execute(&mut **tx) + .await + .map_sql_err()?; + refresh_global_summary(tx, cutoff_date, now_unix_secs).await?; + Ok(()) +} + +async fn refresh_global_summary( + tx: &mut sqlx::Transaction<'_, Sqlite>, + cutoff_date: i64, + now_unix_secs: i64, +) -> Result<(), DataLayerError> { + let existing_id: Option = + sqlx::query_scalar("SELECT id FROM stats_summary ORDER BY created_at, id LIMIT 1") + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + let summary_id = existing_id.unwrap_or_else(|| stats_id("stats-summary")); + sqlx::query( + r#" +INSERT INTO stats_summary ( + id, cutoff_date, all_time_requests, all_time_success_requests, + all_time_error_requests, all_time_input_tokens, all_time_output_tokens, + all_time_cache_creation_tokens, all_time_cache_read_tokens, all_time_cost, + all_time_actual_cost, total_users, active_users, total_api_keys, + active_api_keys, created_at, updated_at +) +SELECT ?, ?, COALESCE(SUM(total_requests), 0), COALESCE(SUM(success_requests), 0), + COALESCE(SUM(error_requests), 0), COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), COALESCE(SUM(cache_creation_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), COALESCE(SUM(total_cost), 0), + COALESCE(SUM(actual_total_cost), 0), + (SELECT COUNT(*) FROM users), + (SELECT COUNT(*) FROM users WHERE is_active <> 0), + (SELECT COUNT(*) FROM api_keys), + (SELECT COUNT(*) FROM api_keys WHERE is_active <> 0), ?, ? +FROM stats_daily +WHERE "date" < ? +ON CONFLICT (id) DO UPDATE SET + cutoff_date = excluded.cutoff_date, + all_time_requests = excluded.all_time_requests, + all_time_success_requests = excluded.all_time_success_requests, + all_time_error_requests = excluded.all_time_error_requests, + all_time_input_tokens = excluded.all_time_input_tokens, + all_time_output_tokens = excluded.all_time_output_tokens, + all_time_cache_creation_tokens = excluded.all_time_cache_creation_tokens, + all_time_cache_read_tokens = excluded.all_time_cache_read_tokens, + all_time_cost = excluded.all_time_cost, + all_time_actual_cost = excluded.all_time_actual_cost, + total_users = excluded.total_users, active_users = excluded.active_users, + total_api_keys = excluded.total_api_keys, active_api_keys = excluded.active_api_keys, + updated_at = excluded.updated_at +"#, + ) + .bind(summary_id) + .bind(cutoff_date) + .bind(now_unix_secs) + .bind(now_unix_secs) + .bind(cutoff_date) + .execute(&mut **tx) + .await + .map_sql_err()?; + Ok(()) +} diff --git a/crates/aether-data/runtime/src/backend/stats_common.rs b/crates/aether-data/runtime/src/backend/stats_common.rs index ed15bba78..0881d580c 100644 --- a/crates/aether-data/runtime/src/backend/stats_common.rs +++ b/crates/aether-data/runtime/src/backend/stats_common.rs @@ -7,12 +7,6 @@ pub(crate) fn unix_secs(value: DateTime) -> i64 { value.timestamp().max(0) } -pub(crate) fn unix_ms(value: i64) -> Result { - value.checked_mul(1000).ok_or_else(|| { - DataLayerError::InvalidInput(format!("timestamp overflow while converting {value} to ms")) - }) -} - pub(crate) fn utc_from_unix_secs( value: i64, field_name: &str, diff --git a/crates/aether-data/runtime/src/backend/system.rs b/crates/aether-data/runtime/src/backend/system.rs index 0aa073675..c3be27999 100644 --- a/crates/aether-data/runtime/src/backend/system.rs +++ b/crates/aether-data/runtime/src/backend/system.rs @@ -70,10 +70,12 @@ const ADMIN_STATS_PURGE_TABLES: &[&str] = &[ ]; const ADMIN_USAGE_CHILD_TABLES: &[&str] = &[ + "usage_counter_deltas", "usage_body_blobs", "usage_http_audits", "usage_routing_snapshots", "usage_settlement_snapshots", + "user_model_usage_counts", ]; const USAGE_BODY_FIELD_COLUMNS: &[&str] = &[ diff --git a/crates/aether-data/runtime/src/backend/system/mysql.rs b/crates/aether-data/runtime/src/backend/system/mysql.rs index 3fa3cd634..15d4122b2 100644 --- a/crates/aether-data/runtime/src/backend/system/mysql.rs +++ b/crates/aether-data/runtime/src/backend/system/mysql.rs @@ -855,6 +855,22 @@ SET provider_id = NULL, WHERE provider_id IS NOT NULL OR provider_endpoint_id IS NOT NULL OR provider_api_key_id IS NOT NULL +"#, + summary, + ) + .await?; + mysql_execute_if_table( + tx, + "usage_routing_snapshots", + "usage_routing_provider_refs_cleared", + r#" +UPDATE usage_routing_snapshots +SET selected_provider_id = NULL, + selected_endpoint_id = NULL, + selected_provider_api_key_id = NULL +WHERE selected_provider_id IS NOT NULL + OR selected_endpoint_id IS NOT NULL + OR selected_provider_api_key_id IS NOT NULL "#, summary, ) @@ -952,6 +968,14 @@ WHERE request_count <> 0 summary, ) .await?; + mysql_execute_if_table( + tx, + "global_models", + "global_model_usage_stats_reset", + "UPDATE global_models SET usage_count = 0 WHERE usage_count <> 0", + summary, + ) + .await?; } AdminSystemPurgeTarget::AuditLogs => { mysql_delete_table(tx, "audit_logs", summary).await?; diff --git a/crates/aether-data/runtime/src/backend/system/postgres.rs b/crates/aether-data/runtime/src/backend/system/postgres.rs index 8f641449a..b6596e880 100644 --- a/crates/aether-data/runtime/src/backend/system/postgres.rs +++ b/crates/aether-data/runtime/src/backend/system/postgres.rs @@ -460,6 +460,22 @@ SET provider_id = NULL, WHERE provider_id IS NOT NULL OR provider_endpoint_id IS NOT NULL OR provider_api_key_id IS NOT NULL +"#, + summary, + ) + .await?; + pg_execute_if_table( + tx, + "usage_routing_snapshots", + "usage_routing_provider_refs_cleared", + r#" +UPDATE public.usage_routing_snapshots +SET selected_provider_id = NULL, + selected_endpoint_id = NULL, + selected_provider_api_key_id = NULL +WHERE selected_provider_id IS NOT NULL + OR selected_endpoint_id IS NOT NULL + OR selected_provider_api_key_id IS NOT NULL "#, summary, ) @@ -559,6 +575,14 @@ WHERE request_count <> 0 summary, ) .await?; + pg_execute_if_table( + tx, + "global_models", + "global_model_usage_stats_reset", + "UPDATE public.global_models SET usage_count = 0 WHERE usage_count <> 0", + summary, + ) + .await?; } AdminSystemPurgeTarget::AuditLogs => { pg_delete_table(tx, "audit_logs", summary).await?; diff --git a/crates/aether-data/runtime/src/backend/system/sqlite.rs b/crates/aether-data/runtime/src/backend/system/sqlite.rs index 9e500b0ad..033bd38bd 100644 --- a/crates/aether-data/runtime/src/backend/system/sqlite.rs +++ b/crates/aether-data/runtime/src/backend/system/sqlite.rs @@ -601,6 +601,22 @@ SET provider_id = NULL, WHERE provider_id IS NOT NULL OR provider_endpoint_id IS NOT NULL OR provider_api_key_id IS NOT NULL +"#, + summary, + ) + .await?; + sqlite_execute_if_table( + tx, + "usage_routing_snapshots", + "usage_routing_provider_refs_cleared", + r#" +UPDATE usage_routing_snapshots +SET selected_provider_id = NULL, + selected_endpoint_id = NULL, + selected_provider_api_key_id = NULL +WHERE selected_provider_id IS NOT NULL + OR selected_endpoint_id IS NOT NULL + OR selected_provider_api_key_id IS NOT NULL "#, summary, ) @@ -698,6 +714,14 @@ WHERE request_count <> 0 summary, ) .await?; + sqlite_execute_if_table( + tx, + "global_models", + "global_model_usage_stats_reset", + "UPDATE global_models SET usage_count = 0 WHERE usage_count <> 0", + summary, + ) + .await?; } AdminSystemPurgeTarget::AuditLogs => { sqlite_delete_table(tx, "audit_logs", summary).await?; @@ -1250,3 +1274,132 @@ pub(super) fn map_admin_system_stats( .max(0) as u64, }) } + +#[cfg(test)] +mod tests { + use super::{purge_sqlite_admin_system_data, AdminSystemPurgeSummary, AdminSystemPurgeTarget}; + + #[tokio::test] + async fn usage_purge_removes_pending_counters_and_resets_model_usage() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + crate::lifecycle::migrate::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + sqlx::raw_sql( + r#" +INSERT INTO users (id, email, username, created_at, updated_at) +VALUES ('purge-user', 'purge@example.com', 'purge-user', 1, 1); +INSERT INTO global_models (id, name, usage_count, created_at, updated_at) +VALUES ('purge-model', 'purge-model', 7, 1, 1); +INSERT INTO "usage" ( + request_id, user_id, provider_name, model, status, billing_status, + created_at_unix_ms, updated_at_unix_secs +) VALUES ('purge-request', 'purge-user', 'provider', 'purge-model', 'completed', 'settled', 1, 1); +INSERT INTO usage_counter_deltas ( + id, request_id, kind, target_id, request_count_delta, created_at +) VALUES ('purge-delta', 'purge-request', 'model', 'purge-model', 1, 1); +INSERT INTO user_model_usage_counts ( + id, user_id, model, usage_count, created_at, updated_at +) VALUES ('purge-user-model', 'purge-user', 'purge-model', 7, 1, 1); +"#, + ) + .execute(&pool) + .await + .expect("usage purge fixtures should insert"); + + let mut tx = pool.begin().await.expect("purge transaction should begin"); + let mut summary = AdminSystemPurgeSummary::default(); + purge_sqlite_admin_system_data(&mut tx, AdminSystemPurgeTarget::Usage, &mut summary) + .await + .expect("usage purge should succeed"); + tx.commit().await.expect("purge transaction should commit"); + + let usage_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM \"usage\"") + .fetch_one(&pool) + .await + .expect("usage count should load"); + let delta_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas") + .fetch_one(&pool) + .await + .expect("counter delta count should load"); + let user_model_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM user_model_usage_counts") + .fetch_one(&pool) + .await + .expect("user model count should load"); + let model_usage_count: i64 = + sqlx::query_scalar("SELECT usage_count FROM global_models WHERE id = 'purge-model'") + .fetch_one(&pool) + .await + .expect("global model usage count should load"); + + assert_eq!(usage_count, 0); + assert_eq!(delta_count, 0); + assert_eq!(user_model_count, 0); + assert_eq!(model_usage_count, 0); + assert_eq!(summary.affected.get("usage_counter_deltas"), Some(&1)); + assert_eq!(summary.affected.get("user_model_usage_counts"), Some(&1)); + assert_eq!( + summary.affected.get("global_model_usage_stats_reset"), + Some(&1) + ); + } + + #[tokio::test] + async fn config_purge_clears_canonical_routing_provider_refs() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + crate::lifecycle::migrate::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + sqlx::raw_sql( + r#" +INSERT INTO "usage" ( + request_id, provider_name, model, status, billing_status, + created_at_unix_ms, updated_at_unix_secs +) VALUES ('routing-purge-request', 'provider', 'model', 'completed', 'settled', 1, 1); +INSERT INTO usage_routing_snapshots ( + request_id, selected_provider_id, selected_endpoint_id, + selected_provider_api_key_id, created_at, updated_at +) VALUES ( + 'routing-purge-request', 'provider-1', 'endpoint-1', 'key-1', 1, 1 +); +"#, + ) + .execute(&pool) + .await + .expect("routing purge fixture should insert"); + + let mut tx = pool.begin().await.expect("purge transaction should begin"); + let mut summary = AdminSystemPurgeSummary::default(); + purge_sqlite_admin_system_data(&mut tx, AdminSystemPurgeTarget::Config, &mut summary) + .await + .expect("config purge should succeed"); + tx.commit().await.expect("purge transaction should commit"); + + let refs = sqlx::query_as::<_, (Option, Option, Option)>( + r#" +SELECT selected_provider_id, selected_endpoint_id, selected_provider_api_key_id +FROM usage_routing_snapshots +WHERE request_id = 'routing-purge-request' +"#, + ) + .fetch_one(&pool) + .await + .expect("routing refs should load"); + assert_eq!(refs, (None, None, None)); + assert_eq!( + summary.affected.get("usage_routing_provider_refs_cleared"), + Some(&1) + ); + } +} diff --git a/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs b/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs index 6d6719bbb..4741f4db5 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs @@ -1,16 +1,253 @@ -use sqlx::migrate::MigrateError; -use tracing::info; +use std::collections::{HashMap, HashSet}; + +use sqlx::{ + migrate::{Migrate, MigrateError, Migrator}, + query, Connection, MySqlConnection, Row, +}; +use tracing::{error, info, warn}; use super::types::PendingBackfillInfo; use crate::driver::mysql::MysqlPool; -pub async fn run_backfills(_pool: &MysqlPool) -> Result<(), MigrateError> { - info!("mysql database backfills are up to date"); +static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/mysql"); + +const ENSURE_SCHEMA_BACKFILLS_TABLE_SQL: &str = r#" +CREATE TABLE IF NOT EXISTS schema_backfills ( + version BIGINT NOT NULL, + description TEXT NOT NULL, + success BOOLEAN NOT NULL DEFAULT TRUE, + checksum BLOB NOT NULL, + execution_time BIGINT NOT NULL DEFAULT 0, + applied_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + PRIMARY KEY (version) +) +"#; +const LIST_APPLIED_BACKFILLS_SQL: &str = r#" +SELECT version, checksum +FROM schema_backfills +WHERE success IS TRUE +ORDER BY version ASC +"#; +const INSERT_APPLIED_BACKFILL_SQL: &str = r#" +INSERT INTO schema_backfills ( + version, + description, + success, + checksum, + execution_time, + applied_at +) VALUES ( + ?, + ?, + TRUE, + ?, + ?, + CURRENT_TIMESTAMP(6) +) +ON DUPLICATE KEY UPDATE version = schema_backfills.version +"#; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AppliedBackfill { + version: i64, + checksum: Vec, +} + +pub async fn run_backfills(pool: &MysqlPool) -> Result<(), MigrateError> { + let mut conn = pool.acquire().await?; + + if BACKFILL_MIGRATOR.locking { + conn.lock().await?; + } + + let result = run_backfills_locked(&mut conn).await; + + if BACKFILL_MIGRATOR.locking { + match conn.unlock().await { + Ok(()) => {} + Err(unlock_error) if result.is_ok() => return Err(unlock_error), + Err(unlock_error) => { + warn!( + error = %unlock_error, + "mysql database backfill lock release failed after backfill error" + ); + } + } + } + + result +} + +pub async fn pending_backfills(pool: &MysqlPool) -> Result, MigrateError> { + let mut conn = pool.acquire().await?; + pending_backfills_locked(&mut conn).await +} + +async fn run_backfills_locked(conn: &mut MySqlConnection) -> Result<(), MigrateError> { + ensure_schema_backfills_table(conn).await?; + + let applied_backfills = list_applied_backfills(conn).await?; + validate_applied_backfills(&applied_backfills)?; + + let applied_by_version: HashMap<_, _> = applied_backfills + .iter() + .map(|backfill| (backfill.version, backfill)) + .collect(); + let pending_backfills: Vec<_> = BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + .filter(|backfill| !applied_by_version.contains_key(&backfill.version)) + .collect(); + + if pending_backfills.is_empty() { + info!( + driver = "mysql", + pending_backfills = 0, + "database backfills already up to date" + ); + return Ok(()); + } + + info!( + driver = "mysql", + pending_backfills = pending_backfills.len(), + "database backfills pending" + ); + + for (index, backfill) in pending_backfills.iter().enumerate() { + let current = index + 1; + let total = pending_backfills.len(); + info!( + driver = "mysql", + current, + total, + version = backfill.version, + description = %backfill.description, + "applying database backfill" + ); + + let mut tx = conn.begin().await?; + let started_at = std::time::Instant::now(); + sqlx::raw_sql(&backfill.sql).execute(&mut *tx).await?; + let elapsed_ms = i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX); + query(INSERT_APPLIED_BACKFILL_SQL) + .bind(backfill.version) + .bind(backfill.description.as_ref()) + .bind(backfill.checksum.as_ref()) + .bind(elapsed_ms) + .execute(&mut *tx) + .await?; + tx.commit().await?; + + info!( + driver = "mysql", + current, + total, + version = backfill.version, + description = %backfill.description, + elapsed_ms, + "applied database backfill" + ); + } + + info!( + driver = "mysql", + pending_backfills = 0, + "database backfills complete" + ); Ok(()) } -pub async fn pending_backfills( - _pool: &MysqlPool, +async fn pending_backfills_locked( + conn: &mut MySqlConnection, ) -> Result, MigrateError> { - Ok(Vec::new()) + ensure_schema_backfills_table(conn).await?; + let applied_backfills = list_applied_backfills(conn).await?; + validate_applied_backfills(&applied_backfills)?; + Ok(pending_backfills_from_applied(&applied_backfills)) +} + +async fn ensure_schema_backfills_table(conn: &mut MySqlConnection) -> Result<(), MigrateError> { + query(ENSURE_SCHEMA_BACKFILLS_TABLE_SQL) + .execute(&mut *conn) + .await?; + Ok(()) +} + +async fn list_applied_backfills( + conn: &mut MySqlConnection, +) -> Result, MigrateError> { + let rows = query(LIST_APPLIED_BACKFILLS_SQL) + .fetch_all(&mut *conn) + .await?; + rows.into_iter() + .map(|row| { + Ok(AppliedBackfill { + version: row.try_get("version")?, + checksum: row.try_get("checksum")?, + }) + }) + .collect::, sqlx::Error>>() + .map_err(MigrateError::from) +} + +fn validate_applied_backfills(applied_backfills: &[AppliedBackfill]) -> Result<(), MigrateError> { + if BACKFILL_MIGRATOR.ignore_missing { + return Ok(()); + } + + let known_versions: HashSet<_> = BACKFILL_MIGRATOR + .iter() + .map(|backfill| backfill.version) + .collect(); + for applied_backfill in applied_backfills { + if !known_versions.contains(&applied_backfill.version) { + error!( + driver = "mysql", + version = applied_backfill.version, + "applied database backfill is missing from embedded backfills" + ); + return Err(MigrateError::VersionMissing(applied_backfill.version)); + } + } + + for backfill in BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + { + let Some(applied) = applied_backfills + .iter() + .find(|applied| applied.version == backfill.version) + else { + continue; + }; + if backfill.checksum != applied.checksum { + warn!( + driver = "mysql", + version = backfill.version, + description = %backfill.description, + "applied database backfill checksum differs from embedded backfill; skipping strict enforcement" + ); + } + } + + Ok(()) +} + +fn pending_backfills_from_applied( + applied_backfills: &[AppliedBackfill], +) -> Vec { + let applied_versions: HashSet<_> = applied_backfills + .iter() + .map(|backfill| backfill.version) + .collect(); + BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + .filter(|backfill| !applied_versions.contains(&backfill.version)) + .map(|backfill| PendingBackfillInfo { + version: backfill.version, + description: backfill.description.to_string(), + }) + .collect() } diff --git a/crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs b/crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs index 9eee1125c..6e8a17780 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs @@ -1,16 +1,254 @@ -use sqlx::migrate::MigrateError; -use tracing::info; +use std::collections::{HashMap, HashSet}; + +use sqlx::{ + migrate::{Migrate, MigrateError, Migrator}, + query, Connection, Row, SqliteConnection, +}; +use tracing::{error, info, warn}; use super::types::PendingBackfillInfo; use crate::driver::sqlite::SqlitePool; -pub async fn run_backfills(_pool: &SqlitePool) -> Result<(), MigrateError> { - info!("sqlite database backfills are up to date"); - Ok(()) +static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/sqlite"); + +const ENSURE_SCHEMA_BACKFILLS_TABLE_SQL: &str = r#" +CREATE TABLE IF NOT EXISTS schema_backfills ( + version INTEGER NOT NULL PRIMARY KEY, + description TEXT NOT NULL, + success INTEGER NOT NULL DEFAULT 1, + checksum BLOB NOT NULL, + execution_time INTEGER NOT NULL DEFAULT 0, + applied_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s', 'now') AS INTEGER)) +) +"#; +const LIST_APPLIED_BACKFILLS_SQL: &str = r#" +SELECT version, checksum +FROM schema_backfills +WHERE success = 1 +ORDER BY version ASC +"#; +const INSERT_APPLIED_BACKFILL_SQL: &str = r#" +INSERT INTO schema_backfills ( + version, + description, + success, + checksum, + execution_time, + applied_at +) VALUES ( + ?, + ?, + 1, + ?, + ?, + CAST(strftime('%s', 'now') AS INTEGER) +) +ON CONFLICT(version) DO NOTHING +"#; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AppliedBackfill { + version: i64, + checksum: Vec, +} + +pub async fn run_backfills(pool: &SqlitePool) -> Result<(), MigrateError> { + let mut conn = pool.acquire().await?; + + if BACKFILL_MIGRATOR.locking { + conn.lock().await?; + } + + let result = run_backfills_locked(&mut conn).await; + + if BACKFILL_MIGRATOR.locking { + match conn.unlock().await { + Ok(()) => {} + Err(unlock_error) if result.is_ok() => return Err(unlock_error), + Err(unlock_error) => { + warn!( + error = %unlock_error, + "sqlite database backfill lock release failed after backfill error" + ); + } + } + } + + result } pub async fn pending_backfills( - _pool: &SqlitePool, + pool: &SqlitePool, ) -> Result, MigrateError> { - Ok(Vec::new()) + let mut conn = pool.acquire().await?; + pending_backfills_locked(&mut conn).await +} + +async fn run_backfills_locked(conn: &mut SqliteConnection) -> Result<(), MigrateError> { + ensure_schema_backfills_table(conn).await?; + + let applied_backfills = list_applied_backfills(conn).await?; + validate_applied_backfills(&applied_backfills)?; + + let applied_by_version: HashMap<_, _> = applied_backfills + .iter() + .map(|backfill| (backfill.version, backfill)) + .collect(); + let pending_backfills: Vec<_> = BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + .filter(|backfill| !applied_by_version.contains_key(&backfill.version)) + .collect(); + + if pending_backfills.is_empty() { + info!( + driver = "sqlite", + pending_backfills = 0, + "database backfills already up to date" + ); + return Ok(()); + } + + info!( + driver = "sqlite", + pending_backfills = pending_backfills.len(), + "database backfills pending" + ); + + for (index, backfill) in pending_backfills.iter().enumerate() { + let current = index + 1; + let total = pending_backfills.len(); + info!( + driver = "sqlite", + current, + total, + version = backfill.version, + description = %backfill.description, + "applying database backfill" + ); + + let mut tx = conn.begin().await?; + let started_at = std::time::Instant::now(); + sqlx::raw_sql(&backfill.sql).execute(&mut *tx).await?; + let elapsed_ms = i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX); + query(INSERT_APPLIED_BACKFILL_SQL) + .bind(backfill.version) + .bind(backfill.description.as_ref()) + .bind(backfill.checksum.as_ref()) + .bind(elapsed_ms) + .execute(&mut *tx) + .await?; + tx.commit().await?; + + info!( + driver = "sqlite", + current, + total, + version = backfill.version, + description = %backfill.description, + elapsed_ms, + "applied database backfill" + ); + } + + info!( + driver = "sqlite", + pending_backfills = 0, + "database backfills complete" + ); + Ok(()) +} + +async fn pending_backfills_locked( + conn: &mut SqliteConnection, +) -> Result, MigrateError> { + ensure_schema_backfills_table(conn).await?; + let applied_backfills = list_applied_backfills(conn).await?; + validate_applied_backfills(&applied_backfills)?; + Ok(pending_backfills_from_applied(&applied_backfills)) +} + +async fn ensure_schema_backfills_table(conn: &mut SqliteConnection) -> Result<(), MigrateError> { + query(ENSURE_SCHEMA_BACKFILLS_TABLE_SQL) + .execute(&mut *conn) + .await?; + Ok(()) +} + +async fn list_applied_backfills( + conn: &mut SqliteConnection, +) -> Result, MigrateError> { + let rows = query(LIST_APPLIED_BACKFILLS_SQL) + .fetch_all(&mut *conn) + .await?; + rows.into_iter() + .map(|row| { + Ok(AppliedBackfill { + version: row.try_get("version")?, + checksum: row.try_get("checksum")?, + }) + }) + .collect::, sqlx::Error>>() + .map_err(MigrateError::from) +} + +fn validate_applied_backfills(applied_backfills: &[AppliedBackfill]) -> Result<(), MigrateError> { + if BACKFILL_MIGRATOR.ignore_missing { + return Ok(()); + } + + let known_versions: HashSet<_> = BACKFILL_MIGRATOR + .iter() + .map(|backfill| backfill.version) + .collect(); + for applied_backfill in applied_backfills { + if !known_versions.contains(&applied_backfill.version) { + error!( + driver = "sqlite", + version = applied_backfill.version, + "applied database backfill is missing from embedded backfills" + ); + return Err(MigrateError::VersionMissing(applied_backfill.version)); + } + } + + for backfill in BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + { + let Some(applied) = applied_backfills + .iter() + .find(|applied| applied.version == backfill.version) + else { + continue; + }; + if backfill.checksum != applied.checksum { + warn!( + driver = "sqlite", + version = backfill.version, + description = %backfill.description, + "applied database backfill checksum differs from embedded backfill; skipping strict enforcement" + ); + } + } + + Ok(()) +} + +fn pending_backfills_from_applied( + applied_backfills: &[AppliedBackfill], +) -> Vec { + let applied_versions: HashSet<_> = applied_backfills + .iter() + .map(|backfill| backfill.version) + .collect(); + BACKFILL_MIGRATOR + .iter() + .filter(|backfill| backfill.migration_type.is_up_migration()) + .filter(|backfill| !applied_versions.contains(&backfill.version)) + .map(|backfill| PendingBackfillInfo { + version: backfill.version, + description: backfill.description.to_string(), + }) + .collect() } diff --git a/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs b/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs index f2a4ed888..0dd156a3c 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs @@ -4,15 +4,14 @@ use std::{ time::{Duration, Instant}, }; -use sqlx::{query, query_scalar, Connection, PgConnection, PgPool}; +use sqlx::{query, query_as, query_scalar, Connection, PgConnection, PgPool}; use super::{ pending_backfills, pending_backfills_from_applied, pending_mysql_backfills, pending_sqlite_backfills, run_backfills, run_mysql_backfills, run_sqlite_backfills, AppliedBackfill, }; -use crate::lifecycle::migrate::prepare_database_for_startup; -use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; +use crate::lifecycle::migrate::{prepare_database_for_startup, run_sqlite_migrations}; const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION: i64 = 20260517012000; const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL: &str = @@ -88,44 +87,558 @@ fn corrected_legacy_backfill_is_not_requeued_after_application() { } #[tokio::test] -async fn mysql_backfills_are_empty_until_driver_specific_backfills_exist() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), +async fn mysql_backfills_apply_portable_repairs_when_url_is_set() { + let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + eprintln!("skipping mysql backfill test because AETHER_TEST_MYSQL_URL is unset"); + return; + }; + + let pool = sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&database_url) + .await + .expect("mysql backfill test pool should connect"); + let mut conn = pool + .acquire() + .await + .expect("mysql backfill test connection should acquire"); + sqlx::raw_sql( + r#" +CREATE TEMPORARY TABLE schema_backfills ( + version BIGINT PRIMARY KEY, + description TEXT NOT NULL, + success BOOLEAN NOT NULL, + checksum BLOB NOT NULL, + execution_time BIGINT NOT NULL, + applied_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) +); +CREATE TEMPORARY TABLE api_keys ( + id VARCHAR(64) PRIMARY KEY, + total_requests BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + total_cost_usd DOUBLE NOT NULL DEFAULT 0, + last_used_at BIGINT +); +CREATE TEMPORARY TABLE provider_api_keys ( + id VARCHAR(64) PRIMARY KEY, + total_tokens BIGINT NOT NULL DEFAULT 0 +); +CREATE TEMPORARY TABLE global_models ( + id VARCHAR(64) PRIMARY KEY, + name VARCHAR(255) NOT NULL, + usage_count BIGINT NOT NULL DEFAULT 0, + updated_at BIGINT NOT NULL +); +CREATE TEMPORARY TABLE providers ( + id VARCHAR(64) PRIMARY KEY, + enabled BOOLEAN NOT NULL, + is_active BOOLEAN NOT NULL +); +CREATE TEMPORARY TABLE provider_endpoints ( + id VARCHAR(64) PRIMARY KEY, + enabled BOOLEAN NOT NULL, + is_active BOOLEAN NOT NULL +); +CREATE TEMPORARY TABLE models ( + id VARCHAR(64) PRIMARY KEY, + enabled BOOLEAN NOT NULL, + is_active BOOLEAN NOT NULL +); +CREATE TEMPORARY TABLE `usage` ( + request_id VARCHAR(128) PRIMARY KEY, + api_key_id VARCHAR(64), + provider_api_key_id VARCHAR(64), + model VARCHAR(255), + status VARCHAR(64) NOT NULL, + total_tokens BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_input_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_input_tokens_5m BIGINT NOT NULL DEFAULT 0, + cache_creation_input_tokens_1h BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_5m_input_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_ephemeral_1h_input_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_input_tokens BIGINT NOT NULL DEFAULT 0, + endpoint_api_format VARCHAR(64), + api_format VARCHAR(64), + total_cost_usd DOUBLE NOT NULL DEFAULT 0, + created_at BIGINT, + created_at_unix_ms BIGINT NOT NULL DEFAULT 0, + updated_at_unix_secs BIGINT NOT NULL DEFAULT 0 +); +CREATE TEMPORARY TABLE usage_settlement_snapshots ( + request_id VARCHAR(128) PRIMARY KEY, + billing_effective_input_tokens BIGINT, + billing_output_tokens BIGINT, + billing_cache_creation_tokens BIGINT, + billing_cache_creation_5m_tokens BIGINT, + billing_cache_creation_1h_tokens BIGINT, + billing_cache_read_tokens BIGINT, + billing_total_input_context BIGINT +); +INSERT INTO api_keys (id, total_requests, total_tokens, total_cost_usd) +VALUES ('mysql-backfill-api-key', 77, 7777, 77.0); +INSERT INTO provider_api_keys (id, total_tokens) +VALUES ('mysql-backfill-provider-key', 7777); +INSERT INTO global_models (id, name, usage_count, updated_at) +VALUES ('mysql-backfill-model', 'gpt-portable', 77, 1); +INSERT INTO providers (id, enabled, is_active) +VALUES ('mysql-backfill-provider', TRUE, FALSE); +INSERT INTO provider_endpoints (id, enabled, is_active) +VALUES ('mysql-backfill-endpoint', TRUE, FALSE); +INSERT INTO models (id, enabled, is_active) +VALUES ('mysql-backfill-provider-model', TRUE, FALSE); +INSERT INTO `usage` ( + request_id, + api_key_id, + provider_api_key_id, + model, + status, + total_tokens, + input_tokens, + output_tokens, + cache_read_input_tokens, + api_format, + total_cost_usd, + created_at, + created_at_unix_ms, + updated_at_unix_secs +) VALUES + ( + 'mysql-backfill-completed', + 'mysql-backfill-api-key', + 'mysql-backfill-provider-key', + 'gpt-portable', + 'completed', + 0, + 120, + 30, + 20, + 'openai', + 1.25, + 1714979289, + 1714979289, + 1714979289 + ), + ( + 'mysql-backfill-pending', + 'mysql-backfill-api-key', + 'mysql-backfill-provider-key', + 'gpt-portable', + 'pending', + 777, + 700, + 77, + 0, + 'openai', + 0.25, + 1714979349, + 1714979349, + 1714979349 ); +INSERT INTO usage_settlement_snapshots ( + request_id, + billing_effective_input_tokens, + billing_output_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens +) VALUES ('mysql-backfill-completed', 100, 30, 10, 20); +"#, + ) + .execute(&mut *conn) + .await + .expect("mysql temporary backfill schema should initialize"); + drop(conn); + + let pending_versions = pending_mysql_backfills(&pool) + .await + .expect("mysql pending backfills should load") + .into_iter() + .map(|item| item.version) + .collect::>(); assert_eq!( - pending_mysql_backfills(&pool) - .await - .expect("mysql pending backfills should load"), - Vec::new() + pending_versions, + vec![ + 20260422120000, + 20260505120000, + 20260517012000, + 20260716010000 + ] ); + run_mysql_backfills(&pool) .await - .expect("mysql backfills should no-op"); + .expect("mysql backfills should apply"); + assert!(pending_mysql_backfills(&pool) + .await + .expect("mysql pending backfills should reload") + .is_empty()); + + let api_key_stats: (i64, i64, f64, Option) = query_as( + "SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = 'mysql-backfill-api-key'", + ) + .fetch_one(&pool) + .await + .expect("mysql api key backfill result should load"); + assert_eq!(api_key_stats, (2, 160, 1.5, Some(1714979349))); + let provider_total_tokens: i64 = query_scalar( + "SELECT total_tokens FROM provider_api_keys WHERE id = 'mysql-backfill-provider-key'", + ) + .fetch_one(&pool) + .await + .expect("mysql provider key total should load"); + assert_eq!(provider_total_tokens, 160); + let global_usage_count: i64 = + query_scalar("SELECT usage_count FROM global_models WHERE id = 'mysql-backfill-model'") + .fetch_one(&pool) + .await + .expect("mysql global model count should load"); + assert_eq!(global_usage_count, 1); + for table in ["providers", "provider_endpoints", "models"] { + let enabled: bool = query_scalar(&format!( + "SELECT enabled FROM {table} WHERE is_active = FALSE" + )) + .fetch_one(&pool) + .await + .unwrap_or_else(|error| panic!("mysql {table} legacy flag should load: {error}")); + assert!(!enabled, "mysql {table}.enabled should follow is_active"); + } } #[tokio::test] -async fn sqlite_backfills_are_empty_until_driver_specific_backfills_exist() { - let config = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - "sqlite::memory:", - SqlPoolConfig::default(), +async fn sqlite_backfills_apply_portable_repairs_and_record_versions() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite backfill test pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite schema should migrate"); + + query( + r#" +INSERT INTO api_keys ( + id, user_id, key_hash, total_requests, total_tokens, total_cost_usd, created_at, updated_at +) VALUES ( + 'sqlite-backfill-api-key', 'sqlite-backfill-user', 'sqlite-backfill-hash', + 77, 7777, 77.0, 1, 1 +) +"#, ) - .expect("sqlite config should build"); - let pool = crate::driver::sqlite::SqlitePoolFactory::new(config) - .expect("sqlite factory should build") - .connect_lazy() - .expect("sqlite pool should build"); + .execute(&pool) + .await + .expect("sqlite api key fixture should insert"); + query( + r#" +INSERT INTO provider_api_keys ( + id, provider_id, name, total_tokens, created_at, updated_at +) VALUES ( + 'sqlite-backfill-provider-key', 'sqlite-backfill-provider', 'Portable key', 7777, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite provider key fixture should insert"); + query( + r#" +INSERT INTO global_models ( + id, name, display_name, usage_count, created_at, updated_at +) VALUES ( + 'sqlite-backfill-model', 'gpt-portable', 'GPT Portable', 77, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite global model fixture should insert"); + query( + r#" +INSERT INTO providers ( + id, name, provider_type, enabled, is_active, created_at, updated_at +) VALUES ( + 'sqlite-backfill-provider', 'SQLite Backfill Provider', 'openai', 1, 0, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite provider flag fixture should insert"); + query( + r#" +INSERT INTO provider_endpoints ( + id, provider_id, name, base_url, enabled, is_active, created_at, updated_at +) VALUES ( + 'sqlite-backfill-endpoint', 'sqlite-backfill-provider', 'Default', + 'https://example.invalid', 1, 0, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite provider endpoint flag fixture should insert"); + query( + r#" +INSERT INTO models ( + id, provider_id, provider_model_name, enabled, is_active, created_at, updated_at +) VALUES ( + 'sqlite-backfill-provider-model', 'sqlite-backfill-provider', 'gpt-portable', + 1, 0, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite model flag fixture should insert"); + query( + r#" +INSERT INTO "usage" ( + request_id, + api_key_id, + provider_api_key_id, + model, + status, + total_tokens, + input_tokens, + output_tokens, + cache_read_input_tokens, + api_format, + total_cost_usd, + created_at, + created_at_unix_ms, + updated_at_unix_secs +) VALUES + ( + 'sqlite-backfill-completed', + 'sqlite-backfill-api-key', + 'sqlite-backfill-provider-key', + 'gpt-portable', + 'completed', + 0, + 120, + 30, + 20, + 'openai', + 1.25, + 1714979289, + 1714979289, + 1714979289 + ), + ( + 'sqlite-backfill-pending', + 'sqlite-backfill-api-key', + 'sqlite-backfill-provider-key', + 'gpt-portable', + 'pending', + 777, + 700, + 77, + 0, + 'openai', + 0.25, + 1714979349, + 1714979349, + 1714979349 + ) +"#, + ) + .execute(&pool) + .await + .expect("sqlite usage fixtures should insert"); + query( + r#" +INSERT INTO usage_settlement_snapshots ( + request_id, + billing_status, + billing_effective_input_tokens, + billing_output_tokens, + billing_cache_creation_tokens, + billing_cache_read_tokens, + created_at, + updated_at +) VALUES ( + 'sqlite-backfill-completed', 'settled', 100, 30, 10, 20, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite settlement fixture should insert"); + + let pending_versions = pending_sqlite_backfills(&pool) + .await + .expect("sqlite pending backfills should load") + .into_iter() + .map(|item| item.version) + .collect::>(); assert_eq!( - pending_sqlite_backfills(&pool) - .await - .expect("sqlite pending backfills should load"), - Vec::new() + pending_versions, + vec![ + 20260422120000, + 20260505120000, + 20260517012000, + 20260716010000 + ] ); + run_sqlite_backfills(&pool) .await - .expect("sqlite backfills should no-op"); + .expect("sqlite backfills should apply"); + assert!(pending_sqlite_backfills(&pool) + .await + .expect("sqlite pending backfills should reload") + .is_empty()); + + let applied_versions: Vec = + query_scalar("SELECT version FROM schema_backfills ORDER BY version") + .fetch_all(&pool) + .await + .expect("sqlite applied backfill versions should load"); + assert_eq!( + applied_versions, + vec![ + 20260422120000, + 20260505120000, + 20260517012000, + 20260716010000 + ] + ); + let api_key_stats: (i64, i64, f64, Option) = query_as( + "SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = 'sqlite-backfill-api-key'", + ) + .fetch_one(&pool) + .await + .expect("sqlite api key backfill result should load"); + assert_eq!(api_key_stats, (2, 160, 1.5, Some(1714979349))); + let provider_total_tokens: i64 = query_scalar( + "SELECT total_tokens FROM provider_api_keys WHERE id = 'sqlite-backfill-provider-key'", + ) + .fetch_one(&pool) + .await + .expect("sqlite provider key total should load"); + assert_eq!(provider_total_tokens, 160); + let global_usage_count: i64 = + query_scalar("SELECT usage_count FROM global_models WHERE id = 'sqlite-backfill-model'") + .fetch_one(&pool) + .await + .expect("sqlite global model count should load"); + assert_eq!(global_usage_count, 1); + for table in ["providers", "provider_endpoints", "models"] { + let enabled: i64 = + query_scalar(&format!("SELECT enabled FROM {table} WHERE is_active = 0")) + .fetch_one(&pool) + .await + .unwrap_or_else(|error| panic!("sqlite {table} legacy flag should load: {error}")); + assert_eq!(enabled, 0, "sqlite {table}.enabled should follow is_active"); + } + + run_sqlite_backfills(&pool) + .await + .expect("sqlite backfills should be idempotent"); + let applied_count: i64 = query_scalar("SELECT COUNT(*) FROM schema_backfills") + .fetch_one(&pool) + .await + .expect("sqlite applied backfill count should load"); + assert_eq!(applied_count, 4); + + query("UPDATE schema_backfills SET checksum = X'00' WHERE version = 20260422120000") + .execute(&pool) + .await + .expect("sqlite checksum compatibility fixture should update"); + assert!(pending_sqlite_backfills(&pool) + .await + .expect("checksum drift should retain the postgres compatibility policy") + .is_empty()); + + query( + r#" +INSERT INTO schema_backfills ( + version, description, success, checksum, execution_time +) VALUES ( + 99999999999999, 'missing embedded backfill', 1, X'', 0 +) +"#, + ) + .execute(&pool) + .await + .expect("unknown sqlite backfill fixture should insert"); + let error = pending_sqlite_backfills(&pool) + .await + .expect_err("unknown applied sqlite backfill should fail validation"); + assert!(matches!( + error, + sqlx::migrate::MigrateError::VersionMissing(99999999999999) + )); +} + +#[tokio::test] +async fn sqlite_backfill_sql_and_version_record_commit_atomically() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite backfill transaction test pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite schema should migrate"); + query( + r#" +INSERT INTO global_models ( + id, name, display_name, usage_count, created_at, updated_at +) VALUES ( + 'sqlite-backfill-rollback-model', 'rollback-model', 'Rollback Model', 77, 1, 1 +) +"#, + ) + .execute(&pool) + .await + .expect("sqlite rollback global model fixture should insert"); + query( + r#" +CREATE TRIGGER reject_global_model_backfill +BEFORE UPDATE OF usage_count ON global_models +BEGIN + SELECT RAISE(ABORT, 'forced global model backfill failure'); +END +"#, + ) + .execute(&pool) + .await + .expect("sqlite rollback trigger should create"); + + run_sqlite_backfills(&pool) + .await + .expect_err("forced sqlite backfill failure should propagate"); + let applied_versions: Vec = + query_scalar("SELECT version FROM schema_backfills ORDER BY version") + .fetch_all(&pool) + .await + .expect("sqlite partial applied versions should load"); + assert_eq!(applied_versions, vec![20260422120000]); + let usage_count: i64 = query_scalar( + "SELECT usage_count FROM global_models WHERE id = 'sqlite-backfill-rollback-model'", + ) + .fetch_one(&pool) + .await + .expect("sqlite rolled back global model should load"); + assert_eq!(usage_count, 77); + + query("DROP TRIGGER reject_global_model_backfill") + .execute(&pool) + .await + .expect("sqlite rollback trigger should drop"); + run_sqlite_backfills(&pool) + .await + .expect("sqlite backfills should resume after the failed transaction"); + let applied_count: i64 = query_scalar("SELECT COUNT(*) FROM schema_backfills") + .fetch_one(&pool) + .await + .expect("sqlite resumed applied backfill count should load"); + assert_eq!(applied_count, 4); } #[derive(Debug)] diff --git a/crates/aether-data/runtime/src/lifecycle/export.rs b/crates/aether-data/runtime/src/lifecycle/export.rs index 36c741d92..dd7d07522 100644 --- a/crates/aether-data/runtime/src/lifecycle/export.rs +++ b/crates/aether-data/runtime/src/lifecycle/export.rs @@ -3,6 +3,8 @@ use std::collections::{BTreeMap, BTreeSet}; #[cfg(all(feature = "postgres", feature = "sqlite"))] use futures_util::TryStreamExt; use serde_json::Value; +#[cfg(all(feature = "postgres", feature = "sqlite"))] +use sqlx::Acquire; use sqlx::Row; #[cfg(any(feature = "mysql", feature = "sqlite"))] use sqlx::{Column, TypeInfo, ValueRef}; @@ -41,7 +43,8 @@ use postgres::{ #[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] use postgres::normalize_postgres_import_payload; -pub const EXPORT_FORMAT_VERSION: u32 = 1; +pub const EXPORT_FORMAT_VERSION: u32 = 2; +const MIN_SUPPORTED_EXPORT_FORMAT_VERSION: u32 = 1; #[derive( Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize, @@ -65,6 +68,7 @@ pub enum ExportDomain { Wallets, Usage, Billing, + Auxiliary, } impl ExportDomain { @@ -87,10 +91,296 @@ impl ExportDomain { Self::Wallets => "wallets", Self::Usage => "usage", Self::Billing => "billing", + Self::Auxiliary => "auxiliary", } } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct AuxiliaryTable { + name: &'static str, + primary_key: &'static [&'static str], +} + +const AUXILIARY_TABLES: &[AuxiliaryTable] = &[ + AuxiliaryTable { + name: "audit_logs", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "announcements", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "announcement_reads", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "management_tokens", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "user_preferences", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "user_sessions", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "ldap_configs", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "pool_member_scores", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "api_key_provider_mappings", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "provider_usage_tracking", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "gemini_file_mappings", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "routing_groups", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "routing_group_versions", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "routing_group_bindings", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "proxy_node_events", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "proxy_node_metrics_1m", + primary_key: &["node_id", "bucket_start_unix_secs"], + }, + AuxiliaryTable { + name: "proxy_node_metrics_1h", + primary_key: &["node_id", "bucket_start_unix_secs"], + }, + AuxiliaryTable { + name: "user_invite_codes", + primary_key: &["user_id"], + }, + AuxiliaryTable { + name: "user_referrals", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "referral_rewards", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "payment_gateway_configs", + primary_key: &["provider"], + }, + AuxiliaryTable { + name: "billing_plans", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "user_plan_entitlements", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "entitlement_usage_ledgers", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "request_candidates", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "video_tasks", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "usage_body_blobs", + primary_key: &["body_ref"], + }, + AuxiliaryTable { + name: "usage_http_audits", + primary_key: &["request_id"], + }, + AuxiliaryTable { + name: "usage_routing_snapshots", + primary_key: &["request_id"], + }, + AuxiliaryTable { + name: "usage_counter_deltas", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "background_task_runs", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "background_task_events", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_hourly", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_summary", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_hourly_user", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_hourly_user_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "user_model_usage_counts", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_hourly_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_hourly_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_api_key", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_error", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_summary", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_api_format", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_model_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_model_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_cost_savings", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_cost_savings_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_cost_savings_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_daily_cost_savings_model_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_cost_savings", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_cost_savings_provider", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_cost_savings_model", + primary_key: &["id"], + }, + AuxiliaryTable { + name: "stats_user_daily_cost_savings_model_provider", + primary_key: &["id"], + }, +]; + +fn auxiliary_table(table_name: &str) -> Result { + AUXILIARY_TABLES + .iter() + .copied() + .find(|table| table.name == table_name) + .ok_or_else(|| { + DataLayerError::InvalidInput(format!( + "unsupported auxiliary export table '{table_name}'" + )) + }) +} + +fn auxiliary_row_id(table: AuxiliaryTable, payload: &Value) -> Result { + let object = payload.as_object().ok_or_else(|| { + DataLayerError::UnexpectedValue(format!( + "auxiliary export row in table '{}' is not a JSON object", + table.name + )) + })?; + let key = table + .primary_key + .iter() + .map(|column| { + object + .get(*column) + .filter(|value| !value.is_null()) + .cloned() + .ok_or_else(|| { + DataLayerError::UnexpectedValue(format!( + "auxiliary export row in table '{}' has null or missing primary key column '{}'", + table.name, column + )) + }) + }) + .collect::, _>>()?; + let encoded = serde_json::to_string(&key) + .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; + Ok(format!("{}:{encoded}", table.name)) +} + #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct DataExportManifest { pub format_version: u32, @@ -229,11 +519,167 @@ const USAGE_REQUEST_BODY_DETAIL_COLUMNS: &[&str] = &[ "client_response_body_compressed", ]; +const USAGE_HTTP_BODY_DETAIL_COLUMNS: &[&str] = &[ + "request_body_ref", + "provider_request_body_ref", + "response_body_ref", + "client_response_body_ref", + "request_body_state", + "provider_request_body_state", + "response_body_state", + "client_response_body_state", + "body_capture_mode", +]; + #[cfg(all(feature = "postgres", feature = "sqlite"))] -const REQUEST_BODY_DETAIL_TABLES: &[&str] = &["usage_body_blobs", "usage_http_audits"]; +const REQUEST_BODY_DETAIL_TABLES: &[&str] = &["usage_body_blobs"]; #[cfg(all(feature = "postgres", feature = "sqlite"))] const LIFECYCLE_TABLES: &[&str] = &["_sqlx_migrations", "schema_backfills"]; +fn import_column_stores_timestamp(column_name: &str) -> bool { + column_name.ends_with("_at") + || column_name.ends_with("_unix_secs") + || column_name.ends_with("_unix_ms") + || column_name.ends_with("_date") + || matches!( + column_name, + "start_time" | "end_time" | "window_start" | "window_end" | "hour_utc" | "date" + ) +} + +fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool { + if !column_name.ends_with("_unix_ms") { + return false; + } + + // This legacy field is named `_unix_ms`, but every repository and API path + // has always stored and consumed it as Unix seconds. + let relation_name = table_name + .rsplit('.') + .next() + .unwrap_or(table_name) + .trim_matches(['"', '`']); + !(relation_name == "usage" && column_name == "created_at_unix_ms") +} + +fn normalize_imported_integer_timestamp( + driver_name: &str, + table_name: &str, + column_name: &str, + value: &Value, +) -> Result, DataLayerError> { + let invalid = || { + DataLayerError::InvalidInput(format!( + "{driver_name} import timestamp column '{column_name}' must contain an integer or supported datetime" + )) + }; + + let timestamp = match value { + Value::Null => return Ok(None), + Value::Number(value) => value + .as_i64() + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + .ok_or_else(invalid)?, + Value::String(value) => { + if let Ok(timestamp) = value.trim().parse::() { + timestamp + } else { + let datetime = parse_imported_datetime(value).ok_or_else(invalid)?; + if import_timestamp_uses_millis(table_name, column_name) { + datetime.timestamp_millis() + } else { + datetime.timestamp() + } + } + } + Value::Bool(_) | Value::Array(_) | Value::Object(_) => return Err(invalid()), + }; + Ok(Some(timestamp)) +} + +fn parse_imported_datetime(value: &str) -> Option> { + let value = value.trim(); + if let Ok(datetime) = chrono::DateTime::parse_from_rfc3339(value) { + return Some(datetime.with_timezone(&chrono::Utc)); + } + if let Ok(datetime) = chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S%.f%:z") { + return Some(datetime.with_timezone(&chrono::Utc)); + } + for format in ["%Y-%m-%d %H:%M:%S%.f", "%Y-%m-%dT%H:%M:%S%.f"] { + if let Ok(datetime) = chrono::NaiveDateTime::parse_from_str(value, format) { + return Some(datetime.and_utc()); + } + } + chrono::NaiveDate::parse_from_str(value, "%Y-%m-%d") + .ok() + .and_then(|date| date.and_hms_opt(0, 0, 0)) + .map(|datetime| datetime.and_utc()) +} + +#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +fn normalize_imported_binary( + driver_name: &str, + column_name: &str, + value: &Value, +) -> Result>, DataLayerError> { + let invalid = |detail: &str| { + DataLayerError::InvalidInput(format!( + "{driver_name} import binary column '{column_name}' {detail}" + )) + }; + match value { + Value::Null => Ok(None), + Value::Array(values) => values + .iter() + .map(|value| { + value + .as_u64() + .and_then(|value| u8::try_from(value).ok()) + .ok_or_else(|| invalid("contains a non-byte array value")) + }) + .collect::, _>>() + .map(Some), + Value::String(value) => { + let encoded = value + .trim() + .strip_prefix("\\x") + .ok_or_else(|| invalid("must use PostgreSQL \\x hex encoding"))?; + if !encoded.len().is_multiple_of(2) { + return Err(invalid("contains odd-length hex data")); + } + let mut bytes = Vec::with_capacity(encoded.len() / 2); + for index in (0..encoded.len()).step_by(2) { + let byte = u8::from_str_radix(&encoded[index..index + 2], 16).map_err(|err| { + invalid(&format!( + "contains invalid hex data at byte {}: {err}", + index / 2 + )) + })?; + bytes.push(byte); + } + Ok(Some(bytes)) + } + Value::Bool(_) | Value::Number(_) | Value::Object(_) => { + Err(invalid("must contain a byte array or PostgreSQL hex value")) + } + } +} + +#[cfg(feature = "postgres")] +fn postgres_bytea_json_value(column_name: &str, value: &Value) -> Result { + let Some(bytes) = normalize_imported_binary("postgres", column_name, value)? else { + return Ok(Value::Null); + }; + let mut encoded = String::with_capacity(2 + bytes.len() * 2); + encoded.push_str("\\x"); + for byte in bytes { + use std::fmt::Write as _; + write!(&mut encoded, "{byte:02x}") + .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; + } + Ok(Value::String(encoded)) +} + pub fn encode_jsonl(records: &[DataExportRecord]) -> Result { validate_export_records(records)?; @@ -300,10 +746,12 @@ pub fn validate_export_records(records: &[DataExportRecord]) -> Result<(), DataL "export JSONL must start with a manifest record".to_string(), )); }; - if manifest.format_version != EXPORT_FORMAT_VERSION { + if !(MIN_SUPPORTED_EXPORT_FORMAT_VERSION..=EXPORT_FORMAT_VERSION) + .contains(&manifest.format_version) + { return Err(DataLayerError::InvalidInput(format!( - "unsupported export format version {}; expected {}", - manifest.format_version, EXPORT_FORMAT_VERSION + "unsupported export format version {}; supported versions are {} through {}", + manifest.format_version, MIN_SUPPORTED_EXPORT_FORMAT_VERSION, EXPORT_FORMAT_VERSION ))); } @@ -369,6 +817,7 @@ pub fn sqlite_core_export_domains() -> Vec { ExportDomain::Wallets, ExportDomain::Usage, ExportDomain::Billing, + ExportDomain::Auxiliary, ] } @@ -490,23 +939,39 @@ pub async fn copy_database_records( import_database_jsonl(target, &encode_jsonl(&records)?).await } -fn omit_request_body_details_from_records(records: &mut [DataExportRecord]) { - for record in records { +fn omit_request_body_details_from_records(records: &mut Vec) { + records.retain_mut(|record| { let DataExportRecord::Row { - domain: ExportDomain::Usage, - payload, - .. + domain, payload, .. } = record else { - continue; + return true; }; - - if let Some(object) = payload.as_object_mut() { - for column_name in USAGE_REQUEST_BODY_DETAIL_COLUMNS { - object.remove(*column_name); + let Some(object) = payload.as_object_mut() else { + return true; + }; + match *domain { + ExportDomain::Usage => { + for column_name in USAGE_REQUEST_BODY_DETAIL_COLUMNS { + object.remove(*column_name); + } } + ExportDomain::Auxiliary + if object.get("__table").and_then(Value::as_str) == Some("usage_body_blobs") => + { + return false; + } + ExportDomain::Auxiliary + if object.get("__table").and_then(Value::as_str) == Some("usage_http_audits") => + { + for column_name in USAGE_HTTP_BODY_DETAIL_COLUMNS { + object.remove(*column_name); + } + } + _ => {} } - } + true + }); } #[cfg(all(feature = "postgres", feature = "sqlite"))] @@ -522,24 +987,24 @@ async fn copy_postgres_to_sqlite_from_target_schema( crate::driver::postgres::PostgresPoolFactory::new(source.to_postgres_config()?)? .connect_lazy()?; let sqlite_pool = crate::driver::sqlite::SqlitePoolFactory::new(target)?.connect_lazy()?; + let mut postgres_tx = postgres_pool.begin().await.map_sql_err()?; + sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY") + .execute(&mut *postgres_tx) + .await + .map_sql_err()?; - let source_tables = load_postgres_public_table_names(&postgres_pool).await?; + let source_tables = load_postgres_public_table_names(&mut postgres_tx).await?; let target_tables = load_sqlite_copy_table_names(&sqlite_pool).await?; ensure_no_nonempty_source_tables_outside_target_schema( - &postgres_pool, + &mut postgres_tx, &source_tables, &target_tables, options, ) .await?; - let mut imported = 0usize; - sqlx::raw_sql("PRAGMA foreign_keys = OFF") - .execute(&sqlite_pool) - .await - .map_sql_err()?; - + let mut table_plans = Vec::new(); for table_name in target_tables { if copy_table_is_lifecycle(&table_name) || copy_table_is_sqlite_internal(&table_name) @@ -550,7 +1015,7 @@ async fn copy_postgres_to_sqlite_from_target_schema( } let table_plan = build_postgres_sqlite_copy_table_plan( - &postgres_pool, + &mut postgres_tx, &sqlite_pool, &table_name, options, @@ -559,22 +1024,39 @@ async fn copy_postgres_to_sqlite_from_target_schema( if table_plan.columns.is_empty() { continue; } - imported = imported.saturating_add( - copy_postgres_sqlite_table(&postgres_pool, &sqlite_pool, &table_plan).await?, - ); + table_plans.push(table_plan); } - sqlx::raw_sql("PRAGMA foreign_keys = ON") - .execute(&sqlite_pool) + let mut connection = sqlite_pool.acquire().await.map_sql_err()?; + sqlx::raw_sql("PRAGMA foreign_keys = OFF") + .execute(&mut *connection) .await .map_sql_err()?; - ensure_sqlite_foreign_key_check_passes(&sqlite_pool).await?; + let copy_result = async { + let mut tx = connection.begin().await.map_sql_err()?; + let mut imported = 0usize; + for table_plan in &table_plans { + imported = imported.saturating_add( + copy_postgres_sqlite_table(&mut postgres_tx, &mut tx, table_plan).await?, + ); + } + ensure_sqlite_foreign_key_check_passes(&mut tx).await?; + tx.commit().await.map_sql_err()?; + Ok::<_, DataLayerError>(imported) + } + .await; + sqlx::raw_sql("PRAGMA foreign_keys = ON") + .execute(&mut *connection) + .await + .map_sql_err()?; + let imported = copy_result?; + postgres_tx.commit().await.map_sql_err()?; Ok(imported) } #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn ensure_no_nonempty_source_tables_outside_target_schema( - postgres_pool: &crate::driver::postgres::PostgresPool, + postgres_tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, source_tables: &BTreeSet, target_tables: &BTreeSet, options: DataCopyOptions, @@ -587,7 +1069,7 @@ async fn ensure_no_nonempty_source_tables_outside_target_schema( { continue; } - if postgres_public_table_has_rows(postgres_pool, table_name).await? { + if postgres_public_table_has_rows(postgres_tx, table_name).await? { missing.push(table_name.clone()); } } @@ -603,15 +1085,15 @@ async fn ensure_no_nonempty_source_tables_outside_target_schema( #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn build_postgres_sqlite_copy_table_plan( - postgres_pool: &crate::driver::postgres::PostgresPool, + postgres_tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, sqlite_pool: &crate::driver::sqlite::SqlitePool, table_name: &str, options: DataCopyOptions, ) -> Result { let sqlite_columns = load_sqlite_copy_columns(sqlite_pool, table_name).await?; let postgres_columns = - load_postgres_import_columns(postgres_pool, &format!("public.{table_name}")).await?; - let source_has_rows = postgres_public_table_has_rows(postgres_pool, table_name).await?; + load_postgres_import_columns(&mut **postgres_tx, &format!("public.{table_name}")).await?; + let source_has_rows = postgres_public_table_has_rows(postgres_tx, table_name).await?; let mut columns = Vec::new(); for sqlite_column in sqlite_columns { @@ -621,6 +1103,12 @@ async fn build_postgres_sqlite_copy_table_plan( { continue; } + if options.omit_request_body_details + && table_name == "usage_http_audits" + && USAGE_HTTP_BODY_DETAIL_COLUMNS.contains(&sqlite_column.name.as_str()) + { + continue; + } if let Some(postgres_column) = postgres_columns.get(&sqlite_column.name) { columns.push(SchemaCopyColumn { @@ -652,13 +1140,13 @@ async fn build_postgres_sqlite_copy_table_plan( #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn copy_postgres_sqlite_table( - postgres_pool: &crate::driver::postgres::PostgresPool, - sqlite_pool: &crate::driver::sqlite::SqlitePool, + postgres_tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + sqlite_tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, table: &SchemaCopyTable, ) -> Result { let source_sql = postgres_schema_copy_select_sql(table)?; let target_sql = sqlite_schema_copy_insert_sql(table)?; - let mut rows = sqlx::query(&source_sql).fetch(postgres_pool); + let mut rows = sqlx::query(&source_sql).fetch(&mut **postgres_tx); let mut imported = 0usize; while let Some(row) = rows.try_next().await.map_sql_err()? { @@ -679,7 +1167,7 @@ async fn copy_postgres_sqlite_table( })?; query = bind_sqlite_copy_value(query, value, &column.sqlite)?; } - query.execute(sqlite_pool).await.map_sql_err()?; + query.execute(&mut **sqlite_tx).await.map_sql_err()?; imported = imported.saturating_add(1); } @@ -694,7 +1182,7 @@ fn postgres_schema_copy_select_sql(table: &SchemaCopyTable) -> Result Result Result, DataLayerError> { let column_sql = format!("t.{}", postgres_quote_identifier(&column.sqlite.name)?); @@ -755,7 +1244,7 @@ fn postgres_schema_copy_override_expr( } else { column_sql.clone() }; - let multiplier = if sqlite_copy_column_stores_unix_millis(&column.sqlite.name) { + let multiplier = if import_timestamp_uses_millis(table_name, &column.sqlite.name) { " * 1000" } else { "" @@ -778,14 +1267,46 @@ fn sqlite_schema_copy_insert_sql(table: &SchemaCopyTable) -> Result, _>>()? .join(", "); let placeholder_sql = vec!["?"; table.columns.len()].join(", "); + let mut primary_key = table + .columns + .iter() + .filter(|column| column.sqlite.primary_key_position > 0) + .collect::>(); + primary_key.sort_by_key(|column| column.sqlite.primary_key_position); + if primary_key.is_empty() { + return Ok(format!( + "INSERT INTO {table_sql} ({column_sql}) VALUES ({placeholder_sql})" + )); + } + + let conflict_columns = primary_key + .iter() + .map(|column| sqlite_quote_identifier(&column.sqlite.name)) + .collect::, _>>()? + .join(", "); + let update_sql = table + .columns + .iter() + .filter(|column| column.sqlite.primary_key_position == 0) + .map(|column| { + let quoted = sqlite_quote_identifier(&column.sqlite.name)?; + Ok(format!("{quoted} = excluded.{quoted}")) + }) + .collect::, DataLayerError>>()? + .join(", "); + let conflict_sql = if update_sql.is_empty() { + format!("ON CONFLICT ({conflict_columns}) DO NOTHING") + } else { + format!("ON CONFLICT ({conflict_columns}) DO UPDATE SET {update_sql}") + }; Ok(format!( - "INSERT OR REPLACE INTO {table_sql} ({column_sql}) VALUES ({placeholder_sql})" + "INSERT INTO {table_sql} ({column_sql}) VALUES ({placeholder_sql}) {conflict_sql}" )) } #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn load_postgres_public_table_names( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, ) -> Result, DataLayerError> { let rows = sqlx::query( r#" @@ -796,7 +1317,7 @@ WHERE table_schema = 'public' ORDER BY table_name "#, ) - .fetch_all(pool) + .fetch_all(&mut **tx) .await .map_sql_err()?; @@ -872,24 +1393,24 @@ async fn load_sqlite_copy_columns( #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn postgres_public_table_has_rows( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, table_name: &str, ) -> Result { let table_sql = format!("public.{}", postgres_quote_identifier(table_name)?); sqlx::query_scalar::<_, bool>(&format!( "SELECT EXISTS (SELECT 1 FROM {table_sql} LIMIT 1)" )) - .fetch_one(pool) + .fetch_one(&mut **tx) .await .map_sql_err() } #[cfg(all(feature = "postgres", feature = "sqlite"))] async fn ensure_sqlite_foreign_key_check_passes( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, ) -> Result<(), DataLayerError> { let rows = sqlx::query("PRAGMA foreign_key_check") - .fetch_all(pool) + .fetch_all(&mut **tx) .await .map_sql_err()?; if rows.is_empty() { @@ -935,11 +1456,6 @@ fn sqlite_copy_column_is_required(column: &SqliteCopyColumn) -> bool { (column.not_null || column.primary_key_position > 0) && !column.has_default } -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn sqlite_copy_column_stores_unix_millis(column_name: &str) -> bool { - column_name.ends_with("_unix_ms") -} - #[cfg(all(feature = "postgres", feature = "sqlite"))] fn sqlite_copy_affinity(column: &SqliteCopyColumn) -> SqliteCopyAffinity { let declared_type = column.declared_type.to_ascii_uppercase(); @@ -962,7 +1478,7 @@ fn sqlite_copy_affinity(column: &SqliteCopyColumn) -> SqliteCopyAffinity { } } -#[cfg(all(feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn is_postgres_bytea_column(column: &PostgresImportColumn) -> bool { column.data_type == "bytea" || column.udt_name == "bytea" } @@ -1194,7 +1710,17 @@ fn filter_import_payload( for (column_name, value) in object { if target_columns.contains(column_name) { filtered.insert(column_name.clone(), value.clone()); + continue; } + if value.is_null() { + continue; + } + return Err(DataLayerError::InvalidInput(format!( + "{} export row '{}' contains column '{}' that does not exist in {driver_name} table '{table_name}'", + domain.as_str(), + row.id, + column_name + ))); } if filtered.is_empty() { @@ -1208,7 +1734,6 @@ fn filter_import_payload( Ok(filtered) } -#[cfg(any(feature = "mysql", feature = "sqlite"))] fn payload_with_table(payload: Value, table_name: &str) -> Result { let mut object = payload.as_object().cloned().ok_or_else(|| { DataLayerError::UnexpectedValue("export row payload must be a JSON object".to_string()) @@ -1218,7 +1743,6 @@ fn payload_with_table(payload: Value, table_name: &str) -> Result, diff --git a/crates/aether-data/runtime/src/lifecycle/export/mysql.rs b/crates/aether-data/runtime/src/lifecycle/export/mysql.rs index e52522d5c..b92215126 100644 --- a/crates/aether-data/runtime/src/lifecycle/export/mysql.rs +++ b/crates/aether-data/runtime/src/lifecycle/export/mysql.rs @@ -1,4 +1,12 @@ use super::*; +use sqlx::Acquire; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct MysqlImportColumns { + names: ImportColumnNames, + data_types: BTreeMap, + primary_key: Vec, +} pub async fn export_mysql_core_jsonl( pool: &crate::driver::mysql::MysqlPool, @@ -12,6 +20,12 @@ pub async fn export_mysql_jsonl( domains: Vec, created_at_unix_secs: u64, ) -> Result { + let mut connection = pool.acquire().await.map_sql_err()?; + sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ") + .execute(&mut *connection) + .await + .map_sql_err()?; + let mut tx = connection.begin().await.map_sql_err()?; let manifest = DataExportManifest::new( created_at_unix_secs, Some(DatabaseDriver::Mysql), @@ -20,24 +34,29 @@ pub async fn export_mysql_jsonl( let mut records = vec![DataExportRecord::manifest(manifest)]; for domain in domains { + if domain == ExportDomain::Auxiliary { + export_mysql_auxiliary_records(&mut tx, &mut records).await?; + continue; + } if domain == ExportDomain::Billing { - export_mysql_billing_records(pool, &mut records).await?; + export_mysql_billing_records(&mut tx, &mut records).await?; continue; } if domain == ExportDomain::Wallets { - export_mysql_wallet_records(pool, &mut records).await?; + export_mysql_wallet_records(&mut tx, &mut records).await?; continue; } let (table_name, id_column) = mysql_domain_table(domain)?; let order_by = export_order_by(domain, id_column); let sql = format!("SELECT * FROM {table_name} ORDER BY {order_by}"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut *tx).await.map_sql_err()?; for row in rows { let id = mysql_export_row_id(domain, &row, id_column)?; records.push(DataExportRecord::row(domain, id, mysql_row_payload(&row)?)); } } + tx.commit().await.map_sql_err()?; encode_jsonl(&records) } @@ -53,31 +72,40 @@ pub async fn import_mysql_plan( pool: &crate::driver::mysql::MysqlPool, plan: &DataImportPlan, ) -> Result { + let mut tx = pool.begin().await.map_sql_err()?; let mut imported = 0usize; - let mut column_cache = BTreeMap::::new(); + let mut column_cache = BTreeMap::::new(); for domain in &plan.manifest.domains { + if *domain == ExportDomain::Auxiliary { + for row in plan.rows(*domain) { + import_mysql_auxiliary_row(&mut tx, row, &mut column_cache).await?; + imported = imported.saturating_add(1); + } + continue; + } if *domain == ExportDomain::Billing { for row in plan.rows(*domain) { - import_mysql_billing_row(pool, row, &mut column_cache).await?; + import_mysql_billing_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; } if *domain == ExportDomain::Wallets { for row in plan.rows(*domain) { - import_mysql_wallet_row(pool, row, &mut column_cache).await?; + import_mysql_wallet_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; } let (table_name, _id_column) = mysql_domain_table(*domain)?; let target_columns = - mysql_import_columns_cached(pool, &mut column_cache, table_name).await?; + mysql_import_columns_cached(&mut tx, &mut column_cache, table_name).await?; for row in plan.rows(*domain) { - import_mysql_row(pool, table_name, *domain, row, &target_columns).await?; + import_mysql_row(&mut tx, table_name, *domain, row, &target_columns).await?; imported = imported.saturating_add(1); } } + tx.commit().await.map_sql_err()?; Ok(imported) } @@ -106,9 +134,42 @@ fn mysql_domain_table( ExportDomain::Billing => Err(DataLayerError::InvalidInput( "mysql billing export uses multiple tables and must be handled as a domain".to_string(), )), + ExportDomain::Auxiliary => Err(DataLayerError::InvalidInput( + "mysql auxiliary export uses multiple tables and must be handled as a domain" + .to_string(), + )), } } +async fn export_mysql_auxiliary_records( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + records: &mut Vec, +) -> Result<(), DataLayerError> { + for table in AUXILIARY_TABLES { + let table_sql = mysql_quote_identifier(table.name)?; + let order_sql = table + .primary_key + .iter() + .map(|column| mysql_quote_identifier(column).map(|column| format!("{column} ASC"))) + .collect::, _>>()? + .join(", "); + let rows = sqlx::query(&format!("SELECT * FROM {table_sql} ORDER BY {order_sql}")) + .fetch_all(&mut **tx) + .await + .map_sql_err()?; + for row in rows { + let payload = mysql_row_payload(&row)?; + let id = auxiliary_row_id(*table, &payload)?; + records.push(DataExportRecord::row( + ExportDomain::Auxiliary, + id, + payload_with_table(payload, table.name)?, + )); + } + } + Ok(()) +} + fn mysql_export_row_id( domain: ExportDomain, row: &sqlx::mysql::MySqlRow, @@ -139,7 +200,7 @@ fn mysql_required_export_text( } async fn export_mysql_billing_records( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, id_column) in [ @@ -148,7 +209,7 @@ async fn export_mysql_billing_records( ("usage_settlement_snapshots", "request_id"), ] { let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) @@ -169,12 +230,12 @@ async fn export_mysql_billing_records( } async fn export_mysql_wallet_records( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, id_column) in mysql_wallet_tables() { let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) @@ -195,53 +256,115 @@ async fn export_mysql_wallet_records( } async fn import_mysql_row( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, table_name: &str, domain: ExportDomain, row: &ExportRow, - target_columns: &ImportColumnNames, + target_columns: &MysqlImportColumns, ) -> Result<(), DataLayerError> { - let object = filter_import_payload("mysql", table_name, domain, row, target_columns)?; + let object = filter_import_payload("mysql", table_name, domain, row, &target_columns.names)?; let columns = object.keys().map(String::as_str).collect::>(); + for primary_key in &target_columns.primary_key { + if object.get(primary_key).is_none_or(Value::is_null) { + return Err(DataLayerError::InvalidInput(format!( + "{} export row '{}' is missing non-null primary key column '{}' for mysql table '{}'", + domain.as_str(), + row.id, + primary_key, + table_name + ))); + } + } + + let primary_key_predicate = target_columns + .primary_key + .iter() + .map(|column| mysql_quote_identifier(column).map(|column| format!("{column} = ?"))) + .collect::, _>>()? + .join(" AND "); + let lock_sql = + format!("SELECT 1 FROM {table_name} WHERE {primary_key_predicate} LIMIT 1 FOR UPDATE"); + let mut lock_query = sqlx::query(&lock_sql); + for column in &target_columns.primary_key { + lock_query = + bind_mysql_import_column(lock_query, &object, target_columns, table_name, column)?; + } + let exists = lock_query + .fetch_optional(&mut **tx) + .await + .map_sql_err()? + .is_some(); + + if exists { + let update_columns = columns + .iter() + .copied() + .filter(|column| !target_columns.primary_key.iter().any(|key| key == column)) + .collect::>(); + if update_columns.is_empty() { + return Ok(()); + } + let update_sql = update_columns + .iter() + .map(|column| mysql_quote_identifier(column).map(|column| format!("{column} = ?"))) + .collect::, _>>()? + .join(", "); + let sql = format!("UPDATE {table_name} SET {update_sql} WHERE {primary_key_predicate}"); + let mut query = sqlx::query(&sql); + for column in update_columns { + query = bind_mysql_import_column(query, &object, target_columns, table_name, column)?; + } + for column in &target_columns.primary_key { + query = bind_mysql_import_column(query, &object, target_columns, table_name, column)?; + } + query.execute(&mut **tx).await.map_sql_err()?; + return Ok(()); + } + let column_sql = columns .iter() .map(|column| mysql_quote_identifier(column)) .collect::, _>>()? .join(", "); let placeholder_sql = vec!["?"; columns.len()].join(", "); - let update_sql = columns - .iter() - .map(|column| { - let quoted = mysql_quote_identifier(column)?; - Ok(format!("{quoted} = VALUES({quoted})")) - }) - .collect::, DataLayerError>>()? - .join(", "); - let sql = format!( - "INSERT INTO {table_name} ({column_sql}) VALUES ({placeholder_sql}) ON DUPLICATE KEY UPDATE {update_sql}" - ); + let sql = format!("INSERT INTO {table_name} ({column_sql}) VALUES ({placeholder_sql})"); let mut query = sqlx::query(&sql); for column in columns { - let value = object - .get(column) - .expect("column name came from payload object keys"); - query = bind_mysql_json_value(query, value)?; + query = bind_mysql_import_column(query, &object, target_columns, table_name, column)?; } - query.execute(pool).await.map_sql_err()?; + query.execute(&mut **tx).await.map_sql_err()?; Ok(()) } +fn bind_mysql_import_column<'q>( + query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, + object: &'q serde_json::Map, + target_columns: &MysqlImportColumns, + table_name: &str, + column: &str, +) -> Result, DataLayerError> { + let value = object + .get(column) + .expect("column name came from payload object keys"); + let data_type = target_columns + .data_types + .get(column) + .map(String::as_str) + .unwrap_or_default(); + bind_mysql_import_value(query, value, table_name, column, data_type) +} + async fn import_mysql_billing_row( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, row: &ExportRow, - column_cache: &mut BTreeMap, + column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (table_name, payload) = billing_payload_table(row)?; let table_name = mysql_billing_table_name(&table_name)?; - let target_columns = mysql_import_columns_cached(pool, column_cache, table_name).await?; + let target_columns = mysql_import_columns_cached(tx, column_cache, table_name).await?; import_mysql_row( - pool, + tx, table_name, ExportDomain::Billing, &ExportRow { @@ -253,6 +376,27 @@ async fn import_mysql_billing_row( .await } +async fn import_mysql_auxiliary_row( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + row: &ExportRow, + column_cache: &mut BTreeMap, +) -> Result<(), DataLayerError> { + let (table_name, payload) = domain_payload_table(row, "auxiliary", None)?; + let table = auxiliary_table(&table_name)?; + let target_columns = mysql_import_columns_cached(tx, column_cache, table.name).await?; + import_mysql_row( + tx, + table.name, + ExportDomain::Auxiliary, + &ExportRow { + id: row.id.clone(), + payload, + }, + &target_columns, + ) + .await +} + fn mysql_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { match table_name { "billing_rules" => Ok("billing_rules"), @@ -265,15 +409,15 @@ fn mysql_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerE } async fn import_mysql_wallet_row( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, row: &ExportRow, - column_cache: &mut BTreeMap, + column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (table_name, payload) = domain_payload_table(row, "wallet", Some("wallets"))?; let table_name = mysql_wallet_table_name(&table_name)?; - let target_columns = mysql_import_columns_cached(pool, column_cache, table_name).await?; + let target_columns = mysql_import_columns_cached(tx, column_cache, table_name).await?; import_mysql_row( - pool, + tx, table_name, ExportDomain::Wallets, &ExportRow { @@ -311,51 +455,130 @@ fn mysql_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerEr } async fn mysql_import_columns_cached( - pool: &crate::driver::mysql::MysqlPool, - cache: &mut BTreeMap, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + cache: &mut BTreeMap, table_name: &str, -) -> Result { +) -> Result { if let Some(columns) = cache.get(table_name) { return Ok(columns.clone()); } - let columns = load_mysql_import_columns(pool, table_name).await?; + let columns = load_mysql_import_columns(tx, table_name).await?; cache.insert(table_name.to_string(), columns.clone()); Ok(columns) } async fn load_mysql_import_columns( - pool: &crate::driver::mysql::MysqlPool, + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, table_name: &str, -) -> Result { +) -> Result { let relation_name = table_name.trim_matches('`'); let rows = sqlx::query( r#" -SELECT COLUMN_NAME AS column_name +SELECT + COLUMN_NAME AS column_name, + DATA_TYPE AS data_type, + COLUMN_KEY AS column_key, + ORDINAL_POSITION AS ordinal_position FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = ? "#, ) .bind(relation_name) - .fetch_all(pool) + .fetch_all(&mut **tx) .await .map_sql_err()?; - let mut columns = ImportColumnNames::new(); + let mut columns = MysqlImportColumns::default(); + let mut primary_key = BTreeMap::new(); for row in rows { - columns.insert(row.try_get::("column_name").map_sql_err()?); + let name = row.try_get::("column_name").map_sql_err()?; + let data_type = row + .try_get::("data_type") + .map_sql_err()? + .to_ascii_lowercase(); + columns.names.insert(name.clone()); + columns.data_types.insert(name.clone(), data_type); + if row + .try_get::("column_key") + .map_sql_err()? + .eq_ignore_ascii_case("PRI") + { + primary_key.insert( + row.try_get::("ordinal_position").map_sql_err()?, + name, + ); + } } - if columns.is_empty() { + if columns.names.is_empty() { return Err(DataLayerError::UnexpectedValue(format!( "mysql import target table '{table_name}' has no visible columns" ))); } + if primary_key.is_empty() { + return Err(DataLayerError::UnexpectedValue(format!( + "mysql import target table '{table_name}' has no primary key" + ))); + } + columns.primary_key = primary_key.into_values().collect(); Ok(columns) } +fn bind_mysql_import_value<'q>( + query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, + json_value: &'q Value, + table_name: &str, + column_name: &str, + data_type: &str, +) -> Result, DataLayerError> { + if matches!( + data_type, + "binary" | "varbinary" | "blob" | "tinyblob" | "mediumblob" | "longblob" + ) { + return match normalize_imported_binary("mysql", column_name, json_value)? { + Some(bytes) => Ok(query.bind(bytes)), + None => Ok(query.bind(Option::>::None)), + }; + } + if matches!(data_type, "decimal" | "numeric") { + return match normalize_mysql_decimal_value(column_name, json_value)? { + Some(value) => Ok(query.bind(value)), + None => Ok(query.bind(Option::::None)), + }; + } + let has_integer_type = matches!( + data_type, + "tinyint" | "smallint" | "mediumint" | "int" | "integer" | "bigint" + ); + if !has_integer_type || !import_column_stores_timestamp(column_name) { + return bind_mysql_json_value(query, json_value); + } + + match normalize_imported_integer_timestamp("mysql", table_name, column_name, json_value)? { + Some(timestamp) => Ok(query.bind(timestamp)), + None => Ok(query.bind(Option::::None)), + } +} + +fn normalize_mysql_decimal_value( + column_name: &str, + value: &Value, +) -> Result, DataLayerError> { + match value { + Value::Null => Ok(None), + Value::Number(value) => Ok(Some(value.to_string())), + Value::String(value) => Ok(Some(value.clone())), + Value::Bool(_) | Value::Array(_) | Value::Object(_) => { + Err(DataLayerError::InvalidInput(format!( + "mysql decimal import column '{column_name}' must contain a number or numeric string" + ))) + } + } +} + fn mysql_quote_identifier(identifier: &str) -> Result { if identifier.trim().is_empty() { return Err(DataLayerError::InvalidInput( @@ -454,3 +677,30 @@ fn mysql_value_to_json(row: &sqlx::mysql::MySqlRow, index: usize) -> Result, created_at_unix_secs: u64, ) -> Result { + let mut tx = pool.begin().await.map_sql_err()?; + sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY") + .execute(&mut *tx) + .await + .map_sql_err()?; let manifest = DataExportManifest::new( created_at_unix_secs, Some(DatabaseDriver::Postgres), @@ -20,12 +25,16 @@ pub async fn export_postgres_jsonl( let mut records = vec![DataExportRecord::manifest(manifest)]; for domain in domains { + if domain == ExportDomain::Auxiliary { + export_postgres_auxiliary_records(&mut tx, &mut records).await?; + continue; + } if domain == ExportDomain::Billing { - export_postgres_billing_records(pool, &mut records).await?; + export_postgres_billing_records(&mut tx, &mut records).await?; continue; } if domain == ExportDomain::Wallets { - export_postgres_wallet_records(pool, &mut records).await?; + export_postgres_wallet_records(&mut tx, &mut records).await?; continue; } let (table_name, id_column) = postgres_domain_table(domain)?; @@ -34,7 +43,7 @@ pub async fn export_postgres_jsonl( let sql = format!( "SELECT {export_id_sql} AS export_id, to_jsonb(t) AS payload FROM {table_name} AS t ORDER BY {order_by}" ); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut *tx).await.map_sql_err()?; for row in rows { let id = row.try_get::("export_id").map_sql_err()?; let payload = row.try_get::("payload").map_sql_err()?; @@ -42,6 +51,7 @@ pub async fn export_postgres_jsonl( } } + tx.commit().await.map_sql_err()?; encode_jsonl(&records) } @@ -57,19 +67,27 @@ pub async fn import_postgres_plan( pool: &crate::driver::postgres::PostgresPool, plan: &DataImportPlan, ) -> Result { + let mut tx = pool.begin().await.map_sql_err()?; let mut imported = 0usize; let mut column_cache = BTreeMap::::new(); for domain in &plan.manifest.domains { + if *domain == ExportDomain::Auxiliary { + for row in plan.rows(*domain) { + import_postgres_auxiliary_row(&mut tx, row, &mut column_cache).await?; + imported = imported.saturating_add(1); + } + continue; + } if *domain == ExportDomain::Billing { for row in plan.rows(*domain) { - import_postgres_billing_row(pool, row, &mut column_cache).await?; + import_postgres_billing_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; } if *domain == ExportDomain::Wallets { for row in plan.rows(*domain) { - import_postgres_wallet_row(pool, row, &mut column_cache).await?; + import_postgres_wallet_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; @@ -81,10 +99,10 @@ pub async fn import_postgres_plan( continue; } let target_columns = - postgres_import_columns_cached(pool, &mut column_cache, table_name).await?; + postgres_import_columns_cached(&mut tx, &mut column_cache, table_name).await?; for row in rows { import_postgres_row( - pool, + &mut tx, table_name, &conflict_columns, *domain, @@ -95,9 +113,52 @@ pub async fn import_postgres_plan( imported = imported.saturating_add(1); } } + if !plan.rows(ExportDomain::Auxiliary).is_empty() { + reset_postgres_auxiliary_sequences(&mut tx).await?; + } + tx.commit().await.map_sql_err()?; Ok(imported) } +async fn reset_postgres_auxiliary_sequences( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, +) -> Result<(), DataLayerError> { + for table in AUXILIARY_TABLES { + let [primary_key] = table.primary_key else { + continue; + }; + let relation_name = format!("public.{}", table.name); + let sequence = + sqlx::query_scalar::<_, Option>("SELECT pg_get_serial_sequence($1, $2)") + .bind(&relation_name) + .bind(*primary_key) + .fetch_one(&mut **tx) + .await + .map_sql_err()?; + let Some(sequence) = sequence else { + continue; + }; + + let table_sql = postgres_quote_identifier(table.name)?; + let primary_key_sql = postgres_quote_identifier(primary_key)?; + let maximum = sqlx::query_scalar::<_, Option>(&format!( + "SELECT MAX({primary_key_sql})::bigint FROM public.{table_sql}" + )) + .fetch_one(&mut **tx) + .await + .map_sql_err()?; + let (value, is_called) = maximum.map_or((1_i64, false), |value| (value, true)); + sqlx::query("SELECT setval($1::regclass, $2, $3)") + .bind(sequence) + .bind(value) + .bind(is_called) + .execute(&mut **tx) + .await + .map_sql_err()?; + } + Ok(()) +} + fn postgres_domain_table( domain: ExportDomain, ) -> Result<(&'static str, &'static str), DataLayerError> { @@ -125,9 +186,44 @@ fn postgres_domain_table( "postgres billing export uses multiple tables and must be handled as a domain" .to_string(), )), + ExportDomain::Auxiliary => Err(DataLayerError::InvalidInput( + "postgres auxiliary export uses multiple tables and must be handled as a domain" + .to_string(), + )), } } +async fn export_postgres_auxiliary_records( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + records: &mut Vec, +) -> Result<(), DataLayerError> { + for table in AUXILIARY_TABLES { + let table_sql = postgres_quote_identifier(table.name)?; + let order_sql = table + .primary_key + .iter() + .map(|column| postgres_quote_identifier(column).map(|column| format!("{column} ASC"))) + .collect::, _>>()? + .join(", "); + let rows = sqlx::query(&format!( + "SELECT to_jsonb(t) AS payload FROM public.{table_sql} AS t ORDER BY {order_sql}" + )) + .fetch_all(&mut **tx) + .await + .map_sql_err()?; + for row in rows { + let payload = row.try_get::("payload").map_sql_err()?; + let id = auxiliary_row_id(*table, &payload)?; + records.push(DataExportRecord::row( + ExportDomain::Auxiliary, + id, + payload_with_table(payload, table.name)?, + )); + } + } + Ok(()) +} + fn postgres_export_id_sql(domain: ExportDomain, id_column: &str) -> String { if domain == ExportDomain::UserGroupMembers { "group_id::text || ':' || user_id::text".to_string() @@ -145,7 +241,7 @@ fn postgres_conflict_columns(domain: ExportDomain, id_column: &str) -> Vec<&str> } async fn postgres_import_columns_cached( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, cache: &mut BTreeMap, table_name: &str, ) -> Result { @@ -153,13 +249,13 @@ async fn postgres_import_columns_cached( return Ok(columns.clone()); } - let columns = load_postgres_import_columns(pool, table_name).await?; + let columns = load_postgres_import_columns(&mut **tx, table_name).await?; cache.insert(table_name.to_string(), columns.clone()); Ok(columns) } -pub(super) async fn load_postgres_import_columns( - pool: &crate::driver::postgres::PostgresPool, +pub(super) async fn load_postgres_import_columns<'e>( + executor: impl sqlx::Executor<'e, Database = sqlx::Postgres>, table_name: &str, ) -> Result { let (schema_name, relation_name) = postgres_table_parts(table_name)?; @@ -173,7 +269,7 @@ WHERE table_schema = $1 ) .bind(schema_name) .bind(relation_name) - .fetch_all(pool) + .fetch_all(executor) .await .map_sql_err()?; @@ -223,7 +319,7 @@ fn postgres_table_parts(table_name: &str) -> Result<(&str, &str), DataLayerError } async fn export_postgres_billing_records( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, export_table, id_column) in [ @@ -238,7 +334,7 @@ async fn export_postgres_billing_records( let sql = format!( "SELECT {id_column}::text AS export_id, to_jsonb(t) || jsonb_build_object('__table', '{export_table}') AS payload FROM {table_name} AS t ORDER BY {id_column} ASC" ); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row.try_get::("export_id").map_sql_err()?; let payload = row.try_get::("payload").map_sql_err()?; @@ -253,14 +349,14 @@ async fn export_postgres_billing_records( } async fn export_postgres_wallet_records( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, export_table, id_column) in postgres_wallet_tables() { let sql = format!( "SELECT {id_column}::text AS export_id, to_jsonb(t) || jsonb_build_object('__table', '{export_table}') AS payload FROM {table_name} AS t ORDER BY {id_column} ASC" ); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row.try_get::("export_id").map_sql_err()?; let payload = row.try_get::("payload").map_sql_err()?; @@ -275,7 +371,7 @@ async fn export_postgres_wallet_records( } async fn import_postgres_row( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, table_name: &str, conflict_columns: &[&str], domain: ExportDomain, @@ -316,7 +412,7 @@ async fn import_postgres_row( sqlx::query(&sql) .bind(&payload) - .execute(pool) + .execute(&mut **tx) .await .map_sql_err()?; Ok(()) @@ -351,7 +447,7 @@ pub(super) fn normalize_postgres_import_payload( } normalized.insert( column_name.clone(), - normalize_postgres_import_value(column_name, target_column, value)?, + normalize_postgres_import_value(table_name, column_name, target_column, value)?, ); continue; } @@ -380,6 +476,7 @@ pub(super) fn normalize_postgres_import_payload( } fn normalize_postgres_import_value( + table_name: &str, column_name: &str, target_column: &PostgresImportColumn, value: &Value, @@ -392,7 +489,10 @@ fn normalize_postgres_import_value( return normalize_postgres_boolean_value(column_name, value); } if is_postgres_timestamp_column(target_column) { - return normalize_postgres_timestamp_value(column_name, value); + return normalize_postgres_timestamp_value(table_name, column_name, value); + } + if is_postgres_bytea_column(target_column) { + return postgres_bytea_json_value(column_name, value); } if is_postgres_json_column(target_column) { return normalize_postgres_json_value(value); @@ -450,6 +550,7 @@ fn normalize_postgres_boolean_value( } fn normalize_postgres_timestamp_value( + table_name: &str, column_name: &str, value: &Value, ) -> Result { @@ -465,7 +566,7 @@ fn normalize_postgres_timestamp_value( ))); }; - let datetime = if column_name.ends_with("_unix_ms") + let datetime = if import_timestamp_uses_millis(table_name, column_name) || timestamp >= 100_000_000_000 || timestamp <= -100_000_000_000 { @@ -497,17 +598,17 @@ fn normalize_postgres_json_value(value: &Value) -> Result } async fn import_postgres_billing_row( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: &ExportRow, column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (export_table_name, payload) = billing_payload_table(row)?; - let table_name = postgres_billing_table_name(&export_table_name)?; - let target_columns = postgres_import_columns_cached(pool, column_cache, table_name).await?; + let (table_name, conflict_column) = postgres_billing_table_name(&export_table_name)?; + let target_columns = postgres_import_columns_cached(tx, column_cache, table_name).await?; import_postgres_row( - pool, + tx, table_name, - &["id"], + &[conflict_column], ExportDomain::Billing, &ExportRow { id: row.id.clone(), @@ -518,27 +619,66 @@ async fn import_postgres_billing_row( .await } -fn postgres_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { +async fn import_postgres_auxiliary_row( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + row: &ExportRow, + column_cache: &mut BTreeMap, +) -> Result<(), DataLayerError> { + let (table_name, payload) = domain_payload_table(row, "auxiliary", None)?; + let table = auxiliary_table(&table_name)?; + let target_table = format!("public.{}", postgres_quote_identifier(table.name)?); + let target_columns = postgres_import_columns_cached(tx, column_cache, &target_table).await?; + import_postgres_row( + tx, + &target_table, + table.primary_key, + ExportDomain::Auxiliary, + &ExportRow { + id: row.id.clone(), + payload, + }, + &target_columns, + ) + .await +} + +fn postgres_billing_table_name( + table_name: &str, +) -> Result<(&'static str, &'static str), DataLayerError> { match table_name { - "billing_rules" => Ok("public.billing_rules"), - "dimension_collectors" => Ok("public.dimension_collectors"), - "usage_settlement_snapshots" => Ok("public.usage_settlement_snapshots"), + "billing_rules" => Ok(("public.billing_rules", "id")), + "dimension_collectors" => Ok(("public.dimension_collectors", "id")), + "usage_settlement_snapshots" => Ok(("public.usage_settlement_snapshots", "request_id")), other => Err(DataLayerError::InvalidInput(format!( "unsupported postgres billing export table '{other}'" ))), } } +#[cfg(test)] +mod billing_table_tests { + use super::postgres_billing_table_name; + + #[test] + fn settlement_snapshot_import_uses_request_id_conflict_key() { + assert_eq!( + postgres_billing_table_name("usage_settlement_snapshots") + .expect("settlement snapshot table should be supported"), + ("public.usage_settlement_snapshots", "request_id") + ); + } +} + async fn import_postgres_wallet_row( - pool: &crate::driver::postgres::PostgresPool, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: &ExportRow, column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (export_table_name, payload) = domain_payload_table(row, "wallet", Some("wallets"))?; let (table_name, id_column) = postgres_wallet_table_name(&export_table_name)?; - let target_columns = postgres_import_columns_cached(pool, column_cache, table_name).await?; + let target_columns = postgres_import_columns_cached(tx, column_cache, table_name).await?; import_postgres_row( - pool, + tx, table_name, &[id_column], ExportDomain::Wallets, diff --git a/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs b/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs index 372f02aa5..95e729ead 100644 --- a/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs +++ b/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs @@ -1,5 +1,12 @@ use super::*; +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct SqliteImportColumns { + names: ImportColumnNames, + declared_types: BTreeMap, + primary_key: Vec, +} + pub async fn export_sqlite_core_jsonl( pool: &crate::driver::sqlite::SqlitePool, created_at_unix_secs: u64, @@ -12,6 +19,7 @@ pub async fn export_sqlite_jsonl( domains: Vec, created_at_unix_secs: u64, ) -> Result { + let mut tx = pool.begin().await.map_sql_err()?; let manifest = DataExportManifest::new( created_at_unix_secs, Some(DatabaseDriver::Sqlite), @@ -20,24 +28,29 @@ pub async fn export_sqlite_jsonl( let mut records = vec![DataExportRecord::manifest(manifest)]; for domain in domains { + if domain == ExportDomain::Auxiliary { + export_sqlite_auxiliary_records(&mut tx, &mut records).await?; + continue; + } if domain == ExportDomain::Billing { - export_sqlite_billing_records(pool, &mut records).await?; + export_sqlite_billing_records(&mut tx, &mut records).await?; continue; } if domain == ExportDomain::Wallets { - export_sqlite_wallet_records(pool, &mut records).await?; + export_sqlite_wallet_records(&mut tx, &mut records).await?; continue; } let (table_name, id_column) = sqlite_domain_table(domain)?; let order_by = export_order_by(domain, id_column); let sql = format!("SELECT * FROM {table_name} ORDER BY {order_by}"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut *tx).await.map_sql_err()?; for row in rows { let id = sqlite_export_row_id(domain, &row, id_column)?; records.push(DataExportRecord::row(domain, id, sqlite_row_payload(&row)?)); } } + tx.commit().await.map_sql_err()?; encode_jsonl(&records) } @@ -53,31 +66,40 @@ pub async fn import_sqlite_plan( pool: &crate::driver::sqlite::SqlitePool, plan: &DataImportPlan, ) -> Result { + let mut tx = pool.begin().await.map_sql_err()?; let mut imported = 0usize; - let mut column_cache = BTreeMap::::new(); + let mut column_cache = BTreeMap::::new(); for domain in &plan.manifest.domains { + if *domain == ExportDomain::Auxiliary { + for row in plan.rows(*domain) { + import_sqlite_auxiliary_row(&mut tx, row, &mut column_cache).await?; + imported = imported.saturating_add(1); + } + continue; + } if *domain == ExportDomain::Billing { for row in plan.rows(*domain) { - import_sqlite_billing_row(pool, row, &mut column_cache).await?; + import_sqlite_billing_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; } if *domain == ExportDomain::Wallets { for row in plan.rows(*domain) { - import_sqlite_wallet_row(pool, row, &mut column_cache).await?; + import_sqlite_wallet_row(&mut tx, row, &mut column_cache).await?; imported = imported.saturating_add(1); } continue; } let (table_name, _id_column) = sqlite_domain_table(*domain)?; let target_columns = - sqlite_import_columns_cached(pool, &mut column_cache, table_name).await?; + sqlite_import_columns_cached(&mut tx, &mut column_cache, table_name).await?; for row in plan.rows(*domain) { - import_sqlite_row(pool, table_name, *domain, row, &target_columns).await?; + import_sqlite_row(&mut tx, table_name, *domain, row, &target_columns).await?; imported = imported.saturating_add(1); } } + tx.commit().await.map_sql_err()?; Ok(imported) } @@ -107,9 +129,42 @@ fn sqlite_domain_table( "sqlite billing export uses multiple tables and must be handled as a domain" .to_string(), )), + ExportDomain::Auxiliary => Err(DataLayerError::InvalidInput( + "sqlite auxiliary export uses multiple tables and must be handled as a domain" + .to_string(), + )), } } +async fn export_sqlite_auxiliary_records( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + records: &mut Vec, +) -> Result<(), DataLayerError> { + for table in AUXILIARY_TABLES { + let table_sql = sqlite_quote_identifier(table.name)?; + let order_sql = table + .primary_key + .iter() + .map(|column| sqlite_quote_identifier(column).map(|column| format!("{column} ASC"))) + .collect::, _>>()? + .join(", "); + let rows = sqlx::query(&format!("SELECT * FROM {table_sql} ORDER BY {order_sql}")) + .fetch_all(&mut **tx) + .await + .map_sql_err()?; + for row in rows { + let payload = sqlite_row_payload(&row)?; + let id = auxiliary_row_id(*table, &payload)?; + records.push(DataExportRecord::row( + ExportDomain::Auxiliary, + id, + payload_with_table(payload, table.name)?, + )); + } + } + Ok(()) +} + fn sqlite_export_row_id( domain: ExportDomain, row: &sqlx::sqlite::SqliteRow, @@ -140,7 +195,7 @@ fn sqlite_required_export_text( } async fn export_sqlite_billing_records( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, records: &mut Vec, ) -> Result<(), DataLayerError> { for table_name in [ @@ -154,7 +209,7 @@ async fn export_sqlite_billing_records( "id" }; let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) @@ -175,12 +230,12 @@ async fn export_sqlite_billing_records( } async fn export_sqlite_wallet_records( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, records: &mut Vec, ) -> Result<(), DataLayerError> { for (table_name, id_column) in sqlite_wallet_tables() { let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; for row in rows { let id = row .try_get::, _>(id_column) @@ -201,13 +256,13 @@ async fn export_sqlite_wallet_records( } async fn import_sqlite_row( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, table_name: &str, domain: ExportDomain, row: &ExportRow, - target_columns: &ImportColumnNames, + target_columns: &SqliteImportColumns, ) -> Result<(), DataLayerError> { - let object = filter_import_payload("sqlite", table_name, domain, row, target_columns)?; + let object = filter_import_payload("sqlite", table_name, domain, row, &target_columns.names)?; let columns = object.keys().map(String::as_str).collect::>(); let column_sql = columns @@ -216,29 +271,55 @@ async fn import_sqlite_row( .collect::, _>>()? .join(", "); let placeholder_sql = vec!["?"; columns.len()].join(", "); - let sql = - format!("INSERT OR REPLACE INTO {table_name} ({column_sql}) VALUES ({placeholder_sql})"); + let conflict_columns = target_columns + .primary_key + .iter() + .map(|column| sqlite_quote_identifier(column)) + .collect::, _>>()? + .join(", "); + let update_sql = columns + .iter() + .filter(|column| !target_columns.primary_key.iter().any(|key| key == *column)) + .map(|column| { + let quoted = sqlite_quote_identifier(column)?; + Ok(format!("{quoted} = excluded.{quoted}")) + }) + .collect::, DataLayerError>>()? + .join(", "); + let conflict_sql = if update_sql.is_empty() { + format!("ON CONFLICT ({conflict_columns}) DO NOTHING") + } else { + format!("ON CONFLICT ({conflict_columns}) DO UPDATE SET {update_sql}") + }; + let sql = format!( + "INSERT INTO {table_name} ({column_sql}) VALUES ({placeholder_sql}) {conflict_sql}" + ); let mut query = sqlx::query(&sql); for column in columns { let value = object .get(column) .expect("column name came from payload object keys"); - query = bind_sqlite_json_value(query, value)?; + let declared_type = target_columns + .declared_types + .get(column) + .map(String::as_str) + .unwrap_or_default(); + query = bind_sqlite_import_value(query, value, table_name, column, declared_type)?; } - query.execute(pool).await.map_sql_err()?; + query.execute(&mut **tx).await.map_sql_err()?; Ok(()) } async fn import_sqlite_billing_row( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, row: &ExportRow, - column_cache: &mut BTreeMap, + column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (table_name, payload) = billing_payload_table(row)?; let table_name = sqlite_billing_table_name(&table_name)?; - let target_columns = sqlite_import_columns_cached(pool, column_cache, table_name).await?; + let target_columns = sqlite_import_columns_cached(tx, column_cache, table_name).await?; import_sqlite_row( - pool, + tx, table_name, ExportDomain::Billing, &ExportRow { @@ -250,6 +331,27 @@ async fn import_sqlite_billing_row( .await } +async fn import_sqlite_auxiliary_row( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + row: &ExportRow, + column_cache: &mut BTreeMap, +) -> Result<(), DataLayerError> { + let (table_name, payload) = domain_payload_table(row, "auxiliary", None)?; + let table = auxiliary_table(&table_name)?; + let target_columns = sqlite_import_columns_cached(tx, column_cache, table.name).await?; + import_sqlite_row( + tx, + table.name, + ExportDomain::Auxiliary, + &ExportRow { + id: row.id.clone(), + payload, + }, + &target_columns, + ) + .await +} + fn sqlite_billing_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { match table_name { "billing_rules" => Ok("billing_rules"), @@ -262,15 +364,15 @@ fn sqlite_billing_table_name(table_name: &str) -> Result<&'static str, DataLayer } async fn import_sqlite_wallet_row( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, row: &ExportRow, - column_cache: &mut BTreeMap, + column_cache: &mut BTreeMap, ) -> Result<(), DataLayerError> { let (table_name, payload) = domain_payload_table(row, "wallet", Some("wallets"))?; let table_name = sqlite_wallet_table_name(&table_name)?; - let target_columns = sqlite_import_columns_cached(pool, column_cache, table_name).await?; + let target_columns = sqlite_import_columns_cached(tx, column_cache, table_name).await?; import_sqlite_row( - pool, + tx, table_name, ExportDomain::Wallets, &ExportRow { @@ -308,39 +410,81 @@ fn sqlite_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerE } async fn sqlite_import_columns_cached( - pool: &crate::driver::sqlite::SqlitePool, - cache: &mut BTreeMap, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + cache: &mut BTreeMap, table_name: &str, -) -> Result { +) -> Result { if let Some(columns) = cache.get(table_name) { return Ok(columns.clone()); } - let columns = load_sqlite_import_columns(pool, table_name).await?; + let columns = load_sqlite_import_columns(tx, table_name).await?; cache.insert(table_name.to_string(), columns.clone()); Ok(columns) } async fn load_sqlite_import_columns( - pool: &crate::driver::sqlite::SqlitePool, + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, table_name: &str, -) -> Result { +) -> Result { let sql = format!("PRAGMA table_info({table_name})"); - let rows = sqlx::query(&sql).fetch_all(pool).await.map_sql_err()?; - let mut columns = ImportColumnNames::new(); + let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; + let mut columns = SqliteImportColumns::default(); + let mut primary_key = BTreeMap::new(); for row in rows { - columns.insert(row.try_get::("name").map_sql_err()?); + let name = row.try_get::("name").map_sql_err()?; + let declared_type = row + .try_get::, _>("type") + .map_sql_err()? + .unwrap_or_default(); + columns.names.insert(name.clone()); + columns.declared_types.insert(name.clone(), declared_type); + let primary_key_position = row.try_get::("pk").map_sql_err()?; + if primary_key_position > 0 { + primary_key.insert(primary_key_position, name); + } } - if columns.is_empty() { + if columns.names.is_empty() { return Err(DataLayerError::UnexpectedValue(format!( "sqlite import target table '{table_name}' has no visible columns" ))); } + if primary_key.is_empty() { + return Err(DataLayerError::UnexpectedValue(format!( + "sqlite import target table '{table_name}' has no primary key" + ))); + } + columns.primary_key = primary_key.into_values().collect(); Ok(columns) } +fn bind_sqlite_import_value<'q>( + query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>, + json_value: &'q Value, + table_name: &str, + column_name: &str, + declared_type: &str, +) -> Result>, DataLayerError> +{ + if declared_type.to_ascii_uppercase().contains("BLOB") { + return match normalize_imported_binary("sqlite", column_name, json_value)? { + Some(bytes) => Ok(query.bind(bytes)), + None => Ok(query.bind(Option::>::None)), + }; + } + let has_integer_affinity = declared_type.to_ascii_uppercase().contains("INT"); + if !has_integer_affinity || !import_column_stores_timestamp(column_name) { + return bind_sqlite_json_value(query, json_value); + } + + match normalize_imported_integer_timestamp("sqlite", table_name, column_name, json_value)? { + Some(timestamp) => Ok(query.bind(timestamp)), + None => Ok(query.bind(Option::::None)), + } +} + fn sqlite_row_payload(row: &sqlx::sqlite::SqliteRow) -> Result { let mut object = serde_json::Map::new(); for (index, column) in row.columns().iter().enumerate() { diff --git a/crates/aether-data/runtime/src/lifecycle/export/tests.rs b/crates/aether-data/runtime/src/lifecycle/export/tests.rs index 6fe7e44f1..08e02f618 100644 --- a/crates/aether-data/runtime/src/lifecycle/export/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/export/tests.rs @@ -1,14 +1,16 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use serde_json::json; use super::{ build_import_plan, decode_jsonl, encode_jsonl, export_mysql_core_jsonl, export_mysql_jsonl, - export_postgres_core_jsonl, export_sqlite_core_jsonl, import_mysql_jsonl, - import_postgres_jsonl, import_sqlite_jsonl, mysql_core_export_domains, - normalize_postgres_import_payload, postgres_core_export_domains, sqlite_core_export_domains, - DataExportManifest, DataExportRecord, DataImportPlan, ExportDomain, ExportRow, - PostgresImportColumn, + export_postgres_core_jsonl, export_sqlite_core_jsonl, filter_import_payload, + import_mysql_jsonl, import_postgres_jsonl, import_sqlite_jsonl, mysql_core_export_domains, + normalize_imported_binary, normalize_imported_integer_timestamp, + normalize_postgres_import_payload, postgres_bytea_json_value, postgres_core_export_domains, + sqlite_core_export_domains, sqlite_schema_copy_insert_sql, DataExportManifest, + DataExportRecord, DataImportPlan, ExportDomain, ExportRow, PostgresImportColumn, + SchemaCopyColumn, SchemaCopyTable, SqliteCopyColumn, AUXILIARY_TABLES, }; use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory}; use crate::lifecycle::migrate::{ @@ -64,6 +66,81 @@ fn jsonl_round_trips_manifest_and_domain_rows() { fn core_export_domains_match_across_sql_drivers() { assert_eq!(sqlite_core_export_domains(), mysql_core_export_domains()); assert_eq!(sqlite_core_export_domains(), postgres_core_export_domains()); + assert!(sqlite_core_export_domains().contains(&ExportDomain::Auxiliary)); +} + +#[tokio::test] +async fn sqlite_core_export_covers_every_portable_table() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + let schema_tables = sqlx::query_scalar::<_, String>( + r#" +SELECT name +FROM sqlite_master +WHERE type = 'table' + AND name NOT LIKE 'sqlite_%' + AND name NOT IN ('_sqlx_migrations', 'schema_backfills') +ORDER BY name +"#, + ) + .fetch_all(&pool) + .await + .expect("sqlite schema tables should load") + .into_iter() + .collect::>(); + + let mut exported_tables = [ + "users", + "api_keys", + "providers", + "provider_api_keys", + "provider_endpoints", + "global_models", + "models", + "auth_modules", + "oauth_providers", + "user_oauth_links", + "user_groups", + "user_group_members", + "proxy_nodes", + "system_configs", + "usage", + "wallets", + "wallet_transactions", + "wallet_daily_usage_ledgers", + "payment_orders", + "payment_callbacks", + "refund_requests", + "redeem_code_batches", + "redeem_codes", + "billing_rules", + "dimension_collectors", + "usage_settlement_snapshots", + ] + .into_iter() + .map(str::to_string) + .collect::>(); + exported_tables.extend(AUXILIARY_TABLES.iter().map(|table| table.name.to_string())); + + assert_eq!(schema_tables, exported_tables); +} + +#[test] +fn version_one_exports_remain_importable_after_full_export_expansion() { + let records = decode_jsonl( + r#"{"record_type":"manifest","manifest":{"format_version":1,"created_at_unix_secs":1,"source_driver":null,"domains":["users"]}} +{"record_type":"row","domain":"users","id":"user-1","payload":{"id":"user-1"}}"#, + ) + .expect("version one exports should remain supported"); + + assert_eq!(records.len(), 2); } #[test] @@ -171,6 +248,70 @@ fn postgres_import_payload_normalizes_sqlite_values_for_target_columns() { assert!(!normalized.contains_key("legacy_nullable")); } +#[test] +fn cross_driver_timestamp_normalization_preserves_usage_second_contract() { + assert_eq!( + normalize_imported_integer_timestamp( + "sqlite", + r#""usage""#, + "created_at_unix_ms", + &json!("1970-01-01T00:00:01.234900Z"), + ) + .expect("usage timestamp should normalize"), + Some(1), + ); + assert_eq!( + normalize_imported_integer_timestamp( + "mysql", + "request_candidates", + "created_at_unix_ms", + &json!("1970-01-01T00:00:01.234900Z"), + ) + .expect("millisecond timestamp should normalize"), + Some(1_234), + ); + + let target_columns = BTreeMap::from([( + "created_at_unix_ms".to_string(), + postgres_column("timestamp with time zone", "timestamptz"), + )]); + let row = ExportRow { + id: "usage-1".to_string(), + payload: json!({ "created_at_unix_ms": 1_700_000_000 }), + }; + let normalized = normalize_postgres_import_payload( + "public.usage", + ExportDomain::Usage, + &row, + &target_columns, + ) + .expect("postgres usage timestamp should normalize"); + assert_eq!( + normalized["created_at_unix_ms"], + json!("2023-11-14T22:13:20+00:00") + ); +} + +#[test] +fn cross_driver_binary_normalization_preserves_raw_bytes() { + assert_eq!( + normalize_imported_binary("sqlite", "payload_gzip", &json!([0, 1, 127, 255])) + .expect("byte array should normalize"), + Some(vec![0, 1, 127, 255]), + ); + assert_eq!( + normalize_imported_binary("mysql", "payload_gzip", &json!("\\x00017fff")) + .expect("postgres hex should normalize"), + Some(vec![0, 1, 127, 255]), + ); + assert!(normalize_imported_binary("sqlite", "payload_gzip", &json!([256])).is_err()); + assert_eq!( + postgres_bytea_json_value("payload_gzip", &json!([0, 1, 127, 255])) + .expect("postgres bytea should normalize"), + json!("\\x00017fff"), + ); +} + #[test] fn postgres_import_payload_rejects_non_null_unknown_columns() { let target_columns = BTreeMap::from([( @@ -197,6 +338,95 @@ fn postgres_import_payload_rejects_non_null_unknown_columns() { assert!(err.to_string().contains("does not exist")); } +#[test] +fn mysql_and_sqlite_import_payloads_reject_non_null_unknown_columns() { + let target_columns = BTreeSet::from(["id".to_string()]); + let row = ExportRow { + id: "user-1".to_string(), + payload: json!({ + "id": "user-1", + "legacy_nullable": null, + "unexpected_column": "value" + }), + }; + + for driver_name in ["mysql", "sqlite"] { + let err = filter_import_payload( + driver_name, + "users", + ExportDomain::Users, + &row, + &target_columns, + ) + .expect_err("non-null unknown columns should fail"); + + assert!(err.to_string().contains("unexpected_column")); + assert!(err.to_string().contains("does not exist")); + assert!(err.to_string().contains(driver_name)); + } +} + +#[test] +fn mysql_and_sqlite_import_payloads_ignore_unknown_null_columns() { + let target_columns = BTreeSet::from(["id".to_string()]); + let row = ExportRow { + id: "user-1".to_string(), + payload: json!({ + "id": "user-1", + "legacy_nullable": null + }), + }; + + let filtered = filter_import_payload( + "sqlite", + "users", + ExportDomain::Users, + &row, + &target_columns, + ) + .expect("unknown null columns should remain backward compatible"); + + assert_eq!( + filtered, + serde_json::Map::from_iter([("id".to_string(), json!("user-1"))]) + ); +} + +#[test] +fn postgres_to_sqlite_copy_uses_primary_key_upsert_instead_of_replace() { + let table = SchemaCopyTable { + table_name: "usage".to_string(), + columns: vec![ + SchemaCopyColumn { + sqlite: SqliteCopyColumn { + name: "request_id".to_string(), + declared_type: "TEXT".to_string(), + not_null: true, + has_default: false, + primary_key_position: 1, + }, + postgres: postgres_column("character varying", "varchar"), + }, + SchemaCopyColumn { + sqlite: SqliteCopyColumn { + name: "status".to_string(), + declared_type: "TEXT".to_string(), + not_null: true, + has_default: false, + primary_key_position: 0, + }, + postgres: postgres_column("character varying", "varchar"), + }, + ], + }; + + let sql = sqlite_schema_copy_insert_sql(&table).expect("copy SQL should build"); + + assert!(!sql.contains("OR REPLACE")); + assert!(sql.contains("ON CONFLICT (\"request_id\") DO UPDATE SET")); + assert!(sql.contains("\"status\" = excluded.\"status\"")); +} + fn postgres_column(data_type: &str, udt_name: &str) -> PostgresImportColumn { PostgresImportColumn { data_type: data_type.to_ascii_lowercase(), @@ -215,6 +445,183 @@ fn postgres_not_null_default_column(data_type: &str, udt_name: &str) -> Postgres } } +#[tokio::test] +async fn sqlite_import_rejects_non_integer_timestamp_values() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + for invalid_value in [ + json!("not-a-timestamp"), + json!(1.5), + json!(true), + json!({"unexpected": "object"}), + ] { + let encoded = encode_jsonl(&[ + DataExportRecord::manifest(DataExportManifest::new( + 1_700_000_000, + Some(DatabaseDriver::Postgres), + vec![ExportDomain::GlobalModels], + )), + DataExportRecord::row( + ExportDomain::GlobalModels, + "invalid-timestamp", + json!({ + "id": "invalid-timestamp", + "name": "invalid-timestamp", + "created_at": invalid_value, + "updated_at": 1 + }), + ), + ]) + .expect("invalid timestamp fixture should encode"); + + let err = import_sqlite_jsonl(&pool, &encoded) + .await + .expect_err("non-integer timestamp should be rejected"); + assert!(err.to_string().contains( + "timestamp column 'created_at' must contain an integer or supported datetime" + )); + } +} + +#[tokio::test] +async fn sqlite_import_updates_parent_without_cascading_child_rows() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + sqlx::query("PRAGMA foreign_keys = ON") + .execute(&pool) + .await + .expect("foreign keys should be enabled"); + sqlx::raw_sql( + r#" +INSERT INTO users (id, email, username, created_at, updated_at) +VALUES ('import-user', 'import@example.test', 'import-user', 1, 1); +INSERT INTO user_groups ( + id, name, normalized_name, description, priority, + allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, rate_limit_mode, + created_at, updated_at +) +VALUES ( + 'import-group', 'Before', 'import-group', 'preserve-me', 0, + 'inherit', 'inherit', 'inherit', 'inherit', 1, 1 +); +INSERT INTO user_group_members (group_id, user_id, created_at) +VALUES ('import-group', 'import-user', 1); +"#, + ) + .execute(&pool) + .await + .expect("parent and child fixtures should insert"); + + let encoded = encode_jsonl(&[ + DataExportRecord::manifest(DataExportManifest::new( + 1_700_000_000, + Some(DatabaseDriver::Postgres), + vec![ExportDomain::UserGroups], + )), + DataExportRecord::row( + ExportDomain::UserGroups, + "import-group", + json!({ + "id": "import-group", + "name": "After", + "normalized_name": "import-group", + "priority": 10, + "allowed_providers_mode": "inherit", + "allowed_api_formats_mode": "inherit", + "allowed_models_mode": "inherit", + "rate_limit_mode": "inherit", + "created_at": 1, + "updated_at": 2 + }), + ), + ]) + .expect("group export should encode"); + + assert_eq!( + import_sqlite_jsonl(&pool, &encoded) + .await + .expect("group import should update in place"), + 1 + ); + let group = sqlx::query_as::<_, (String, String)>( + "SELECT name, description FROM user_groups WHERE id = 'import-group'", + ) + .fetch_one(&pool) + .await + .expect("updated group should load"); + assert_eq!(group, ("After".to_string(), "preserve-me".to_string())); + let member_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM user_group_members WHERE group_id = 'import-group'", + ) + .fetch_one(&pool) + .await + .expect("group member count should load"); + assert_eq!(member_count, 1); +} + +#[tokio::test] +async fn sqlite_import_rolls_back_rows_after_late_failure() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let encoded = encode_jsonl(&[ + DataExportRecord::manifest(DataExportManifest::new( + 1_700_000_000, + Some(DatabaseDriver::Postgres), + vec![ExportDomain::GlobalModels], + )), + DataExportRecord::row( + ExportDomain::GlobalModels, + "rollback-valid", + json!({ + "id": "rollback-valid", + "name": "rollback-valid", + "created_at": 1, + "updated_at": 1 + }), + ), + DataExportRecord::row( + ExportDomain::GlobalModels, + "rollback-invalid", + json!({ + "id": "rollback-invalid", + "name": "rollback-invalid", + "created_at": "invalid-timestamp", + "updated_at": 1 + }), + ), + ]) + .expect("rollback fixture should encode"); + + import_sqlite_jsonl(&pool, &encoded) + .await + .expect_err("late invalid row should fail the import"); + let count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM global_models WHERE id LIKE 'rollback-%'") + .fetch_one(&pool) + .await + .expect("rolled back row count should load"); + assert_eq!(count, 0); +} + #[tokio::test] async fn sqlite_core_export_reads_migrated_database_rows() { let pool = sqlx::sqlite::SqlitePoolOptions::new() @@ -243,7 +650,7 @@ VALUES ('provider-key-1', 'provider-1', 'Provider Key', 'ciphertext-provider', ' INSERT INTO provider_endpoints (id, provider_id, name, base_url, created_at, updated_at) VALUES ('endpoint-1', 'provider-1', 'Primary', 'https://example.test', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); INSERT INTO global_models (id, name, created_at, updated_at) -VALUES ('global-model-1', 'gpt-test', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); +VALUES ('global-model-1', 'gpt-test', '1970-01-01T00:00:01Z', '1970-01-01 00:00:02.123456'); INSERT INTO models (id, provider_id, global_model_id, provider_model_name, created_at, updated_at) VALUES ('model-1', 'provider-1', 'global-model-1', 'gpt-test', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); INSERT INTO billing_rules (id, global_model_id, name, task_type, expression, variables, dimension_mappings, is_enabled, created_at, updated_at) @@ -255,7 +662,21 @@ VALUES ('config-1', 'billing.enabled', 'true', '1970-01-01T00:00:01Z', '1970-01- INSERT INTO wallets (id, user_id, created_at, updated_at) VALUES ('wallet-1', 'user-1', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); INSERT INTO "usage" (request_id, id, user_id, provider_name, model, status, billing_status, created_at_unix_ms, updated_at_unix_secs) -VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'completed', 'settled', 1, 2); +VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'completed', 'settled', '1970-01-01T00:00:01.234900Z', 2); +INSERT INTO audit_logs (id, event_type, description, request_id, created_at) +VALUES ('audit-1', 'request.completed', 'Exported audit', 'request-1', '1970-01-01T00:00:02Z'); +INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip, created_at, updated_at) +VALUES ('body-ref-1', 'request-1', 'request', X'00117FFF', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); +INSERT INTO usage_http_audits (request_id, request_body_ref, request_body_state, body_capture_mode, created_at, updated_at) +VALUES ('request-1', 'body-ref-1', 'captured', 'full', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); +INSERT INTO usage_routing_snapshots ( + request_id, candidate_id, candidate_index, selected_provider_id, + selected_endpoint_id, selected_provider_api_key_id, created_at, updated_at +) +VALUES ( + 'request-1', 'candidate-1', 2, 'provider-1', + 'endpoint-1', 'provider-key-1', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z' +); "#, ) .execute(&pool) @@ -304,6 +725,21 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet import_plan.rows(ExportDomain::Billing)[0].payload["dimension_mappings"]["input"], "input_tokens" ); + assert!(import_plan + .rows(ExportDomain::Auxiliary) + .iter() + .any(|row| row.payload["__table"] == "audit_logs" && row.payload["id"] == "audit-1")); + assert!(import_plan + .rows(ExportDomain::Auxiliary) + .iter() + .any(|row| row.payload["__table"] == "usage_body_blobs" + && row.payload["payload_gzip"] == json!([0, 17, 127, 255]))); + assert!(import_plan + .rows(ExportDomain::Auxiliary) + .iter() + .any(|row| row.payload["__table"] == "usage_routing_snapshots" + && row.payload["candidate_id"] == "candidate-1" + && row.payload["selected_provider_id"] == "provider-1")); let target_pool = sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) @@ -316,7 +752,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet let imported = import_sqlite_jsonl(&target_pool, &encoded) .await .expect("sqlite import should load exported rows"); - assert_eq!(imported, 16); + assert_eq!(imported, 20); let imported_api_key = sqlx::query_as::<_, (String,)>("SELECT key_encrypted FROM api_keys WHERE id = 'api-key-1'") @@ -325,13 +761,31 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet .expect("imported api key should load"); assert_eq!(imported_api_key.0, "ciphertext-1"); - let imported_usage = sqlx::query_as::<_, (String,)>( - "SELECT request_id FROM \"usage\" WHERE request_id = 'request-1'", + let imported_usage = sqlx::query_as::<_, (String, i64, String)>( + "SELECT request_id, created_at_unix_ms, typeof(created_at_unix_ms) FROM \"usage\" WHERE request_id = 'request-1'", ) .fetch_one(&target_pool) .await .expect("imported usage should load"); - assert_eq!(imported_usage.0, "request-1"); + assert_eq!( + imported_usage, + ("request-1".to_string(), 1, "integer".to_string()) + ); + + let imported_global_model_timestamps = sqlx::query_as::<_, (i64, i64, String, String)>( + r#" +SELECT created_at, updated_at, typeof(created_at), typeof(updated_at) +FROM global_models +WHERE id = 'global-model-1' +"#, + ) + .fetch_one(&target_pool) + .await + .expect("imported global model timestamps should decode as integers"); + assert_eq!( + imported_global_model_timestamps, + (1, 2, "integer".to_string(), "integer".to_string()) + ); let imported_group_member = sqlx::query_as::<_, (String, String)>( "SELECT group_id, user_id FROM user_group_members WHERE group_id = 'group-1' AND user_id = 'user-1'", @@ -350,6 +804,29 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet .expect("imported billing rule should load"); assert_eq!(imported_billing_rule.0, "input_tokens * 0.01"); + let imported_body: Vec = sqlx::query_scalar( + "SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = 'body-ref-1'", + ) + .fetch_one(&target_pool) + .await + .expect("imported body blob should load"); + assert_eq!(imported_body, vec![0, 17, 127, 255]); + + let imported_routing = sqlx::query_as::<_, (String, i64, String)>( + r#" +SELECT candidate_id, candidate_index, selected_provider_id +FROM usage_routing_snapshots +WHERE request_id = 'request-1' +"#, + ) + .fetch_one(&target_pool) + .await + .expect("imported routing snapshot should load"); + assert_eq!( + imported_routing, + ("candidate-1".to_string(), 2, "provider-1".to_string()) + ); + if let Some(database_url) = std::env::var("AETHER_TEST_POSTGRES_URL") .ok() .filter(|value| !value.trim().is_empty()) @@ -375,7 +852,7 @@ VALUES ('request-1', 'request-1', 'user-1', 'Provider One', 'gpt-test', 'complet let imported = import_postgres_jsonl(&postgres_pool, &encoded) .await .expect("postgres import should load exported rows"); - assert_eq!(imported, 16); + assert_eq!(imported, 20); let imported_api_key = sqlx::query_as::<_, (String,)>( "SELECT key_encrypted FROM api_keys WHERE id = 'api-key-1'", @@ -616,6 +1093,17 @@ async fn postgres_core_export_reads_migrated_database_rows_when_url_is_set() { .await .expect("imported sqlite api key should load"); assert_eq!(imported_api_key.0, "ciphertext-1"); + let imported_global_model_timestamps = sqlx::query_as::<_, (i64, i64, String, String)>( + "SELECT created_at, updated_at, typeof(created_at), typeof(updated_at) FROM global_models WHERE id = ?", + ) + .bind(&global_model_id) + .fetch_one(&target_pool) + .await + .expect("imported sqlite global model timestamps should decode as integers"); + assert_eq!( + imported_global_model_timestamps, + (1, 2, "integer".to_string(), "integer".to_string()) + ); let imported_group_member = sqlx::query_as::<_, (String, String)>( "SELECT group_id, user_id FROM user_group_members WHERE group_id = ? AND user_id = ?", ) diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs index 5961c0b54..0b75fa6db 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs @@ -453,15 +453,173 @@ fn create_table_names(sql: &str) -> BTreeSet { let trimmed = line.trim_start(); let table_part = trimmed .strip_prefix("CREATE TABLE IF NOT EXISTS public.") - .or_else(|| trimmed.strip_prefix("CREATE TABLE IF NOT EXISTS "))?; + .or_else(|| trimmed.strip_prefix("CREATE TABLE IF NOT EXISTS ")) + .or_else(|| trimmed.strip_prefix("CREATE TABLE public.")) + .or_else(|| trimmed.strip_prefix("CREATE TABLE "))?; let table_name = table_part .split(|ch: char| ch.is_ascii_whitespace() || ch == '(') .next()?; - Some(table_name.trim_matches('"').to_string()) + Some( + table_name + .trim_matches(|ch| ch == '"' || ch == '`') + .to_string(), + ) }) .collect() } +#[test] +fn portable_driver_migrations_create_the_postgres_table_set() { + let mut postgres_tables = POSTGRES_MIGRATOR + .iter() + .filter(|migration| migration.migration_type.is_up_migration()) + .flat_map(|migration| create_table_names(migration.sql.as_ref())) + .collect::>(); + postgres_tables.remove("schema_backfills"); + + let mysql_tables = super::mysql::MIGRATOR + .iter() + .filter(|migration| migration.migration_type.is_up_migration()) + .flat_map(|migration| create_table_names(migration.sql.as_ref())) + .collect::>(); + let sqlite_tables = super::sqlite::MIGRATOR + .iter() + .filter(|migration| migration.migration_type.is_up_migration()) + .flat_map(|migration| create_table_names(migration.sql.as_ref())) + .collect::>(); + + assert_eq!(mysql_tables, postgres_tables, "MySQL table set drifted"); + assert_eq!(sqlite_tables, postgres_tables, "SQLite table set drifted"); +} + +#[tokio::test] +async fn migrated_sqlite_columns_match_the_generated_logical_schema() { + const GENERATED_SQLITE_SCHEMA: &[&str] = &[ + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/001_identity.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/002_provider_catalog.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/003_auth_config.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/004_proxy_nodes.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/005_wallet_billing.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/006_usage.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/007_stats.sql" + )), + include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/schema/generated/sqlite/baseline/008_background_tasks.sql" + )), + ]; + + let migrated = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("migrated sqlite pool should connect"); + super::run_sqlite_migrations(&migrated) + .await + .expect("sqlite migrations should run"); + + let generated = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("generated sqlite pool should connect"); + for source in GENERATED_SQLITE_SCHEMA { + sqlx::raw_sql(source) + .execute(&generated) + .await + .expect("generated sqlite schema fragment should run"); + } + + let migrated_tables = sqlite_portable_table_names(&migrated).await; + let generated_tables = sqlite_portable_table_names(&generated).await; + assert_eq!(migrated_tables, generated_tables); + + for table in generated_tables { + let migrated_columns = sqlite_table_column_names(&migrated, &table).await; + let generated_columns = sqlite_table_column_names(&generated, &table).await; + assert_eq!( + migrated_columns, generated_columns, + "SQLite migration columns drifted for table {table}" + ); + } + + let migrated_indexes = sqlite_named_index_names(&migrated).await; + let generated_indexes = sqlite_named_index_names(&generated).await; + let missing_indexes = generated_indexes + .difference(&migrated_indexes) + .cloned() + .collect::>(); + assert!( + missing_indexes.is_empty(), + "SQLite migrations are missing generated logical indexes: {missing_indexes:?}" + ); +} + +async fn sqlite_portable_table_names(pool: &SqlitePool) -> BTreeSet { + query_scalar::<_, String>( + r#" +SELECT name +FROM sqlite_master +WHERE type = 'table' + AND name NOT LIKE 'sqlite_%' + AND name NOT IN ('_sqlx_migrations', 'schema_backfills') +ORDER BY name +"#, + ) + .fetch_all(pool) + .await + .expect("sqlite table names should load") + .into_iter() + .collect() +} + +async fn sqlite_table_column_names(pool: &SqlitePool, table: &str) -> BTreeSet { + query_scalar::<_, String>("SELECT name FROM pragma_table_info(?) ORDER BY cid") + .bind(table) + .fetch_all(pool) + .await + .expect("sqlite table columns should load") + .into_iter() + .collect() +} + +async fn sqlite_named_index_names(pool: &SqlitePool) -> BTreeSet { + query_scalar::<_, String>( + r#" +SELECT name +FROM sqlite_master +WHERE type = 'index' + AND sql IS NOT NULL +ORDER BY name +"#, + ) + .fetch_all(pool) + .await + .expect("sqlite named indexes should load") + .into_iter() + .collect() +} + #[test] fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role() { assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("'audit_admin'")); @@ -862,6 +1020,9 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260527000000, 20260528000000, 20260528020000, + 20260725010000, + 20260725020000, + 20260725030000, ] ); assert_eq!( @@ -889,10 +1050,304 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260527000000, 20260528000000, 20260528020000, + 20260725000000, + 20260725010000, + 20260725020000, + 20260725030000, + 20260725040000, ] ); } +#[tokio::test] +async fn sqlite_imported_timestamp_migration_normalizes_text_storage() { + let pool = SqlitePool::connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + super::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + query( + r#" +INSERT INTO global_models (id, name, created_at, updated_at) +VALUES + ('timestamp-rfc3339', 'timestamp-rfc3339', '1970-01-01T00:00:01Z', '1970-01-01T08:00:02+08:00'), + ('timestamp-sqlalchemy', 'timestamp-sqlalchemy', '1970-01-01 00:00:03.123456', '1970-01-01 00:00:04.987654'), + ('timestamp-integer', 'timestamp-integer', 5, 6); +"#, + ) + .execute(&pool) + .await + .expect("timestamp fixtures should insert"); + query( + r#" +INSERT INTO "usage" (request_id, created_at_unix_ms, updated_at_unix_secs) +VALUES ('timestamp-usage', '1970-01-01T00:00:01.234900Z', '1970-01-01T00:00:02Z'); +"#, + ) + .execute(&pool) + .await + .expect("usage timestamp fixture should insert"); + + let migration = super::sqlite::MIGRATOR + .iter() + .find(|migration| migration.version == 20260725000000) + .expect("timestamp normalization migration should be embedded"); + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect("timestamp normalization migration should apply"); + + let rows = sqlx::query_as::<_, (String, i64, i64, String, String)>( + r#" +SELECT id, created_at, updated_at, typeof(created_at), typeof(updated_at) +FROM global_models +WHERE id LIKE 'timestamp-%' +ORDER BY id +"#, + ) + .fetch_all(&pool) + .await + .expect("normalized timestamps should decode as integers"); + + assert_eq!( + rows, + vec![ + ( + "timestamp-integer".to_string(), + 5, + 6, + "integer".to_string(), + "integer".to_string(), + ), + ( + "timestamp-rfc3339".to_string(), + 1, + 2, + "integer".to_string(), + "integer".to_string(), + ), + ( + "timestamp-sqlalchemy".to_string(), + 3, + 4, + "integer".to_string(), + "integer".to_string(), + ), + ] + ); + + let usage_timestamps = sqlx::query_as::<_, (i64, i64, String, String)>( + r#" +SELECT created_at_unix_ms, updated_at_unix_secs, + typeof(created_at_unix_ms), typeof(updated_at_unix_secs) +FROM "usage" +WHERE request_id = 'timestamp-usage' +"#, + ) + .fetch_one(&pool) + .await + .expect("normalized usage timestamps should decode as integers"); + assert_eq!( + usage_timestamps, + (1, 2, "integer".to_string(), "integer".to_string()) + ); +} + +#[tokio::test] +async fn sqlite_imported_timestamp_migration_rejects_non_integer_storage() { + let pool = SqlitePool::connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + super::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + query( + r#" +INSERT INTO global_models (id, name, created_at, updated_at) +VALUES ('timestamp-invalid', 'timestamp-invalid', 1.5, 1); +"#, + ) + .execute(&pool) + .await + .expect("non-integer timestamp fixture should insert"); + + let migration = super::sqlite::MIGRATOR + .iter() + .find(|migration| migration.version == 20260725000000) + .expect("timestamp normalization migration should be embedded"); + let err = sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect_err("non-integer timestamp should fail the migration"); + assert!(err + .to_string() + .contains("imported_timestamp_storage_must_be_integer")); +} + +#[tokio::test] +async fn sqlite_remaining_timestamp_migration_repairs_other_repository_domains() { + let pool = SqlitePool::connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + super::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + query( + r#" +INSERT INTO users (id, email, username, auth_source, created_at, updated_at) +VALUES ('timestamp-user', 'timestamp@example.com', 'timestamp-user', 'local', 1, 1); + +INSERT INTO audit_logs (id, event_type, description, created_at) +VALUES ('timestamp-audit', 'test', 'test', '1970-01-01T00:00:01Z'); + +INSERT INTO request_candidates ( + id, request_id, candidate_index, status, created_at, started_at, finished_at +) VALUES ( + 'timestamp-candidate', 'timestamp-request', 0, 'success', + '1970-01-01T00:00:02Z', '1970-01-01T00:00:03Z', '1970-01-01T00:00:04Z' +); + +INSERT INTO stats_daily (id, date, created_at, updated_at) +VALUES ( + 'timestamp-stats', '1970-01-02', + '1970-01-01T00:00:05Z', '1970-01-01T00:00:06Z' +); + +INSERT INTO user_sessions ( + id, user_id, client_device_id, refresh_token_hash, + last_seen_at, expires_at, created_at, updated_at +) VALUES ( + 'timestamp-session', 'timestamp-user', 'device', 'hash', + '1970-01-01T00:00:07Z', '1970-01-01T00:00:08Z', + '1970-01-01T00:00:09Z', '1970-01-01T00:00:10Z' +); +"#, + ) + .execute(&pool) + .await + .expect("remaining timestamp fixtures should insert"); + + let migration = super::sqlite::MIGRATOR + .iter() + .find(|migration| migration.version == 20260725040000) + .expect("remaining timestamp migration should be embedded"); + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect("remaining timestamp migration should apply"); + + let audit = sqlx::query_as::<_, (i64, String)>( + "SELECT created_at, typeof(created_at) FROM audit_logs WHERE id = 'timestamp-audit'", + ) + .fetch_one(&pool) + .await + .expect("normalized audit timestamp should load"); + assert_eq!(audit, (1, "integer".to_string())); + + let candidate = sqlx::query_as::<_, (i64, i64, i64, String, String, String)>( + r#" +SELECT created_at, started_at, finished_at, + typeof(created_at), typeof(started_at), typeof(finished_at) +FROM request_candidates +WHERE id = 'timestamp-candidate' +"#, + ) + .fetch_one(&pool) + .await + .expect("normalized candidate timestamps should load"); + assert_eq!( + candidate, + ( + 2, + 3, + 4, + "integer".to_string(), + "integer".to_string(), + "integer".to_string(), + ) + ); + + let stats = sqlx::query_as::<_, (i64, i64, i64, String, String, String)>( + r#" +SELECT date, created_at, updated_at, + typeof(date), typeof(created_at), typeof(updated_at) +FROM stats_daily +WHERE id = 'timestamp-stats' +"#, + ) + .fetch_one(&pool) + .await + .expect("normalized stats timestamps should load"); + assert_eq!( + stats, + ( + 86_400, + 5, + 6, + "integer".to_string(), + "integer".to_string(), + "integer".to_string(), + ) + ); + + let session = sqlx::query_as::<_, (i64, i64, i64, i64, String, String, String, String)>( + r#" +SELECT last_seen_at, expires_at, created_at, updated_at, + typeof(last_seen_at), typeof(expires_at), typeof(created_at), typeof(updated_at) +FROM user_sessions +WHERE id = 'timestamp-session' +"#, + ) + .fetch_one(&pool) + .await + .expect("normalized session timestamps should load"); + assert_eq!( + session, + ( + 7, + 8, + 9, + 10, + "integer".to_string(), + "integer".to_string(), + "integer".to_string(), + "integer".to_string(), + ) + ); +} + +#[tokio::test] +async fn sqlite_remaining_timestamp_migration_rejects_invalid_storage() { + let pool = SqlitePool::connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + super::run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + query( + r#" +INSERT INTO audit_logs (id, event_type, description, created_at) +VALUES ('timestamp-invalid-audit', 'test', 'test', 1.5); +"#, + ) + .execute(&pool) + .await + .expect("invalid timestamp fixture should insert"); + + let migration = super::sqlite::MIGRATOR + .iter() + .find(|migration| migration.version == 20260725040000) + .expect("remaining timestamp migration should be embedded"); + let err = sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect_err("invalid remaining timestamp should fail the migration"); + assert!(err.to_string().contains("invalid_count = 0")); +} + #[tokio::test] async fn endpoint_api_root_migration_moves_v1_from_stored_default_paths() { let pool = SqlitePool::connect("sqlite::memory:") diff --git a/crates/aether-data/runtime/src/repository/usage/mod.rs b/crates/aether-data/runtime/src/repository/usage/mod.rs index 97f5d3ce0..8b630b92b 100644 --- a/crates/aether-data/runtime/src/repository/usage/mod.rs +++ b/crates/aether-data/runtime/src/repository/usage/mod.rs @@ -1,358 +1,3 @@ -#[cfg(feature = "mysql")] -macro_rules! impl_materialized_usage_read_repository { - ($repository:ty) => { - #[async_trait::async_trait] - impl $crate::repository::usage::UsageReadRepository for $repository { - async fn find_by_id( - &self, - id: &str, - ) -> Result< - Option<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::find_by_id(&repository, id).await - } - - async fn list_by_ids( - &self, - ids: &[String], - ) -> Result< - Vec<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_by_ids(&repository, ids).await - } - - async fn find_by_request_id( - &self, - request_id: &str, - ) -> Result< - Option<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::find_by_request_id(&repository, request_id).await - } - - async fn resolve_body_ref( - &self, - body_ref: &str, - ) -> Result, $crate::DataLayerError> { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::resolve_body_ref(&repository, body_ref).await - } - - async fn list_usage_audits( - &self, - query: &$crate::repository::usage::UsageAuditListQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_usage_audits(&repository, query).await - } - - async fn count_usage_audits( - &self, - query: &$crate::repository::usage::UsageAuditListQuery, - ) -> Result { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::count_usage_audits(&repository, query).await - } - - async fn list_usage_audits_by_keyword_search( - &self, - query: &$crate::repository::usage::UsageAuditKeywordSearchQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_usage_audits_by_keyword_search(&repository, query).await - } - - async fn count_usage_audits_by_keyword_search( - &self, - query: &$crate::repository::usage::UsageAuditKeywordSearchQuery, - ) -> Result { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::count_usage_audits_by_keyword_search(&repository, query).await - } - - async fn aggregate_usage_audits( - &self, - query: &$crate::repository::usage::UsageAuditAggregationQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageAuditAggregation>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::aggregate_usage_audits(&repository, query).await - } - - async fn summarize_usage_audits( - &self, - query: &$crate::repository::usage::UsageAuditSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageAuditSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_audits(&repository, query).await - } - - async fn summarize_usage_totals_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, $crate::DataLayerError> - { - <$repository>::summarize_usage_totals_by_user_ids(self, user_ids).await - } - - async fn summarize_usage_cache_hit_summary( - &self, - query: &$crate::repository::usage::UsageCacheHitSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageCacheHitSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_cache_hit_summary(&repository, query).await - } - - async fn summarize_usage_settled_cost( - &self, - query: &$crate::repository::usage::UsageSettledCostSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageSettledCostSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_settled_cost(&repository, query).await - } - - async fn summarize_usage_cache_affinity_hit_summary( - &self, - query: &$crate::repository::usage::UsageCacheAffinityHitSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageCacheAffinityHitSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_cache_affinity_hit_summary(&repository, query).await - } - - async fn list_usage_cache_affinity_intervals( - &self, - query: &$crate::repository::usage::UsageCacheAffinityIntervalQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageCacheAffinityIntervalRow>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_usage_cache_affinity_intervals(&repository, query).await - } - - async fn summarize_dashboard_usage( - &self, - query: &$crate::repository::usage::UsageDashboardSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageDashboardSummary, - $crate::DataLayerError, - > { - <$repository>::summarize_dashboard_usage(self, query).await - } - - async fn list_dashboard_daily_breakdown( - &self, - query: &$crate::repository::usage::UsageDashboardDailyBreakdownQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageDashboardDailyBreakdownRow>, - $crate::DataLayerError, - > { - <$repository>::list_dashboard_daily_breakdown(self, query).await - } - - async fn summarize_dashboard_provider_counts( - &self, - query: &$crate::repository::usage::UsageDashboardProviderCountsQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageDashboardProviderCount>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_dashboard_provider_counts(&repository, query).await - } - - async fn summarize_usage_breakdown( - &self, - query: &$crate::repository::usage::UsageBreakdownSummaryQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageBreakdownSummaryRow>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_breakdown(&repository, query).await - } - - async fn count_monitoring_usage_errors( - &self, - query: &$crate::repository::usage::UsageMonitoringErrorCountQuery, - ) -> Result { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::count_monitoring_usage_errors(&repository, query).await - } - - async fn list_monitoring_usage_errors( - &self, - query: &$crate::repository::usage::UsageMonitoringErrorListQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_monitoring_usage_errors(&repository, query).await - } - - async fn summarize_usage_error_distribution( - &self, - query: &$crate::repository::usage::UsageErrorDistributionQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageErrorDistributionRow>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_error_distribution(&repository, query).await - } - - async fn summarize_usage_performance_percentiles( - &self, - query: &$crate::repository::usage::UsagePerformancePercentilesQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsagePerformancePercentilesRow>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_performance_percentiles(&repository, query).await - } - - async fn summarize_usage_provider_performance( - &self, - query: &$crate::repository::usage::UsageProviderPerformanceQuery, - ) -> Result< - $crate::repository::usage::StoredUsageProviderPerformance, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_provider_performance(&repository, query).await - } - - async fn summarize_usage_cost_savings( - &self, - query: &$crate::repository::usage::UsageCostSavingsSummaryQuery, - ) -> Result< - $crate::repository::usage::StoredUsageCostSavingsSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_cost_savings(&repository, query).await - } - - async fn summarize_usage_time_series( - &self, - query: &$crate::repository::usage::UsageTimeSeriesQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageTimeSeriesBucket>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_time_series(&repository, query).await - } - - async fn summarize_usage_leaderboard( - &self, - query: &$crate::repository::usage::UsageLeaderboardQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageLeaderboardSummary>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_leaderboard(&repository, query).await - } - - async fn list_recent_usage_audits( - &self, - user_id: Option<&str>, - limit: usize, - ) -> Result< - Vec<$crate::repository::usage::StoredRequestUsageAudit>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_recent_usage_audits(&repository, user_id, limit).await - } - - async fn summarize_total_tokens_by_api_key_ids( - &self, - api_key_ids: &[String], - ) -> Result, $crate::DataLayerError> { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_total_tokens_by_api_key_ids(&repository, api_key_ids).await - } - - async fn summarize_usage_by_provider_api_key_ids( - &self, - provider_api_key_ids: &[String], - ) -> Result< - std::collections::BTreeMap< - String, - $crate::repository::usage::StoredProviderApiKeyUsageSummary, - >, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_by_provider_api_key_ids(&repository, provider_api_key_ids).await - } - - async fn summarize_usage_by_provider_api_key_windows( - &self, - requests: &[$crate::repository::usage::ProviderApiKeyWindowUsageRequest], - ) -> Result< - Vec<$crate::repository::usage::StoredProviderApiKeyWindowUsageSummary>, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_by_provider_api_key_windows(&repository, requests).await - } - - async fn summarize_provider_usage_since( - &self, - provider_id: &str, - since_unix_secs: u64, - ) -> Result< - $crate::repository::usage::StoredProviderUsageSummary, - $crate::DataLayerError, - > { - let repository = self.materialize_read_model().await?; - <$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_provider_usage_since(&repository, provider_id, since_unix_secs).await - } - - async fn summarize_usage_daily_heatmap( - &self, - query: &$crate::repository::usage::UsageDailyHeatmapQuery, - ) -> Result< - Vec<$crate::repository::usage::StoredUsageDailySummary>, - $crate::DataLayerError, - > { - <$repository>::summarize_usage_daily_heatmap(self, query).await - } - } - }; -} - mod memory; #[cfg(feature = "mysql")] mod mysql; diff --git a/crates/aether-data/runtime/src/repository/usage/mysql.rs b/crates/aether-data/runtime/src/repository/usage/mysql.rs index cb57ad207..74bd34c29 100644 --- a/crates/aether-data/runtime/src/repository/usage/mysql.rs +++ b/crates/aether-data/runtime/src/repository/usage/mysql.rs @@ -1,7 +1,23 @@ +use std::collections::BTreeMap; + use aether_data_contracts::repository::usage::{ - StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardSummary, - StoredUsageUserTotals, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, - UsageDashboardSummaryQuery, UsageReadRepository, + ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary, + StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit, + StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, + StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, + StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary, + StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, + StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, + StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, + StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, + UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, + UsageAuditSummaryQuery, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, + UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery, + UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, + UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardQuery, + UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, + UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, + UsageSettledCostSummaryQuery, UsageTimeSeriesQuery, }; use super::InMemoryUsageReadRepository; @@ -9,30 +25,116 @@ use crate::driver::mysql::MysqlPool; use crate::DataLayerError; pub use aether_data_mysql::MysqlUsageWriteRepository; +use aether_data_mysql::{MysqlUsageReadFilter, MysqlUsageStorage}; #[derive(Debug, Clone)] pub struct MysqlUsageReadRepository { - storage: aether_data_mysql::MysqlUsageStorage, + storage: MysqlUsageStorage, } impl MysqlUsageReadRepository { pub fn new(pool: MysqlPool) -> Self { Self { - storage: aether_data_mysql::MysqlUsageStorage::new(pool), + storage: MysqlUsageStorage::new(pool), } } - async fn materialize_read_model(&self) -> Result { + async fn materialize_read_model( + &self, + filter: MysqlUsageReadFilter, + ) -> Result { Ok(InMemoryUsageReadRepository::seed( - self.storage.load_usage_records().await?, + self.storage.load_usage_records_in_range(&filter).await?, )) } - async fn summarize_usage_daily_heatmap( + fn range(created_from_unix_secs: u64, created_until_unix_secs: u64) -> MysqlUsageReadFilter { + MysqlUsageReadFilter::new(created_from_unix_secs, created_until_unix_secs) + } +} + +#[async_trait::async_trait] +impl UsageReadRepository for MysqlUsageReadRepository { + async fn find_by_id( &self, - query: &UsageDailyHeatmapQuery, - ) -> Result, DataLayerError> { - self.storage.summarize_usage_daily_heatmap(query).await + id: &str, + ) -> Result, DataLayerError> { + self.storage.find_by_id(id).await + } + + async fn list_by_ids( + &self, + ids: &[String], + ) -> Result, DataLayerError> { + self.storage.list_by_ids(ids).await + } + + async fn find_by_request_id( + &self, + request_id: &str, + ) -> Result, DataLayerError> { + self.storage.find_by_request_id(request_id).await + } + + async fn resolve_body_ref( + &self, + body_ref: &str, + ) -> Result, DataLayerError> { + self.storage.resolve_body_ref(body_ref).await + } + + async fn list_usage_audits( + &self, + query: &UsageAuditListQuery, + ) -> Result, DataLayerError> { + self.storage.list_usage_audits(query).await + } + + async fn count_usage_audits(&self, query: &UsageAuditListQuery) -> Result { + self.storage.count_usage_audits(query).await + } + + async fn list_usage_audits_by_keyword_search( + &self, + query: &UsageAuditKeywordSearchQuery, + ) -> Result, DataLayerError> { + self.storage + .list_usage_audits_by_keyword_search(query) + .await + } + + async fn count_usage_audits_by_keyword_search( + &self, + query: &UsageAuditKeywordSearchQuery, + ) -> Result { + self.storage + .count_usage_audits_by_keyword_search(query) + .await + } + + async fn aggregate_usage_audits( + &self, + query: &UsageAuditAggregationQuery, + ) -> Result, DataLayerError> { + let repository = self + .materialize_read_model(Self::range( + query.created_from_unix_secs, + query.created_until_unix_secs, + )) + .await?; + repository.aggregate_usage_audits(query).await + } + + async fn summarize_usage_audits( + &self, + query: &UsageAuditSummaryQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_provider_name(query.provider_name.as_deref()) + .with_model(query.model.as_deref()); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_audits(query).await } async fn summarize_usage_totals_by_user_ids( @@ -44,10 +146,60 @@ impl MysqlUsageReadRepository { .await } + async fn summarize_usage_cache_hit_summary( + &self, + query: &UsageCacheHitSummaryQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_cache_hit_summary(query).await + } + + async fn summarize_usage_settled_cost( + &self, + query: &UsageSettledCostSummaryQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_api_key_id(query.api_key_id.as_deref()); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_settled_cost(query).await + } + + async fn summarize_usage_cache_affinity_hit_summary( + &self, + query: &UsageCacheAffinityHitSummaryQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_api_key_id(query.api_key_id.as_deref()) + .completed_only(); + let repository = self.materialize_read_model(filter).await?; + repository + .summarize_usage_cache_affinity_hit_summary(query) + .await + } + + async fn list_usage_cache_affinity_intervals( + &self, + query: &UsageCacheAffinityIntervalQuery, + ) -> Result, DataLayerError> { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_api_key_id(query.api_key_id.as_deref()) + .completed_only(); + let repository = self.materialize_read_model(filter).await?; + repository.list_usage_cache_affinity_intervals(query).await + } + async fn summarize_dashboard_usage( &self, query: &UsageDashboardSummaryQuery, ) -> Result { + if query.created_from_unix_secs >= query.created_until_unix_secs { + return Ok(StoredUsageDashboardSummary::default()); + } if let Some(summary) = self .storage .summarize_dashboard_usage_from_daily_aggregates(query) @@ -55,7 +207,10 @@ impl MysqlUsageReadRepository { { return Ok(summary); } - let repository = self.materialize_read_model().await?; + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; repository.summarize_dashboard_usage(query).await } @@ -63,6 +218,9 @@ impl MysqlUsageReadRepository { &self, query: &UsageDashboardDailyBreakdownQuery, ) -> Result, DataLayerError> { + if query.created_from_unix_secs >= query.created_until_unix_secs { + return Ok(Vec::new()); + } let rows = self .storage .list_dashboard_daily_breakdown_from_daily_aggregates(query) @@ -70,9 +228,230 @@ impl MysqlUsageReadRepository { if !rows.is_empty() { return Ok(rows); } - let repository = self.materialize_read_model().await?; + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; repository.list_dashboard_daily_breakdown(query).await } + + async fn summarize_dashboard_provider_counts( + &self, + query: &UsageDashboardProviderCountsQuery, + ) -> Result, DataLayerError> { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_dashboard_provider_counts(query).await + } + + async fn summarize_usage_breakdown( + &self, + query: &UsageBreakdownSummaryQuery, + ) -> Result, DataLayerError> { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_provider_name(query.provider_name.as_deref()) + .with_model(query.model.as_deref()) + .with_api_format(query.api_format.as_deref()) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_breakdown(query).await + } + + async fn count_monitoring_usage_errors( + &self, + query: &UsageMonitoringErrorCountQuery, + ) -> Result { + self.storage.count_monitoring_usage_errors(query).await + } + + async fn list_monitoring_usage_errors( + &self, + query: &UsageMonitoringErrorListQuery, + ) -> Result, DataLayerError> { + self.storage.list_monitoring_usage_errors(query).await + } + + async fn summarize_usage_error_distribution( + &self, + query: &UsageErrorDistributionQuery, + ) -> Result, DataLayerError> { + let repository = self + .materialize_read_model(Self::range( + query.created_from_unix_secs, + query.created_until_unix_secs, + )) + .await?; + repository.summarize_usage_error_distribution(query).await + } + + async fn summarize_usage_performance_percentiles( + &self, + query: &UsagePerformancePercentilesQuery, + ) -> Result, DataLayerError> { + let repository = self + .materialize_read_model( + Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .completed_only(), + ) + .await?; + repository + .summarize_usage_performance_percentiles(query) + .await + } + + async fn summarize_usage_provider_performance( + &self, + query: &UsageProviderPerformanceQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_provider_id(query.provider_id.as_deref()) + .with_model(query.model.as_deref()) + .with_api_format(query.api_format.as_deref()) + .with_endpoint_kind(query.endpoint_kind.as_deref()) + .with_is_stream(query.is_stream) + .with_has_format_conversion(query.has_format_conversion) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_provider_performance(query).await + } + + async fn summarize_usage_cost_savings( + &self, + query: &UsageCostSavingsSummaryQuery, + ) -> Result { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_provider_name(query.provider_name.as_deref()) + .with_model(query.model.as_deref()); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_cost_savings(query).await + } + + async fn summarize_usage_time_series( + &self, + query: &UsageTimeSeriesQuery, + ) -> Result, DataLayerError> { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_provider_name(query.provider_name.as_deref()) + .with_model(query.model.as_deref()); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_time_series(query).await + } + + async fn summarize_usage_leaderboard( + &self, + query: &UsageLeaderboardQuery, + ) -> Result, DataLayerError> { + let filter = Self::range(query.created_from_unix_secs, query.created_until_unix_secs) + .with_user_id(query.user_id.as_deref()) + .with_provider_name(query.provider_name.as_deref()) + .with_model(query.model.as_deref()) + .finalized_only(); + let repository = self.materialize_read_model(filter).await?; + repository.summarize_usage_leaderboard(query).await + } + + async fn list_recent_usage_audits( + &self, + user_id: Option<&str>, + limit: usize, + ) -> Result, DataLayerError> { + self.storage.list_recent_usage_audits(user_id, limit).await + } + + async fn summarize_total_tokens_by_api_key_ids( + &self, + api_key_ids: &[String], + ) -> Result, DataLayerError> { + let repository = InMemoryUsageReadRepository::seed( + self.storage + .load_usage_records_by_api_key_ids(api_key_ids) + .await?, + ); + repository + .summarize_total_tokens_by_api_key_ids(api_key_ids) + .await + } + + async fn summarize_usage_by_provider_api_key_ids( + &self, + provider_api_key_ids: &[String], + ) -> Result, DataLayerError> { + let repository = InMemoryUsageReadRepository::seed( + self.storage + .load_usage_records_by_provider_api_key_ids(provider_api_key_ids) + .await?, + ); + repository + .summarize_usage_by_provider_api_key_ids(provider_api_key_ids) + .await + } + + async fn summarize_usage_by_provider_api_key_windows( + &self, + requests: &[ProviderApiKeyWindowUsageRequest], + ) -> Result, DataLayerError> { + let repository = InMemoryUsageReadRepository::seed( + self.storage + .load_usage_records_by_provider_api_key_windows(requests) + .await?, + ); + repository + .summarize_usage_by_provider_api_key_windows(requests) + .await + } + + async fn summarize_provider_usage_since( + &self, + provider_id: &str, + since_unix_secs: u64, + ) -> Result { + let repository = InMemoryUsageReadRepository::seed( + self.storage + .load_usage_records_for_provider_since(provider_id, since_unix_secs) + .await?, + ); + repository + .summarize_provider_usage_since(provider_id, since_unix_secs) + .await + } + + async fn summarize_usage_daily_heatmap( + &self, + query: &UsageDailyHeatmapQuery, + ) -> Result, DataLayerError> { + self.storage.summarize_usage_daily_heatmap(query).await + } + + async fn read_usage_counter_health( + &self, + ) -> Result + { + self.storage.read_usage_counter_health().await + } + + async fn read_usage_counter_pending_health( + &self, + ) -> Result< + aether_data_contracts::repository::usage::UsageCounterPendingHealthSnapshot, + DataLayerError, + > { + self.storage.read_usage_counter_pending_health().await + } } -impl_materialized_usage_read_repository!(MysqlUsageReadRepository); +#[cfg(test)] +mod tests { + #[test] + fn mysql_usage_reads_do_not_restore_the_unconditional_full_table_loader() { + let source = include_str!("mysql.rs"); + let forbidden = ["load_usage_", "records()"].concat(); + assert!(!source.contains(&forbidden)); + assert!(source.contains("load_usage_records_in_range")); + assert!(source.contains("MysqlUsageReadFilter::new")); + } +}