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")); + } +}