From 2281f2b754551c3a2cb3586486e9c733cb54ab60 Mon Sep 17 00:00:00 2001 From: elky Date: Mon, 7 Sep 2026 00:09:42 +0800 Subject: [PATCH] refactor(data): remove MySQL and SQLite support Use PostgreSQL as the only database backend across runtime, schema tooling, installation, Compose, and CI. Update regression tests and reject removed drivers explicitly. --- .env.example | 10 +- .github/workflows/rust-ci.yml | 167 +- Cargo.lock | 42 +- Cargo.toml | 4 - Makefile | 12 - README.md | 40 +- apps/aether-gateway/Cargo.toml | 3 +- .../src/control/auth/resolution.rs | 47 +- apps/aether-gateway/src/data/config.rs | 52 +- apps/aether-gateway/src/data/state/core.rs | 11 +- apps/aether-gateway/src/data/tests.rs | 34 - .../handlers/admin/request/system/import.rs | 11 +- .../handlers/public/support/user_me_usage.rs | 2 +- apps/aether-gateway/src/main.rs | 370 +- .../src/tests/architecture/admin_provider.rs | 6 +- .../src/tests/architecture/mod.rs | 28 +- .../src/tests/architecture/sql_and_data.rs | 854 +-- .../src/tests/architecture/workspace_tiers.rs | 133 +- crates/aether-data/adapters/mysql/Cargo.toml | 24 - .../migrations/20260403000000_baseline.sql | 1051 --- ...20000_add_management_token_permissions.sql | 3 - .../20260508000000_add_proxy_node_metrics.sql | 42 - .../20260509000000_add_background_tasks.sql | 37 - .../20260509120000_add_user_groups.sql | 103 - ...0000_normalize_empty_user_policy_modes.sql | 50 - ...exclude_admins_from_default_user_group.sql | 12 - .../20260511130000_add_pool_member_scores.sql | 29 - .../20260512000000_add_feature_settings.sql | 5 - ...260512090000_add_payment_gateway_plans.sql | 79 - ..._add_billing_plan_purchase_limit_scope.sql | 2 - .../20260516000000_add_routing_profiles.sql | 41 - ...0260518000000_add_usage_counter_deltas.sql | 32 - ...ferrals_privacy_required_announcements.sql | 62 - ...0519120000_add_oauth_provider_icon_url.sql | 1 - ...ovider_api_key_default_pool_sort_index.sql | 15 - .../20260520000000_add_api_key_ip_rules.sql | 2 - ...10000_preserve_usage_provider_identity.sql | 10 - ...24000000_add_scheduler_hotspot_indexes.sql | 31 - ...00_add_provider_key_name_hotspot_index.sql | 15 - ...528000000_normalize_endpoint_api_roots.sql | 132 - ...528020000_add_gateway_pressure_indexes.sql | 26 - ...5010000_add_cross_driver_schema_parity.sql | 230 - ...260725020000_add_advanced_stats_parity.sql | 381 - ...0000_add_usage_routing_snapshot_parity.sql | 69 - ...0000_repair_routing_default_uniqueness.sql | 115 - ...000_cleanup_duplicate_worker_boot_runs.sql | 22 - ...0814000000_add_usage_cost_reservations.sql | 35 - ...815000000_add_usage_request_admissions.sql | 22 - ...sage_cost_reservation_user_foreign_key.sql | 12 - ...age_request_admission_user_foreign_key.sql | 10 - ..._migrate_legacy_codex_live_permissions.sql | 203 - ...force_payment_gateway_order_uniqueness.sql | 55 - ...260821130000_add_user_security_version.sql | 5 - ...40000_expand_proxy_password_ciphertext.sql | 3 - ...7050000_anonymize_deleted_user_history.sql | 93 - ...31000000_enforce_ldap_config_singleton.sql | 13 - ...10000_add_proxy_node_tunnel_generation.sql | 9 - ...enforce_proxy_node_endpoint_uniqueness.sql | 10 - ..._usage_counter_delta_tunnel_generation.sql | 2 - ...903000000_add_routing_group_sort_order.sql | 3 - .../adapters/mysql/src/announcements.rs | 365 - .../aether-data/adapters/mysql/src/audit.rs | 277 - crates/aether-data/adapters/mysql/src/auth.rs | 1498 ---- .../adapters/mysql/src/auth_modules.rs | 439 -- .../adapters/mysql/src/background_tasks.rs | 450 -- .../aether-data/adapters/mysql/src/billing.rs | 1716 ----- .../adapters/mysql/src/candidate_selection.rs | 1405 ---- .../adapters/mysql/src/candidates.rs | 1227 --- .../aether-data/adapters/mysql/src/error.rs | 11 - .../mysql/src/gemini_file_mappings.rs | 581 -- .../adapters/mysql/src/global_models.rs | 1394 ---- crates/aether-data/adapters/mysql/src/lib.rs | 53 - .../adapters/mysql/src/management_tokens.rs | 786 -- .../adapters/mysql/src/migrations.rs | 274 - .../adapters/mysql/src/oauth_providers.rs | 491 -- crates/aether-data/adapters/mysql/src/pool.rs | 184 - .../adapters/mysql/src/pool_scores.rs | 693 -- .../adapters/mysql/src/provider_catalog.rs | 3780 --------- .../adapters/mysql/src/proxy_nodes.rs | 2197 ------ .../aether-data/adapters/mysql/src/quota.rs | 174 - .../adapters/mysql/src/routing_profiles.rs | 518 -- .../adapters/mysql/src/settlement.rs | 1604 ---- .../aether-data/adapters/mysql/src/usage.rs | 1873 ----- .../adapters/mysql/src/usage/cleanup.rs | 865 --- .../adapters/mysql/src/usage/counters.rs | 1210 --- .../adapters/mysql/src/usage/http_capture.rs | 759 -- .../adapters/mysql/src/usage/read.rs | 852 --- .../adapters/mysql/src/usage/snapshots.rs | 1043 --- .../adapters/mysql/src/usage/tests.rs | 1338 ---- .../aether-data/adapters/mysql/src/users.rs | 3354 -------- .../adapters/mysql/src/video_tasks.rs | 1050 --- .../aether-data/adapters/mysql/src/wallet.rs | 6778 ----------------- .../adapters/mysql/src/wallet/tests.rs | 630 -- .../adapters/postgres/src/quota.rs | 12 +- crates/aether-data/adapters/sqlite/Cargo.toml | 24 - .../migrations/20260403000000_baseline.sql | 1074 --- ...20000_add_management_token_permissions.sql | 3 - .../20260508000000_add_proxy_node_metrics.sql | 46 - .../20260509000000_add_background_tasks.sql | 43 - .../20260509120000_add_user_groups.sql | 101 - ...0000_normalize_empty_user_policy_modes.sql | 29 - ...exclude_admins_from_default_user_group.sql | 14 - .../20260511130000_add_pool_member_scores.sql | 37 - .../20260512000000_add_feature_settings.sql | 2 - ...260512090000_add_payment_gateway_plans.sql | 85 - ..._add_billing_plan_purchase_limit_scope.sql | 2 - .../20260516000000_add_routing_profiles.sql | 45 - ...0260518000000_add_usage_counter_deltas.sql | 34 - ...ferrals_privacy_required_announcements.sql | 69 - ...0519120000_add_oauth_provider_icon_url.sql | 1 - ...ovider_api_key_default_pool_sort_index.sql | 2 - .../20260520000000_add_api_key_ip_rules.sql | 1 - ...10000_preserve_usage_provider_identity.sql | 18 - ...24000000_add_scheduler_hotspot_indexes.sql | 16 - ...00_add_provider_key_name_hotspot_index.sql | 2 - ...528000000_normalize_endpoint_api_roots.sql | 177 - ...528020000_add_gateway_pressure_indexes.sql | 26 - ...0_normalize_imported_timestamp_storage.sql | 784 -- ...5010000_add_cross_driver_schema_parity.sql | 270 - ...260725020000_add_advanced_stats_parity.sql | 447 -- ...0000_add_usage_routing_snapshot_parity.sql | 70 - ...e_remaining_imported_timestamp_storage.sql | 821 -- ...0000_repair_routing_default_uniqueness.sql | 38 - ...000_cleanup_duplicate_worker_boot_runs.sql | 22 - ...0814000000_add_usage_cost_reservations.sql | 34 - ...815000000_add_usage_request_admissions.sql | 21 - ...000_add_usage_policy_user_foreign_keys.sql | 117 - ..._migrate_legacy_codex_live_permissions.sql | 190 - ...force_payment_gateway_order_uniqueness.sql | 19 - ...260821130000_add_user_security_version.sql | 5 - ...7050000_anonymize_deleted_user_history.sql | 157 - ...31000000_enforce_ldap_config_singleton.sql | 11 - ...10000_add_proxy_node_tunnel_generation.sql | 20 - ...enforce_proxy_node_endpoint_uniqueness.sql | 9 - ..._usage_counter_delta_tunnel_generation.sql | 2 - ...903000000_add_routing_group_sort_order.sql | 4 - .../adapters/sqlite/src/announcements.rs | 480 -- .../aether-data/adapters/sqlite/src/audit.rs | 405 - .../aether-data/adapters/sqlite/src/auth.rs | 2199 ------ .../adapters/sqlite/src/auth_modules.rs | 597 -- .../adapters/sqlite/src/background_tasks.rs | 537 -- .../adapters/sqlite/src/billing.rs | 2144 ------ .../sqlite/src/candidate_selection.rs | 1989 ----- .../adapters/sqlite/src/candidates.rs | 1294 ---- .../aether-data/adapters/sqlite/src/error.rs | 11 - .../sqlite/src/gemini_file_mappings.rs | 667 -- .../adapters/sqlite/src/global_models.rs | 1486 ---- crates/aether-data/adapters/sqlite/src/lib.rs | 75 - .../adapters/sqlite/src/management_tokens.rs | 1112 --- .../adapters/sqlite/src/migrations.rs | 995 --- .../adapters/sqlite/src/oauth_providers.rs | 653 -- .../aether-data/adapters/sqlite/src/pool.rs | 159 - .../adapters/sqlite/src/pool_scores.rs | 901 --- .../adapters/sqlite/src/provider_catalog.rs | 5131 ------------- .../adapters/sqlite/src/proxy_nodes.rs | 2931 ------- .../aether-data/adapters/sqlite/src/quota.rs | 217 - .../adapters/sqlite/src/routing_profiles.rs | 856 --- .../adapters/sqlite/src/settlement.rs | 2156 ------ .../aether-data/adapters/sqlite/src/usage.rs | 5260 ------------- .../adapters/sqlite/src/usage/cleanup.rs | 841 -- .../adapters/sqlite/src/usage/counters.rs | 1603 ---- .../adapters/sqlite/src/usage/http_capture.rs | 745 -- .../adapters/sqlite/src/usage/snapshots.rs | 1047 --- .../adapters/sqlite/src/usage/tests.rs | 2711 ------- .../aether-data/adapters/sqlite/src/users.rs | 5336 ------------- .../adapters/sqlite/src/video_tasks.rs | 1168 --- .../aether-data/adapters/sqlite/src/wallet.rs | 6698 ---------------- .../adapters/sqlite/src/wallet/tests.rs | 5478 ------------- crates/aether-data/contracts/src/database.rs | 99 +- crates/aether-data/contracts/src/lib.rs | 5 +- crates/aether-data/query/Cargo.toml | 2 +- crates/aether-data/query/src/lib.rs | 118 +- crates/aether-data/runtime/Cargo.toml | 6 +- crates/aether-data/runtime/README.md | 71 +- ...422120000_backfill_api_key_usage_stats.sql | 36 - ...20000_rebuild_global_model_usage_count.sql | 15 - ...12000_sync_legacy_enabled_active_flags.sql | 11 - ...0_rebuild_canonical_usage_token_totals.sql | 229 - .../runtime/backfills/mysql/README.md | 2 - ...422120000_backfill_api_key_usage_stats.sql | 28 - ...20000_rebuild_global_model_usage_count.sql | 9 - ...12000_sync_legacy_enabled_active_flags.sql | 11 - ...0_rebuild_canonical_usage_token_totals.sql | 225 - .../runtime/backfills/sqlite/README.md | 2 - crates/aether-data/runtime/schema/README.md | 21 +- .../runtime/schema/compose_schema.sh | 49 - .../drivers/mysql/baseline/001_identity.sql | 167 - .../mysql/baseline/002_provider_catalog.sql | 326 - .../mysql/baseline/003_auth_config.sql | 71 - .../mysql/baseline/004_proxy_nodes.sql | 39 - .../mysql/baseline/005_wallet_billing.sql | 183 - .../drivers/mysql/baseline/006_usage.sql | 265 - .../drivers/mysql/baseline/manifest.txt | 6 - .../drivers/sqlite/baseline/001_identity.sql | 167 - .../sqlite/baseline/002_provider_catalog.sql | 330 - .../sqlite/baseline/003_auth_config.sql | 71 - .../sqlite/baseline/004_proxy_nodes.sql | 39 - .../sqlite/baseline/005_wallet_billing.sql | 202 - .../drivers/sqlite/baseline/006_usage.sql | 265 - .../drivers/sqlite/baseline/manifest.txt | 6 - .../runtime/schema/generated/README.md | 2 +- .../generated/mysql/baseline/001_identity.sql | 222 - .../mysql/baseline/002_provider_catalog.sql | 431 -- .../mysql/baseline/003_auth_config.sql | 83 - .../mysql/baseline/004_proxy_nodes.sql | 88 - .../mysql/baseline/005_wallet_billing.sql | 330 - .../generated/mysql/baseline/006_usage.sql | 282 - .../generated/mysql/baseline/007_stats.sql | 599 -- .../mysql/baseline/008_background_tasks.sql | 42 - .../generated/mysql/baseline/manifest.txt | 11 - .../sqlite/baseline/001_identity.sql | 213 - .../sqlite/baseline/002_provider_catalog.sql | 415 - .../sqlite/baseline/003_auth_config.sql | 78 - .../sqlite/baseline/004_proxy_nodes.sql | 86 - .../sqlite/baseline/005_wallet_billing.sql | 315 - .../generated/sqlite/baseline/006_usage.sql | 274 - .../generated/sqlite/baseline/007_stats.sql | 572 -- .../sqlite/baseline/008_background_tasks.sql | 40 - .../generated/sqlite/baseline/manifest.txt | 11 - .../schema/logical/005_wallet_billing.toml | 3 - .../runtime/src/backend/maintenance.rs | 203 +- .../runtime/src/backend/maintenance/mysql.rs | 67 - .../runtime/src/backend/maintenance/sqlite.rs | 34 - crates/aether-data/runtime/src/backend/mod.rs | 219 +- .../aether-data/runtime/src/backend/mysql.rs | 703 -- .../aether-data/runtime/src/backend/read.rs | 163 +- .../runtime/src/backend/referrals.rs | 3232 +------- .../aether-data/runtime/src/backend/sqlite.rs | 1207 --- .../aether-data/runtime/src/backend/stats.rs | 4 - .../runtime/src/backend/stats/mysql.rs | 831 -- .../src/backend/stats/mysql/advanced.rs | 973 --- .../runtime/src/backend/stats/sqlite.rs | 871 --- .../src/backend/stats/sqlite/advanced.rs | 986 --- .../runtime/src/backend/stats_common.rs | 27 - .../aether-data/runtime/src/backend/system.rs | 27 - .../runtime/src/backend/system/mysql.rs | 1317 ---- .../runtime/src/backend/system/sqlite.rs | 1438 ---- .../aether-data/runtime/src/backend/wallet.rs | 41 - .../runtime/src/backend/wallet/mysql.rs | 155 - .../runtime/src/backend/wallet/sqlite.rs | 156 - .../aether-data/runtime/src/backend/write.rs | 142 +- crates/aether-data/runtime/src/config.rs | 7 +- crates/aether-data/runtime/src/driver/mod.rs | 7 +- .../aether-data/runtime/src/driver/mysql.rs | 3 - .../aether-data/runtime/src/driver/sqlite.rs | 3 - crates/aether-data/runtime/src/lib.rs | 6 +- .../runtime/src/lifecycle/backfill.rs | 16 +- .../runtime/src/lifecycle/backfill/mysql.rs | 263 - .../runtime/src/lifecycle/backfill/sqlite.rs | 265 - .../runtime/src/lifecycle/backfill/tests.rs | 590 +- .../runtime/src/lifecycle/export.rs | 963 +-- .../runtime/src/lifecycle/export/mysql.rs | 914 --- .../runtime/src/lifecycle/export/sqlite.rs | 772 -- .../runtime/src/lifecycle/export/tests.rs | 1581 +--- .../runtime/src/lifecycle/migrate.rs | 47 +- .../runtime/src/lifecycle/migrate/mysql.rs | 5 - .../runtime/src/lifecycle/migrate/postgres.rs | 2 +- .../runtime/src/lifecycle/migrate/sqlite.rs | 5 - .../runtime/src/lifecycle/migrate/tests.rs | 1091 +-- .../src/repository/announcements/mod.rs | 4 - .../runtime/src/repository/audit.rs | 4 - .../runtime/src/repository/auth/mod.rs | 4 - .../src/repository/auth_modules/mod.rs | 4 - .../src/repository/background_tasks/mod.rs | 4 - .../runtime/src/repository/billing/mod.rs | 4 - .../src/repository/candidate_selection/mod.rs | 4 - .../runtime/src/repository/candidates/mod.rs | 4 - .../repository/gemini_file_mappings/mod.rs | 12 - .../src/repository/global_models/mod.rs | 4 - .../src/repository/management_tokens/mod.rs | 4 - .../src/repository/oauth_providers/mod.rs | 4 - .../runtime/src/repository/pool_scores/mod.rs | 4 - .../src/repository/provider_catalog/mod.rs | 4 - .../runtime/src/repository/proxy_nodes/mod.rs | 4 - .../runtime/src/repository/quota/mod.rs | 4 - .../src/repository/routing_profiles/mod.rs | 4 - .../runtime/src/repository/settlement/mod.rs | 4 - .../runtime/src/repository/usage/mod.rs | 6 - .../runtime/src/repository/usage/mysql.rs | 457 -- .../runtime/src/repository/users/mod.rs | 4 - .../runtime/src/repository/video_tasks/mod.rs | 4 - .../runtime/src/repository/wallet/mod.rs | 4 - .../runtime/tests/public_entrypoints.rs | 4 +- .../schema/src/bin/aether-schema.rs | 6 +- crates/aether-data/schema/src/dialect/mod.rs | 66 +- .../aether-data/schema/src/dialect/mysql.rs | 117 - .../schema/src/dialect/postgres.rs | 5 +- .../aether-data/schema/src/dialect/sqlite.rs | 92 - crates/aether-data/schema/src/lib.rs | 76 +- .../tests/responses_websocket_e2e.rs | 62 +- docker-compose.release-local.yml | 48 +- docker-compose.single-node.yml | 77 +- docker-compose.yml | 27 +- generate_keys.sh | 4 - install.sh | 188 +- tests/compose_database_config_test.py | 68 + ...install_container_runtime_security_test.sh | 112 +- tests/release_supply_chain_test.sh | 4 +- 298 files changed, 793 insertions(+), 134804 deletions(-) delete mode 100644 crates/aether-data/adapters/mysql/Cargo.toml delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260403000000_baseline.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260507120000_add_management_token_permissions.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260508000000_add_proxy_node_metrics.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260509000000_add_background_tasks.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260509120000_add_user_groups.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260510120000_normalize_empty_user_policy_modes.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260511120000_exclude_admins_from_default_user_group.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260511130000_add_pool_member_scores.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260512000000_add_feature_settings.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260512090000_add_payment_gateway_plans.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260516000000_add_routing_profiles.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260518000000_add_usage_counter_deltas.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260519000000_add_referrals_privacy_required_announcements.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260519120000_add_oauth_provider_icon_url.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260520000000_add_api_key_ip_rules.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260520010000_preserve_usage_provider_identity.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260524000000_add_scheduler_hotspot_indexes.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260527000000_add_provider_key_name_hotspot_index.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260528000000_normalize_endpoint_api_roots.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260528020000_add_gateway_pressure_indexes.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260727000000_repair_routing_default_uniqueness.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260814000000_add_usage_cost_reservations.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260815000000_add_usage_request_admissions.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260816000000_add_usage_cost_reservation_user_foreign_key.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260817000000_add_usage_request_admission_user_foreign_key.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260821130000_add_user_security_version.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260827040000_expand_proxy_password_ciphertext.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260827050000_anonymize_deleted_user_history.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260831000000_enforce_ldap_config_singleton.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260831010000_add_proxy_node_tunnel_generation.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql delete mode 100644 crates/aether-data/adapters/mysql/migrations/20260903000000_add_routing_group_sort_order.sql delete mode 100644 crates/aether-data/adapters/mysql/src/announcements.rs delete mode 100644 crates/aether-data/adapters/mysql/src/audit.rs delete mode 100644 crates/aether-data/adapters/mysql/src/auth.rs delete mode 100644 crates/aether-data/adapters/mysql/src/auth_modules.rs delete mode 100644 crates/aether-data/adapters/mysql/src/background_tasks.rs delete mode 100644 crates/aether-data/adapters/mysql/src/billing.rs delete mode 100644 crates/aether-data/adapters/mysql/src/candidate_selection.rs delete mode 100644 crates/aether-data/adapters/mysql/src/candidates.rs delete mode 100644 crates/aether-data/adapters/mysql/src/error.rs delete mode 100644 crates/aether-data/adapters/mysql/src/gemini_file_mappings.rs delete mode 100644 crates/aether-data/adapters/mysql/src/global_models.rs delete mode 100644 crates/aether-data/adapters/mysql/src/lib.rs delete mode 100644 crates/aether-data/adapters/mysql/src/management_tokens.rs delete mode 100644 crates/aether-data/adapters/mysql/src/migrations.rs delete mode 100644 crates/aether-data/adapters/mysql/src/oauth_providers.rs delete mode 100644 crates/aether-data/adapters/mysql/src/pool.rs delete mode 100644 crates/aether-data/adapters/mysql/src/pool_scores.rs delete mode 100644 crates/aether-data/adapters/mysql/src/provider_catalog.rs delete mode 100644 crates/aether-data/adapters/mysql/src/proxy_nodes.rs delete mode 100644 crates/aether-data/adapters/mysql/src/quota.rs delete mode 100644 crates/aether-data/adapters/mysql/src/routing_profiles.rs delete mode 100644 crates/aether-data/adapters/mysql/src/settlement.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/cleanup.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/counters.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/http_capture.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/read.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/snapshots.rs delete mode 100644 crates/aether-data/adapters/mysql/src/usage/tests.rs delete mode 100644 crates/aether-data/adapters/mysql/src/users.rs delete mode 100644 crates/aether-data/adapters/mysql/src/video_tasks.rs delete mode 100644 crates/aether-data/adapters/mysql/src/wallet.rs delete mode 100644 crates/aether-data/adapters/mysql/src/wallet/tests.rs delete mode 100644 crates/aether-data/adapters/sqlite/Cargo.toml delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260403000000_baseline.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260507120000_add_management_token_permissions.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260508000000_add_proxy_node_metrics.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260509000000_add_background_tasks.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260509120000_add_user_groups.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260510120000_normalize_empty_user_policy_modes.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260511120000_exclude_admins_from_default_user_group.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260511130000_add_pool_member_scores.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260512000000_add_feature_settings.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260512090000_add_payment_gateway_plans.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260516000000_add_routing_profiles.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260518000000_add_usage_counter_deltas.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260519000000_add_referrals_privacy_required_announcements.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260519120000_add_oauth_provider_icon_url.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260520000000_add_api_key_ip_rules.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260520010000_preserve_usage_provider_identity.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260524000000_add_scheduler_hotspot_indexes.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260527000000_add_provider_key_name_hotspot_index.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260528000000_normalize_endpoint_api_roots.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260528020000_add_gateway_pressure_indexes.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260727000000_repair_routing_default_uniqueness.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260814000000_add_usage_cost_reservations.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260815000000_add_usage_request_admissions.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260816000000_add_usage_policy_user_foreign_keys.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260821130000_add_user_security_version.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260827050000_anonymize_deleted_user_history.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260831000000_enforce_ldap_config_singleton.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260831010000_add_proxy_node_tunnel_generation.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql delete mode 100644 crates/aether-data/adapters/sqlite/migrations/20260903000000_add_routing_group_sort_order.sql delete mode 100644 crates/aether-data/adapters/sqlite/src/announcements.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/audit.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/auth.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/auth_modules.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/background_tasks.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/billing.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/candidate_selection.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/candidates.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/error.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/gemini_file_mappings.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/global_models.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/lib.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/management_tokens.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/migrations.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/oauth_providers.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/pool.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/pool_scores.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/provider_catalog.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/proxy_nodes.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/quota.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/routing_profiles.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/settlement.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage/cleanup.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage/counters.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage/http_capture.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage/snapshots.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/usage/tests.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/users.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/video_tasks.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/wallet.rs delete mode 100644 crates/aether-data/adapters/sqlite/src/wallet/tests.rs delete mode 100644 crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql delete mode 100644 crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql delete mode 100644 crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql delete mode 100644 crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql delete mode 100644 crates/aether-data/runtime/backfills/mysql/README.md delete mode 100644 crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql delete mode 100644 crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql delete mode 100644 crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql delete mode 100644 crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql delete mode 100644 crates/aether-data/runtime/backfills/sqlite/README.md delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/001_identity.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/002_provider_catalog.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/003_auth_config.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/004_proxy_nodes.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/005_wallet_billing.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/006_usage.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/mysql/baseline/manifest.txt delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/001_identity.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/002_provider_catalog.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/003_auth_config.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/004_proxy_nodes.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/005_wallet_billing.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/006_usage.sql delete mode 100644 crates/aether-data/runtime/schema/drivers/sqlite/baseline/manifest.txt delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/002_provider_catalog.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/003_auth_config.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/004_proxy_nodes.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/005_wallet_billing.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/008_background_tasks.sql delete mode 100644 crates/aether-data/runtime/schema/generated/mysql/baseline/manifest.txt delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/002_provider_catalog.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/003_auth_config.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/004_proxy_nodes.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/005_wallet_billing.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/008_background_tasks.sql delete mode 100644 crates/aether-data/runtime/schema/generated/sqlite/baseline/manifest.txt delete mode 100644 crates/aether-data/runtime/src/backend/maintenance/mysql.rs delete mode 100644 crates/aether-data/runtime/src/backend/maintenance/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/backend/mysql.rs delete mode 100644 crates/aether-data/runtime/src/backend/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/backend/stats/mysql.rs delete mode 100644 crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs delete mode 100644 crates/aether-data/runtime/src/backend/stats/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs delete mode 100644 crates/aether-data/runtime/src/backend/stats_common.rs delete mode 100644 crates/aether-data/runtime/src/backend/system/mysql.rs delete mode 100644 crates/aether-data/runtime/src/backend/system/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/backend/wallet/mysql.rs delete mode 100644 crates/aether-data/runtime/src/backend/wallet/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/driver/mysql.rs delete mode 100644 crates/aether-data/runtime/src/driver/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/export/mysql.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/export/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/migrate/mysql.rs delete mode 100644 crates/aether-data/runtime/src/lifecycle/migrate/sqlite.rs delete mode 100644 crates/aether-data/runtime/src/repository/usage/mysql.rs delete mode 100644 crates/aether-data/schema/src/dialect/mysql.rs delete mode 100644 crates/aether-data/schema/src/dialect/sqlite.rs create mode 100644 tests/compose_database_config_test.py diff --git a/.env.example b/.env.example index c50116a15..e9aba1428 100644 --- a/.env.example +++ b/.env.example @@ -15,8 +15,8 @@ APP_PORT=8084 # APP_IMAGE=ghcr.io/fawney19/aether:beta # APP_IMAGE=ghcr.io/fawney19/aether:0.7.0-rc.1 -# Compose 应用容器的非 root 数字身份。Single Node 的 ./data 必须归该身份所有; -# install.sh 会自动写入安装用户的 UID/GID 并迁移旧数据目录。 +# Compose 应用容器的非 root 数字身份。 +# install.sh 会自动写入安装用户的 UID/GID。 AETHER_CONTAINER_UID=65532 AETHER_CONTAINER_GID=65532 @@ -35,7 +35,7 @@ RUST_LOG=aether_gateway=info # AUTH_REFRESH_COOKIE_SAMESITE=None # AUTH_REFRESH_COOKIE_SECURE=true -# 数据库配置 +# 数据库配置(仅支持 PostgreSQL) DB_HOST=localhost DB_PORT=5432 DB_USER=postgres @@ -47,10 +47,6 @@ REDIS_HOST=localhost REDIS_PORT=6379 REDIS_PASSWORD= -# 可选 MySQL profile 的应用用户与 root 密码 -MYSQL_PASSWORD= -MYSQL_ROOT_PASSWORD= - # JWT密钥(使用 ./generate_keys.sh 生成) # 用于用户登录 token 签名,更换后所有用户需重新登录 JWT_SECRET_KEY=change-this-to-a-secure-random-string diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index d5563c019..cdbd699a4 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -20,6 +20,9 @@ on: - "Dockerfile.app" - "docker-compose.yml" - "docker-compose.single-node.yml" + - "docker-compose.local.yml" + - "docker-compose.release-local.yml" + - "tests/compose_database_config_test.py" - "tests/install_*_test.sh" - "tests/deploy_*_test.sh" - "tests/update_*_test.sh" @@ -45,6 +48,9 @@ on: - "Dockerfile.app" - "docker-compose.yml" - "docker-compose.single-node.yml" + - "docker-compose.local.yml" + - "docker-compose.release-local.yml" + - "tests/compose_database_config_test.py" - "tests/install_*_test.sh" - "tests/deploy_*_test.sh" - "tests/update_*_test.sh" @@ -79,6 +85,7 @@ jobs: - name: Run installer and supply-chain fixtures shell: bash run: | + python3 tests/compose_database_config_test.py bash tests/deploy_state_safety_test.sh bash tests/install_archive_safety_test.sh bash tests/install_container_runtime_security_test.sh @@ -252,6 +259,9 @@ jobs: - name: Install nextest uses: taiki-e/install-action@d5f9268ff7620505a81ada10ddf18cdd72240185 # nextest + - name: Expose PostgreSQL test binaries + run: pg_config --bindir >> "$GITHUB_PATH" + - name: Test lib env: RUSTC_WRAPPER: sccache @@ -299,10 +309,14 @@ jobs: - name: Install nextest uses: taiki-e/install-action@d5f9268ff7620505a81ada10ddf18cdd72240185 # nextest + - name: Expose PostgreSQL test binaries + run: pg_config --bindir >> "$GITHUB_PATH" + - name: Test env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: "true" + AETHER_REQUIRE_LOCAL_POSTGRES_TESTS: "true" run: cargo nextest run -p aether-data - name: Show sccache stats @@ -320,8 +334,6 @@ jobs: matrix: feature: - postgres - - mysql - - sqlite - all-drivers steps: - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5 @@ -396,8 +408,6 @@ jobs: matrix: package: - aether-data-postgres - - aether-data-mysql - - aether-data-sqlite steps: - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5 @@ -447,6 +457,9 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 + - name: Expose PostgreSQL test binaries + run: pg_config --bindir >> "$GITHUB_PATH" + - name: Test scenario binaries and end-to-end suites env: RUSTC_WRAPPER: sccache @@ -484,40 +497,6 @@ jobs: exit 1 fi - data_db_smoke_sqlite: - name: Data DB Smoke (SQLite) - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5 - - - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - - name: Show Rust toolchain - run: rustup show active-toolchain - - - name: Rust cache - uses: Swatinem/rust-cache@49a0bdc70d2e1b713ca9e2869b211fcce03d3c1c # v2 - with: - shared-key: rust-ci-${{ runner.os }} - workspaces: . -> target - - - name: Setup sccache - uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - - - name: Run SQLite data smoke tests - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - run: cargo test -p aether-data --all-features sqlite --lib - - - name: Show sccache stats - if: always() - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - run: sccache --show-stats - data_db_smoke_postgres: name: Data DB Smoke (Postgres) runs-on: ubuntu-latest @@ -589,112 +568,6 @@ jobs: AETHER_TEST_POSTGRES_URL: postgres://aether:aether@127.0.0.1:5432/aether_test run: cargo test -p aether-data --all-features postgres_core_export_reads_migrated_database_rows_when_url_is_set --lib -- --nocapture - - name: Run SQLite-to-Postgres import smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_POSTGRES_URL: postgres://aether:aether@127.0.0.1:5432/aether_test - run: cargo test -p aether-data --all-features sqlite_core_export_reads_migrated_database_rows --lib -- --nocapture - - - name: Show sccache stats - if: always() - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - run: sccache --show-stats - - data_db_smoke_mysql: - name: Data DB Smoke (MySQL) - runs-on: ubuntu-latest - services: - mysql: - image: mysql:8.0 - env: - MYSQL_DATABASE: aether_test - MYSQL_USER: aether - MYSQL_PASSWORD: aether - MYSQL_ROOT_PASSWORD: aether_root - ports: - - 3306:3306 - options: >- - --health-cmd="mysqladmin ping -h 127.0.0.1 -uaether -paether --silent" - --health-interval=5s - --health-timeout=5s - --health-retries=20 - steps: - - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5 - - - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable - - - name: Show Rust toolchain - run: rustup show active-toolchain - - - name: Rust cache - uses: Swatinem/rust-cache@49a0bdc70d2e1b713ca9e2869b211fcce03d3c1c # v2 - with: - shared-key: rust-ci-${{ runner.os }} - workspaces: . -> target - - - name: Setup sccache - uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad # v0.0.9 - - - name: Run MySQL migration smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data --all-features mysql_migrations_create_core_config_tables_when_url_is_set --lib -- --nocapture - - - name: Run MySQL usage write smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data-mysql mysql_usage_write_repository_upserts_and_flushes_counters_when_url_is_set --lib -- --nocapture - - - name: Run MySQL usage read smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data-mysql mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set --lib -- --nocapture - - - name: Run MySQL provider catalog smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data-mysql mysql_provider_catalog_repository_round_trips_when_url_is_set --lib -- --nocapture - - - name: Run MySQL core export smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data --all-features mysql_core_export_reads_migrated_database_rows_when_url_is_set --lib -- --nocapture - - - name: Run MySQL wallet read smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data-mysql mysql_wallet_read_repository_reads_wallet_contract_views --lib -- --nocapture - - - name: Run MySQL wallet daily usage aggregation smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data --all-features mysql_wallet_daily_usage_aggregation_uses_settlement_wallets_when_url_is_set --lib -- --nocapture - - - name: Run MySQL stats aggregation smoke test - env: - RUSTC_WRAPPER: sccache - SCCACHE_GHA_ENABLED: "true" - AETHER_TEST_MYSQL_URL: mysql://aether:aether@127.0.0.1:3306/aether_test - run: cargo test -p aether-data --all-features mysql_stats_aggregation_runs_after_mysql_migrations_when_url_is_set --lib -- --nocapture - - name: Show sccache stats if: always() env: @@ -706,16 +579,12 @@ jobs: name: Data DB Smoke runs-on: ubuntu-latest needs: - - data_db_smoke_sqlite - data_db_smoke_postgres - - data_db_smoke_mysql if: ${{ always() }} steps: - name: Verify database smoke jobs run: | - if [ "${{ needs.data_db_smoke_sqlite.result }}" != "success" ] || \ - [ "${{ needs.data_db_smoke_postgres.result }}" != "success" ] || \ - [ "${{ needs.data_db_smoke_mysql.result }}" != "success" ]; then + if [ "${{ needs.data_db_smoke_postgres.result }}" != "success" ]; then echo "Data DB smoke failed" exit 1 fi diff --git a/Cargo.lock b/Cargo.lock index 95e1f8529..e3568af0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -166,9 +166,7 @@ dependencies = [ "aether-ai-formats", "aether-cache", "aether-data-contracts", - "aether-data-mysql", "aether-data-postgres", - "aether-data-sqlite", "aether-wallet", "async-trait", "chrono", @@ -207,25 +205,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "aether-data-mysql" -version = "0.1.0" -dependencies = [ - "aether-ai-formats", - "aether-data-contracts", - "aether-data-query", - "async-trait", - "chrono", - "chrono-tz", - "flate2", - "serde_json", - "sha2", - "sqlx", - "tokio", - "tracing", - "uuid", -] - [[package]] name = "aether-data-postgres" version = "0.1.0" @@ -263,25 +242,6 @@ dependencies = [ "toml", ] -[[package]] -name = "aether-data-sqlite" -version = "0.1.0" -dependencies = [ - "aether-ai-formats", - "aether-data-contracts", - "aether-data-query", - "async-trait", - "chrono", - "chrono-tz", - "flate2", - "serde_json", - "sha2", - "sqlx", - "tokio", - "tracing", - "uuid", -] - [[package]] name = "aether-dispatch-core" version = "0.1.0" @@ -325,6 +285,7 @@ dependencies = [ "aether-scheduler-core", "aether-task-runtime", "aether-test-support", + "aether-testkit", "aether-usage-runtime", "aether-video-tasks-core", "aether-wallet", @@ -3038,7 +2999,6 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" dependencies = [ - "cc", "pkg-config", "vcpkg", ] diff --git a/Cargo.toml b/Cargo.toml index 7763dffe8..96272433b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,8 +11,6 @@ members = [ "crates/aether-routing-core", "crates/aether-data/contracts", "crates/aether-data/adapters/postgres", - "crates/aether-data/adapters/mysql", - "crates/aether-data/adapters/sqlite", "crates/aether-data/query", "crates/aether-data/schema", "crates/aether-dispatch-core", @@ -66,8 +64,6 @@ aether-provider-pool = { path = "crates/aether-provider/pool" } aether-routing-core = { path = "crates/aether-routing-core" } aether-data-contracts = { path = "crates/aether-data/contracts" } aether-data-postgres = { path = "crates/aether-data/adapters/postgres" } -aether-data-mysql = { path = "crates/aether-data/adapters/mysql" } -aether-data-sqlite = { path = "crates/aether-data/adapters/sqlite" } aether-data-query = { path = "crates/aether-data/query" } aether-data-schema = { path = "crates/aether-data/schema" } aether-dispatch-core = { path = "crates/aether-dispatch-core" } diff --git a/Makefile b/Makefile index 3c3de9d1b..c586e684a 100644 --- a/Makefile +++ b/Makefile @@ -36,15 +36,6 @@ lowercase() { printf '%s' "$$1" | tr '[:upper:]' '[:lower:]' } -dev_uses_sqlite_database() { - local driver - local url - driver="$$(lowercase "$${AETHER_DATABASE_DRIVER:-}")" - url="$${AETHER_DATABASE_URL:-$${DATABASE_URL:-}}" - - [[ "$${driver}" == "sqlite" || "$${url}" == sqlite:* ]] -} - dev_uses_postgres_database() { local driver local url @@ -68,9 +59,6 @@ dev_uses_redis_runtime() { if [[ "$${backend}" == "redis" ]]; then return 0 fi - if dev_uses_sqlite_database; then - return 1 - fi return 0 } diff --git a/README.md b/README.md index 425cd6a30..d59539842 100644 --- a/README.md +++ b/README.md @@ -46,27 +46,19 @@ cd Aether cp .env.example .env # .env 包含数据库、JWT 和数据加密密钥,先限制为仅当前用户可读写 chmod 600 .env -# 生成 JWT / 加密 / Postgres / Redis / MySQL 独立随机密钥,并填入 .env +# 生成 JWT / 加密 / Postgres / Redis 独立随机密钥,并填入 .env ./generate_keys.sh # 编辑 .env 设置 ADMIN_PASSWORD # 3. 首次部署 / 更新 (从以下部署形态任选其一) # Postgres + Redis (推荐) docker compose pull && docker compose up -d -# Single Node:默认容器身份为 65532:65532,先停止旧容器并检查/迁移 SQLite bind 目录 -docker compose -f docker-compose.single-node.yml stop app -mkdir -p data -test -z "$(find data ! -type d ! -type f -print -quit)" || { echo "data 中存在 symlink/FIFO/socket/device,拒绝迁移" >&2; exit 1; } -test -z "$(find data -type f -links +1 -print -quit)" || { echo "data 中存在硬链接,拒绝迁移" >&2; exit 1; } -sudo chown -R -P 65532:65532 ./data -sudo find data -type d -exec chmod 0700 {} + -sudo find data -type f -exec chmod 0600 {} + +# Single Node:同样使用 PostgreSQL + Redis,无需挂载本地数据库文件 docker compose -f docker-compose.single-node.yml pull && docker compose -f docker-compose.single-node.yml up -d ``` -应用镜像默认以固定非 root 身份 `65532:65532` 运行;Compose 同时移除全部 Linux capabilities、禁止提权、启用只读根文件系统,并仅提供带 `nosuid,nodev,noexec` 的 `/tmp` 临时文件系统。若宿主机不适合使用固定 UID/GID,可在 `.env` 中把 `AETHER_CONTAINER_UID` / `AETHER_CONTAINER_GID` 改成其他非零数字身份,并让 Single Node 的 `./data` 归该身份所有。`install.sh --mode compose-single-node` 会按安装用户自动生成这两个值;使用 `sudo` 运行时会采用原调用用户身份,并安全迁移已有 SQLite 数据。 +应用镜像默认以固定非 root 身份 `65532:65532` 运行;Compose 移除全部 Linux capabilities、禁止提权、启用只读根文件系统,并提供带 `nosuid,nodev,noexec` 的 `/tmp`。如需使用其他身份,可在 `.env` 中设置非零的 `AETHER_CONTAINER_UID` / `AETHER_CONTAINER_GID`。数据库使用独立 PostgreSQL 容器和 named volume,不再需要调整应用数据库目录的权限。 -从旧版 root 容器升级 Single Node 时,必须先停止旧 `app` 容器,再在第一次启动新版 Compose 前完成一次数据目录迁移;安装器检测到容器仍在运行会拒绝迁移,避免并发改写造成检查竞态。停止容器后用 `sudo` 重新执行一键安装器会自动处理;非 root 安装器发现旧数据所有权不匹配时会拒绝启动并提示迁移,不会放宽目录权限。手工部署且仍使用默认身份时执行上面的检查、`chown` 和 `find ... chmod` 命令即可。迁移只改变 `./data` 的所有权和权限,不会删除数据库、WAL 或备份文件。 ### 一键更新 @@ -82,6 +74,8 @@ Docker Compose 部署后,可在部署目录直接执行: ./update.sh --mode single-node ``` +现在仅支持 PostgreSQL。标准和单节点 Docker Compose 均部署 PostgreSQL + Redis;原生 systemd / launchd 安装需要显式提供 PostgreSQL `DATABASE_URL`,例如 `DATABASE_URL=postgresql://user:password@host:5432/aether`。旧数据库不会自动迁移或清空。升级时保留原有 PostgreSQL 密码、`JWT_SECRET_KEY` 和 `ENCRYPTION_KEY`,不要重新生成整个 `.env`。 + 仓库自带的 Docker Compose 默认把应用日志输出到容器 `stdout/stderr`,直接用 `docker compose logs -f app` 查看,并由 Docker 轮转日志,避免非 root 用户被宿主机日志目录权限拖垮启动。如果你确实需要文件日志,需要在 compose 里把 `AETHER_LOG_DESTINATION` 改成 `file|both`,额外挂载目录到 `/opt/aether/logs`,并让它归 `.env` 中配置的容器 UID/GID 所有;只读根文件系统不会阻止显式可写挂载。 管理后台右上角“版本信息”会检测新版本。Docker Compose 部署只提示版本,实际更新继续执行 `./update.sh`;systemd / launchd / 二进制部署才使用后台自更新,流程是下载对应平台的 GitHub Release 包、强制校验 `SHA256SUMS`、解压到 `/opt/aether/releases/`,再切换 `/opt/aether/current` 并退出进程,交给 systemd / launchd 拉起新版本。 @@ -96,13 +90,13 @@ gh attestation verify "aether-${TAG}-linux-amd64.tar.gz" \ --bundle AETHER_RELEASE_PROVENANCE.sigstore.json ``` -`docker-compose.yml` 中的官方 PostgreSQL、Redis 和 MySQL 镜像均固定到多架构 OCI index digest。升级这些依赖时应在发布变更中显式更新 digest,避免同名 tag 在无人审查的情况下改变部署内容。 +`docker-compose.yml` 中的官方 PostgreSQL 和 Redis 镜像均固定到多架构 OCI index digest。升级这些依赖时应在发布变更中显式更新 digest,避免同名 tag 在无人审查的情况下改变部署内容。 正式发布到 GHCR 和 Docker Hub 的多架构 Aether 镜像也带有同一 GitHub Actions OIDC / Sigstore provenance;生产 `Dockerfile.app` 的 BusyBox 与 Distroless 基础镜像同样固定到多架构 OCI index digest。 源码或本地构建版本不会启用后台在线更新,请继续使用源码更新流程。Docker Compose 用户如果希望“容器重建后也保持镜像层面的新版本”,仍建议定期运行 `./update.sh` 拉取并重建 app 镜像。服务器访问 GitHub 需要代理时,可设置 `AETHER_UPDATE_PROXY_URL`,也兼容 `UPDATE_PROXY_URL`、`HTTPS_PROXY`、`ALL_PROXY`、`HTTP_PROXY` 以及 `NO_PROXY`。共享出口触发 GitHub API 限流时,可设置只读 `AETHER_UPDATE_GITHUB_TOKEN`,也兼容 `GITHUB_TOKEN` / `GH_TOKEN`。下载总超时默认 600 秒,连续无响应/无数据默认 30 秒,可通过 `AETHER_UPDATE_DOWNLOAD_TIMEOUT_SECS` 和 `AETHER_UPDATE_DOWNLOAD_IDLE_TIMEOUT_SECS` 调整。 -标准 Docker Compose 使用 Docker named volumes 存放 Postgres/Redis/MySQL 数据;Single Node 使用部署目录下的 `./data` 存放 SQLite 数据。 +标准和 Single Node Docker Compose 均使用 Docker named volume 存放 PostgreSQL 数据。 如果是本地源码构建镜像的部署,继续使用: @@ -132,25 +126,27 @@ docker compose -f docker-compose.release-local.yml down -v - `AETHER_RELEASE_LOCAL_PORT`:本地联调端口,默认 `18085` - `LOCAL_RELEASE_APP_IMAGE`:本地联调镜像名,默认 `aether-app:release-local` -### 一键安装(默认 Single Node:Linux systemd / macOS launchd + SQLite) +### 一键安装(PostgreSQL + Redis) ```bash git clone https://github.com/fawney19/Aether.git cd Aether -curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash +curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash -s -- --mode compose ``` +原生 Linux systemd / macOS launchd 安装需先准备 PostgreSQL,将连接串通过 `DATABASE_URL` 传给安装进程,并选择 `--mode single-node`;不再自动创建本地数据库文件。 + ### Nightly(每日 main 构建) Nightly workflow 每天从 `main` 的固定 commit 构建并发布滚动的 GitHub Release `nightly`,同时推送多架构 GHCR 镜像 `ghcr.io/fawney19/aether:nightly`。Nightly 是预发布版本,适合验证最新代码,不保证与正式版相同的稳定性。滚动 Release 需要仓库保持关闭 GitHub Release immutability。 -安装最新 nightly(Linux systemd / macOS launchd + SQLite): +安装最新 nightly(PostgreSQL + Redis): ```bash -curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash -s -- --channel nightly +curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/main/install.sh | sudo bash -s -- --mode compose --channel nightly ``` -Docker Compose 用户可在部署目录的 `.env` 中设置 `APP_IMAGE=ghcr.io/fawney19/aether:nightly`,然后运行 `./update.sh` 获取下一次 nightly。二进制方式可重新执行上述安装命令升级;当前管理后台的在线更新列表只跟踪正式版/RC/Beta,不会自动提示下一次 nightly。 +Docker Compose 用户可在部署目录的 `.env` 中设置 `APP_IMAGE=ghcr.io/fawney19/aether:nightly`,然后运行 `./update.sh` 获取下一次 nightly。二进制部署请沿用已有 PostgreSQL 环境配置,并使用 `--mode single-node --channel nightly` 重新运行安装脚本升级;当前管理后台的在线更新列表只跟踪正式版/RC/Beta,不会自动提示下一次 nightly。 ## 本地开发 @@ -195,8 +191,8 @@ Aether Tunnel 是配套的正向代理节点,部署在海外 VPS 上,为墙 ## 环境变量 - `APP_PORT`:`aether-gateway` 唯一监听端口,固定绑定 `0.0.0.0:${APP_PORT}` -- `DATABASE_URL`:数据库连接串;SQLite 例如 `sqlite:///opt/aether/data/aether.db`,Postgres 例如 `postgresql://postgres:aether@postgres:5432/aether` -- `AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS` / `AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS`:数据库连接池手动覆盖值;未配置时 SQLite 固定 `1/1`,Postgres/MySQL 按每核 `4` 条自动推导,总池范围为 `32-100`。该预算按进程计算,多实例部署应按数据库连接上限显式分配 +- `DATABASE_URL`:PostgreSQL 连接串,例如 `postgresql://USER:PASSWORD@HOST:5432/aether` +- `AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS` / `AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS`:数据库连接池手动覆盖值;未配置时 PostgreSQL 按每核 `4` 条自动推导,总池范围为 `32-100`。该预算按进程计算,多实例部署应按数据库连接上限显式分配 - `AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS`:单实例请求并发上限;未配置时按 CPU 自动推导(基础范围 `512-65536`),低文件描述符预算时会进一步下调 - `AETHER_GATEWAY_REQUEST_BODY_BUFFER_BUDGET_MB`:单实例同时读取和解压请求体的加权内存预算,默认 `256MB` - `AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS`:可选的请求体完整读取超时;默认或显式设为 `0` 时关闭,非零值限制在 `1000-600000ms` @@ -209,7 +205,7 @@ Aether Tunnel 是配套的正向代理节点,部署在海外 VPS 上,为墙 - `AETHER_GATEWAY_SECURITY_CACHE_TTL_MS`:IP 黑白名单本地缓存时间,默认 `1000ms`,写操作会主动失效相关缓存 - `AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB`:PII 恢复同步响应缓冲上限,默认 `64MB`;显式设为 `0` 表示不再收紧默认值,但仍受 `256MB` 安全硬上限约束 - `REDIS_URL`:Redis 连接串;仅 Postgres + Redis 的 Docker Compose 部署需要配置 -- `AETHER_RUNTIME_BACKEND=memory|redis`:运行时缓存/协调后端。SQLite 默认用 `memory`,不会连接 Redis;多节点部署和需要跨 gateway 重启恢复 OpenAI Responses continuation history 的部署必须使用共享 Redis +- `AETHER_RUNTIME_BACKEND=memory|redis`:运行时缓存/协调后端。配置 Redis 时使用 `redis`,否则使用 `memory`;多节点部署和需要跨 gateway 重启恢复 OpenAI Responses continuation history 的部署必须使用共享 Redis - `AETHER_GATEWAY_DATABASE_MODE=auto|verify-only`:数据库启动策略,默认 `auto`,自动完成挂起的 schema migration 和 backfill;`verify-only` 仅检查并在数据库落后时拒绝启动 - `AETHER_GATEWAY_AUTO_PREPARE_DATABASE`:旧版兼容开关;新配置请使用 `AETHER_GATEWAY_DATABASE_MODE` - `JWT_SECRET_KEY` / `ENCRYPTION_KEY`:认证和敏感数据加密所需密钥 @@ -218,7 +214,7 @@ Aether Tunnel 是配套的正向代理节点,部署在海外 VPS 上,为墙 - `ADMIN_USERNAME` / `ADMIN_PASSWORD` / `ADMIN_EMAIL`:首次启动时自举首个本地管理员;`install.sh` 会提示输入管理员密码 - `CORS_ORIGINS` / `CORS_ALLOW_CREDENTIALS`:前端跨域来源控制;如果要跨域带登录 Cookie,`CORS_ORIGINS` 不能写 `*` - `RUST_LOG`:Rust 日志过滤,例如 `aether_gateway=info`、`aether_gateway=debug,sqlx=warn` -- `DB_PASSWORD` / `REDIS_PASSWORD` / `MYSQL_PASSWORD` / `MYSQL_ROOT_PASSWORD`:Docker Compose 后端密码,首次安装时分别随机生成;手工部署必须替换示例占位值,不要互相复用 +- `DB_PASSWORD` / `REDIS_PASSWORD`:Docker Compose 后端密码,首次安装时分别随机生成;手工部署必须替换示例占位值,不要互相复用 ### S3 备份离线恢复 diff --git a/apps/aether-gateway/Cargo.toml b/apps/aether-gateway/Cargo.toml index ce44803eb..f664afcff 100644 --- a/apps/aether-gateway/Cargo.toml +++ b/apps/aether-gateway/Cargo.toml @@ -81,7 +81,7 @@ sha1 = "0.10" sha2 = { workspace = true, features = ["oid"] } socket2.workspace = true tar.workspace = true -sqlx = { workspace = true, features = ["postgres", "mysql", "sqlite", "migrate"] } +sqlx = { workspace = true, features = ["postgres", "migrate"] } sysinfo = "0.32" thiserror.workspace = true tokio.workspace = true @@ -102,6 +102,7 @@ tikv-jemallocator = { version = "0.6", optional = true } tikv-jemalloc-sys = { version = "0.6", optional = true } [dev-dependencies] +aether-testkit = { workspace = true, features = ["postgres"] } aether-test-support.workspace = true aws-lc-rs.workspace = true tracing-subscriber.workspace = true diff --git a/apps/aether-gateway/src/control/auth/resolution.rs b/apps/aether-gateway/src/control/auth/resolution.rs index c16ac7a8a..51efbd5da 100644 --- a/apps/aether-gateway/src/control/auth/resolution.rs +++ b/apps/aether-gateway/src/control/auth/resolution.rs @@ -1499,7 +1499,6 @@ fn get_cached_auth_context_with_age( #[cfg(test)] mod tests { - use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -1595,40 +1594,24 @@ mod tests { .expect("endpoint should build") } - struct SqliteAuthConfigNodes { + struct PostgresAuthConfigNodes { first: AppState, second: AppState, - _database: TemporarySqliteDatabase, + _database: aether_testkit::ManagedPostgresServer, } - struct TemporarySqliteDatabase(PathBuf); - - impl Drop for TemporarySqliteDatabase { - fn drop(&mut self) { - for suffix in ["", "-shm", "-wal"] { - let mut path = self.0.as_os_str().to_os_string(); - path.push(suffix); - let _ = std::fs::remove_file(PathBuf::from(path)); - } - } - } - - async fn sqlite_auth_config_nodes( + async fn postgres_auth_config_nodes( auth_repository: Arc, - ) -> SqliteAuthConfigNodes { - let database_path = std::env::temp_dir().join(format!( - "aether-auth-config-{}.sqlite", - uuid::Uuid::new_v4() - )); + ) -> PostgresAuthConfigNodes { + let server = aether_testkit::ManagedPostgresServer::start() + .await + .expect("temporary PostgreSQL should start"); let mut pool = SqlPoolConfig::default(); pool.min_connections = 0; pool.max_connections = 4; - let database = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - format!("sqlite://{}", database_path.display()), - pool, - ) - .expect("sqlite config should build"); + let database = + SqlDatabaseConfig::new(DatabaseDriver::Postgres, server.database_url(), pool) + .expect("postgres config should build"); let config = GatewayDataConfig::from_database_config(database); let first_data = GatewayDataState::from_config(config.clone()) .expect("first data state should build") @@ -1637,27 +1620,27 @@ mod tests { assert!(first_data .run_database_migrations() .await - .expect("sqlite migrations should run")); + .expect("postgres migrations should run")); let second_data = GatewayDataState::from_config(config) .expect("second data state should build") .with_auth_api_key_reader(auth_repository) .without_wallet_reader_for_tests(); - SqliteAuthConfigNodes { + PostgresAuthConfigNodes { first: AppState::new() .expect("first app state should build") .with_data_state_for_tests(first_data), second: AppState::new() .expect("second app state should build") .with_data_state_for_tests(second_data), - _database: TemporarySqliteDatabase(database_path), + _database: server, } } #[tokio::test] async fn strong_system_config_read_bypasses_app_and_data_caches() { let nodes = - sqlite_auth_config_nodes(Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([]))) + postgres_auth_config_nodes(Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([]))) .await; let key = format!("test.auth.strong-read.{}", uuid::Uuid::new_v4()); let old_value = serde_json::json!({"version": "old"}); @@ -2627,7 +2610,7 @@ mod tests { let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( None, snapshot, )])); - let nodes = sqlite_auth_config_nodes(auth_repository.clone()).await; + let nodes = postgres_auth_config_nodes(auth_repository.clone()).await; nodes .first .auth_context_cache diff --git a/apps/aether-gateway/src/data/config.rs b/apps/aether-gateway/src/data/config.rs index c694b8c26..8ee136633 100644 --- a/apps/aether-gateway/src/data/config.rs +++ b/apps/aether-gateway/src/data/config.rs @@ -113,10 +113,7 @@ impl GatewayDataConfig { return (self.clone(), None); }; let total_max = database.pool.max_connections; - if total_max < 2 - || configured_background_max == Some(0) - || is_private_sqlite_memory_database(database) - { + if total_max < 2 || configured_background_max == Some(0) { return (self.clone(), None); } @@ -152,14 +149,9 @@ impl GatewayDataConfig { } } -fn is_private_sqlite_memory_database(database: &aether_data::SqlDatabaseConfig) -> bool { - database.driver == aether_data::DatabaseDriver::Sqlite - && matches!(database.url.trim(), "sqlite::memory:" | "sqlite://:memory:") -} - #[cfg(test)] mod tests { - use super::GatewayDataConfig; + use super::{GatewayDataConfig, PostgresPoolConfig}; use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; #[test] @@ -226,7 +218,10 @@ mod tests { #[test] fn runtime_pool_split_can_be_disabled_or_degrade_for_single_connection() { - let mut database = SqlDatabaseConfig::sqlite_default(); + let mut database = SqlDatabaseConfig::from_postgres_config(PostgresPoolConfig { + database_url: "postgres://localhost/aether".to_string(), + ..Default::default() + }); database.pool.max_connections = 1; let config = GatewayDataConfig::from_database_config(database); assert!(config @@ -234,7 +229,10 @@ mod tests { .1 .is_none()); - let mut database = SqlDatabaseConfig::sqlite_default(); + let mut database = SqlDatabaseConfig::from_postgres_config(PostgresPoolConfig { + database_url: "postgres://localhost/aether".to_string(), + ..Default::default() + }); database.pool.max_connections = 8; let config = GatewayDataConfig::from_database_config(database); assert!(config @@ -242,34 +240,4 @@ mod tests { .1 .is_none()); } - - #[test] - fn runtime_pool_split_keeps_private_sqlite_memory_database_in_one_pool() { - for url in ["sqlite::memory:", "sqlite://:memory:"] { - let config = GatewayDataConfig::from_database_config( - SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - url, - SqlPoolConfig { - min_connections: 1, - max_connections: 8, - ..SqlPoolConfig::default() - }, - ) - .expect("sqlite memory database config should be valid"), - ); - - let (foreground, background) = config.split_runtime_pools_with_background_max(Some(2)); - - assert!(background.is_none(), "private SQLite URL {url} was split"); - assert_eq!( - foreground - .database() - .expect("foreground database") - .pool - .max_connections, - 8 - ); - } - } } diff --git a/apps/aether-gateway/src/data/state/core.rs b/apps/aether-gateway/src/data/state/core.rs index c41e4593c..91ad12463 100644 --- a/apps/aether-gateway/src/data/state/core.rs +++ b/apps/aether-gateway/src/data/state/core.rs @@ -1116,10 +1116,7 @@ impl GatewayDataState { } fn database_driver_supports_usage_counter_flush(driver: Option) -> bool { - matches!( - driver, - Some(DatabaseDriver::Postgres | DatabaseDriver::Mysql | DatabaseDriver::Sqlite) - ) + matches!(driver, Some(DatabaseDriver::Postgres)) } #[cfg(test)] @@ -1132,12 +1129,6 @@ mod usage_counter_flush_backend_tests { 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/data/tests.rs b/apps/aether-gateway/src/data/tests.rs index de857c64a..aabe78d89 100644 --- a/apps/aether-gateway/src/data/tests.rs +++ b/apps/aether-gateway/src/data/tests.rs @@ -254,40 +254,6 @@ async fn app_state_wires_gateway_data_state_from_config() { assert!(state.data.has_video_task_reader()); } -#[tokio::test] -async fn app_state_prepares_sqlite_database_startup() -> Result<(), Box> { - let mut pool = SqlPoolConfig::default(); - pool.min_connections = 0; - pool.max_connections = 1; - let database = SqlDatabaseConfig::new(DatabaseDriver::Sqlite, "sqlite::memory:", pool)?; - let state = - AppState::new()?.with_data_config(GatewayDataConfig::from_database_config(database))?; - - let pending = state - .prepare_database_for_startup() - .await? - .expect("sqlite database should expose migration state"); - assert!( - !pending.is_empty(), - "fresh sqlite gateway databases should report pending migrations" - ); - - assert!( - state.run_database_migrations().await?, - "sqlite gateway database should run migrations" - ); - let pending = state - .prepare_database_for_startup() - .await? - .expect("sqlite database should expose migration state"); - assert!( - pending.is_empty(), - "sqlite gateway databases should be current after migrations" - ); - - Ok(()) -} - #[tokio::test] async fn data_state_checks_user_uniqueness_through_user_reader() { let user = StoredUserAuthRecord::new( diff --git a/apps/aether-gateway/src/handlers/admin/request/system/import.rs b/apps/aether-gateway/src/handlers/admin/request/system/import.rs index 0a9bf2759..a63e08ea8 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/import.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/import.rs @@ -8631,7 +8631,7 @@ enum WalletOwner<'a> { mod tests { use std::sync::Arc; - use aether_data::repository::pool_scores::SqlitePoolMemberScoreRepository; + use aether_data::repository::pool_scores::PostgresPoolMemberScoreRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogKey, StoredProviderCatalogProvider, @@ -9349,12 +9349,11 @@ mod tests { .await .expect("a disabled score writer remains an allowed no-op"); - let pool = sqlx::sqlite::SqlitePoolOptions::new() + let pool = sqlx::postgres::PgPoolOptions::new() .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("sqlite pool should connect"); - let score_repository = Arc::new(SqlitePoolMemberScoreRepository::new(pool.clone())); + .connect_lazy("postgres://localhost/aether") + .expect("postgres pool should build"); + let score_repository = Arc::new(PostgresPoolMemberScoreRepository::new(pool.clone())); pool.close().await; let app = AppState::new() .expect("app state should build") diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index 3cca499df..5ca327b2d 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -32,7 +32,7 @@ use super::{ const USERS_ME_USAGE_DATA_UNAVAILABLE_DETAIL: &str = "用户用量数据暂不可用"; // The active-usage endpoint accepts an explicit list of request IDs. Keep -// this list bounded before it reaches the repository layer: SQLite/MySQL +// this list bounded before it reaches the repository layer: the database // expand every value into a bind parameter, while PostgreSQL still has to // materialize the complete array. Request IDs are normally UUIDs, but a // generous per-item bound preserves compatibility with provider-generated diff --git a/apps/aether-gateway/src/main.rs b/apps/aether-gateway/src/main.rs index 571dee507..2341a493e 100644 --- a/apps/aether-gateway/src/main.rs +++ b/apps/aether-gateway/src/main.rs @@ -120,7 +120,7 @@ use aether_data::lifecycle::export::{ copy_database_records, export_database_jsonl, import_database_jsonl, DataCopyOptions, ExportDomain, MAX_JSONL_INPUT_BYTES, }; -use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL}; +use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; use aether_gateway::{ attach_static_frontend, build_router_with_state, prewarm_direct_h2c_sender_cache_from_env_for_startup, set_gateway_frontdoor_app_port, AppState, @@ -192,8 +192,6 @@ impl DeploymentTopologyArg { #[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)] enum DatabaseDriverArg { - Sqlite, - Mysql, Postgres, } @@ -270,8 +268,6 @@ impl From for ExportDomain { impl From for DatabaseDriver { fn from(value: DatabaseDriverArg) -> Self { match value { - DatabaseDriverArg::Sqlite => DatabaseDriver::Sqlite, - DatabaseDriverArg::Mysql => DatabaseDriver::Mysql, DatabaseDriverArg::Postgres => DatabaseDriver::Postgres, } } @@ -398,7 +394,6 @@ const DEFAULT_SQL_POOL_ACQUIRE_TIMEOUT_MS: u64 = 10_000; const DEFAULT_SQL_POOL_IDLE_TIMEOUT_MS: u64 = 30_000; const DEFAULT_SQL_POOL_MAX_LIFETIME_MS: u64 = 30 * 60_000; const DEFAULT_SQL_POOL_STATEMENT_CACHE_CAPACITY: usize = 100; -const DEFAULT_SQLITE_POOL_MAX_CONNECTIONS: u32 = 1; // Per-process default for server SQL backends. Keep this below common // database server max_connections defaults; operators can override with // AETHER_GATEWAY_DATA_POSTGRES_{MIN,MAX}_CONNECTIONS after sizing the DB. @@ -552,9 +547,6 @@ fn usage_queue_worker_database_cap( let Some(database) = database else { return MAX_USAGE_QUEUE_WORKERS; }; - if database.driver == DatabaseDriver::Sqlite { - return 1; - } let max_connections = database.pool.max_connections.max(1) as usize; // An isolated pool is already a dedicated background budget. Applying the shared-pool @@ -583,9 +575,6 @@ fn usage_worker_record_concurrency_database_cap( database_is_isolated: bool, ) -> Option { let database = database?; - if database.driver == DatabaseDriver::Sqlite { - return Some(1); - } let max_connections = database.pool.max_connections.max(1) as usize; // The isolated background pool has already been carved out of the foreground pool. Keep one @@ -662,8 +651,7 @@ fn automatic_sql_pool_config_for_parallelism( parallelism: u32, ) -> SqlPoolConfig { let (min_connections, max_connections) = match driver { - DatabaseDriver::Sqlite => (1, DEFAULT_SQLITE_POOL_MAX_CONNECTIONS), - DatabaseDriver::Mysql | DatabaseDriver::Postgres => { + DatabaseDriver::Postgres => { let cpu_count = parallelism.max(1); let max_connections = cpu_count .saturating_mul(AUTO_SERVER_SQL_POOL_CONNECTIONS_PER_CPU) @@ -854,7 +842,7 @@ impl GatewayDataArgs { statement_cache_capacity: self .postgres_statement_cache_capacity .unwrap_or(auto.statement_cache_capacity), - require_ssl: driver != DatabaseDriver::Sqlite && self.postgres_require_ssl, + require_ssl: self.postgres_require_ssl, } } @@ -931,10 +919,6 @@ fn resolve_database_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), } } @@ -1378,7 +1362,7 @@ struct DataCopyArgs { source_url: String, /// Permit a cleartext source connection for a non-loopback database. - /// Leave unset to require TLS for remote MySQL/Postgres URLs. + /// Leave unset to require TLS for remote Postgres URLs. #[arg(long)] source_allow_insecure: bool, @@ -1389,7 +1373,7 @@ struct DataCopyArgs { target_url: String, /// Permit a cleartext target connection for a non-loopback database. - /// Leave unset to require TLS for remote MySQL/Postgres URLs. + /// Leave unset to require TLS for remote Postgres URLs. #[arg(long)] target_allow_insecure: bool, @@ -1654,7 +1638,7 @@ impl Args { fn effective_runtime_backend( &self, - database: Option<&SqlDatabaseConfig>, + _database: Option<&SqlDatabaseConfig>, data_redis_url: Option<&str>, ) -> RuntimeBackendArg { if let Some(runtime_backend) = self.runtime_backend { @@ -1665,9 +1649,6 @@ impl Args { if matches!(self.deployment_topology, DeploymentTopologyArg::MultiNode) { return RuntimeBackendArg::Redis; } - if database.is_some_and(|database| database.driver == DatabaseDriver::Sqlite) { - return RuntimeBackendArg::Memory; - } if data_redis_url.is_some() { RuntimeBackendArg::Redis } else { @@ -2039,13 +2020,6 @@ fn validate_deployment_topology( )); } - if database.is_some_and(|database| database.driver == DatabaseDriver::Sqlite) { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "AETHER_DATABASE_DRIVER=sqlite is only valid for single-node deployment", - )); - } - if args .video_task_store_path .as_deref() @@ -3135,7 +3109,7 @@ fn copy_database_url_is_literal_loopback( ) -> Result { // Parse with the same SQLx driver that will open the pool. This preserves // query-parameter overrides such as PostgreSQL `host`/`hostaddr` and - // MySQL/PostgreSQL Unix `socket` paths, which URL authority inspection + // PostgreSQL Unix `socket` paths, which URL authority inspection // alone would miss. match driver { DatabaseDriver::Postgres => { @@ -3150,19 +3124,6 @@ fn copy_database_url_is_literal_loopback( Ok(options.get_socket().is_some() || copy_database_host_is_literal_loopback(options.get_host())) } - DatabaseDriver::Mysql => { - let options = url - .parse::() - .map_err(|error| { - io::Error::new( - io::ErrorKind::InvalidInput, - format!("{label} database URL is invalid: {error}"), - ) - })?; - Ok(options.get_socket().is_some() - || copy_database_host_is_literal_loopback(options.get_host())) - } - DatabaseDriver::Sqlite => Ok(false), } } @@ -3185,12 +3146,8 @@ fn copy_database_config( // while every named/remote SQL host defaults to an encrypted connection. // `allow_insecure` is deliberately endpoint-specific so a local source // does not silently downgrade a remote target (or vice versa). - let literal_loopback = if driver == DatabaseDriver::Sqlite { - false - } else { - copy_database_url_is_literal_loopback(driver, url, label)? - }; - let require_ssl = driver != DatabaseDriver::Sqlite && !allow_insecure && !literal_loopback; + let literal_loopback = copy_database_url_is_literal_loopback(driver, url, label)?; + let require_ssl = !allow_insecure && !literal_loopback; Ok(SqlDatabaseConfig::new( driver, url, @@ -3591,8 +3548,6 @@ mod tests { fn test_database(driver: DatabaseDriver, max_connections: u32) -> SqlDatabaseConfig { let url = match driver { - DatabaseDriver::Sqlite => "sqlite://./data/aether.db", - DatabaseDriver::Mysql => "mysql://root:root@localhost/aether", DatabaseDriver::Postgres => "postgres://postgres:postgres@localhost/aether", }; let max_connections = max_connections.max(1); @@ -3608,19 +3563,6 @@ mod tests { .expect("test database config should build") } - fn temporary_sqlite_args(label: &str) -> (Args, std::path::PathBuf) { - let mut args = test_args(); - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .expect("clock should be available") - .as_nanos(); - let database_path = - std::env::temp_dir().join(format!("aether-{label}-{}-{nonce}.db", std::process::id())); - args.data.database_driver = Some(DatabaseDriverArg::Sqlite); - args.data.database_url = Some(format!("sqlite://{}", database_path.display())); - (args, database_path) - } - #[test] fn resolves_healthcheck_url_from_app_port() { assert_eq!( @@ -3851,16 +3793,16 @@ mod tests { let before = Args::try_parse_from([ "aether-gateway", "--database-driver", - "sqlite", + "postgres", "--database-url", - "sqlite:///tmp/before.db", + "postgres://localhost/before", "db", "status", ]) .expect("database arguments before db should parse"); assert_eq!( before.data.database_url.as_deref(), - Some("sqlite:///tmp/before.db") + Some("postgres://localhost/before") ); let after = Args::try_parse_from([ @@ -3868,57 +3810,17 @@ mod tests { "db", "prepare", "--database-driver", - "sqlite", + "postgres", "--database-url", - "sqlite:///tmp/after.db", + "postgres://localhost/after", ]) .expect("database arguments after db prepare should parse"); assert_eq!( after.data.database_url.as_deref(), - Some("sqlite:///tmp/after.db") + Some("postgres://localhost/after") ); } - #[test] - fn gateway_data_pool_auto_sizes_sqlite_to_single_connection() { - let mut args = test_args(); - args.data.database_driver = Some(DatabaseDriverArg::Sqlite); - args.data.database_url = Some("sqlite://./data/aether.db".to_string()); - - let database = args - .data - .effective_sql_database_config() - .expect("sqlite database config should build"); - - assert_eq!(database.driver, DatabaseDriver::Sqlite); - assert_eq!(database.pool.min_connections, 1); - 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( @@ -4262,27 +4164,11 @@ mod tests { assert_eq!(workers, 10); } - #[test] - fn gateway_usage_queue_workers_auto_uses_single_worker_for_sqlite() { - let database = test_database(DatabaseDriver::Sqlite, 1); - - let workers = automatic_usage_queue_workers_for_parallelism( - 16, - NodeRoleArg::All, - Some(5_000), - None, - Some(&database), - false, - ); - - assert_eq!(workers, 1); - } - #[test] fn gateway_data_pool_explicit_values_override_auto_sizing() { let mut args = test_args(); - args.data.database_driver = Some(DatabaseDriverArg::Sqlite); - args.data.database_url = Some("sqlite://./data/aether.db".to_string()); + args.data.database_driver = Some(DatabaseDriverArg::Postgres); + args.data.database_url = Some("postgres://localhost/aether".to_string()); args.data.postgres_min_connections = Some(2); args.data.postgres_max_connections = Some(8); args.data.postgres_acquire_timeout_ms = Some(2_000); @@ -4290,7 +4176,7 @@ mod tests { let database = args .data .effective_sql_database_config() - .expect("sqlite database config should build"); + .expect("postgres database config should build"); assert_eq!(database.pool.min_connections, 2); assert_eq!(database.pool.max_connections, 8); @@ -4329,27 +4215,11 @@ mod tests { assert_eq!(database.pool.max_connections, 128); } - #[test] - fn sqlite_database_defaults_to_memory_runtime_backend() { - let args = test_args(); - let database = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - "sqlite://./data/aether.db".to_string(), - SqlPoolConfig::default(), - ) - .expect("sqlite config should build"); - - assert_eq!( - args.effective_runtime_backend(Some(&database), Some("redis://127.0.0.1/0")), - RuntimeBackendArg::Memory - ); - } - #[test] fn memory_runtime_data_config_keeps_redis_out_of_data_layer() { let mut args = test_args(); - args.data.database_driver = Some(DatabaseDriverArg::Sqlite); - args.data.database_url = Some("sqlite://./data/aether.db".to_string()); + args.data.database_driver = Some(DatabaseDriverArg::Postgres); + args.data.database_url = Some("postgres://localhost/aether".to_string()); args.data.redis_url = Some("redis://127.0.0.1/0".to_string()); let config = args.data.to_config(); @@ -4359,7 +4229,7 @@ mod tests { .database() .expect("database should be configured") .driver, - DatabaseDriver::Sqlite + DatabaseDriver::Postgres ); } @@ -4394,15 +4264,6 @@ mod tests { .expect("remote postgres config should build"); assert!(remote_postgres.pool.require_ssl); - let remote_mysql = copy_database_config( - DatabaseDriverArg::Mysql, - "mysql://user:pass@192.0.2.10:3306/aether", - "target", - false, - ) - .expect("remote mysql config should build"); - assert!(remote_mysql.pool.require_ssl); - for url in [ "postgres://user:pass@localhost/aether", "postgres://user:pass@127.42.17.9/aether", @@ -4425,15 +4286,6 @@ mod tests { ) .expect("explicit insecure opt-out should build"); assert!(!explicitly_insecure.pool.require_ssl); - - let sqlite = copy_database_config( - DatabaseDriverArg::Sqlite, - "sqlite://./data/aether.db", - "target", - false, - ) - .expect("sqlite config should build"); - assert!(!sqlite.pool.require_ssl); } #[test] @@ -4498,9 +4350,9 @@ mod tests { "postgres://user:pass@db.example/aether", "--source-allow-insecure", "--target-driver", - "mysql", + "postgres", "--target-url", - "mysql://user:pass@db.example/aether", + "postgres://user:pass@db.example/aether", ]) .expect("copy command should parse endpoint-specific TLS flags"); @@ -4677,82 +4529,6 @@ mod tests { ); } - #[test] - fn mysql_database_with_redis_defaults_to_redis_runtime_backend() { - let args = test_args(); - let database = SqlDatabaseConfig::new( - DatabaseDriver::Mysql, - "mysql://aether:aether@localhost:3306/aether".to_string(), - SqlPoolConfig::default(), - ) - .expect("mysql config should build"); - - assert_eq!( - args.effective_runtime_backend(Some(&database), Some("redis://127.0.0.1/0")), - RuntimeBackendArg::Redis - ); - } - - #[test] - fn sqlite_database_allows_explicit_redis_runtime_backend_when_redis_is_configured() { - let mut args = test_args(); - args.runtime_backend = Some(RuntimeBackendArg::Redis); - let database = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - "sqlite://./data/aether.db".to_string(), - SqlPoolConfig::default(), - ) - .expect("sqlite config should build"); - - assert_eq!( - args.effective_runtime_backend(Some(&database), Some("redis://127.0.0.1/0")), - RuntimeBackendArg::Redis - ); - super::validate_deployment_topology( - &args, - Some(&database), - Some("redis://127.0.0.1/0"), - RuntimeBackendArg::Redis, - ) - .expect("single-node sqlite should allow explicit redis runtime"); - } - - #[test] - fn single_node_sqlite_without_redis_allows_memory_runtime_backend() { - let args = test_args(); - let database = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - "sqlite://./data/aether.db".to_string(), - SqlPoolConfig::default(), - ) - .expect("sqlite config should build"); - - super::validate_deployment_topology( - &args, - Some(&database), - None, - RuntimeBackendArg::Memory, - ) - .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(); @@ -4799,29 +4575,6 @@ mod tests { assert!(error.to_string().contains("REDIS_URL")); } - #[test] - fn multi_node_rejects_sqlite_database_backend() { - let mut args = test_args(); - args.deployment_topology = DeploymentTopologyArg::MultiNode; - args.node_role = NodeRoleArg::Frontdoor; - let database = SqlDatabaseConfig::new( - DatabaseDriver::Sqlite, - "sqlite://./data/aether.db".to_string(), - SqlPoolConfig::default(), - ) - .expect("sqlite config should build"); - - let error = super::validate_deployment_topology( - &args, - Some(&database), - Some("redis://127.0.0.1/0"), - RuntimeBackendArg::Redis, - ) - .expect_err("multi-node sqlite should be rejected"); - assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); - assert!(error.to_string().contains("AETHER_DATABASE_DRIVER=sqlite")); - } - #[test] fn pending_schema_error_mentions_database_prepare_command() { let error = pending_schema_error(2, 20260413020000, "squash usage schema split"); @@ -4871,63 +4624,6 @@ mod tests { .expect("disabled data backend should not block startup"); } - #[tokio::test] - async fn verify_only_does_not_prepare_fresh_sqlite_database() { - let (args, database_path) = temporary_sqlite_args("verify-only"); - let state = AppState::new() - .expect("state should build") - .with_data_config(args.data.to_config()) - .expect("sqlite state should build"); - let pending_before = state - .pending_database_migrations() - .await - .expect("pending migrations should load") - .expect("sqlite should expose migration state"); - assert!(!pending_before.is_empty()); - - let error = - super::prepare_database_startup_requirements(&state, DatabaseModeArg::VerifyOnly) - .await - .expect_err("verify-only should reject a fresh database"); - assert!(error.to_string().contains("aether-gateway db prepare")); - - let pending_after = state - .pending_database_migrations() - .await - .expect("pending migrations should reload") - .expect("sqlite should expose migration state"); - assert_eq!(pending_after, pending_before); - drop(state); - let _ = std::fs::remove_file(database_path); - } - - #[tokio::test] - async fn auto_mode_prepares_fresh_sqlite_database() { - let (args, database_path) = temporary_sqlite_args("auto-prepare"); - let state = AppState::new() - .expect("state should build") - .with_data_config(args.data.to_config()) - .expect("sqlite state should build"); - - super::prepare_database_startup_requirements(&state, DatabaseModeArg::Auto) - .await - .expect("auto mode should prepare a fresh database"); - assert!(state - .pending_database_migrations() - .await - .expect("pending migrations should load") - .expect("sqlite should expose migration state") - .is_empty()); - assert!(state - .pending_database_backfills() - .await - .expect("pending backfills should load") - .expect("sqlite should expose backfill state") - .is_empty()); - drop(state); - let _ = std::fs::remove_file(database_path); - } - #[tokio::test] async fn database_prepare_requires_database_url() { let data = test_args().data; @@ -4974,26 +4670,6 @@ mod tests { assert!(message.contains("--apply-backfills")); } - #[tokio::test] - async fn explicit_backfills_are_noop_for_sqlite_database() { - let mut args = test_args(); - let database_path = std::env::temp_dir().join(format!( - "aether-sqlite-backfill-noop-{}-{}.db", - std::process::id(), - crate::current_unix_secs().expect("clock should be available") - )); - args.data.database_driver = Some(DatabaseDriverArg::Sqlite); - args.data.database_url = Some(format!("sqlite://{}", database_path.display())); - - super::run_explicit_migrations(&args) - .await - .expect("sqlite migrations should run before backfills"); - super::run_explicit_backfills(&args) - .await - .expect("sqlite backfills should be an explicit no-op"); - let _ = std::fs::remove_file(database_path); - } - #[tokio::test] async fn first_request_gate_closes_a_connection_that_never_reaches_service() { let gate = super::GatewayFirstRequestGate::new(); diff --git a/apps/aether-gateway/src/tests/architecture/admin_provider.rs b/apps/aether-gateway/src/tests/architecture/admin_provider.rs index 14039106c..9c27d423d 100644 --- a/apps/aether-gateway/src/tests/architecture/admin_provider.rs +++ b/apps/aether-gateway/src/tests/architecture/admin_provider.rs @@ -137,11 +137,7 @@ fn postgres_provider_cleanup_preserves_usage_history() { #[test] fn provider_cleanup_keeps_common_backends_in_sync() { - for path in [ - "crates/aether-data/adapters/postgres/src/provider_catalog.rs", - "crates/aether-data/adapters/mysql/src/provider_catalog.rs", - "crates/aether-data/adapters/sqlite/src/provider_catalog.rs", - ] { + for path in ["crates/aether-data/adapters/postgres/src/provider_catalog.rs"] { let source = read_workspace_file(path); for required in [ "UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id =", diff --git a/apps/aether-gateway/src/tests/architecture/mod.rs b/apps/aether-gateway/src/tests/architecture/mod.rs index 40398252d..374e1074b 100644 --- a/apps/aether-gateway/src/tests/architecture/mod.rs +++ b/apps/aether-gateway/src/tests/architecture/mod.rs @@ -36,7 +36,7 @@ pub(super) fn assert_no_sqlx_queries(root_relative_path: &str) { let source = fs::read_to_string(&path).expect("source file should be readable"); let hits = patterns .iter() - .filter(|pattern| source.contains(**pattern)) + .filter(|pattern| source_contains_sql_pattern(&source, pattern)) .copied() .collect::>(); if hits.is_empty() { @@ -54,6 +54,32 @@ pub(super) fn assert_no_sqlx_queries(root_relative_path: &str) { ); } +fn source_contains_sql_pattern(source: &str, pattern: &str) -> bool { + if pattern == "PostgresPool" { + source + .split(|character: char| !character.is_ascii_alphanumeric() && character != '_') + .any(|identifier| identifier == pattern) + } else { + source.contains(pattern) + } +} + +#[test] +fn sql_pool_scan_distinguishes_pool_types_from_repository_names() { + assert!(source_contains_sql_pattern( + "pool: PostgresPool", + "PostgresPool" + )); + assert!(source_contains_sql_pattern( + "PostgresPool::connect(url)", + "PostgresPool" + )); + assert!(!source_contains_sql_pattern( + "PostgresPoolMemberScoreRepository::new(pool)", + "PostgresPool" + )); +} + pub(super) fn assert_no_sensitive_log_patterns(root_relative_path: &str, patterns: &[&str]) { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join(root_relative_path); let mut files = Vec::new(); 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 b7d7baba2..781b6f93f 100644 --- a/apps/aether-gateway/src/tests/architecture/sql_and_data.rs +++ b/apps/aether-gateway/src/tests/architecture/sql_and_data.rs @@ -56,11 +56,7 @@ fn aether_data_bootstrap_snapshot_is_built_from_schema_sources() { #[test] fn aether_data_backend_pool_modules_do_not_own_maintenance_sql() { - for path in [ - "crates/aether-data/runtime/src/backend/postgres.rs", - "crates/aether-data/runtime/src/backend/mysql.rs", - "crates/aether-data/runtime/src/backend/sqlite.rs", - ] { + for path in ["crates/aether-data/runtime/src/backend/postgres.rs"] { let source = read_workspace_file(path); let production = production_source(&source); for forbidden in [ @@ -86,17 +82,9 @@ fn aether_data_backend_pool_modules_do_not_own_maintenance_sql() { let maintenance = read_workspace_file("crates/aether-data/runtime/src/backend/maintenance.rs"); for pattern in [ "Self::Postgres(postgres) => postgres.run_table_maintenance(table_names).await", - "Self::Mysql(mysql) => mysql.run_table_maintenance(table_names).await", - "Self::Sqlite(sqlite) => sqlite.run_table_maintenance(table_names).await", "Self::Postgres(postgres) => postgres.aggregate_wallet_daily_usage(input).await", - "Self::Mysql(mysql) => mysql.aggregate_wallet_daily_usage(input).await", - "Self::Sqlite(sqlite) => sqlite.aggregate_wallet_daily_usage(input).await", "Self::Postgres(postgres) => postgres.aggregate_stats_hourly(input).await", - "Self::Mysql(mysql) => mysql.aggregate_stats_hourly(input).await", - "Self::Sqlite(sqlite) => sqlite.aggregate_stats_hourly(input).await", "Self::Postgres(postgres) => postgres.aggregate_stats_daily(input).await", - "Self::Mysql(mysql) => mysql.aggregate_stats_daily(input).await", - "Self::Sqlite(sqlite) => sqlite.aggregate_stats_daily(input).await", ] { assert!( maintenance.contains(pattern), @@ -108,7 +96,7 @@ fn aether_data_backend_pool_modules_do_not_own_maintenance_sql() { #[test] fn wallet_maintenance_sql_is_partitioned_by_driver() { let facade = read_workspace_file("crates/aether-data/runtime/src/backend/wallet.rs"); - for module in ["mod postgres;", "mod mysql;", "mod sqlite;"] { + for module in ["mod postgres;"] { assert!( facade.contains(module), "wallet facade should declare {module}" @@ -117,8 +105,6 @@ fn wallet_maintenance_sql_is_partitioned_by_driver() { for forbidden in [ "sqlx::", "PostgresBackend", - "MysqlBackend", - "SqliteBackend", "SELECT ", "INSERT INTO ", "DELETE FROM ", @@ -129,11 +115,7 @@ fn wallet_maintenance_sql_is_partitioned_by_driver() { ); } - for (driver, backend) in [ - ("postgres", "PostgresBackend"), - ("mysql", "MysqlBackend"), - ("sqlite", "SqliteBackend"), - ] { + for (driver, backend) in [("postgres", "PostgresBackend")] { let path = format!("crates/aether-data/runtime/src/backend/wallet/{driver}.rs"); let source = read_workspace_file(&path); assert!(source.contains(&format!("impl {backend}"))); @@ -145,7 +127,7 @@ fn wallet_maintenance_sql_is_partitioned_by_driver() { #[test] fn table_maintenance_is_partitioned_for_each_driver() { let facade = read_workspace_file("crates/aether-data/runtime/src/backend/maintenance.rs"); - for module in ["mod postgres;", "mod mysql;", "mod sqlite;"] { + for module in ["mod postgres;"] { assert!( facade.contains(module), "maintenance facade should declare {module}" @@ -153,8 +135,6 @@ fn table_maintenance_is_partitioned_for_each_driver() { } for forbidden in [ "impl PostgresBackend", - "impl MysqlBackend", - "impl SqliteBackend", "VACUUM ANALYZE", "ANALYZE TABLE", "PRAGMA optimize", @@ -171,21 +151,12 @@ fn table_maintenance_is_partitioned_for_each_driver() { assert!(postgres.contains("impl PostgresBackend")); assert!(postgres.contains("postgres_observability_snapshot")); assert!(postgres.contains("VACUUM ANALYZE")); - - let mysql = read_workspace_file("crates/aether-data/runtime/src/backend/maintenance/mysql.rs"); - assert!(mysql.contains("impl MysqlBackend")); - assert!(mysql.contains("ANALYZE TABLE")); - - let sqlite = - read_workspace_file("crates/aether-data/runtime/src/backend/maintenance/sqlite.rs"); - assert!(sqlite.contains("impl SqliteBackend")); - assert!(sqlite.contains("PRAGMA optimize")); } #[test] fn system_driver_database_operations_are_partitioned() { let facade = read_workspace_file("crates/aether-data/runtime/src/backend/system.rs"); - for module in ["mod postgres;", "mod mysql;", "mod sqlite;"] { + for module in ["mod postgres;"] { assert!( facade.contains(module), "system facade should declare {module}" @@ -193,41 +164,17 @@ fn system_driver_database_operations_are_partitioned() { } for forbidden in [ "impl PostgresBackend", - "impl MysqlBackend", - "impl SqliteBackend", "fn map_postgres_stats_daily_aggregate(", "fn map_postgres_stats_user_daily_aggregate(", - "fn map_mysql_stats_daily_aggregate(", - "fn map_mysql_stats_user_daily_aggregate(", - "fn map_sqlite_stats_daily_aggregate(", - "fn map_sqlite_stats_user_daily_aggregate(", "async fn export_postgres_admin_system_usage_aggregates(", - "async fn export_mysql_admin_system_usage_aggregates(", - "async fn export_sqlite_admin_system_usage_aggregates(", "async fn import_postgres_admin_system_usage_aggregates(", - "async fn import_mysql_admin_system_usage_aggregates(", - "async fn import_sqlite_admin_system_usage_aggregates(", "async fn pg_delete_table(", "async fn pg_execute_if_table(", "async fn pg_execute_batch_if_table(", "async fn pg_table_exists(", - "async fn sqlite_delete_table(", - "async fn sqlite_execute_if_table(", - "async fn sqlite_execute_batch_if_table(", - "async fn sqlite_table_exists(", - "async fn mysql_delete_table(", - "async fn mysql_execute_if_table(", - "async fn mysql_execute_batch_if_table(", - "async fn mysql_table_exists(", "async fn purge_postgres_admin_system_data(", "async fn purge_postgres_non_admin_users(", "async fn purge_postgres_request_bodies_batch(", - "async fn purge_mysql_admin_system_data(", - "async fn purge_mysql_non_admin_users(", - "async fn purge_mysql_request_bodies_batch(", - "async fn purge_sqlite_admin_system_data(", - "async fn purge_sqlite_non_admin_users(", - "async fn purge_sqlite_request_bodies_batch(", ] { assert!( !facade.contains(forbidden), @@ -235,14 +182,13 @@ fn system_driver_database_operations_are_partitioned() { ); } - for driver in ["postgres", "mysql", "sqlite"] { + for driver in ["postgres"] { let source = read_workspace_file(&format!( "crates/aether-data/runtime/src/backend/system/{driver}.rs" )); let backend = match driver { "postgres" => "Postgres", - "mysql" => "Mysql", - "sqlite" => "Sqlite", + _ => unreachable!(), }; for required in [ @@ -275,38 +221,6 @@ fn system_driver_database_operations_are_partitioned() { ); } } - if driver == "sqlite" { - for required in [ - "sqlite_delete_table", - "sqlite_execute_if_table", - "sqlite_execute_batch_if_table", - "sqlite_table_exists", - "purge_sqlite_admin_system_data", - "purge_sqlite_non_admin_users", - "purge_sqlite_request_bodies_batch", - ] { - assert!( - source.contains(required), - "system/sqlite.rs should own {required}" - ); - } - } - if driver == "mysql" { - for required in [ - "mysql_delete_table", - "mysql_execute_if_table", - "mysql_execute_batch_if_table", - "mysql_table_exists", - "purge_mysql_admin_system_data", - "purge_mysql_non_admin_users", - "purge_mysql_request_bodies_batch", - ] { - assert!( - source.contains(required), - "system/mysql.rs should own {required}" - ); - } - } } } @@ -322,8 +236,6 @@ fn audit_repository_database_operations_are_partitioned() { for forbidden in [ "sqlx::", "PostgresPool", - "MysqlPool", - "SqlitePool", "impl AuditLogReadRepository", "SELECT ", "DELETE FROM ", @@ -347,7 +259,7 @@ fn audit_repository_database_operations_are_partitioned() { "aether-data-contracts audit.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !types.contains(forbidden), "aether-data-contracts audit.rs should remain driver-independent from {forbidden}" @@ -367,23 +279,11 @@ fn audit_repository_database_operations_are_partitioned() { ); } - for (driver, repository, row_mapper) in [ - ( - "postgres", - "PostgresAuditLogReadRepository", - "map_postgres_admin_audit_log_row", - ), - ( - "mysql", - "MysqlAuditLogReadRepository", - "map_mysql_admin_audit_log_row", - ), - ( - "sqlite", - "SqliteAuditLogReadRepository", - "map_sqlite_admin_audit_log_row", - ), - ] { + for (driver, repository, row_mapper) in [( + "postgres", + "PostgresAuditLogReadRepository", + "map_postgres_admin_audit_log_row", + )] { let source = read_workspace_file(&format!( "crates/aether-data/adapters/{driver}/src/audit.rs" )); @@ -411,14 +311,7 @@ fn audit_repository_database_operations_are_partitioned() { fn auth_module_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/auth_modules/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "auth module facade should not own contracts or driver code via {forbidden}" @@ -438,33 +331,19 @@ fn auth_module_repositories_are_owned_by_contracts_and_driver_adapters() { "aether-data-contracts auth_modules.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "auth module contracts should remain driver-independent from {forbidden}" ); } - for (feature, adapter, read_repository, write_repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxAuthModuleReadRepository", - "SqlxAuthModuleRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlAuthModuleReadRepository", - "MysqlAuthModuleRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteAuthModuleReadRepository", - "SqliteAuthModuleRepository", - ), - ] { + for (feature, adapter, read_repository, write_repository) in [( + "postgres", + "aether_data_postgres", + "SqlxAuthModuleReadRepository", + "SqlxAuthModuleRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "auth module facade should preserve the {feature} feature boundary" @@ -509,14 +388,7 @@ fn auth_module_repositories_are_owned_by_contracts_and_driver_adapters() { fn announcement_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/announcements/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "announcement facade should not own contracts or driver code via {forbidden}" @@ -538,26 +410,18 @@ fn announcement_repositories_are_owned_by_contracts_and_driver_adapters() { "aether-data-contracts announcements.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "announcement contracts should remain driver-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxAnnouncementReadRepository", - ), - ("mysql", "aether_data_mysql", "MysqlAnnouncementRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteAnnouncementRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxAnnouncementReadRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "announcement facade should preserve the {feature} feature boundary" @@ -599,14 +463,7 @@ fn announcement_repositories_are_owned_by_contracts_and_driver_adapters() { fn oauth_provider_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/oauth_providers/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "oauth provider facade should not own contracts or driver code via {forbidden}" @@ -628,26 +485,18 @@ fn oauth_provider_repositories_are_owned_by_contracts_and_driver_adapters() { "aether-data-contracts oauth_providers.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "oauth provider contracts should remain driver-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxOAuthProviderRepository", - ), - ("mysql", "aether_data_mysql", "MysqlOAuthProviderRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteOAuthProviderRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxOAuthProviderRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "oauth provider facade should preserve the {feature} feature boundary" @@ -691,14 +540,7 @@ fn oauth_provider_repositories_are_owned_by_contracts_and_driver_adapters() { fn management_token_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/management_tokens/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "management token facade should not own contracts or driver code via {forbidden}" @@ -720,30 +562,18 @@ fn management_token_repositories_are_owned_by_contracts_and_driver_adapters() { "aether-data-contracts management_tokens.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "management token contracts should remain driver-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxManagementTokenRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlManagementTokenRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteManagementTokenRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxManagementTokenRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "management token facade should preserve the {feature} feature boundary" @@ -788,13 +618,7 @@ fn gemini_file_mapping_repositories_are_owned_by_contracts_and_driver_adapters() let facade = read_workspace_file( "crates/aether-data/runtime/src/repository/gemini_file_mappings/mod.rs", ); - for forbidden in [ - "sqlx::", - "QueryBuilder", - "PgPool", - "MySqlPool", - "SqlitePool", - ] { + for forbidden in ["sqlx::", "QueryBuilder", "PgPool"] { assert!( !facade.contains(forbidden), "gemini file mapping facade should not own driver code via {forbidden}" @@ -819,30 +643,18 @@ fn gemini_file_mapping_repositories_are_owned_by_contracts_and_driver_adapters() "aether-data-contracts gemini_file_mappings.rs should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "gemini file mapping contracts should remain driver-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxGeminiFileMappingRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlGeminiFileMappingRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteGeminiFileMappingRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxGeminiFileMappingRepository", + )] { assert!( facade.contains(&format!("pub use {adapter}::{repository};")), "gemini file mapping facade should re-export {feature} repository from its adapter" @@ -881,13 +693,7 @@ fn gemini_file_mapping_repositories_are_owned_by_contracts_and_driver_adapters() fn video_task_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/video_tasks/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "video task facade should not own driver code via {forbidden}" @@ -908,23 +714,11 @@ fn video_task_repositories_are_owned_by_driver_adapters() { ); } - for (feature, adapter, repositories) in [ - ( - "postgres", - "aether_data_postgres", - vec!["SqlxVideoTaskReadRepository", "SqlxVideoTaskRepository"], - ), - ( - "mysql", - "aether_data_mysql", - vec!["MysqlVideoTaskRepository"], - ), - ( - "sqlite", - "aether_data_sqlite", - vec!["SqliteVideoTaskRepository"], - ), - ] { + for (feature, adapter, repositories) in [( + "postgres", + "aether_data_postgres", + vec!["SqlxVideoTaskReadRepository", "SqlxVideoTaskRepository"], + )] { for repository in repositories { assert!( facade.contains(repository) && facade.contains(adapter), @@ -952,15 +746,7 @@ fn video_task_repositories_are_owned_by_driver_adapters() { fn candidate_selection_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/candidate_selection/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "PgPool", - "MySqlPool", - "SqlitePool", - ] { + for forbidden in ["mod postgres;", "sqlx::", "PgPool"] { assert!( !facade.contains(forbidden), "candidate selection facade should not own driver code via {forbidden}" @@ -982,23 +768,11 @@ fn candidate_selection_repositories_are_owned_by_driver_adapters() { ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxMinimalCandidateSelectionReadRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlMinimalCandidateSelectionReadRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteMinimalCandidateSelectionReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxMinimalCandidateSelectionReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/candidate_selection.rs" @@ -1026,15 +800,7 @@ fn candidate_selection_repositories_are_owned_by_driver_adapters() { #[test] fn request_candidate_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/candidates/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "PgPool", - "MySqlPool", - "SqlitePool", - ] { + for forbidden in ["mod postgres;", "sqlx::", "PgPool"] { assert!( !facade.contains(forbidden), "request candidate facade should not own driver code via {forbidden}" @@ -1056,23 +822,11 @@ fn request_candidate_repositories_are_owned_by_driver_adapters() { ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxRequestCandidateReadRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlRequestCandidateRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteRequestCandidateRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxRequestCandidateReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/candidates.rs" @@ -1101,13 +855,7 @@ fn request_candidate_repositories_are_owned_by_driver_adapters() { #[test] fn billing_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/billing/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "billing facade should not own driver code via {forbidden}" @@ -1129,19 +877,11 @@ fn billing_repositories_are_owned_by_driver_adapters() { ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxBillingReadRepository", - ), - ("mysql", "aether_data_mysql", "MysqlBillingReadRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteBillingReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxBillingReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/billing.rs" @@ -1169,13 +909,7 @@ fn billing_repositories_are_owned_by_driver_adapters() { #[test] fn settlement_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/settlement/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "settlement facade should not own driver code via {forbidden}" @@ -1197,15 +931,11 @@ fn settlement_repositories_are_owned_by_contracts_and_driver_adapters() { ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxSettlementRepository", - ), - ("mysql", "aether_data_mysql", "MysqlSettlementRepository"), - ("sqlite", "aether_data_sqlite", "SqliteSettlementRepository"), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxSettlementRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/settlement.rs" @@ -1234,14 +964,7 @@ fn settlement_repositories_are_owned_by_contracts_and_driver_adapters() { fn proxy_node_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/proxy_nodes/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "proxy node facade should not own contracts or driver code via {forbidden}" @@ -1263,26 +986,18 @@ fn proxy_node_repositories_are_owned_by_contracts_and_driver_adapters() { "proxy node contracts should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool", "tracing::"] { + for forbidden in ["sqlx::", "PgPool", "tracing::"] { assert!( !contracts.contains(forbidden), "proxy node contracts should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxProxyNodeRepository", - ), - ("mysql", "aether_data_mysql", "MysqlProxyNodeReadRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteProxyNodeReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxProxyNodeRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/proxy_nodes.rs" @@ -1318,8 +1033,6 @@ fn global_model_repositories_are_owned_by_contracts_and_driver_adapters() { read_workspace_file("crates/aether-data/runtime/src/repository/global_models/mod.rs"); for forbidden in [ "mod postgres;", - "mod mysql;", - "mod sqlite;", "sqlx::", "QueryBuilder", "EMBEDDING_API_FORMATS", @@ -1356,30 +1069,18 @@ fn global_model_repositories_are_owned_by_contracts_and_driver_adapters() { "global model snapshot policy should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool", "RwLock"] { + for forbidden in ["sqlx::", "PgPool", "RwLock"] { assert!( !snapshot.contains(forbidden), "global model snapshot policy should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxGlobalModelReadRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlGlobalModelReadRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteGlobalModelReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxGlobalModelReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/global_models.rs" @@ -1408,14 +1109,7 @@ fn global_model_repositories_are_owned_by_contracts_and_driver_adapters() { #[test] fn auth_api_key_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/auth/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "auth facade should not own contracts or driver code via {forbidden}" @@ -1434,30 +1128,18 @@ fn auth_api_key_repositories_are_owned_by_contracts_and_driver_adapters() { "auth contracts should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "auth contracts should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxAuthApiKeySnapshotReadRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlAuthApiKeyReadRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteAuthApiKeyReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxAuthApiKeySnapshotReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/auth.rs" @@ -1490,14 +1172,7 @@ fn auth_api_key_repositories_are_owned_by_contracts_and_driver_adapters() { #[test] fn user_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/users/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "users facade should not own contracts or driver code via {forbidden}" @@ -1516,18 +1191,16 @@ fn user_repositories_are_owned_by_contracts_and_driver_adapters() { "user contracts should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["sqlx::", "PgPool"] { assert!( !contracts.contains(forbidden), "user contracts should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ("postgres", "aether_data_postgres", "SqlxUserReadRepository"), - ("mysql", "aether_data_mysql", "MysqlUserReadRepository"), - ("sqlite", "aether_data_sqlite", "SqliteUserReadRepository"), - ] { + for (feature, adapter, repository) in + [("postgres", "aether_data_postgres", "SqlxUserReadRepository")] + { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/users.rs" @@ -1560,13 +1233,7 @@ fn user_repositories_are_owned_by_contracts_and_driver_adapters() { fn provider_catalog_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/provider_catalog/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "provider catalog facade should not own driver code via {forbidden}" @@ -1600,30 +1267,18 @@ fn provider_catalog_repositories_are_owned_by_contracts_and_driver_adapters() { "provider catalog snapshot policy should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool", "RwLock"] { + for forbidden in ["sqlx::", "PgPool", "RwLock"] { assert!( !snapshot.contains(forbidden), "provider catalog snapshot should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxProviderCatalogReadRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlProviderCatalogReadRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteProviderCatalogReadRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxProviderCatalogReadRepository", + )] { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/provider_catalog.rs" @@ -1652,14 +1307,7 @@ fn provider_catalog_repositories_are_owned_by_contracts_and_driver_adapters() { #[test] fn wallet_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/wallet/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "mod types;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "wallet facade should not own contracts or driver code via {forbidden}" @@ -1691,18 +1339,16 @@ fn wallet_repositories_are_owned_by_contracts_and_driver_adapters() { "wallet snapshot policy should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool", "RwLock"] { + for forbidden in ["sqlx::", "PgPool", "RwLock"] { assert!( !snapshot.contains(forbidden), "wallet snapshot should remain infrastructure-independent from {forbidden}" ); } - for (feature, adapter, repository) in [ - ("postgres", "aether_data_postgres", "SqlxWalletRepository"), - ("mysql", "aether_data_mysql", "MysqlWalletReadRepository"), - ("sqlite", "aether_data_sqlite", "SqliteWalletReadRepository"), - ] { + for (feature, adapter, repository) in + [("postgres", "aether_data_postgres", "SqlxWalletRepository")] + { assert!(facade.contains(adapter) && facade.contains(repository)); let source = read_workspace_file(&format!( "crates/aether-data/adapters/{feature}/src/wallet.rs" @@ -1737,7 +1383,6 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/usage/mod.rs"); for forbidden in [ "mod postgres;", - "mod sqlite;", "sqlx::", "QueryBuilder", "pub(crate) struct ApiKeyUsageDelta", @@ -1748,17 +1393,6 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() { "usage facade should not own driver or pure policy code via {forbidden}" ); } - let mysql_facade = - read_workspace_file("crates/aether-data/runtime/src/repository/usage/mysql.rs"); - assert!(mysql_facade.contains("use aether_data_mysql::")); - assert!(mysql_facade.contains("MysqlUsageStorage")); - for forbidden in ["sqlx::query", "FROM `usage`", "INSERT INTO `usage`"] { - assert!( - !mysql_facade.contains(forbidden), - "MySQL usage facade should only adapt shared read policy, not own SQL via {forbidden}" - ); - } - let contracts = read_workspace_file("crates/aether-data/contracts/src/repository/usage/types.rs"); for required in [ @@ -1784,7 +1418,7 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() { "usage policy should own {required}" ); } - for forbidden in ["sqlx::", "PgPool", "MySqlPool", "SqlitePool", "RwLock"] { + for forbidden in ["sqlx::", "PgPool", "RwLock"] { assert!( !policy.contains(forbidden), "usage policy should remain infrastructure-independent from {forbidden}" @@ -1803,36 +1437,10 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() { "PostgreSQL usage adapter should own {required}" ); } - let mysql = read_workspace_file("crates/aether-data/adapters/mysql/src/usage.rs"); - for required in [ - "pub struct MysqlUsageStorage", - "pub struct MysqlUsageWriteRepository", - "impl UsageWriteRepository for MysqlUsageWriteRepository", - "sqlx::", - ] { - assert!( - mysql.contains(required), - "MySQL usage adapter should own {required}" - ); - } - let sqlite = read_workspace_file("crates/aether-data/adapters/sqlite/src/usage.rs"); - for required in [ - "pub struct SqliteUsageReadRepository", - "pub struct SqliteUsageWriteRepository", - "impl UsageReadRepository for SqliteUsageReadRepository", - "impl UsageWriteRepository for SqliteUsageWriteRepository", - "sqlx::", - ] { - assert!( - sqlite.contains(required), - "SQLite usage adapter should own {required}" - ); - } for path in [ "crates/aether-data/runtime/src/repository/usage/postgres/mod.rs", "crates/aether-data/runtime/src/repository/usage/postgres/cleanup.rs", - "crates/aether-data/runtime/src/repository/usage/sqlite.rs", ] { assert!( !workspace_file_exists(path), @@ -1853,26 +1461,13 @@ fn usage_repositories_are_owned_by_contracts_and_driver_adapters() { #[test] fn lifecycle_backfills_are_partitioned_by_driver() { let facade = read_workspace_file("crates/aether-data/runtime/src/lifecycle/backfill.rs"); - for module in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "mod tests;", - ] { + for module in ["mod postgres;", "mod types;", "mod tests;"] { assert!( facade.contains(module), "backfill facade should declare {module}" ); } - for forbidden in [ - "sqlx::", - "PgPool", - "MysqlPool", - "SqlitePool", - "BACKFILL_MIGRATOR", - "CREATE TABLE", - ] { + for forbidden in ["sqlx::", "PgPool", "BACKFILL_MIGRATOR", "CREATE TABLE"] { assert!( !facade.contains(forbidden), "backfill facade should not own driver operations via {forbidden}" @@ -1896,47 +1491,12 @@ fn lifecycle_backfills_are_partitioned_by_driver() { "backfill/postgres.rs should own {required}" ); } - - for (driver, pool) in [("mysql", "MysqlPool"), ("sqlite", "SqlitePool")] { - let source = read_workspace_file(&format!( - "crates/aether-data/runtime/src/lifecycle/backfill/{driver}.rs" - )); - 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", "PgConnection", "crate::driver::postgres"] { - assert!( - !source.contains(forbidden), - "backfill/{driver}.rs should not depend on PostgreSQL via {forbidden}" - ); - } - } } #[test] fn lifecycle_migrations_are_partitioned_by_driver() { let facade = read_workspace_file("crates/aether-data/runtime/src/lifecycle/migrate.rs"); - for module in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "mod types;", - "mod tests;", - ] { + for module in ["mod postgres;", "mod types;", "mod tests;"] { assert!( facade.contains(module), "migration facade should declare {module}" @@ -1961,7 +1521,7 @@ fn lifecycle_migrations_are_partitioned_by_driver() { types.contains(required), "migrate/types.rs should own {required}" ); - for forbidden in ["PgPool", "MySqlPool", "SqlitePool"] { + for forbidden in ["PgPool"] { assert!( !types.contains(forbidden), "migrate/types.rs should remain driver-independent from {forbidden}" @@ -1996,47 +1556,6 @@ fn lifecycle_migrations_are_partitioned_by_driver() { ); } - for (driver, pool, adapter_module, adapter_dir) in [ - ( - "mysql", - "MySqlPool", - "aether_data_mysql", - "aether-data-mysql", - ), - ( - "sqlite", - "SqlitePool", - "aether_data_sqlite", - "aether-data-sqlite", - ), - ] { - let source = read_workspace_file(&format!( - "crates/aether-data/runtime/src/lifecycle/migrate/{driver}.rs" - )); - assert!( - source.contains(&format!("pub(super) use {adapter_module}")), - "migrate/{driver}.rs should remain an adapter compatibility facade" - ); - assert!(!source.contains("sqlx::migrate!")); - let adapter_source = read_workspace_file(&format!( - "crates/aether-data/adapters/{driver}/src/migrations.rs" - )); - assert!(adapter_source.contains("sqlx::migrate!(\"./migrations\")")); - assert!(workspace_file_exists(&format!( - "crates/aether-data/adapters/{driver}/migrations/20260403000000_baseline.sql" - ))); - for required in [ - "sqlx::migrate!".to_string(), - pool.to_string(), - "pub async fn run_migrations".to_string(), - "pub async fn pending_migrations".to_string(), - ] { - assert!( - adapter_source.contains(&required), - "{adapter_dir}/src/migrations.rs should own {required}" - ); - } - } assert!( !workspace_file_exists("crates/aether-data/runtime/migrations"), "driver migration SQL must be owned by adapter crates" @@ -2047,25 +1566,17 @@ fn lifecycle_migrations_are_partitioned_by_driver() { fn routing_profile_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/routing_profiles/mod.rs"); - for forbidden in ["mod postgres;", "mod mysql;", "mod sqlite;", "sqlx::"] { + for forbidden in ["mod postgres;", "sqlx::"] { assert!( !facade.contains(forbidden), "routing profile facade should not own driver implementation via {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "PostgresRoutingGroupRepository", - ), - ("mysql", "aether_data_mysql", "MysqlRoutingGroupRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteRoutingGroupRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "PostgresRoutingGroupRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "routing profile facade should preserve the {feature} feature boundary" @@ -2096,31 +1607,17 @@ fn routing_profile_repositories_are_owned_by_driver_adapters() { #[test] fn provider_quota_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/quota/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "SelectQuery", - ] { + for forbidden in ["mod postgres;", "sqlx::", "SelectQuery"] { assert!( !facade.contains(forbidden), "provider quota facade should not own driver implementation via {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxProviderQuotaRepository", - ), - ("mysql", "aether_data_mysql", "MysqlProviderQuotaRepository"), - ( - "sqlite", - "aether_data_sqlite", - "SqliteProviderQuotaRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxProviderQuotaRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "provider quota facade should preserve the {feature} feature boundary" @@ -2152,35 +1649,17 @@ fn provider_quota_repositories_are_owned_by_driver_adapters() { fn pool_score_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/pool_scores/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "pool score facade should not own driver implementation via {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "PostgresPoolMemberScoreRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlPoolMemberScoreRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqlitePoolMemberScoreRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "PostgresPoolMemberScoreRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "pool score facade should preserve the {feature} feature boundary" @@ -2212,35 +1691,17 @@ fn pool_score_repositories_are_owned_by_driver_adapters() { fn background_task_repositories_are_owned_by_driver_adapters() { let facade = read_workspace_file("crates/aether-data/runtime/src/repository/background_tasks/mod.rs"); - for forbidden in [ - "mod postgres;", - "mod mysql;", - "mod sqlite;", - "sqlx::", - "QueryBuilder", - ] { + for forbidden in ["mod postgres;", "sqlx::", "QueryBuilder"] { assert!( !facade.contains(forbidden), "background task facade should not own driver implementation via {forbidden}" ); } - for (feature, adapter, repository) in [ - ( - "postgres", - "aether_data_postgres", - "SqlxBackgroundTaskRepository", - ), - ( - "mysql", - "aether_data_mysql", - "MysqlBackgroundTaskRepository", - ), - ( - "sqlite", - "aether_data_sqlite", - "SqliteBackgroundTaskRepository", - ), - ] { + for (feature, adapter, repository) in [( + "postgres", + "aether_data_postgres", + "SqlxBackgroundTaskRepository", + )] { assert!( facade.contains(&format!("#[cfg(feature = \"{feature}\")]")), "background task facade should preserve the {feature} feature boundary" @@ -2272,7 +1733,7 @@ fn background_task_repositories_are_owned_by_driver_adapters() { #[test] fn lifecycle_driver_export_operations_are_partitioned() { let facade = read_workspace_file("crates/aether-data/runtime/src/lifecycle/export.rs"); - for required in ["mod postgres;", "mod mysql;", "mod sqlite;", "mod tests;"] { + for required in ["mod postgres;", "mod tests;"] { assert!( facade.contains(required), "export facade should declare {required}" @@ -2280,11 +1741,7 @@ fn lifecycle_driver_export_operations_are_partitioned() { } assert!(!facade.contains("mod tests {")); - for (driver, implementation_marker) in [ - ("postgres", "fn normalize_postgres_import_payload("), - ("mysql", "fn mysql_row_payload("), - ("sqlite", "fn sqlite_row_payload("), - ] { + for (driver, implementation_marker) in [("postgres", "fn normalize_postgres_import_payload(")] { for required in [ format!("pub use {driver}::"), format!("export_{driver}_core_jsonl"), @@ -2330,9 +1787,7 @@ fn lifecycle_driver_export_operations_are_partitioned() { let tests = read_workspace_file("crates/aether-data/runtime/src/lifecycle/export/tests.rs"); assert!(tests.contains("jsonl_round_trips_manifest_and_domain_rows")); - assert!(tests.contains("postgres_import_payload_normalizes_sqlite_values_for_target_columns")); - assert!(tests.contains("sqlite_core_export_reads_migrated_database_rows")); - assert!(tests.contains("mysql_core_export_reads_migrated_database_rows_when_url_is_set")); + assert!(tests.contains("postgres_import_payload_normalizes_imported_values_for_target_columns")); } #[test] @@ -2375,12 +1830,7 @@ fn testkit_does_not_copy_aether_business_schema_sql() { #[test] fn gateway_main_keeps_database_export_import_driver_selection_in_data_layer() { let main_rs = read_workspace_file("apps/aether-gateway/src/main.rs"); - for forbidden in [ - "PostgresPoolFactory", - "MysqlPoolFactory", - "SqlitePoolFactory", - "to_postgres_config()", - ] { + for forbidden in ["PostgresPoolFactory", "to_postgres_config()"] { assert!( !main_rs.contains(forbidden), "main.rs should delegate database export/import driver selection to aether-data instead of {forbidden}" @@ -2455,11 +1905,7 @@ fn gateway_system_config_types_are_owned_by_aether_data() { } let data_backends = read_workspace_file("crates/aether-data/runtime/src/backend/maintenance.rs"); - for pattern in [ - "postgres.list_system_config_entries().await", - "mysql.list_system_config_entries().await", - "sqlite.list_system_config_entries().await", - ] { + for pattern in ["postgres.list_system_config_entries().await"] { assert!( data_backends.contains(pattern), "aether-data backends should own driver-specific system config dispatch {pattern}" diff --git a/apps/aether-gateway/src/tests/architecture/workspace_tiers.rs b/apps/aether-gateway/src/tests/architecture/workspace_tiers.rs index a9d4dd099..5f808e06f 100644 --- a/apps/aether-gateway/src/tests/architecture/workspace_tiers.rs +++ b/apps/aether-gateway/src/tests/architecture/workspace_tiers.rs @@ -40,58 +40,18 @@ fn pure_policy_crates_do_not_depend_on_runtime_adapters() { #[test] fn database_adapters_are_independent_driver_boundaries() { - let adapters = [ - ( - "crates/aether-data/adapters/postgres/Cargo.toml", - "features = [\"postgres\"", - ["features = [\"mysql\"", "features = [\"sqlite\""], - ), - ( - "crates/aether-data/adapters/mysql/Cargo.toml", - "features = [\"mysql\"", - ["features = [\"postgres\"", "features = [\"sqlite\""], - ), - ( - "crates/aether-data/adapters/sqlite/Cargo.toml", - "features = [\"sqlite\"", - ["features = [\"postgres\"", "features = [\"mysql\""], - ), - ]; - - for (manifest_path, expected_driver, other_drivers) in adapters { - let manifest = read_workspace_file(manifest_path); - assert!(manifest.contains("aether-data-contracts.workspace = true")); - assert!(manifest.contains(expected_driver)); - assert_manifest_excludes( - manifest_path, - &[other_drivers[0], other_drivers[1], "aether-gateway", "axum"], - ); - } + let path = "crates/aether-data/adapters/postgres/Cargo.toml"; + let manifest = read_workspace_file(path); + assert!(manifest.contains("aether-data-contracts.workspace = true")); + assert!(manifest.contains("features = [\"postgres\"")); + assert_manifest_excludes(path, &["aether-gateway", "axum"]); } #[test] fn data_facade_preserves_legacy_driver_paths_without_owning_driver_code() { - for (path, adapter) in [ - ( - "crates/aether-data/runtime/src/driver/postgres.rs", - "aether_data_postgres", - ), - ( - "crates/aether-data/runtime/src/driver/mysql.rs", - "aether_data_mysql", - ), - ( - "crates/aether-data/runtime/src/driver/sqlite.rs", - "aether_data_sqlite", - ), - ] { - let source = read_workspace_file(path); - assert!( - source.contains(&format!("pub use {adapter}::*;")), - "{path} should remain a thin compatibility facade" - ); - assert!(!source.contains("sqlx::")); - } + let source = read_workspace_file("crates/aether-data/runtime/src/driver/postgres.rs"); + assert!(source.contains("pub use aether_data_postgres::*;")); + assert!(!source.contains("sqlx::")); } #[test] @@ -183,32 +143,16 @@ fn tunnel_binary_uses_shared_tunnel_boundary_without_gateway_runtime_dependency( fn data_facade_defaults_to_postgres_and_gateway_selects_all_drivers_explicitly() { let data_manifest = read_workspace_file("crates/aether-data/runtime/Cargo.toml"); assert!(data_manifest.contains("default = [\"postgres\"]")); - for dependency in [ - "aether-data-postgres = { workspace = true, optional = true }", - "aether-data-mysql = { workspace = true, optional = true }", - "aether-data-sqlite = { workspace = true, optional = true }", - ] { - assert!( - data_manifest.contains(dependency), - "aether-data should keep {dependency} optional" - ); - } - assert!( - !data_manifest.contains("features = [\"postgres\", \"mysql\", \"sqlite\"\"]"), - "aether-data must not unconditionally enable every sqlx driver" - ); + assert!(data_manifest.contains("aether-data-postgres = { workspace = true, optional = true }")); + assert!(data_manifest.contains("postgres = [\"dep:aether-data-postgres\", \"sqlx/postgres\"]")); + assert!(data_manifest.contains("all-drivers = [\"postgres\"]")); let gateway_manifest = read_workspace_file("apps/aether-gateway/Cargo.toml"); assert!(gateway_manifest .contains("aether-data = { workspace = true, features = [\"all-drivers\"] }")); let data_lib = read_workspace_file("crates/aether-data/runtime/src/lib.rs"); - for backend in ["PostgresBackend", "MysqlBackend", "SqliteBackend"] { - assert!( - data_lib.contains(&format!("pub use backend::{backend};")), - "aether-data should expose enabled backends symmetrically at its facade root" - ); - } + assert!(data_lib.contains("pub use backend::PostgresBackend;")); } #[test] @@ -219,58 +163,11 @@ fn data_query_helpers_belong_to_adapters_not_the_runtime_facade() { "aether-data should not keep a direct query-helper dependency after SQL repositories move to adapters" ); - for adapter_manifest in [ - "crates/aether-data/adapters/postgres/Cargo.toml", - "crates/aether-data/adapters/mysql/Cargo.toml", - "crates/aether-data/adapters/sqlite/Cargo.toml", - ] { - let manifest = read_workspace_file(adapter_manifest); - assert!( - manifest.contains("aether-data-query.workspace = true"), - "{adapter_manifest} should own its query-helper dependency" - ); - } + let manifest = read_workspace_file("crates/aether-data/adapters/postgres/Cargo.toml"); + assert!(manifest.contains("aether-data-query.workspace = true")); let query_helpers = read_workspace_file("crates/aether-data/query/src/lib.rs"); - for dialect in ["Postgres", "MySql", "Sqlite"] { - assert!( - query_helpers.contains(dialect), - "aether-data-query should render the {dialect} dialect" - ); - } -} - -#[test] -fn sql_adapters_centralize_error_mapping_boilerplate() { - for (adapter, driver) in [ - ("aether-data-mysql", "mysql"), - ("aether-data-sqlite", "sqlite"), - ] { - let root = format!("crates/aether-data/adapters/{driver}/src"); - let files = collect_workspace_rust_files(&root); - let trait_owners = files - .iter() - .filter(|path| { - std::fs::read_to_string(path) - .expect("adapter source should be readable") - .contains("trait SqlResultExt") - }) - .collect::>(); - - assert_eq!( - trait_owners.len(), - 1, - "{adapter} should have exactly one SqlResultExt owner, found: {trait_owners:?}" - ); - assert_eq!( - trait_owners[0].file_name().and_then(|name| name.to_str()), - Some("error.rs"), - "{adapter} should keep SQL error mapping in src/error.rs" - ); - - let lib = read_workspace_file(&format!("crates/aether-data/adapters/{driver}/src/lib.rs")); - assert!(lib.contains("mod error;")); - } + assert!(query_helpers.contains("Postgres")); } #[test] diff --git a/crates/aether-data/adapters/mysql/Cargo.toml b/crates/aether-data/adapters/mysql/Cargo.toml deleted file mode 100644 index 35a96b69b..000000000 --- a/crates/aether-data/adapters/mysql/Cargo.toml +++ /dev/null @@ -1,24 +0,0 @@ -[package] -name = "aether-data-mysql" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true -description = "MySQL repositories, pools, and migrations for Aether" - -[dependencies] -aether-ai-formats.workspace = true -aether-data-contracts.workspace = true -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"] } -tracing.workspace = true -uuid.workspace = true - -[dev-dependencies] -tokio.workspace = true diff --git a/crates/aether-data/adapters/mysql/migrations/20260403000000_baseline.sql b/crates/aether-data/adapters/mysql/migrations/20260403000000_baseline.sql deleted file mode 100644 index 0f91573ad..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260403000000_baseline.sql +++ /dev/null @@ -1,1051 +0,0 @@ -CREATE TABLE IF NOT EXISTS users ( - id VARCHAR(64) PRIMARY KEY, - external_id VARCHAR(255), - email VARCHAR(320), - username VARCHAR(255), - password_hash VARCHAR(255), - role VARCHAR(64), - auth_source VARCHAR(64) NOT NULL DEFAULT 'local', - email_verified TINYINT(1) NOT NULL DEFAULT 0, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_deleted TINYINT(1) NOT NULL DEFAULT 0, - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - model_capability_settings TEXT, - rate_limit INT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - last_login_at BIGINT, - ldap_dn VARCHAR(1024), - ldap_username VARCHAR(255), - UNIQUE KEY users_email_key (email), - UNIQUE KEY users_username_key (username) -); - -CREATE TABLE IF NOT EXISTS api_keys ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - key_hash VARCHAR(255) NOT NULL, - key_encrypted TEXT, - name VARCHAR(255), - key_prefix VARCHAR(64), - status VARCHAR(64) NOT NULL DEFAULT 'active', - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - rate_limit INT DEFAULT 100, - concurrent_limit INT, - force_capabilities TEXT, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_locked TINYINT(1) NOT NULL DEFAULT 0, - is_standalone TINYINT(1) NOT NULL DEFAULT 0, - auto_delete_on_expiry TINYINT(1) NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - metadata TEXT, - expires_at BIGINT, - last_used_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY api_keys_key_hash_key (key_hash), - KEY api_keys_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS audit_logs ( - id VARCHAR(64) PRIMARY KEY, - event_type VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - description TEXT NOT NULL, - ip_address VARCHAR(64), - user_agent VARCHAR(512), - request_id VARCHAR(128), - event_metadata TEXT, - status_code INT, - error_message TEXT, - created_at BIGINT NOT NULL, - KEY audit_logs_created_at_idx (created_at), - KEY audit_logs_event_type_idx (event_type), - KEY audit_logs_request_id_idx (request_id), - KEY audit_logs_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS announcements ( - id VARCHAR(64) PRIMARY KEY, - title VARCHAR(200) NOT NULL, - content TEXT NOT NULL, - `type` VARCHAR(32) NOT NULL DEFAULT 'info', - priority INT NOT NULL DEFAULT 0, - author_id VARCHAR(64), - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_pinned TINYINT(1) NOT NULL DEFAULT 0, - start_time BIGINT, - end_time BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY announcements_author_id_idx (author_id), - KEY announcements_created_at_idx (created_at), - KEY announcements_is_active_idx (is_active) -); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - announcement_id VARCHAR(64) NOT NULL, - read_at BIGINT NOT NULL, - UNIQUE KEY uq_user_announcement (user_id, announcement_id), - KEY announcement_reads_announcement_id_idx (announcement_id), - KEY announcement_reads_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS management_tokens ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - description TEXT, - token_hash VARCHAR(255) NOT NULL, - token_prefix VARCHAR(64), - allowed_ips TEXT, - expires_at BIGINT, - last_used_at BIGINT, - last_used_ip VARCHAR(255), - usage_count BIGINT NOT NULL DEFAULT 0, - is_active TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY management_tokens_token_hash_key (token_hash), - UNIQUE KEY uq_management_tokens_user_name (user_id, name), - KEY management_tokens_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS user_preferences ( - id VARCHAR(36) PRIMARY KEY, - user_id VARCHAR(36) NOT NULL UNIQUE, - avatar_url VARCHAR(500), - bio TEXT, - default_provider_id VARCHAR(36), - theme VARCHAR(20) NOT NULL DEFAULT 'light', - language VARCHAR(10) NOT NULL DEFAULT 'zh-CN', - timezone VARCHAR(50) NOT NULL DEFAULT 'Asia/Shanghai', - email_notifications BOOLEAN NOT NULL DEFAULT TRUE, - usage_alerts BOOLEAN NOT NULL DEFAULT TRUE, - announcement_notifications BOOLEAN NOT NULL DEFAULT TRUE, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - INDEX user_preferences_default_provider_id_idx (default_provider_id), - INDEX user_preferences_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS user_sessions ( - id VARCHAR(36) PRIMARY KEY, - user_id VARCHAR(36) NOT NULL, - client_device_id VARCHAR(128) NOT NULL, - device_label VARCHAR(120), - device_type VARCHAR(20) NOT NULL DEFAULT 'unknown', - browser_name VARCHAR(50), - browser_version VARCHAR(50), - os_name VARCHAR(50), - os_version VARCHAR(50), - device_model VARCHAR(100), - ip_address VARCHAR(45), - user_agent VARCHAR(1000), - client_hints TEXT, - refresh_token_hash VARCHAR(64) NOT NULL, - prev_refresh_token_hash VARCHAR(64), - rotated_at BIGINT, - last_seen_at BIGINT NOT NULL, - expires_at BIGINT NOT NULL, - revoked_at BIGINT, - revoke_reason VARCHAR(100), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - INDEX user_sessions_user_active_idx (user_id, revoked_at, expires_at), - INDEX user_sessions_user_device_idx (user_id, client_device_id) -); -CREATE TABLE IF NOT EXISTS billing_rules ( - id VARCHAR(64) PRIMARY KEY, - global_model_id VARCHAR(64), - model_id VARCHAR(64), - name VARCHAR(255) NOT NULL, - task_type VARCHAR(64) NOT NULL DEFAULT 'chat', - expression TEXT NOT NULL, - variables TEXT NOT NULL, - dimension_mappings TEXT NOT NULL, - is_enabled TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY billing_rules_global_model_task_idx (global_model_id, task_type, is_enabled), - KEY billing_rules_model_task_idx (model_id, task_type, is_enabled) -); - -CREATE TABLE IF NOT EXISTS dimension_collectors ( - id VARCHAR(64) PRIMARY KEY, - api_format VARCHAR(64) NOT NULL, - task_type VARCHAR(64) NOT NULL, - dimension_name VARCHAR(128) NOT NULL, - source_type VARCHAR(64) NOT NULL, - source_path VARCHAR(255), - value_type VARCHAR(64) NOT NULL DEFAULT 'float', - transform_expression TEXT, - default_value VARCHAR(255), - priority INT NOT NULL DEFAULT 0, - is_enabled TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY dimension_collectors_enabled_idx ( - api_format, - task_type, - dimension_name, - priority, - is_enabled - ) -); - -CREATE TABLE IF NOT EXISTS providers ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - description TEXT, - website VARCHAR(500), - provider_type VARCHAR(64) NOT NULL, - billing_type VARCHAR(64), - monthly_quota_usd DOUBLE, - monthly_used_usd DOUBLE, - quota_reset_day INT, - quota_last_reset_at BIGINT, - quota_expires_at BIGINT, - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - priority BIGINT NOT NULL DEFAULT 0, - provider_priority INT NOT NULL DEFAULT 100, - keep_priority_on_conversion TINYINT(1) NOT NULL DEFAULT 0, - enable_format_conversion TINYINT(1) NOT NULL DEFAULT 1, - concurrent_limit INT, - max_retries INT, - proxy TEXT, - request_timeout DOUBLE, - stream_first_byte_timeout DOUBLE, - config TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY providers_name_key (name) -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - api_key TEXT, - encrypted_key TEXT, - auth_type VARCHAR(32) NOT NULL DEFAULT 'api_key', - auth_config TEXT, - note TEXT, - internal_priority INT NOT NULL DEFAULT 50, - capabilities TEXT, - api_formats TEXT, - auth_type_by_format TEXT, - allow_auth_channel_mismatch_formats TEXT, - rate_multipliers TEXT, - global_priority_by_format TEXT, - allowed_models TEXT, - expires_at BIGINT, - cache_ttl_minutes INT NOT NULL DEFAULT 5, - max_probe_interval_minutes INT NOT NULL DEFAULT 32, - proxy TEXT, - fingerprint TEXT, - concurrent_limit INT, - learned_rpm_limit INT, - concurrent_429_count INT NOT NULL DEFAULT 0, - rpm_429_count INT NOT NULL DEFAULT 0, - last_429_at BIGINT, - last_429_type VARCHAR(64), - adjustment_history TEXT, - utilization_samples TEXT, - last_probe_increase_at BIGINT, - last_rpm_peak INT, - request_count BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - success_count BIGINT NOT NULL DEFAULT 0, - error_count BIGINT NOT NULL DEFAULT 0, - total_response_time_ms BIGINT NOT NULL DEFAULT 0, - last_used_at BIGINT, - auto_fetch_models TINYINT(1) NOT NULL DEFAULT 0, - last_models_fetch_at BIGINT, - last_models_fetch_error TEXT, - locked_models TEXT, - model_include_patterns TEXT, - model_exclude_patterns TEXT, - upstream_metadata TEXT, - oauth_invalid_at BIGINT, - oauth_invalid_reason VARCHAR(255), - status_snapshot TEXT, - health_by_format TEXT, - circuit_breaker_by_format TEXT, - status VARCHAR(64) NOT NULL DEFAULT 'active', - is_active TINYINT(1) NOT NULL DEFAULT 1, - weight BIGINT NOT NULL DEFAULT 1, - rpm_limit BIGINT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY provider_api_keys_provider_id_idx (provider_id), - KEY idx_provider_api_keys_provider_default_sort (provider_id, internal_priority, name, id) -); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id VARCHAR(192) PRIMARY KEY, - pool_kind VARCHAR(64) NOT NULL, - pool_id VARCHAR(64) NOT NULL, - member_kind VARCHAR(64) NOT NULL, - member_id VARCHAR(64) NOT NULL, - capability VARCHAR(64) NOT NULL, - scope_kind VARCHAR(64) NOT NULL, - scope_id VARCHAR(128), - score DOUBLE NOT NULL DEFAULT 0, - hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown', - score_version BIGINT NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at BIGINT, - last_scheduled_at BIGINT, - last_success_at BIGINT, - last_failure_at BIGINT, - failure_count BIGINT NOT NULL DEFAULT 0, - last_probe_attempt_at BIGINT, - last_probe_success_at BIGINT, - last_probe_failure_at BIGINT, - probe_failure_count BIGINT NOT NULL DEFAULT 0, - probe_status VARCHAR(64) NOT NULL DEFAULT 'never', - updated_at BIGINT NOT NULL, - KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC), - KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id), - KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at), - KEY pool_member_scores_updated_at_idx (updated_at) -); - -CREATE TABLE IF NOT EXISTS gemini_file_mappings ( - id VARCHAR(64) PRIMARY KEY, - file_name VARCHAR(512) NOT NULL, - key_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - display_name VARCHAR(512), - mime_type VARCHAR(255), - source_hash VARCHAR(128), - created_at BIGINT NOT NULL, - expires_at BIGINT NOT NULL, - UNIQUE KEY gemini_file_mappings_file_name_key (file_name), - KEY gemini_file_mappings_key_id_idx (key_id), - KEY gemini_file_mappings_user_id_idx (user_id), - KEY gemini_file_mappings_expires_at_idx (expires_at), - KEY gemini_file_mappings_source_hash_idx (source_hash) -); - -CREATE TABLE IF NOT EXISTS request_candidates ( - id VARCHAR(64) PRIMARY KEY, - request_id VARCHAR(128) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - username VARCHAR(255), - api_key_name VARCHAR(255), - candidate_index INT NOT NULL, - retry_index INT NOT NULL DEFAULT 0, - provider_id VARCHAR(64), - endpoint_id VARCHAR(64), - key_id VARCHAR(64), - status VARCHAR(32) NOT NULL, - skip_reason TEXT, - is_cached TINYINT(1) NOT NULL DEFAULT 0, - status_code INT, - error_type VARCHAR(128), - error_message TEXT, - latency_ms INT, - concurrent_requests INT, - extra_data TEXT, - required_capabilities TEXT, - created_at BIGINT NOT NULL, - started_at BIGINT, - finished_at BIGINT, - UNIQUE KEY uq_request_candidate_with_retry (request_id, candidate_index, retry_index), - KEY request_candidates_request_id_idx (request_id), - KEY request_candidates_provider_id_idx (provider_id), - KEY request_candidates_endpoint_id_idx (endpoint_id), - KEY request_candidates_status_idx (status), - KEY request_candidates_created_at_idx (created_at), - KEY request_candidates_endpoint_status_created_idx (endpoint_id, status, created_at) -); - -CREATE TABLE IF NOT EXISTS video_tasks ( - id VARCHAR(64) PRIMARY KEY, - short_id VARCHAR(32), - request_id VARCHAR(128) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - username VARCHAR(255), - api_key_name VARCHAR(255), - external_task_id VARCHAR(255), - provider_id VARCHAR(64), - endpoint_id VARCHAR(64), - key_id VARCHAR(64), - client_api_format VARCHAR(128), - provider_api_format VARCHAR(128), - format_converted TINYINT(1) NOT NULL DEFAULT 0, - model VARCHAR(255), - prompt TEXT, - original_request_body TEXT, - duration_seconds INT, - resolution VARCHAR(64), - aspect_ratio VARCHAR(32), - size VARCHAR(64), - status VARCHAR(32) NOT NULL DEFAULT 'pending', - progress_percent INT NOT NULL DEFAULT 0, - progress_message TEXT, - retry_count INT NOT NULL DEFAULT 0, - poll_interval_seconds INT NOT NULL DEFAULT 10, - next_poll_at BIGINT, - poll_count INT NOT NULL DEFAULT 0, - max_poll_count INT NOT NULL DEFAULT 360, - created_at BIGINT NOT NULL, - submitted_at BIGINT, - completed_at BIGINT, - updated_at BIGINT NOT NULL, - error_code VARCHAR(128), - error_message TEXT, - video_url TEXT, - request_metadata TEXT, - UNIQUE KEY video_tasks_short_id_key (short_id), - UNIQUE KEY video_tasks_request_id_key (request_id), - KEY video_tasks_external_id_idx (external_task_id), - KEY video_tasks_next_poll_idx (next_poll_at), - KEY video_tasks_user_status_idx (user_id, status), - KEY video_tasks_api_key_id_idx (api_key_id), - KEY video_tasks_provider_id_idx (provider_id), - KEY video_tasks_endpoint_id_idx (endpoint_id), - KEY video_tasks_key_id_idx (key_id) -); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - base_url TEXT NOT NULL, - api_format VARCHAR(128), - api_family VARCHAR(128), - endpoint_kind VARCHAR(128), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - health_score DOUBLE NOT NULL DEFAULT 1.0, - weight BIGINT NOT NULL DEFAULT 1, - header_rules TEXT, - body_rules TEXT, - max_retries INT, - custom_path TEXT, - metadata TEXT, - config TEXT, - format_acceptance_config TEXT, - proxy TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY provider_endpoints_provider_id_idx (provider_id) -); - -CREATE TABLE IF NOT EXISTS models ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - global_model_id VARCHAR(64), - provider_model_name VARCHAR(255) NOT NULL, - global_model_name VARCHAR(255), - api_format VARCHAR(128), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_available TINYINT(1) NOT NULL DEFAULT 1, - price_per_request DOUBLE, - tiered_pricing TEXT, - supports_vision TINYINT(1), - supports_function_calling TINYINT(1), - supports_streaming TINYINT(1), - supports_extended_thinking TINYINT(1), - supports_image_generation TINYINT(1), - provider_model_mappings TEXT, - config TEXT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY models_provider_id_idx (provider_id) -); - -CREATE TABLE IF NOT EXISTS global_models ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - display_name VARCHAR(255), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - default_price_per_request DOUBLE, - default_tiered_pricing TEXT, - supported_capabilities TEXT, - usage_count BIGINT NOT NULL DEFAULT 0, - config TEXT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY global_models_name_key (name) -); -CREATE TABLE IF NOT EXISTS system_configs ( - id VARCHAR(64) PRIMARY KEY, - `key` VARCHAR(255) NOT NULL, - value TEXT NOT NULL, - description TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY system_configs_key_key (`key`) -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - id VARCHAR(64) PRIMARY KEY, - module_type VARCHAR(128) NOT NULL, - enabled TINYINT(1) NOT NULL DEFAULT 1, - config TEXT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY auth_modules_module_type_key (module_type) -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - provider_type VARCHAR(64) PRIMARY KEY, - display_name VARCHAR(255) NOT NULL, - client_id TEXT NOT NULL, - client_secret_encrypted TEXT, - authorization_url_override VARCHAR(500), - token_url_override VARCHAR(500), - userinfo_url_override VARCHAR(500), - scopes TEXT, - redirect_uri VARCHAR(500) NOT NULL, - frontend_callback_url VARCHAR(500) NOT NULL, - attribute_mapping TEXT, - extra_config TEXT, - is_enabled TINYINT(1) NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - id BIGINT PRIMARY KEY AUTO_INCREMENT, - server_url VARCHAR(255) NOT NULL, - bind_dn TEXT NOT NULL, - bind_password_encrypted TEXT, - base_dn TEXT NOT NULL, - user_search_filter VARCHAR(512) NOT NULL DEFAULT '(uid={username})', - username_attr VARCHAR(50) NOT NULL DEFAULT 'uid', - email_attr VARCHAR(50) NOT NULL DEFAULT 'mail', - display_name_attr VARCHAR(50) NOT NULL DEFAULT 'cn', - is_enabled TINYINT(1) NOT NULL DEFAULT 0, - is_exclusive TINYINT(1) NOT NULL DEFAULT 0, - use_starttls TINYINT(1) NOT NULL DEFAULT 0, - connect_timeout INT NOT NULL DEFAULT 10, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - provider_type VARCHAR(64) NOT NULL, - provider_user_id VARCHAR(255) NOT NULL, - provider_username VARCHAR(255), - provider_email VARCHAR(255), - extra_data TEXT, - linked_at BIGINT NOT NULL, - last_login_at BIGINT, - UNIQUE KEY uq_user_oauth_links_provider_user (provider_type, provider_user_id), - UNIQUE KEY uq_user_oauth_links_user_provider (user_id, provider_type), - KEY user_oauth_links_provider_type_idx (provider_type), - KEY user_oauth_links_user_id_idx (user_id) -); -CREATE TABLE IF NOT EXISTS proxy_nodes ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - ip VARCHAR(512) NOT NULL, - port INT NOT NULL, - region VARCHAR(100), - status VARCHAR(32) NOT NULL DEFAULT 'online', - registered_by VARCHAR(64), - last_heartbeat_at BIGINT, - heartbeat_interval INT NOT NULL DEFAULT 30, - active_connections INT NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - avg_latency_ms DOUBLE, - is_manual TINYINT(1) NOT NULL DEFAULT 0, - proxy_url VARCHAR(500), - proxy_username VARCHAR(255), - proxy_password TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - remote_config TEXT, - config_version INT NOT NULL DEFAULT 0, - hardware_info TEXT, - estimated_max_concurrency INT, - tunnel_mode TINYINT(1) NOT NULL DEFAULT 0, - tunnel_connected TINYINT(1) NOT NULL DEFAULT 0, - tunnel_connected_at BIGINT, - failed_requests BIGINT NOT NULL DEFAULT 0, - dns_failures BIGINT NOT NULL DEFAULT 0, - stream_errors BIGINT NOT NULL DEFAULT 0, - proxy_metadata TEXT -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - id BIGINT PRIMARY KEY AUTO_INCREMENT, - node_id VARCHAR(64) NOT NULL, - event_type VARCHAR(64) NOT NULL, - detail VARCHAR(500), - created_at BIGINT NOT NULL -); -CREATE TABLE IF NOT EXISTS wallets ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - balance DOUBLE NOT NULL DEFAULT 0, - gift_balance DOUBLE NOT NULL DEFAULT 0, - limit_mode VARCHAR(64) NOT NULL DEFAULT 'finite', - currency VARCHAR(16) NOT NULL DEFAULT 'USD', - status VARCHAR(64) NOT NULL DEFAULT 'active', - total_recharged DOUBLE NOT NULL DEFAULT 0, - total_consumed DOUBLE NOT NULL DEFAULT 0, - total_refunded DOUBLE NOT NULL DEFAULT 0, - total_adjusted DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY wallets_user_id_key (user_id), - UNIQUE KEY wallets_api_key_id_key (api_key_id), - KEY wallets_api_key_id_idx (api_key_id), - KEY wallets_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - id VARCHAR(64) PRIMARY KEY, - wallet_id VARCHAR(64) NOT NULL, - category VARCHAR(64) NOT NULL, - reason_code VARCHAR(64) NOT NULL, - amount DOUBLE NOT NULL, - balance_before DOUBLE NOT NULL, - balance_after DOUBLE NOT NULL, - recharge_balance_before DOUBLE NOT NULL, - recharge_balance_after DOUBLE NOT NULL, - gift_balance_before DOUBLE NOT NULL, - gift_balance_after DOUBLE NOT NULL, - link_type VARCHAR(64), - link_id VARCHAR(128), - operator_id VARCHAR(64), - description TEXT, - created_at BIGINT NOT NULL, - KEY idx_wallet_tx_wallet_created (wallet_id, created_at), - KEY idx_wallet_tx_category_created (category, created_at), - KEY idx_wallet_tx_reason_created (reason_code, created_at), - KEY idx_wallet_tx_link (link_type, link_id), - KEY ix_wallet_transactions_operator_id (operator_id) -); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - id VARCHAR(64) PRIMARY KEY, - wallet_id VARCHAR(64) NOT NULL, - billing_date VARCHAR(16) NOT NULL, - billing_timezone VARCHAR(64) NOT NULL, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - first_finalized_at BIGINT, - last_finalized_at BIGINT, - aggregated_at BIGINT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_wallet_daily_usage_wallet_date (wallet_id, billing_timezone, billing_date) -); - -CREATE TABLE IF NOT EXISTS payment_orders ( - id VARCHAR(64) PRIMARY KEY, - order_no VARCHAR(128) NOT NULL, - wallet_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - amount_usd DOUBLE NOT NULL, - pay_amount DOUBLE, - pay_currency VARCHAR(16), - exchange_rate DOUBLE, - refunded_amount_usd DOUBLE NOT NULL DEFAULT 0, - refundable_amount_usd DOUBLE NOT NULL DEFAULT 0, - payment_method VARCHAR(64) NOT NULL, - gateway_order_id VARCHAR(128), - gateway_response TEXT, - status VARCHAR(64) NOT NULL DEFAULT 'pending', - created_at BIGINT NOT NULL, - paid_at BIGINT, - credited_at BIGINT, - expires_at BIGINT, - UNIQUE KEY uq_payment_orders_order_no (order_no), - KEY idx_payment_orders_wallet_created (wallet_id, created_at), - KEY idx_payment_orders_user_created (user_id, created_at), - KEY idx_payment_orders_status (status), - KEY idx_payment_orders_gateway_order_id (gateway_order_id) -); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - id VARCHAR(64) PRIMARY KEY, - payment_order_id VARCHAR(64), - payment_method VARCHAR(64) NOT NULL, - callback_key VARCHAR(128) NOT NULL, - order_no VARCHAR(128), - gateway_order_id VARCHAR(128), - payload_hash VARCHAR(128), - signature_valid TINYINT(1) NOT NULL DEFAULT 0, - status VARCHAR(64) NOT NULL DEFAULT 'received', - payload TEXT, - error_message TEXT, - created_at BIGINT NOT NULL, - processed_at BIGINT, - UNIQUE KEY uq_payment_callbacks_callback_key (callback_key), - KEY idx_payment_callbacks_order (order_no), - KEY idx_payment_callbacks_gateway_order (gateway_order_id), - KEY idx_payment_callbacks_created (created_at), - KEY ix_payment_callbacks_payment_order_id (payment_order_id) -); - -CREATE TABLE IF NOT EXISTS refund_requests ( - id VARCHAR(64) PRIMARY KEY, - refund_no VARCHAR(128) NOT NULL, - wallet_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - payment_order_id VARCHAR(64), - source_type VARCHAR(64) NOT NULL, - source_id VARCHAR(128), - refund_mode VARCHAR(64) NOT NULL, - amount_usd DOUBLE NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'pending_approval', - reason TEXT, - requested_by VARCHAR(64), - approved_by VARCHAR(64), - processed_by VARCHAR(64), - gateway_refund_id VARCHAR(128), - payout_method VARCHAR(64), - payout_reference VARCHAR(255), - payout_proof TEXT, - failure_reason TEXT, - idempotency_key VARCHAR(128), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - processed_at BIGINT, - completed_at BIGINT, - UNIQUE KEY uq_refund_requests_refund_no (refund_no), - UNIQUE KEY uq_refund_requests_idempotency_key (idempotency_key), - KEY idx_refund_wallet_created (wallet_id, created_at), - KEY idx_refund_user_created (user_id, created_at), - KEY idx_refund_status (status), - KEY ix_refund_requests_payment_order_id (payment_order_id), - KEY ix_refund_requests_requested_by (requested_by), - KEY ix_refund_requests_approved_by (approved_by), - KEY ix_refund_requests_processed_by (processed_by) -); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - amount_usd DOUBLE NOT NULL, - currency VARCHAR(16) NOT NULL DEFAULT 'USD', - balance_bucket VARCHAR(64) NOT NULL DEFAULT 'gift', - total_count INT NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'active', - description TEXT, - created_by VARCHAR(64), - expires_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_redeem_code_batches_status (status, created_at) -); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - id VARCHAR(64) PRIMARY KEY, - batch_id VARCHAR(64) NOT NULL, - code_hash VARCHAR(128) NOT NULL, - code_prefix VARCHAR(16) NOT NULL, - code_suffix VARCHAR(16) NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'active', - redeemed_by_user_id VARCHAR(64), - redeemed_wallet_id VARCHAR(64), - redeemed_payment_order_id VARCHAR(64), - redeemed_at BIGINT, - disabled_by VARCHAR(64), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_redeem_codes_code_hash (code_hash), - KEY idx_redeem_codes_batch_created (batch_id, created_at), - KEY idx_redeem_codes_status (status, updated_at), - KEY idx_redeem_codes_redeemed_user (redeemed_by_user_id, redeemed_at), - KEY idx_redeem_codes_redeemed_order (redeemed_payment_order_id) -); -CREATE TABLE IF NOT EXISTS `usage` ( - request_id VARCHAR(128) PRIMARY KEY, - id VARCHAR(128), - user_id VARCHAR(64), - api_key_id VARCHAR(64), - provider_name VARCHAR(255) NOT NULL DEFAULT 'unknown', - model VARCHAR(255) NOT NULL DEFAULT 'unknown', - target_model VARCHAR(255), - provider_id VARCHAR(64), - provider_endpoint_id VARCHAR(64), - provider_api_key_id VARCHAR(64), - request_type VARCHAR(64), - api_format VARCHAR(64), - api_family VARCHAR(64), - endpoint_kind VARCHAR(64), - endpoint_api_format VARCHAR(64), - provider_api_family VARCHAR(64), - provider_endpoint_kind VARCHAR(64), - has_format_conversion TINYINT(1) NOT NULL DEFAULT 0, - is_stream TINYINT(1) NOT NULL DEFAULT 0, - upstream_is_stream TINYINT(1), - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_input_tokens 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, - cache_creation_cost_usd DOUBLE NOT NULL DEFAULT 0, - cache_read_cost_usd DOUBLE NOT NULL DEFAULT 0, - output_price_per_1m DOUBLE, - status_code INT, - error_message TEXT, - error_category VARCHAR(255), - response_time_ms BIGINT, - first_byte_time_ms BIGINT, - wallet_id VARCHAR(64), - status VARCHAR(64) NOT NULL DEFAULT 'completed', - billing_status VARCHAR(64) NOT NULL DEFAULT 'pending', - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - actual_total_cost_usd DOUBLE NOT NULL DEFAULT 0, - request_metadata TEXT, - candidate_id VARCHAR(128), - candidate_index BIGINT, - key_name VARCHAR(255), - planner_kind VARCHAR(64), - route_family VARCHAR(128), - route_kind VARCHAR(128), - execution_path VARCHAR(128), - local_execution_runtime_miss_reason VARCHAR(255), - wallet_balance_before DOUBLE, - wallet_balance_after DOUBLE, - wallet_recharge_balance_before DOUBLE, - wallet_recharge_balance_after DOUBLE, - wallet_gift_balance_before DOUBLE, - wallet_gift_balance_after DOUBLE, - finalized_at BIGINT, - created_at_unix_ms BIGINT NOT NULL DEFAULT 0, - updated_at_unix_secs BIGINT NOT NULL DEFAULT 0, - KEY usage_api_key_id_idx (api_key_id), - KEY usage_billing_status_idx (billing_status), - KEY usage_created_at_idx (created_at_unix_ms), - KEY usage_provider_api_key_id_idx (provider_api_key_id), - KEY usage_provider_id_idx (provider_id), - KEY usage_request_id_idx (request_id), - KEY usage_user_id_idx (user_id), - KEY usage_wallet_id_idx (wallet_id) -); - -CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( - request_id VARCHAR(128) PRIMARY KEY, - billing_status VARCHAR(64) NOT NULL, - wallet_id VARCHAR(64), - wallet_balance_before DOUBLE, - wallet_balance_after DOUBLE, - wallet_recharge_balance_before DOUBLE, - wallet_recharge_balance_after DOUBLE, - wallet_gift_balance_before DOUBLE, - wallet_gift_balance_after DOUBLE, - provider_monthly_used_usd DOUBLE, - finalized_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY usage_settlement_snapshots_billing_status_idx (billing_status), - KEY usage_settlement_snapshots_wallet_id_idx (wallet_id) -); - - -CREATE TABLE IF NOT EXISTS stats_hourly ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - is_complete TINYINT(1) NOT NULL DEFAULT 0, - aggregated_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_hour (hour_utc) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - user_id VARCHAR(64) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_user (hour_utc, user_id) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - user_id VARCHAR(64) NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_user_model (hour_utc, user_id, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_model (hour_utc, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - provider_name VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_provider (hour_utc, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - fallback_count BIGINT NOT NULL DEFAULT 0, - unique_models BIGINT NOT NULL DEFAULT 0, - unique_providers BIGINT NOT NULL DEFAULT 0, - is_complete TINYINT(1) NOT NULL DEFAULT 0, - aggregated_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_date (`date`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_model (`date`, model) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - provider_name VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_provider (`date`, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - id VARCHAR(64) PRIMARY KEY, - api_key_id VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - api_key_name VARCHAR(255), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_api_key (`date`, api_key_id) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - error_category VARCHAR(255) NOT NULL, - provider_name VARCHAR(255), - model VARCHAR(255), - count BIGINT NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_error (`date`, error_category, provider_name, model) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - username VARCHAR(255), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_user_daily (`date`, user_id) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260507120000_add_management_token_permissions.sql b/crates/aether-data/adapters/mysql/migrations/20260507120000_add_management_token_permissions.sql deleted file mode 100644 index 1ef066949..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260507120000_add_management_token_permissions.sql +++ /dev/null @@ -1,3 +0,0 @@ -ALTER TABLE management_tokens - ADD COLUMN permissions TEXT NULL AFTER allowed_ips; - diff --git a/crates/aether-data/adapters/mysql/migrations/20260508000000_add_proxy_node_metrics.sql b/crates/aether-data/adapters/mysql/migrations/20260508000000_add_proxy_node_metrics.sql deleted file mode 100644 index d6449e12a..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260508000000_add_proxy_node_metrics.sql +++ /dev/null @@ -1,42 +0,0 @@ -ALTER TABLE proxy_node_events - ADD COLUMN event_metadata TEXT NULL AFTER detail; - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m ( - node_id VARCHAR(64) NOT NULL, - bucket_start_unix_secs BIGINT NOT NULL, - samples BIGINT NOT NULL DEFAULT 0, - uptime_samples BIGINT NOT NULL DEFAULT 0, - active_connections_sum BIGINT NOT NULL DEFAULT 0, - active_connections_max BIGINT NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum BIGINT NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max BIGINT NOT NULL DEFAULT 0, - connect_errors_delta BIGINT NOT NULL DEFAULT 0, - disconnects_delta BIGINT NOT NULL DEFAULT 0, - error_events_delta BIGINT NOT NULL DEFAULT 0, - ws_in_bytes_delta BIGINT NOT NULL DEFAULT 0, - ws_out_bytes_delta BIGINT NOT NULL DEFAULT 0, - ws_in_frames_delta BIGINT NOT NULL DEFAULT 0, - ws_out_frames_delta BIGINT NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs), - INDEX idx_proxy_node_metrics_1m_bucket_start (bucket_start_unix_secs) -); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h ( - node_id VARCHAR(64) NOT NULL, - bucket_start_unix_secs BIGINT NOT NULL, - samples BIGINT NOT NULL DEFAULT 0, - uptime_samples BIGINT NOT NULL DEFAULT 0, - active_connections_sum BIGINT NOT NULL DEFAULT 0, - active_connections_max BIGINT NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum BIGINT NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max BIGINT NOT NULL DEFAULT 0, - connect_errors_delta BIGINT NOT NULL DEFAULT 0, - disconnects_delta BIGINT NOT NULL DEFAULT 0, - error_events_delta BIGINT NOT NULL DEFAULT 0, - ws_in_bytes_delta BIGINT NOT NULL DEFAULT 0, - ws_out_bytes_delta BIGINT NOT NULL DEFAULT 0, - ws_in_frames_delta BIGINT NOT NULL DEFAULT 0, - ws_out_frames_delta BIGINT NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs), - INDEX idx_proxy_node_metrics_1h_bucket_start (bucket_start_unix_secs) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260509000000_add_background_tasks.sql b/crates/aether-data/adapters/mysql/migrations/20260509000000_add_background_tasks.sql deleted file mode 100644 index ff7884694..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260509000000_add_background_tasks.sql +++ /dev/null @@ -1,37 +0,0 @@ -CREATE TABLE IF NOT EXISTS background_task_runs ( - id VARCHAR(64) PRIMARY KEY, - task_key VARCHAR(200) NOT NULL, - kind VARCHAR(32) NOT NULL, - `trigger` VARCHAR(64) NOT NULL, - status VARCHAR(32) NOT NULL, - attempt INT NOT NULL DEFAULT 0, - max_attempts INT NOT NULL DEFAULT 0, - owner_instance VARCHAR(200), - progress_percent INT NOT NULL DEFAULT 0, - progress_message TEXT, - payload_json JSON, - result_json JSON, - error_message TEXT, - cancel_requested TINYINT(1) NOT NULL DEFAULT 0, - created_by VARCHAR(200), - created_at_unix_secs BIGINT NOT NULL, - started_at_unix_secs BIGINT NULL, - finished_at_unix_secs BIGINT NULL, - updated_at_unix_secs BIGINT NOT NULL, - INDEX idx_background_task_runs_task_key (task_key), - INDEX idx_background_task_runs_status (status), - INDEX idx_background_task_runs_kind (kind), - INDEX idx_background_task_runs_created_at (created_at_unix_secs) -); - -CREATE TABLE IF NOT EXISTS background_task_events ( - id VARCHAR(64) PRIMARY KEY, - run_id VARCHAR(64) NOT NULL, - event_type VARCHAR(64) NOT NULL, - message TEXT NOT NULL, - payload_json JSON, - created_at_unix_secs BIGINT NOT NULL, - INDEX idx_background_task_events_run_id (run_id, created_at_unix_secs), - CONSTRAINT fk_background_task_events_run - FOREIGN KEY (run_id) REFERENCES background_task_runs(id) ON DELETE CASCADE -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260509120000_add_user_groups.sql b/crates/aether-data/adapters/mysql/migrations/20260509120000_add_user_groups.sql deleted file mode 100644 index f329bd653..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260509120000_add_user_groups.sql +++ /dev/null @@ -1,103 +0,0 @@ -ALTER TABLE users - ADD COLUMN allowed_providers_mode VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - ADD COLUMN allowed_api_formats_mode VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - ADD COLUMN allowed_models_mode VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - ADD COLUMN rate_limit_mode VARCHAR(32) NOT NULL DEFAULT 'system'; - -UPDATE users -SET allowed_providers_mode = CASE WHEN allowed_providers IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_providers_mode = 'unrestricted'; - -UPDATE users -SET allowed_api_formats_mode = CASE WHEN allowed_api_formats IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_api_formats_mode = 'unrestricted'; - -UPDATE users -SET allowed_models_mode = CASE WHEN allowed_models IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_models_mode = 'unrestricted'; - -UPDATE users -SET rate_limit_mode = CASE WHEN rate_limit IS NULL THEN 'system' ELSE 'custom' END -WHERE rate_limit_mode = 'system'; - -CREATE TABLE IF NOT EXISTS user_groups ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(100) NOT NULL, - normalized_name VARCHAR(100) NOT NULL, - description TEXT, - priority INT NOT NULL DEFAULT 0, - allowed_providers TEXT, - allowed_providers_mode VARCHAR(32) NOT NULL DEFAULT 'inherit', - allowed_api_formats TEXT, - allowed_api_formats_mode VARCHAR(32) NOT NULL DEFAULT 'inherit', - allowed_models TEXT, - allowed_models_mode VARCHAR(32) NOT NULL DEFAULT 'inherit', - rate_limit INT, - rate_limit_mode VARCHAR(32) NOT NULL DEFAULT 'inherit', - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY user_groups_normalized_name_key (normalized_name), - KEY user_groups_priority_name_idx (priority, name, id) -); - -CREATE TABLE IF NOT EXISTS user_group_members ( - group_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64) NOT NULL, - created_at BIGINT NOT NULL, - PRIMARY KEY (group_id, user_id), - KEY user_group_members_user_id_idx (user_id), - CONSTRAINT user_group_members_group_id_fk - FOREIGN KEY (group_id) REFERENCES user_groups(id) ON DELETE CASCADE, - CONSTRAINT user_group_members_user_id_fk - FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE -); - -INSERT IGNORE 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 ( - '00000000-0000-0000-0000-000000000001', - 'Default', - 'default', - 'Default group for all users', - 0, - 'unrestricted', - 'unrestricted', - 'unrestricted', - 'system', - UNIX_TIMESTAMP(), - UNIX_TIMESTAMP() -); - -INSERT IGNORE INTO system_configs ( - id, - `key`, - value, - description, - created_at, - updated_at -) -VALUES ( - '00000000-0000-0000-0000-000000000002', - 'default_user_group_id', - '"00000000-0000-0000-0000-000000000001"', - 'Default user group', - UNIX_TIMESTAMP(), - UNIX_TIMESTAMP() -); - -INSERT IGNORE INTO user_group_members (group_id, user_id, created_at) -SELECT '00000000-0000-0000-0000-000000000001', id, UNIX_TIMESTAMP() -FROM users -WHERE is_deleted = 0 - AND LOWER(role) <> 'admin'; diff --git a/crates/aether-data/adapters/mysql/migrations/20260510120000_normalize_empty_user_policy_modes.sql b/crates/aether-data/adapters/mysql/migrations/20260510120000_normalize_empty_user_policy_modes.sql deleted file mode 100644 index 9a885d2f7..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260510120000_normalize_empty_user_policy_modes.sql +++ /dev/null @@ -1,50 +0,0 @@ -UPDATE users -SET allowed_providers_mode = 'unrestricted' -WHERE allowed_providers_mode = 'specific' - AND ( - allowed_providers IS NULL - OR TRIM(allowed_providers) = '' - OR CASE - WHEN JSON_VALID(allowed_providers) = 1 THEN - JSON_TYPE(allowed_providers) = 'NULL' - OR ( - JSON_TYPE(allowed_providers) = 'ARRAY' - AND JSON_LENGTH(allowed_providers) = 0 - ) - ELSE FALSE - END - ); - -UPDATE users -SET allowed_api_formats_mode = 'unrestricted' -WHERE allowed_api_formats_mode = 'specific' - AND ( - allowed_api_formats IS NULL - OR TRIM(allowed_api_formats) = '' - OR CASE - WHEN JSON_VALID(allowed_api_formats) = 1 THEN - JSON_TYPE(allowed_api_formats) = 'NULL' - OR ( - JSON_TYPE(allowed_api_formats) = 'ARRAY' - AND JSON_LENGTH(allowed_api_formats) = 0 - ) - ELSE FALSE - END - ); - -UPDATE users -SET allowed_models_mode = 'unrestricted' -WHERE allowed_models_mode = 'specific' - AND ( - allowed_models IS NULL - OR TRIM(allowed_models) = '' - OR CASE - WHEN JSON_VALID(allowed_models) = 1 THEN - JSON_TYPE(allowed_models) = 'NULL' - OR ( - JSON_TYPE(allowed_models) = 'ARRAY' - AND JSON_LENGTH(allowed_models) = 0 - ) - ELSE FALSE - END - ); diff --git a/crates/aether-data/adapters/mysql/migrations/20260511120000_exclude_admins_from_default_user_group.sql b/crates/aether-data/adapters/mysql/migrations/20260511120000_exclude_admins_from_default_user_group.sql deleted file mode 100644 index f0243e140..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260511120000_exclude_admins_from_default_user_group.sql +++ /dev/null @@ -1,12 +0,0 @@ -DELETE member -FROM user_group_members AS member -JOIN users ON users.id = member.user_id -WHERE LOWER(users.role) = 'admin' - AND ( - member.group_id = '00000000-0000-0000-0000-000000000001' - OR member.group_id IN ( - SELECT TRIM(BOTH '"' FROM value) - FROM system_configs - WHERE `key` = 'default_user_group_id' - ) - ); diff --git a/crates/aether-data/adapters/mysql/migrations/20260511130000_add_pool_member_scores.sql b/crates/aether-data/adapters/mysql/migrations/20260511130000_add_pool_member_scores.sql deleted file mode 100644 index b38d194cc..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260511130000_add_pool_member_scores.sql +++ /dev/null @@ -1,29 +0,0 @@ -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id VARCHAR(192) PRIMARY KEY, - pool_kind VARCHAR(64) NOT NULL, - pool_id VARCHAR(64) NOT NULL, - member_kind VARCHAR(64) NOT NULL, - member_id VARCHAR(64) NOT NULL, - capability VARCHAR(64) NOT NULL, - scope_kind VARCHAR(64) NOT NULL, - scope_id VARCHAR(128), - score DOUBLE NOT NULL DEFAULT 0, - hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown', - score_version BIGINT NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at BIGINT, - last_scheduled_at BIGINT, - last_success_at BIGINT, - last_failure_at BIGINT, - failure_count BIGINT NOT NULL DEFAULT 0, - last_probe_attempt_at BIGINT, - last_probe_success_at BIGINT, - last_probe_failure_at BIGINT, - probe_failure_count BIGINT NOT NULL DEFAULT 0, - probe_status VARCHAR(64) NOT NULL DEFAULT 'never', - updated_at BIGINT NOT NULL, - KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC), - KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id), - KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at), - KEY pool_member_scores_updated_at_idx (updated_at) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260512000000_add_feature_settings.sql b/crates/aether-data/adapters/mysql/migrations/20260512000000_add_feature_settings.sql deleted file mode 100644 index 8a7953d12..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260512000000_add_feature_settings.sql +++ /dev/null @@ -1,5 +0,0 @@ -ALTER TABLE users - ADD COLUMN feature_settings TEXT NULL AFTER model_capability_settings; - -ALTER TABLE api_keys - ADD COLUMN feature_settings TEXT NULL AFTER force_capabilities; diff --git a/crates/aether-data/adapters/mysql/migrations/20260512090000_add_payment_gateway_plans.sql b/crates/aether-data/adapters/mysql/migrations/20260512090000_add_payment_gateway_plans.sql deleted file mode 100644 index 036adc8fa..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260512090000_add_payment_gateway_plans.sql +++ /dev/null @@ -1,79 +0,0 @@ -ALTER TABLE payment_orders - ADD COLUMN payment_provider VARCHAR(64), - ADD COLUMN payment_channel VARCHAR(64), - ADD COLUMN order_kind VARCHAR(64) NOT NULL DEFAULT 'wallet_recharge', - ADD COLUMN product_id VARCHAR(64), - ADD COLUMN product_snapshot TEXT, - ADD COLUMN fulfillment_status VARCHAR(64) NOT NULL DEFAULT 'pending', - ADD COLUMN fulfillment_error TEXT; - -CREATE INDEX idx_payment_orders_kind_status - ON payment_orders (order_kind, status); -CREATE INDEX idx_payment_orders_product - ON payment_orders (product_id); - -CREATE TABLE IF NOT EXISTS payment_gateway_configs ( - provider VARCHAR(64) PRIMARY KEY, - enabled BOOLEAN NOT NULL DEFAULT FALSE, - endpoint_url VARCHAR(512) NOT NULL, - callback_base_url VARCHAR(512), - merchant_id VARCHAR(128) NOT NULL, - merchant_key_encrypted TEXT, - pay_currency VARCHAR(16) NOT NULL DEFAULT 'CNY', - usd_exchange_rate DOUBLE NOT NULL DEFAULT 7.2, - min_recharge_usd DOUBLE NOT NULL DEFAULT 1, - channels_json TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS billing_plans ( - id VARCHAR(64) PRIMARY KEY, - title VARCHAR(128) NOT NULL, - description TEXT, - price_amount DOUBLE NOT NULL, - price_currency VARCHAR(16) NOT NULL DEFAULT 'CNY', - duration_unit VARCHAR(32) NOT NULL, - duration_value BIGINT NOT NULL, - enabled BOOLEAN NOT NULL DEFAULT TRUE, - sort_order BIGINT NOT NULL DEFAULT 0, - max_active_per_user BIGINT NOT NULL DEFAULT 1, - entitlements_json TEXT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_billing_plans_enabled_sort (enabled, sort_order) -); - -CREATE TABLE IF NOT EXISTS user_plan_entitlements ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - plan_id VARCHAR(64) NOT NULL, - payment_order_id VARCHAR(64) NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'active', - starts_at BIGINT NOT NULL, - expires_at BIGINT NOT NULL, - entitlements_snapshot TEXT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_user_plan_entitlements_user_active (user_id, status, expires_at), - KEY idx_user_plan_entitlements_order (payment_order_id), - CONSTRAINT user_plan_entitlements_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE, - CONSTRAINT user_plan_entitlements_plan_id_fkey FOREIGN KEY (plan_id) REFERENCES billing_plans(id) ON DELETE RESTRICT, - CONSTRAINT user_plan_entitlements_payment_order_id_fkey FOREIGN KEY (payment_order_id) REFERENCES payment_orders(id) ON DELETE RESTRICT -); - -CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers ( - id VARCHAR(64) PRIMARY KEY, - user_entitlement_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64) NOT NULL, - request_id VARCHAR(128) NOT NULL, - amount_usd DOUBLE NOT NULL, - balance_before DOUBLE NOT NULL, - balance_after DOUBLE NOT NULL, - usage_date VARCHAR(16) NOT NULL, - created_at BIGINT NOT NULL, - UNIQUE KEY uq_entitlement_usage_request (user_entitlement_id, request_id), - KEY idx_entitlement_usage_user_date (user_id, usage_date), - CONSTRAINT entitlement_usage_ledgers_entitlement_fkey FOREIGN KEY (user_entitlement_id) REFERENCES user_plan_entitlements(id) ON DELETE CASCADE, - CONSTRAINT entitlement_usage_ledgers_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql b/crates/aether-data/adapters/mysql/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql deleted file mode 100644 index dc01e5250..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE billing_plans - ADD COLUMN purchase_limit_scope VARCHAR(32) NOT NULL DEFAULT 'active_period'; diff --git a/crates/aether-data/adapters/mysql/migrations/20260516000000_add_routing_profiles.sql b/crates/aether-data/adapters/mysql/migrations/20260516000000_add_routing_profiles.sql deleted file mode 100644 index cc55a0af2..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260516000000_add_routing_profiles.sql +++ /dev/null @@ -1,41 +0,0 @@ -CREATE TABLE IF NOT EXISTS routing_groups ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `description` LONGTEXT, - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_system_default` TINYINT(1) NOT NULL DEFAULT 0, - `config_json` JSON NOT NULL, - `version` BIGINT NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `published_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY routing_groups_name_key (`name`), - KEY routing_groups_system_default_idx (`is_system_default`, `enabled`) -); - -CREATE TABLE IF NOT EXISTS routing_group_bindings ( - `id` VARCHAR(64) NOT NULL, - `group_id` VARCHAR(64) NOT NULL, - `subject_type` VARCHAR(32) NOT NULL, - `subject_id` VARCHAR(64) NOT NULL, - `is_default` TINYINT(1) NOT NULL DEFAULT 0, - `allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY routing_group_bindings_group_id_idx (`group_id`), - KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`) -); - -CREATE TABLE IF NOT EXISTS routing_group_versions ( - `id` VARCHAR(64) NOT NULL, - `group_id` VARCHAR(64) NOT NULL, - `version` BIGINT NOT NULL, - `config_json` JSON NOT NULL, - `created_at` BIGINT NOT NULL, - `created_by` VARCHAR(64), - PRIMARY KEY (`id`), - UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`), - KEY routing_group_versions_group_id_idx (`group_id`) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260518000000_add_usage_counter_deltas.sql b/crates/aether-data/adapters/mysql/migrations/20260518000000_add_usage_counter_deltas.sql deleted file mode 100644 index 17c092e8a..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260518000000_add_usage_counter_deltas.sql +++ /dev/null @@ -1,32 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_counter_deltas ( - `id` VARCHAR(36) NOT NULL, - `request_id` VARCHAR(128) NOT NULL, - `kind` VARCHAR(64) NOT NULL, - `target_id` TEXT NOT NULL, - `request_count_delta` BIGINT NOT NULL DEFAULT 0, - `total_requests_delta` BIGINT NOT NULL DEFAULT 0, - `success_count_delta` BIGINT NOT NULL DEFAULT 0, - `error_count_delta` BIGINT NOT NULL DEFAULT 0, - `dns_failures_delta` BIGINT NOT NULL DEFAULT 0, - `stream_errors_delta` BIGINT NOT NULL DEFAULT 0, - `total_tokens_delta` BIGINT NOT NULL DEFAULT 0, - `total_cost_usd_delta` DOUBLE NOT NULL DEFAULT 0, - `total_response_time_ms_delta` BIGINT NOT NULL DEFAULT 0, - `last_used_at_unix_secs` BIGINT, - `last_used_ip` TEXT, - `candidate_last_used_at_unix_secs` BIGINT, - `removed_last_used_at_unix_secs` BIGINT, - `usage_created_at_unix_secs` BIGINT, - `created_at` BIGINT NOT NULL, - `processed_at` BIGINT, - PRIMARY KEY (`id`), - KEY ix_usage_counter_deltas_unprocessed (`created_at`, `id`), - KEY ix_usage_counter_deltas_processed (`processed_at`, `created_at`, `id`), - KEY ix_usage_counter_deltas_request_kind (`request_id`, `kind`, `target_id`(191)) -); - -CREATE INDEX video_tasks_due_poll_idx - ON video_tasks (status, next_poll_at, updated_at); - -CREATE INDEX idx_entitlement_usage_entitlement_date - ON entitlement_usage_ledgers (user_entitlement_id, usage_date); diff --git a/crates/aether-data/adapters/mysql/migrations/20260519000000_add_referrals_privacy_required_announcements.sql b/crates/aether-data/adapters/mysql/migrations/20260519000000_add_referrals_privacy_required_announcements.sql deleted file mode 100644 index 735c9e782..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260519000000_add_referrals_privacy_required_announcements.sql +++ /dev/null @@ -1,62 +0,0 @@ -ALTER TABLE users - ADD COLUMN privacy_policy_accepted_version VARCHAR(64), - ADD COLUMN privacy_policy_accepted_at BIGINT; - -ALTER TABLE announcements - ADD COLUMN requires_ack BOOLEAN NOT NULL DEFAULT FALSE; - -CREATE TABLE IF NOT EXISTS user_invite_codes ( - user_id VARCHAR(64) PRIMARY KEY, - invite_code VARCHAR(64) NOT NULL UNIQUE, - active BOOLEAN NOT NULL DEFAULT TRUE, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - CONSTRAINT user_invite_codes_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE -); - -CREATE TABLE IF NOT EXISTS user_referrals ( - id VARCHAR(64) PRIMARY KEY, - inviter_user_id VARCHAR(64) NOT NULL, - invitee_user_id VARCHAR(64) NOT NULL UNIQUE, - invite_code_snapshot VARCHAR(64) NOT NULL, - source_json TEXT, - first_paid_order_id VARCHAR(64), - first_paid_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_user_referrals_inviter (inviter_user_id, created_at), - KEY idx_user_referrals_created (created_at), - KEY idx_user_referrals_invite_code (invite_code_snapshot), - CONSTRAINT user_referrals_inviter_user_id_fkey FOREIGN KEY (inviter_user_id) REFERENCES users(id) ON DELETE CASCADE, - CONSTRAINT user_referrals_invitee_user_id_fkey FOREIGN KEY (invitee_user_id) REFERENCES users(id) ON DELETE CASCADE, - CONSTRAINT user_referrals_first_paid_order_fkey FOREIGN KEY (first_paid_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); - -CREATE TABLE IF NOT EXISTS referral_rewards ( - id VARCHAR(64) PRIMARY KEY, - referral_id VARCHAR(64) NOT NULL, - inviter_user_id VARCHAR(64) NOT NULL, - invitee_user_id VARCHAR(64) NOT NULL, - reward_type VARCHAR(32) NOT NULL, - trigger_point VARCHAR(64) NOT NULL, - source_order_id VARCHAR(64), - idempotency_key VARCHAR(128) NOT NULL UNIQUE, - amount_usd DOUBLE NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - wallet_transaction_id VARCHAR(64), - reversed_amount_usd DOUBLE NOT NULL DEFAULT 0, - pending_reversal_amount_usd DOUBLE NOT NULL DEFAULT 0, - failure_reason TEXT, - admin_operator_id VARCHAR(64), - admin_note TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_referral_rewards_inviter_status (inviter_user_id, status, created_at), - KEY idx_referral_rewards_inviter_created (inviter_user_id, created_at), - KEY idx_referral_rewards_created (created_at), - KEY idx_referral_rewards_source_order (source_order_id), - CONSTRAINT referral_rewards_referral_id_fkey FOREIGN KEY (referral_id) REFERENCES user_referrals(id) ON DELETE CASCADE, - CONSTRAINT referral_rewards_inviter_user_id_fkey FOREIGN KEY (inviter_user_id) REFERENCES users(id) ON DELETE CASCADE, - CONSTRAINT referral_rewards_invitee_user_id_fkey FOREIGN KEY (invitee_user_id) REFERENCES users(id) ON DELETE CASCADE, - CONSTRAINT referral_rewards_source_order_fkey FOREIGN KEY (source_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260519120000_add_oauth_provider_icon_url.sql b/crates/aether-data/adapters/mysql/migrations/20260519120000_add_oauth_provider_icon_url.sql deleted file mode 100644 index e9da9b855..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260519120000_add_oauth_provider_icon_url.sql +++ /dev/null @@ -1 +0,0 @@ -ALTER TABLE oauth_providers ADD COLUMN icon_url VARCHAR(500); diff --git a/crates/aether-data/adapters/mysql/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql b/crates/aether-data/adapters/mysql/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql deleted file mode 100644 index 049ead11a..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql +++ /dev/null @@ -1,15 +0,0 @@ -SET @aether_provider_key_default_sort_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'provider_api_keys' - AND index_name = 'idx_provider_api_keys_provider_default_sort' - ) = 0, - 'CREATE INDEX idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id)', - 'DO 0' -); - -PREPARE aether_provider_key_default_sort_index_stmt FROM @aether_provider_key_default_sort_index_sql; -EXECUTE aether_provider_key_default_sort_index_stmt; -DEALLOCATE PREPARE aether_provider_key_default_sort_index_stmt; diff --git a/crates/aether-data/adapters/mysql/migrations/20260520000000_add_api_key_ip_rules.sql b/crates/aether-data/adapters/mysql/migrations/20260520000000_add_api_key_ip_rules.sql deleted file mode 100644 index ca31e5509..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260520000000_add_api_key_ip_rules.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE api_keys -ADD COLUMN ip_rules TEXT NULL AFTER allowed_models; diff --git a/crates/aether-data/adapters/mysql/migrations/20260520010000_preserve_usage_provider_identity.sql b/crates/aether-data/adapters/mysql/migrations/20260520010000_preserve_usage_provider_identity.sql deleted file mode 100644 index 633bf5cbf..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260520010000_preserve_usage_provider_identity.sql +++ /dev/null @@ -1,10 +0,0 @@ --- Usage is a historical fact table. Backfill nullable provider_id snapshots --- from the unique provider name where the catalog row still exists. - -UPDATE `usage` AS usage_rows -JOIN providers - ON providers.name = TRIM(usage_rows.provider_name) -SET usage_rows.provider_id = providers.id -WHERE usage_rows.provider_id IS NULL - AND TRIM(COALESCE(usage_rows.provider_name, '')) <> '' - AND LOWER(TRIM(COALESCE(usage_rows.provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending'); diff --git a/crates/aether-data/adapters/mysql/migrations/20260524000000_add_scheduler_hotspot_indexes.sql b/crates/aether-data/adapters/mysql/migrations/20260524000000_add_scheduler_hotspot_indexes.sql deleted file mode 100644 index bbdfc6d5e..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260524000000_add_scheduler_hotspot_indexes.sql +++ /dev/null @@ -1,31 +0,0 @@ -SET @aether_provider_key_active_priority_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'provider_api_keys' - AND index_name = 'idx_provider_api_keys_provider_active_priority_id' - ) = 0, - 'CREATE INDEX idx_provider_api_keys_provider_active_priority_id ON provider_api_keys (provider_id, is_active, internal_priority, id)', - 'DO 0' -); - -PREPARE aether_provider_key_active_priority_index_stmt FROM @aether_provider_key_active_priority_index_sql; -EXECUTE aether_provider_key_active_priority_index_stmt; -DEALLOCATE PREPARE aether_provider_key_active_priority_index_stmt; - -SET @aether_pool_score_scheduler_rank_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'pool_member_scores' - AND index_name = 'pool_member_scores_scheduler_account_rank_idx' - ) = 0, - 'CREATE INDEX pool_member_scores_scheduler_account_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC, last_ranked_at DESC, member_id, id)', - 'DO 0' -); - -PREPARE aether_pool_score_scheduler_rank_index_stmt FROM @aether_pool_score_scheduler_rank_index_sql; -EXECUTE aether_pool_score_scheduler_rank_index_stmt; -DEALLOCATE PREPARE aether_pool_score_scheduler_rank_index_stmt; diff --git a/crates/aether-data/adapters/mysql/migrations/20260527000000_add_provider_key_name_hotspot_index.sql b/crates/aether-data/adapters/mysql/migrations/20260527000000_add_provider_key_name_hotspot_index.sql deleted file mode 100644 index b60243da1..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260527000000_add_provider_key_name_hotspot_index.sql +++ /dev/null @@ -1,15 +0,0 @@ -SET @aether_provider_key_name_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'provider_api_keys' - AND index_name = 'idx_provider_api_keys_provider_name_id' - ) = 0, - 'CREATE INDEX idx_provider_api_keys_provider_name_id ON provider_api_keys (provider_id, name, id)', - 'DO 0' -); - -PREPARE aether_provider_key_name_index_stmt FROM @aether_provider_key_name_index_sql; -EXECUTE aether_provider_key_name_index_stmt; -DEALLOCATE PREPARE aether_provider_key_name_index_stmt; diff --git a/crates/aether-data/adapters/mysql/migrations/20260528000000_normalize_endpoint_api_roots.sql b/crates/aether-data/adapters/mysql/migrations/20260528000000_normalize_endpoint_api_roots.sql deleted file mode 100644 index 0f14c8b6c..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260528000000_normalize_endpoint_api_roots.sql +++ /dev/null @@ -1,132 +0,0 @@ -UPDATE provider_endpoints e -LEFT JOIN providers p ON p.id = e.provider_id -SET e.base_url = CONCAT( - TRIM(TRAILING '/' FROM SUBSTRING_INDEX(e.base_url, '?', 1)), - CASE - WHEN LOWER(TRIM(e.api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - THEN '/v1beta' - ELSE '/v1' - END, - IF(LOCATE('?', e.base_url) > 0, SUBSTRING(e.base_url, LOCATE('?', e.base_url)), '') -) -WHERE LOWER(TRIM(e.api_format)) IN ( - 'openai:chat', - 'openai:responses', - 'openai:responses:compact', - 'openai:embedding', - 'openai:rerank', - 'openai:image', - 'openai:video', - 'jina:embedding', - 'jina:rerank', - 'claude:messages', - 'gemini:generate_content', - 'gemini:embedding', - 'gemini:video' - ) - AND COALESCE(LOWER(TRIM(p.provider_type)), '') NOT IN ( - 'codex', - 'chatgpt_web', - 'claude_code', - 'kiro', - 'gemini_cli', - 'vertex_ai', - 'antigravity', - 'grok', - 'windsurf' - ) - AND LOWER(TRIM(TRAILING '/' FROM SUBSTRING_INDEX(e.base_url, '?', 1))) NOT REGEXP '/v[0-9]+(beta[0-9]*)?(/|$)' - AND ( - ( - LOWER(TRIM(e.api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) LIKE '/v1beta/%' - ) - OR ( - LOWER(TRIM(e.api_format)) NOT IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) LIKE '/v1/%' - ) - OR COALESCE(TRIM(e.custom_path), '') = '' - ); - -UPDATE provider_endpoints e -LEFT JOIN providers p ON p.id = e.provider_id -SET e.custom_path = CASE - WHEN LOWER(TRIM(e.api_format)) = 'openai:chat' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/chat/completions' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'openai:responses' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/responses' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'openai:responses:compact' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/responses/compact' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'claude:messages' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/messages' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) IN ('openai:embedding', 'jina:embedding') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/embeddings' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) IN ('openai:rerank', 'jina:rerank') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/rerank' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'openai:image' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/images/generations' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'openai:video' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1/videos' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'gemini:generate_content' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1beta/models/{model}:{action}' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'gemini:embedding' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) IN ('/v1beta/models/{model}:embedcontent', '/v1beta/models/{model}:{action}') - THEN NULL - WHEN LOWER(TRIM(e.api_format)) = 'gemini:video' - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) = '/v1beta/models/{model}:predictlongrunning' - THEN NULL - WHEN LOWER(TRIM(e.api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - THEN CONCAT('/', SUBSTRING(TRIM(e.custom_path), 9)) - ELSE CONCAT('/', SUBSTRING(TRIM(e.custom_path), 5)) - END -WHERE LOWER(TRIM(e.api_format)) IN ( - 'openai:chat', - 'openai:responses', - 'openai:responses:compact', - 'openai:embedding', - 'openai:rerank', - 'openai:image', - 'openai:video', - 'jina:embedding', - 'jina:rerank', - 'claude:messages', - 'gemini:generate_content', - 'gemini:embedding', - 'gemini:video' - ) - AND COALESCE(LOWER(TRIM(p.provider_type)), '') NOT IN ( - 'codex', - 'chatgpt_web', - 'claude_code', - 'kiro', - 'gemini_cli', - 'vertex_ai', - 'antigravity', - 'grok', - 'windsurf' - ) - AND ( - ( - LOWER(TRIM(e.api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) LIKE '/v1beta/%' - ) - OR ( - LOWER(TRIM(e.api_format)) NOT IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND LOWER(TRIM(COALESCE(e.custom_path, ''))) LIKE '/v1/%' - ) - ) - AND LOWER(TRIM(TRAILING '/' FROM SUBSTRING_INDEX(e.base_url, '?', 1))) REGEXP '/v[0-9]+(beta[0-9]*)?(/|$)'; - -UPDATE provider_endpoints -SET custom_path = NULL -WHERE custom_path IS NOT NULL - AND TRIM(custom_path) = ''; diff --git a/crates/aether-data/adapters/mysql/migrations/20260528020000_add_gateway_pressure_indexes.sql b/crates/aether-data/adapters/mysql/migrations/20260528020000_add_gateway_pressure_indexes.sql deleted file mode 100644 index a30dcc6b7..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260528020000_add_gateway_pressure_indexes.sql +++ /dev/null @@ -1,26 +0,0 @@ --- High-concurrency gateway read/cleanup paths. --- MySQL stores usage timestamps as unix milliseconds. - -CREATE INDEX idx_usage_created_id_desc - ON `usage` (created_at_unix_ms DESC, id ASC); - -CREATE INDEX idx_usage_user_created_id_desc - ON `usage` (user_id, created_at_unix_ms DESC, id ASC); - -CREATE INDEX idx_usage_api_format_created_id_desc - ON `usage` (api_format, created_at_unix_ms DESC, id ASC); - -CREATE INDEX idx_usage_status_created_id_desc - ON `usage` (status, created_at_unix_ms DESC, id ASC); - -CREATE INDEX idx_request_candidates_provider_created - ON request_candidates (provider_id, created_at DESC, id ASC); - -CREATE INDEX idx_request_candidates_api_key_created - ON request_candidates (api_key_id, created_at ASC, id ASC); - -CREATE INDEX idx_background_task_runs_status_created - ON background_task_runs (status, created_at_unix_secs DESC, updated_at_unix_secs DESC); - -CREATE INDEX idx_background_task_runs_kind_created - ON background_task_runs (kind, created_at_unix_secs DESC, updated_at_unix_secs DESC); 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 deleted file mode 100644 index 6ec4def0a..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260725010000_add_cross_driver_schema_parity.sql +++ /dev/null @@ -1,230 +0,0 @@ --- 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 deleted file mode 100644 index a7b4d84a1..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260725020000_add_advanced_stats_parity.sql +++ /dev/null @@ -1,381 +0,0 @@ -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 deleted file mode 100644 index 9dc553464..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260725030000_add_usage_routing_snapshot_parity.sql +++ /dev/null @@ -1,69 +0,0 @@ -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/migrations/20260727000000_repair_routing_default_uniqueness.sql b/crates/aether-data/adapters/mysql/migrations/20260727000000_repair_routing_default_uniqueness.sql deleted file mode 100644 index 7bc76ba5f..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260727000000_repair_routing_default_uniqueness.sql +++ /dev/null @@ -1,115 +0,0 @@ -UPDATE routing_groups -SET is_system_default = 0 -WHERE id IN ( - SELECT id - FROM ( - SELECT - id, - ROW_NUMBER() OVER (ORDER BY enabled DESC, updated_at DESC, id ASC) AS default_rank - FROM routing_groups - WHERE is_system_default = 1 - ) AS ranked_defaults - WHERE default_rank > 1 -); - -UPDATE routing_group_bindings -SET is_default = 0 -WHERE id IN ( - SELECT id - FROM ( - SELECT - id, - ROW_NUMBER() OVER ( - PARTITION BY subject_type, subject_id - ORDER BY created_at ASC, id ASC - ) AS default_rank - FROM routing_group_bindings - WHERE is_default = 1 - ) AS ranked_defaults - WHERE default_rank > 1 -); - -SET @aether_routing_groups_default_guard_column_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.columns - WHERE table_schema = DATABASE() - AND table_name = 'routing_groups' - AND column_name = 'system_default_unique_guard' - ) = 0, - 'ALTER TABLE routing_groups ADD COLUMN system_default_unique_guard TINYINT GENERATED ALWAYS AS (CASE WHEN is_system_default = 1 THEN 1 ELSE NULL END) VIRTUAL INVISIBLE', - 'DO 0' -); - -PREPARE aether_routing_groups_default_guard_column_stmt - FROM @aether_routing_groups_default_guard_column_sql; -EXECUTE aether_routing_groups_default_guard_column_stmt; -DEALLOCATE PREPARE aether_routing_groups_default_guard_column_stmt; - -SET @aether_routing_bindings_default_type_guard_column_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.columns - WHERE table_schema = DATABASE() - AND table_name = 'routing_group_bindings' - AND column_name = 'default_subject_type_guard' - ) = 0, - 'ALTER TABLE routing_group_bindings ADD COLUMN default_subject_type_guard VARCHAR(32) GENERATED ALWAYS AS (CASE WHEN is_default = 1 THEN subject_type ELSE NULL END) VIRTUAL INVISIBLE', - 'DO 0' -); - -PREPARE aether_routing_bindings_default_type_guard_column_stmt - FROM @aether_routing_bindings_default_type_guard_column_sql; -EXECUTE aether_routing_bindings_default_type_guard_column_stmt; -DEALLOCATE PREPARE aether_routing_bindings_default_type_guard_column_stmt; - -SET @aether_routing_bindings_default_id_guard_column_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.columns - WHERE table_schema = DATABASE() - AND table_name = 'routing_group_bindings' - AND column_name = 'default_subject_id_guard' - ) = 0, - 'ALTER TABLE routing_group_bindings ADD COLUMN default_subject_id_guard VARCHAR(64) GENERATED ALWAYS AS (CASE WHEN is_default = 1 THEN subject_id ELSE NULL END) VIRTUAL INVISIBLE', - 'DO 0' -); - -PREPARE aether_routing_bindings_default_id_guard_column_stmt - FROM @aether_routing_bindings_default_id_guard_column_sql; -EXECUTE aether_routing_bindings_default_id_guard_column_stmt; -DEALLOCATE PREPARE aether_routing_bindings_default_id_guard_column_stmt; - -SET @aether_routing_groups_default_unique_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'routing_groups' - AND index_name = 'routing_groups_one_system_default_key' - ) = 0, - 'CREATE UNIQUE INDEX routing_groups_one_system_default_key ON routing_groups (system_default_unique_guard)', - 'DO 0' -); - -PREPARE aether_routing_groups_default_unique_index_stmt - FROM @aether_routing_groups_default_unique_index_sql; -EXECUTE aether_routing_groups_default_unique_index_stmt; -DEALLOCATE PREPARE aether_routing_groups_default_unique_index_stmt; - -SET @aether_routing_bindings_default_unique_index_sql := IF( - ( - SELECT COUNT(*) - FROM information_schema.statistics - WHERE table_schema = DATABASE() - AND table_name = 'routing_group_bindings' - AND index_name = 'routing_group_bindings_subject_default_key' - ) = 0, - 'CREATE UNIQUE INDEX routing_group_bindings_subject_default_key ON routing_group_bindings (default_subject_type_guard, default_subject_id_guard)', - 'DO 0' -); - -PREPARE aether_routing_bindings_default_unique_index_stmt - FROM @aether_routing_bindings_default_unique_index_sql; -EXECUTE aether_routing_bindings_default_unique_index_stmt; -DEALLOCATE PREPARE aether_routing_bindings_default_unique_index_stmt; diff --git a/crates/aether-data/adapters/mysql/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql b/crates/aether-data/adapters/mysql/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql deleted file mode 100644 index f7d07d142..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql +++ /dev/null @@ -1,22 +0,0 @@ --- Worker supervisors are registered as one logical row per task. Older binaries --- included the ephemeral gateway instance in the row id, leaving a permanently --- running row after every restart. Remove only those system-generated boot rows; --- current workers recreate the stable logical rows after migrations complete. --- The metadata predicate also replaces task-only rows written by early builds of --- this fix that still claimed an instance owner. Delete children explicitly so --- cleanup remains complete after imports performed with FK checks disabled. -DELETE FROM background_task_events -WHERE run_id IN ( - SELECT id - FROM background_task_runs - WHERE id LIKE 'boot:%' - AND owner_instance IS NOT NULL - AND created_by = 'system' - AND progress_message = 'worker booted' -); - -DELETE FROM background_task_runs -WHERE id LIKE 'boot:%' - AND owner_instance IS NOT NULL - AND created_by = 'system' - AND progress_message = 'worker booted'; diff --git a/crates/aether-data/adapters/mysql/migrations/20260814000000_add_usage_cost_reservations.sql b/crates/aether-data/adapters/mysql/migrations/20260814000000_add_usage_cost_reservations.sql deleted file mode 100644 index aa0ac7c88..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260814000000_add_usage_cost_reservations.sql +++ /dev/null @@ -1,35 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_cost_reservations ( - `request_id` VARCHAR(128) NOT NULL, - `subject_id` VARCHAR(128) NOT NULL, - `reservation_token` VARCHAR(128) NOT NULL, - `admitted_at` BIGINT NOT NULL, - `reserved_cost_units` BIGINT NOT NULL, - `actual_cost_units` BIGINT, - `state` VARCHAR(20) NOT NULL, - `reservation_expires_at` BIGINT NOT NULL, - `retain_until` BIGINT NOT NULL, - `finalized_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`reservation_token`), - CONSTRAINT usage_cost_reservations_state_check - CHECK (`state` IN ('reserved', 'finalized', 'released')), - CONSTRAINT usage_cost_reservations_reserved_cost_units_check - CHECK (`reserved_cost_units` >= 0), - CONSTRAINT usage_cost_reservations_actual_cost_units_check - CHECK (`actual_cost_units` IS NULL OR `actual_cost_units` >= 0), - CONSTRAINT usage_cost_reservations_expiry_check - CHECK (`reservation_expires_at` > `admitted_at`), - CONSTRAINT usage_cost_reservations_retention_check - CHECK (`retain_until` >= `reservation_expires_at`), - CONSTRAINT usage_cost_reservations_lifecycle_check CHECK ( - (`state` = 'reserved' AND `actual_cost_units` IS NULL AND `finalized_at` IS NULL) - OR (`state` = 'finalized' AND `actual_cost_units` IS NOT NULL AND `finalized_at` IS NOT NULL) - OR (`state` = 'released' AND `actual_cost_units` IS NOT NULL - AND `actual_cost_units` = 0 AND `finalized_at` IS NOT NULL) - ), - KEY usage_cost_reservations_request_id_idx (`request_id`), - KEY usage_cost_reservations_subject_admitted_at_idx (`subject_id`, `admitted_at`), - KEY usage_cost_reservations_reservation_expires_at_idx (`reservation_expires_at`), - KEY usage_cost_reservations_retain_until_token_idx (`retain_until`, `reservation_token`) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260815000000_add_usage_request_admissions.sql b/crates/aether-data/adapters/mysql/migrations/20260815000000_add_usage_request_admissions.sql deleted file mode 100644 index 7047f7711..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260815000000_add_usage_request_admissions.sql +++ /dev/null @@ -1,22 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_request_admissions ( - `request_id` VARCHAR(128) NOT NULL, - `subject_id` VARCHAR(128) NOT NULL, - `event_token` VARCHAR(128) NOT NULL, - `admitted_at` BIGINT NOT NULL, - `retain_until` BIGINT NOT NULL, - `state` VARCHAR(20) NOT NULL, - `released_at` BIGINT, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`event_token`), - CONSTRAINT usage_request_admissions_retention_check - CHECK (`retain_until` > `admitted_at`), - CONSTRAINT usage_request_admissions_state_check - CHECK (`state` IN ('active', 'released')), - CONSTRAINT usage_request_admissions_lifecycle_check CHECK ( - (`state` = 'active' AND `released_at` IS NULL) - OR (`state` = 'released' AND `released_at` IS NOT NULL - AND `released_at` >= `admitted_at`) - ), - KEY usage_request_admissions_subject_admitted_at_idx (`subject_id`, `admitted_at`), - KEY usage_request_admissions_retain_until_token_idx (`retain_until`, `event_token`) -); diff --git a/crates/aether-data/adapters/mysql/migrations/20260816000000_add_usage_cost_reservation_user_foreign_key.sql b/crates/aether-data/adapters/mysql/migrations/20260816000000_add_usage_cost_reservation_user_foreign_key.sql deleted file mode 100644 index de9f0e7e3..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260816000000_add_usage_cost_reservation_user_foreign_key.sql +++ /dev/null @@ -1,12 +0,0 @@ --- Preserve the already-published cost-ledger migration checksum. This --- follow-up also upgrades databases which ran the feature branch before user --- ownership was enforced. Keep one ALTER TABLE per migration because MySQL --- DDL implicitly commits. -DELETE reservation -FROM usage_cost_reservations AS reservation -LEFT JOIN users AS app_user ON app_user.id = reservation.subject_id -WHERE app_user.id IS NULL; - -ALTER TABLE usage_cost_reservations - ADD CONSTRAINT usage_cost_reservations_subject_id_fkey - FOREIGN KEY (`subject_id`) REFERENCES users (`id`) ON DELETE CASCADE; diff --git a/crates/aether-data/adapters/mysql/migrations/20260817000000_add_usage_request_admission_user_foreign_key.sql b/crates/aether-data/adapters/mysql/migrations/20260817000000_add_usage_request_admission_user_foreign_key.sql deleted file mode 100644 index c294f4b6a..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260817000000_add_usage_request_admission_user_foreign_key.sql +++ /dev/null @@ -1,10 +0,0 @@ --- Split from the cost-ledger foreign key so a MySQL implicit DDL commit cannot --- leave two table changes behind one dirty migration record. -DELETE admission -FROM usage_request_admissions AS admission -LEFT JOIN users AS app_user ON app_user.id = admission.subject_id -WHERE app_user.id IS NULL; - -ALTER TABLE usage_request_admissions - ADD CONSTRAINT usage_request_admissions_subject_id_fkey - FOREIGN KEY (`subject_id`) REFERENCES users (`id`) ON DELETE CASCADE; diff --git a/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql b/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql deleted file mode 100644 index 7f000be64..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql +++ /dev/null @@ -1,203 +0,0 @@ --- Preserve the explicit access granted by the original #741 implementation, --- which represented Codex Live as openai:responses. Runtime permission --- matching remains strict after this one-time data upgrade. - -UPDATE users -SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') -WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' - AND JSON_VALID(allowed_api_formats) - AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('openai:responses'), '$' - ) = 1 - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('codex:live'), '$' - ) = 0; - -UPDATE user_groups -SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') -WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' - AND JSON_VALID(allowed_api_formats) - AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('openai:responses'), '$' - ) = 1 - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('codex:live'), '$' - ) = 0; - -UPDATE api_keys -SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') -WHERE JSON_VALID(allowed_api_formats) - AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('openai:responses'), '$' - ) = 1 - AND JSON_CONTAINS( - IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), - JSON_QUOTE('codex:live'), '$' - ) = 0; - -UPDATE provider_api_keys AS provider_key -INNER JOIN providers AS provider ON provider.id = provider_key.provider_id -SET - provider_key.api_formats = JSON_ARRAY_APPEND(provider_key.api_formats, '$', 'codex:live'), - provider_key.updated_at = UNIX_TIMESTAMP() -WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' - AND JSON_VALID(provider_key.api_formats) - AND JSON_TYPE( - IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, NULL) - ) = 'ARRAY' - AND JSON_CONTAINS( - IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, JSON_ARRAY()), - JSON_QUOTE('openai:responses'), '$' - ) = 1 - AND JSON_CONTAINS( - IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, JSON_ARRAY()), - JSON_QUOTE('codex:live'), '$' - ) = 0; - -UPDATE provider_api_keys AS provider_key -INNER JOIN providers AS provider ON provider.id = provider_key.provider_id -SET - provider_key.auth_type_by_format = JSON_SET( - provider_key.auth_type_by_format, - '$."codex:live"', - JSON_EXTRACT(provider_key.auth_type_by_format, '$."openai:responses"') - ), - provider_key.updated_at = UNIX_TIMESTAMP() -WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' - AND JSON_VALID(provider_key.auth_type_by_format) - AND JSON_TYPE( - IF(JSON_VALID(provider_key.auth_type_by_format), provider_key.auth_type_by_format, NULL) - ) = 'OBJECT' - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.auth_type_by_format), - provider_key.auth_type_by_format, - JSON_OBJECT() - ), - 'one', - '$."openai:responses"' - ) = 1 - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.auth_type_by_format), - provider_key.auth_type_by_format, - JSON_OBJECT() - ), - 'one', - '$."codex:live"' - ) = 0; - -UPDATE provider_api_keys AS provider_key -INNER JOIN providers AS provider ON provider.id = provider_key.provider_id -SET - provider_key.allow_auth_channel_mismatch_formats = JSON_ARRAY_APPEND( - provider_key.allow_auth_channel_mismatch_formats, - '$', - 'codex:live' - ), - provider_key.updated_at = UNIX_TIMESTAMP() -WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' - AND JSON_VALID(provider_key.allow_auth_channel_mismatch_formats) - AND JSON_TYPE( - IF( - JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), - provider_key.allow_auth_channel_mismatch_formats, - NULL - ) - ) = 'ARRAY' - AND JSON_CONTAINS( - IF( - JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), - provider_key.allow_auth_channel_mismatch_formats, - JSON_ARRAY() - ), - JSON_QUOTE('openai:responses'), - '$' - ) = 1 - AND JSON_CONTAINS( - IF( - JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), - provider_key.allow_auth_channel_mismatch_formats, - JSON_ARRAY() - ), - JSON_QUOTE('codex:live'), - '$' - ) = 0; - -UPDATE provider_api_keys AS provider_key -INNER JOIN providers AS provider ON provider.id = provider_key.provider_id -SET - provider_key.rate_multipliers = JSON_SET( - provider_key.rate_multipliers, - '$."codex:live"', - JSON_EXTRACT(provider_key.rate_multipliers, '$."openai:responses"') - ), - provider_key.updated_at = UNIX_TIMESTAMP() -WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' - AND JSON_VALID(provider_key.rate_multipliers) - AND JSON_TYPE( - IF(JSON_VALID(provider_key.rate_multipliers), provider_key.rate_multipliers, NULL) - ) = 'OBJECT' - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.rate_multipliers), - provider_key.rate_multipliers, - JSON_OBJECT() - ), - 'one', - '$."openai:responses"' - ) = 1 - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.rate_multipliers), - provider_key.rate_multipliers, - JSON_OBJECT() - ), - 'one', - '$."codex:live"' - ) = 0; - -UPDATE provider_api_keys AS provider_key -INNER JOIN providers AS provider ON provider.id = provider_key.provider_id -SET - provider_key.global_priority_by_format = JSON_SET( - provider_key.global_priority_by_format, - '$."codex:live"', - JSON_EXTRACT(provider_key.global_priority_by_format, '$."openai:responses"') - ), - provider_key.updated_at = UNIX_TIMESTAMP() -WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' - AND JSON_VALID(provider_key.global_priority_by_format) - AND JSON_TYPE( - IF( - JSON_VALID(provider_key.global_priority_by_format), - provider_key.global_priority_by_format, - NULL - ) - ) = 'OBJECT' - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.global_priority_by_format), - provider_key.global_priority_by_format, - JSON_OBJECT() - ), - 'one', - '$."openai:responses"' - ) = 1 - AND JSON_CONTAINS_PATH( - IF( - JSON_VALID(provider_key.global_priority_by_format), - provider_key.global_priority_by_format, - JSON_OBJECT() - ), - 'one', - '$."codex:live"' - ) = 0; diff --git a/crates/aether-data/adapters/mysql/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql b/crates/aether-data/adapters/mysql/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql deleted file mode 100644 index 434e4c581..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql +++ /dev/null @@ -1,55 +0,0 @@ --- A gateway transaction identifier may repeat across payment methods, but --- must never identify two orders in the same normalized method. MySQL commits --- persistent DDL implicitly, so reject historical conflicts before any --- persistent UPDATE or ALTER TABLE. CREATE/DROP TEMPORARY TABLE do not cause an --- implicit commit, and the leading DROP also makes a same-session retry safe. --- Diagnose with: --- SELECT LOWER(TRIM(payment_method)), --- CONVERT(gateway_order_id USING utf8mb4) COLLATE utf8mb4_0900_bin, --- COUNT(*) --- FROM payment_orders --- WHERE gateway_order_id IS NOT NULL --- GROUP BY LOWER(TRIM(payment_method)), --- CONVERT(gateway_order_id USING utf8mb4) COLLATE utf8mb4_0900_bin --- HAVING COUNT(*) > 1; -DROP TEMPORARY TABLE IF EXISTS aether_payment_gateway_order_uniqueness_preflight; - -CREATE TEMPORARY TABLE aether_payment_gateway_order_uniqueness_preflight ( - conflict_marker TINYINT NOT NULL PRIMARY KEY -); - -INSERT INTO aether_payment_gateway_order_uniqueness_preflight (conflict_marker) -VALUES (1); - --- Inserting the same marker fails on the first conflicting group. The grouping --- mirrors the values and collations used by the normalization and final index: --- payment methods use their existing column collation after LOWER/TRIM, while --- opaque gateway identifiers use MySQL 8's case-sensitive binary collation. -INSERT INTO aether_payment_gateway_order_uniqueness_preflight (conflict_marker) -SELECT 1 -FROM payment_orders -WHERE gateway_order_id IS NOT NULL -GROUP BY - LOWER(TRIM(payment_method)), - CONVERT(gateway_order_id USING utf8mb4) COLLATE utf8mb4_0900_bin -HAVING COUNT(*) > 1 -LIMIT 1; - -DROP TEMPORARY TABLE aether_payment_gateway_order_uniqueness_preflight; - -UPDATE payment_orders -SET payment_method = LOWER(TRIM(payment_method)) -WHERE BINARY payment_method <> BINARY LOWER(TRIM(payment_method)); - -UPDATE payment_callbacks -SET payment_method = LOWER(TRIM(payment_method)) -WHERE BINARY payment_method <> BINARY LOWER(TRIM(payment_method)); - --- Gateway identifiers are opaque and case-sensitive. Changing the column --- collation and adding the unique index in one ALTER avoids a persistent --- intermediate schema if either operation fails. -ALTER TABLE payment_orders - MODIFY COLUMN gateway_order_id VARCHAR(128) - CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin NULL, - ADD UNIQUE INDEX uq_payment_orders_payment_method_gateway_order_id - (payment_method, gateway_order_id); diff --git a/crates/aether-data/adapters/mysql/migrations/20260821130000_add_user_security_version.sql b/crates/aether-data/adapters/mysql/migrations/20260821130000_add_user_security_version.sql deleted file mode 100644 index 9e68683bb..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260821130000_add_user_security_version.sql +++ /dev/null @@ -1,5 +0,0 @@ -ALTER TABLE users - ADD COLUMN security_version BIGINT NOT NULL DEFAULT 0; - -ALTER TABLE user_sessions - ADD COLUMN security_version BIGINT NOT NULL DEFAULT 0; diff --git a/crates/aether-data/adapters/mysql/migrations/20260827040000_expand_proxy_password_ciphertext.sql b/crates/aether-data/adapters/mysql/migrations/20260827040000_expand_proxy_password_ciphertext.sql deleted file mode 100644 index c84e60c0b..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260827040000_expand_proxy_password_ciphertext.sql +++ /dev/null @@ -1,3 +0,0 @@ --- Purpose-bound Fernet envelopes can exceed the former 500-character plaintext limit. -ALTER TABLE proxy_nodes - MODIFY COLUMN proxy_password TEXT NULL; diff --git a/crates/aether-data/adapters/mysql/migrations/20260827050000_anonymize_deleted_user_history.sql b/crates/aether-data/adapters/mysql/migrations/20260827050000_anonymize_deleted_user_history.sql deleted file mode 100644 index f8a2b9cb2..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260827050000_anonymize_deleted_user_history.sql +++ /dev/null @@ -1,93 +0,0 @@ -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'user_plan_entitlements' - AND CONSTRAINT_NAME = 'user_plan_entitlements_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE user_plan_entitlements DROP FOREIGN KEY user_plan_entitlements_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'entitlement_usage_ledgers' - AND CONSTRAINT_NAME = 'entitlement_usage_ledgers_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE entitlement_usage_ledgers DROP FOREIGN KEY entitlement_usage_ledgers_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'user_referrals' - AND CONSTRAINT_NAME = 'user_referrals_inviter_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE user_referrals DROP FOREIGN KEY user_referrals_inviter_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'user_referrals' - AND CONSTRAINT_NAME = 'user_referrals_invitee_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE user_referrals DROP FOREIGN KEY user_referrals_invitee_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'referral_rewards' - AND CONSTRAINT_NAME = 'referral_rewards_inviter_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE referral_rewards DROP FOREIGN KEY referral_rewards_inviter_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - -SET @aether_drop_fact_user_fk_sql := IF( - EXISTS ( - SELECT 1 FROM information_schema.TABLE_CONSTRAINTS - WHERE CONSTRAINT_SCHEMA = DATABASE() - AND TABLE_NAME = 'referral_rewards' - AND CONSTRAINT_NAME = 'referral_rewards_invitee_user_id_fkey' - AND CONSTRAINT_TYPE = 'FOREIGN KEY' - ), - 'ALTER TABLE referral_rewards DROP FOREIGN KEY referral_rewards_invitee_user_id_fkey', - 'DO 0' -); -PREPARE aether_drop_fact_user_fk_stmt FROM @aether_drop_fact_user_fk_sql; -EXECUTE aether_drop_fact_user_fk_stmt; -DEALLOCATE PREPARE aether_drop_fact_user_fk_stmt; - --- Keep existing historical row values unchanged. Runtime deletion paths enforce --- the current anonymization policy for newly deleted users. -SELECT 1; diff --git a/crates/aether-data/adapters/mysql/migrations/20260831000000_enforce_ldap_config_singleton.sql b/crates/aether-data/adapters/mysql/migrations/20260831000000_enforce_ldap_config_singleton.sql deleted file mode 100644 index c8dae1f65..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260831000000_enforce_ldap_config_singleton.sql +++ /dev/null @@ -1,13 +0,0 @@ --- LDAP configuration is a database-wide singleton. Preserve the row selected by the legacy --- reader (the smallest id), remove historical duplicates, and let the database arbitrate --- concurrent first creation. -DELETE FROM ldap_configs -WHERE id <> ( - SELECT keep_id - FROM (SELECT MIN(id) AS keep_id FROM ldap_configs) AS ldap_singleton_keeper -); - -ALTER TABLE ldap_configs - ADD COLUMN singleton_key INT NOT NULL DEFAULT 1, - ADD CONSTRAINT ldap_configs_singleton_key_check CHECK (singleton_key = 1), - ADD UNIQUE KEY ldap_configs_singleton_key_key (singleton_key); diff --git a/crates/aether-data/adapters/mysql/migrations/20260831010000_add_proxy_node_tunnel_generation.sql b/crates/aether-data/adapters/mysql/migrations/20260831010000_add_proxy_node_tunnel_generation.sql deleted file mode 100644 index 5b5d64b6e..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260831010000_add_proxy_node_tunnel_generation.sql +++ /dev/null @@ -1,9 +0,0 @@ -ALTER TABLE proxy_nodes - ADD COLUMN tunnel_generation VARCHAR(64) NULL AFTER id; - -UPDATE proxy_nodes -SET tunnel_generation = UUID() -WHERE tunnel_generation IS NULL OR TRIM(tunnel_generation) = ''; - -ALTER TABLE proxy_nodes - MODIFY COLUMN tunnel_generation VARCHAR(64) NOT NULL; diff --git a/crates/aether-data/adapters/mysql/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql b/crates/aether-data/adapters/mysql/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql deleted file mode 100644 index ce2a46cfc..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql +++ /dev/null @@ -1,10 +0,0 @@ --- A proxy endpoint has one stable node identity across manual and tunnel --- registrations. MySQL 8 performs ALTER TABLE atomically; if historical --- duplicates exist this migration fails without choosing or deleting a row. --- Diagnose with: --- SELECT ip, port, COUNT(*) --- FROM proxy_nodes --- GROUP BY ip, port --- HAVING COUNT(*) > 1; -ALTER TABLE proxy_nodes - ADD UNIQUE INDEX uq_proxy_node_ip_port (ip, port); diff --git a/crates/aether-data/adapters/mysql/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql b/crates/aether-data/adapters/mysql/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql deleted file mode 100644 index 1d0fdf821..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE usage_counter_deltas - ADD COLUMN target_tunnel_generation VARCHAR(64) NULL; diff --git a/crates/aether-data/adapters/mysql/migrations/20260903000000_add_routing_group_sort_order.sql b/crates/aether-data/adapters/mysql/migrations/20260903000000_add_routing_group_sort_order.sql deleted file mode 100644 index f33fd1096..000000000 --- a/crates/aether-data/adapters/mysql/migrations/20260903000000_add_routing_group_sort_order.sql +++ /dev/null @@ -1,3 +0,0 @@ -ALTER TABLE routing_groups - ADD COLUMN sort_order BIGINT NOT NULL DEFAULT 0, - ADD KEY routing_groups_enabled_sort_idx (enabled, sort_order, name, id); diff --git a/crates/aether-data/adapters/mysql/src/announcements.rs b/crates/aether-data/adapters/mysql/src/announcements.rs deleted file mode 100644 index c364eab2f..000000000 --- a/crates/aether-data/adapters/mysql/src/announcements.rs +++ /dev/null @@ -1,365 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; - -use aether_data_contracts::repository::announcements::*; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const ANNOUNCEMENT_SELECT: &str = r#" -SELECT - a.id, - a.title, - a.content, - a.`type` AS type, - a.priority, - a.is_active, - a.is_pinned, - a.requires_ack, - a.author_id, - u.username AS author_username, - a.start_time AS start_time_unix_secs, - a.end_time AS end_time_unix_secs, - a.created_at AS created_at_unix_ms, - a.updated_at AS updated_at_unix_secs -FROM announcements a -LEFT JOIN users u ON u.id = a.author_id -"#; - -#[derive(Debug, Clone)] -pub struct MysqlAnnouncementRepository { - pool: MysqlPool, -} - -impl MysqlAnnouncementRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn reload_by_id( - &self, - announcement_id: &str, - ) -> Result, DataLayerError> { - self.find_by_id(announcement_id).await - } -} - -#[async_trait] -impl AnnouncementReadRepository for MysqlAnnouncementRepository { - async fn find_by_id( - &self, - announcement_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{ANNOUNCEMENT_SELECT} WHERE a.id = ? LIMIT 1")) - .bind(announcement_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_announcement_row).transpose() - } - - async fn list_announcements( - &self, - query: &AnnouncementListQuery, - ) -> Result { - let now_unix_secs = query.now_unix_secs.unwrap_or_else(current_unix_secs); - let total_row = sqlx::query( - r#" -SELECT COUNT(a.id) AS total -FROM announcements a -WHERE ( - NOT ? OR ( - a.is_active = 1 - AND (a.start_time IS NULL OR a.start_time <= ?) - AND (a.end_time IS NULL OR a.end_time >= ?) - ) -) -"#, - ) - .bind(query.active_only) - .bind(now_unix_secs as i64) - .bind(now_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total = total_row.try_get::("total").map_sql_err()?.max(0) as u64; - - let rows = sqlx::query(&format!( - r#" -{ANNOUNCEMENT_SELECT} -WHERE ( - NOT ? OR ( - a.is_active = 1 - AND (a.start_time IS NULL OR a.start_time <= ?) - AND (a.end_time IS NULL OR a.end_time >= ?) - ) -) -ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC -LIMIT ? OFFSET ? -"# - )) - .bind(query.active_only) - .bind(now_unix_secs as i64) - .bind(now_unix_secs as i64) - .bind(query.limit as i64) - .bind(query.offset as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_announcement_row) - .collect::, _>>()?; - - Ok(StoredAnnouncementPage { items, total }) - } - - async fn count_unread_active_announcements( - &self, - user_id: &str, - now_unix_secs: u64, - ) -> Result { - let row = sqlx::query( - r#" -SELECT COUNT(a.id) AS total -FROM announcements a -WHERE a.is_active = 1 - AND (a.start_time IS NULL OR a.start_time <= ?) - AND (a.end_time IS NULL OR a.end_time >= ?) - AND NOT EXISTS ( - SELECT 1 - FROM announcement_reads r - WHERE r.user_id = ? - AND r.announcement_id = a.id - ) -"#, - ) - .bind(now_unix_secs as i64) - .bind(now_unix_secs as i64) - .bind(user_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - async fn list_required_unread_active_announcements( - &self, - user_id: &str, - now_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{ANNOUNCEMENT_SELECT} -WHERE a.is_active = 1 - AND a.requires_ack = 1 - AND (a.start_time IS NULL OR a.start_time <= ?) - AND (a.end_time IS NULL OR a.end_time >= ?) - AND NOT EXISTS ( - SELECT 1 - FROM announcement_reads r - WHERE r.user_id = ? - AND r.announcement_id = a.id - ) -ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC -LIMIT ? -"# - )) - .bind(now_unix_secs as i64) - .bind(now_unix_secs as i64) - .bind(user_id) - .bind(limit as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_announcement_row).collect() - } -} - -#[async_trait] -impl AnnouncementWriteRepository for MysqlAnnouncementRepository { - async fn create_announcement( - &self, - record: CreateAnnouncementRecord, - ) -> Result { - record.validate()?; - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs() as i64; - sqlx::query( - r#" -INSERT INTO announcements ( - id, title, content, `type`, priority, author_id, is_active, is_pinned, - requires_ack, start_time, end_time, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&id) - .bind(record.title) - .bind(record.content) - .bind(record.kind) - .bind(record.priority) - .bind(record.author_id) - .bind(record.is_pinned) - .bind(record.requires_ack) - .bind(optional_i64_from_u64( - record.start_time_unix_secs, - "announcements.start_time", - )?) - .bind(optional_i64_from_u64( - record.end_time_unix_secs, - "announcements.end_time", - )?) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_id(&id) - .await? - .ok_or_else(|| DataLayerError::UnexpectedValue("created announcement missing".into())) - } - - async fn update_announcement( - &self, - record: UpdateAnnouncementRecord, - ) -> Result, DataLayerError> { - record.validate()?; - let id = record.announcement_id; - sqlx::query( - r#" -UPDATE announcements -SET title = COALESCE(?, title), - content = COALESCE(?, content), - `type` = COALESCE(?, `type`), - priority = COALESCE(?, priority), - is_active = COALESCE(?, is_active), - is_pinned = COALESCE(?, is_pinned), - requires_ack = COALESCE(?, requires_ack), - start_time = COALESCE(?, start_time), - end_time = COALESCE(?, end_time), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(record.title) - .bind(record.content) - .bind(record.kind) - .bind(record.priority) - .bind(record.is_active) - .bind(record.is_pinned) - .bind(record.requires_ack) - .bind(optional_i64_from_u64( - record.start_time_unix_secs, - "announcements.start_time", - )?) - .bind(optional_i64_from_u64( - record.end_time_unix_secs, - "announcements.end_time", - )?) - .bind(current_unix_secs() as i64) - .bind(&id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_id(&id).await - } - - async fn delete_announcement(&self, announcement_id: &str) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("DELETE FROM announcement_reads WHERE announcement_id = ?") - .bind(announcement_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let rows_affected = sqlx::query("DELETE FROM announcements WHERE id = ?") - .bind(announcement_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - tx.commit().await.map_sql_err()?; - Ok(rows_affected > 0) - } - - async fn mark_announcement_as_read( - &self, - user_id: &str, - announcement_id: &str, - read_at_unix_secs: u64, - ) -> Result { - let rows_affected = sqlx::query( - r#" -INSERT IGNORE INTO announcement_reads (id, user_id, announcement_id, read_at) -VALUES (?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(user_id) - .bind(announcement_id) - .bind(i64_from_u64( - read_at_unix_secs, - "announcement_reads.read_at", - )?) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| i64_from_u64(value, field_name)) - .transpose() -} - -fn map_announcement_row(row: &MySqlRow) -> Result { - StoredAnnouncement::new( - row.try_get("id").map_sql_err()?, - row.try_get("title").map_sql_err()?, - row.try_get("content").map_sql_err()?, - row.try_get("type").map_sql_err()?, - row.try_get("priority").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_pinned").map_sql_err()?, - row.try_get("requires_ack").map_sql_err()?, - row.try_get("author_id").map_sql_err()?, - row.try_get("author_username").map_sql_err()?, - row.try_get("start_time_unix_secs").map_sql_err()?, - row.try_get("end_time_unix_secs").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - ) -} - -#[cfg(test)] -mod tests { - use super::MysqlAnnouncementRepository; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlAnnouncementRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/audit.rs b/crates/aether-data/adapters/mysql/src/audit.rs deleted file mode 100644 index 2ac958607..000000000 --- a/crates/aether-data/adapters/mysql/src/audit.rs +++ /dev/null @@ -1,277 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; - -use aether_data_contracts::repository::audit::*; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -#[derive(Debug, Clone)] -pub struct MysqlAuditLogReadRepository { - pool: MysqlPool, -} - -impl MysqlAuditLogReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl AuditLogReadRepository for MysqlAuditLogReadRepository { - async fn list_admin_audit_logs( - &self, - query: &AuditLogListQuery, - ) -> Result { - let total = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM audit_logs AS a -LEFT JOIN users AS u ON a.user_id = u.id -WHERE a.created_at >= ? - AND (? IS NULL OR LOWER(u.username) LIKE LOWER(?) ESCAPE '\\') - AND (? IS NULL OR a.event_type = ?) -"#, - ) - .bind(query.cutoff_unix_secs as i64) - .bind(query.username_pattern.as_deref()) - .bind(query.username_pattern.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query( - r#" -SELECT - a.id, - a.event_type, - a.user_id, - u.email AS user_email, - u.username AS user_username, - a.description, - a.ip_address, - a.status_code, - a.error_message, - a.event_metadata AS metadata, - a.created_at -FROM audit_logs AS a -LEFT JOIN users AS u ON a.user_id = u.id -WHERE a.created_at >= ? - AND (? IS NULL OR LOWER(u.username) LIKE LOWER(?) ESCAPE '\\') - AND (? IS NULL OR a.event_type = ?) -ORDER BY a.created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.cutoff_unix_secs as i64) - .bind(query.username_pattern.as_deref()) - .bind(query.username_pattern.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.limit as i64) - .bind(query.offset as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let items = rows - .iter() - .map(map_mysql_admin_audit_log_row) - .collect::, _>>()?; - - Ok(StoredAdminAuditLogPage { - items, - total: total.max(0) as u64, - }) - } - - async fn list_admin_suspicious_activities( - &self, - cutoff_unix_secs: u64, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT id, event_type, user_id, description, ip_address, event_metadata AS metadata, created_at -FROM audit_logs -WHERE created_at >= ? - AND event_type IN (?, ?, ?, ?) -ORDER BY created_at DESC -LIMIT 100 -"#, - ) - .bind(cutoff_unix_secs as i64) - .bind(SUSPICIOUS_EVENT_TYPES[0]) - .bind(SUSPICIOUS_EVENT_TYPES[1]) - .bind(SUSPICIOUS_EVENT_TYPES[2]) - .bind(SUSPICIOUS_EVENT_TYPES[3]) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - rows.iter().map(map_mysql_suspicious_activity_row).collect() - } - - async fn read_admin_user_behavior_event_counts( - &self, - user_id: &str, - cutoff_unix_secs: u64, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT event_type, COUNT(*) AS count -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? -GROUP BY event_type -"#, - ) - .bind(user_id) - .bind(cutoff_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - Ok(rows - .iter() - .filter_map(|row| event_count_from_mysql_row(row).ok()) - .collect()) - } - - async fn list_user_audit_logs( - &self, - user_id: &str, - query: &AuditLogListQuery, - ) -> Result { - let total = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? - AND (? IS NULL OR event_type = ?) -"#, - ) - .bind(user_id) - .bind(query.cutoff_unix_secs as i64) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query( - r#" -SELECT id, event_type, description, ip_address, status_code, created_at -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? - AND (? IS NULL OR event_type = ?) -ORDER BY created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(user_id) - .bind(query.cutoff_unix_secs as i64) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.limit as i64) - .bind(query.offset as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let items = rows - .iter() - .map(map_mysql_user_audit_log_row) - .collect::, _>>()?; - - Ok(StoredUserAuditLogPage { - items, - total: total.max(0) as u64, - }) - } - - async fn delete_audit_logs_before( - &self, - cutoff_unix_secs: u64, - limit: usize, - ) -> Result { - let deleted = sqlx::query( - r#" -DELETE FROM audit_logs -WHERE id IN ( - SELECT id - FROM ( - SELECT id - FROM audit_logs - WHERE created_at < ? - ORDER BY created_at ASC, id ASC - LIMIT ? - ) AS doomed -) -"#, - ) - .bind(cutoff_unix_secs.min(i64::MAX as u64) as i64) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(deleted).unwrap_or(usize::MAX)) - } -} - -fn mysql_created_at_unix_secs(row: &MySqlRow) -> Result { - let value = row.try_get::("created_at").map_sql_err()?; - Ok(value.max(0) as u64) -} - -fn map_mysql_admin_audit_log_row(row: &MySqlRow) -> Result { - Ok(StoredAdminAuditLog { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - user_email: row.try_get("user_email").map_sql_err()?, - user_username: row.try_get("user_username").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - status_code: row.try_get("status_code").map_sql_err()?, - error_message: row.try_get("error_message").map_sql_err()?, - metadata: optional_json_from_text(row.try_get("metadata").map_sql_err()?)?, - created_at_unix_secs: mysql_created_at_unix_secs(row)?, - }) -} - -fn map_mysql_suspicious_activity_row( - row: &MySqlRow, -) -> Result { - Ok(StoredSuspiciousActivity { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - metadata: optional_json_from_text(row.try_get("metadata").map_sql_err()?)?, - created_at_unix_secs: mysql_created_at_unix_secs(row)?, - }) -} - -fn map_mysql_user_audit_log_row(row: &MySqlRow) -> Result { - Ok(StoredUserAuditLog { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - status_code: row.try_get("status_code").map_sql_err()?, - created_at_unix_secs: mysql_created_at_unix_secs(row)?, - }) -} - -fn event_count_from_mysql_row(row: &MySqlRow) -> Result<(String, u64), DataLayerError> { - let event_type = row.try_get("event_type").map_sql_err()?; - let count = row.try_get::("count").map_sql_err()?.max(0) as u64; - Ok((event_type, count)) -} diff --git a/crates/aether-data/adapters/mysql/src/auth.rs b/crates/aether-data/adapters/mysql/src/auth.rs deleted file mode 100644 index 94615d813..000000000 --- a/crates/aether-data/adapters/mysql/src/auth.rs +++ /dev/null @@ -1,1498 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::auth::{ - AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository, - AuthApiKeyWriteRepository, CompareAndSwapAuthApiKeyCiphertext, CreateStandaloneApiKeyRecord, - CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, - StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const SNAPSHOT_COLUMNS: &str = r#" -SELECT - users.id AS user_id, - users.username, - users.email, - users.role AS user_role, - users.auth_source AS user_auth_source, - users.is_active AS user_is_active, - users.is_deleted AS user_is_deleted, - users.rate_limit AS user_rate_limit, - users.allowed_providers AS user_allowed_providers, - users.allowed_api_formats AS user_allowed_api_formats, - users.allowed_models AS user_allowed_models, - api_keys.id AS api_key_id, - api_keys.name AS api_key_name, - api_keys.is_active AS api_key_is_active, - api_keys.is_locked AS api_key_is_locked, - api_keys.is_standalone AS api_key_is_standalone, - api_keys.rate_limit AS api_key_rate_limit, - api_keys.concurrent_limit AS api_key_concurrent_limit, - api_keys.expires_at AS api_key_expires_at_unix_secs, - api_keys.allowed_providers AS api_key_allowed_providers, - api_keys.allowed_api_formats AS api_key_allowed_api_formats, - api_keys.allowed_models AS api_key_allowed_models, - api_keys.ip_rules AS api_key_ip_rules -FROM api_keys -JOIN users ON users.id = api_keys.user_id -"#; - -const EXPORT_COLUMNS: &str = r#" -SELECT - api_keys.user_id, - api_keys.id AS api_key_id, - api_keys.key_hash, - api_keys.key_encrypted, - api_keys.name, - api_keys.allowed_providers, - api_keys.allowed_api_formats, - api_keys.allowed_models, - api_keys.ip_rules, - api_keys.rate_limit, - api_keys.concurrent_limit, - api_keys.force_capabilities, - api_keys.feature_settings, - api_keys.is_active, - api_keys.expires_at AS expires_at_unix_secs, - api_keys.auto_delete_on_expiry, - api_keys.total_requests, - COALESCE(api_keys.total_tokens, 0) AS total_tokens, - COALESCE(api_keys.total_cost_usd, 0) AS total_cost_usd, - api_keys.last_used_at AS last_used_at_unix_secs, - api_keys.created_at AS created_at_unix_secs, - api_keys.updated_at AS updated_at_unix_secs, - api_keys.is_standalone -FROM api_keys -"#; - -const MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL: &[&str] = &[ - "UPDATE request_candidates SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE video_tasks SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE `usage` SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE stats_daily_api_key SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE audit_logs SET description = 'deleted API key event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE api_key_id = ?", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.api_key_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_orders SET gateway_response = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", -]; - -const MYSQL_DELETE_API_KEY_DEPENDENTS_SQL: &[&str] = - &["DELETE FROM api_key_provider_mappings WHERE api_key_id = ?"]; - -#[derive(Debug, Clone)] -pub struct MysqlAuthApiKeyReadRepository { - pool: MysqlPool, -} - -impl MysqlAuthApiKeyReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn fetch_snapshot_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_auth_api_key_snapshot_row).collect() - } - - async fn fetch_export_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_auth_api_key_export_row).collect() - } - - async fn reload_export_by_id( - &self, - api_key_id: &str, - ) -> Result, DataLayerError> { - Ok(self - .list_export_api_keys_by_ids(&[api_key_id.to_string()]) - .await? - .into_iter() - .next()) - } - - async fn create_api_key( - &self, - record: CreateApiKeyInsertRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let owner_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? AND is_deleted = 0 FOR UPDATE") - .bind(&record.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if owner_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query( - r#" -INSERT INTO api_keys ( - id, user_id, key_hash, key_encrypted, name, allowed_providers, - allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit, - force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry, - total_requests, total_tokens, total_cost_usd, is_standalone, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.api_key_id) - .bind(&record.user_id) - .bind(&record.key_hash) - .bind(&record.key_encrypted) - .bind(&record.name) - .bind(json_string_from_string_list( - record.allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(json_string_from_string_list( - record.allowed_api_formats.as_ref(), - "api_keys.allowed_api_formats", - )?) - .bind(json_string_from_string_list( - record.allowed_models.as_ref(), - "api_keys.allowed_models", - )?) - .bind(json_string_from_string_list( - record.ip_rules.as_ref(), - "api_keys.ip_rules", - )?) - .bind(record.rate_limit) - .bind(record.concurrent_limit) - .bind(optional_json_to_string( - &record.force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(optional_json_to_string( - &record.feature_settings, - "api_keys.feature_settings", - )?) - .bind(record.is_active) - .bind(optional_i64_from_u64( - record.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(record.auto_delete_on_expiry) - .bind(i64_from_u64( - record.total_requests, - "api_keys.total_requests", - )?) - .bind(i64_from_u64(record.total_tokens, "api_keys.total_tokens")?) - .bind(record.total_cost_usd) - .bind(record.is_standalone) - .bind(now as i64) - .bind(now as i64) - .execute(&mut *tx) - .await - .map_sql_err()?; - let reload_sql = format!("{EXPORT_COLUMNS}\nWHERE api_keys.id = ?\nLIMIT 1"); - let row = sqlx::query(&reload_sql) - .bind(&record.api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::UnexpectedValue(format!( - "created api_keys row is missing: {}", - record.api_key_id - ))); - }; - let created = map_auth_api_key_export_row(&row)?; - tx.commit().await.map_sql_err()?; - Ok(Some(created)) - } -} - -struct CreateApiKeyInsertRecord { - user_id: String, - api_key_id: String, - key_hash: String, - key_encrypted: Option, - name: Option, - allowed_providers: Option>, - allowed_api_formats: Option>, - allowed_models: Option>, - ip_rules: Option>, - rate_limit: Option, - concurrent_limit: Option, - force_capabilities: Option, - feature_settings: Option, - is_active: bool, - expires_at_unix_secs: Option, - auto_delete_on_expiry: bool, - total_requests: u64, - total_tokens: u64, - total_cost_usd: f64, - is_standalone: bool, -} - -#[async_trait] -impl AuthApiKeyReadRepository for MysqlAuthApiKeyReadRepository { - async fn find_api_key_snapshot( - &self, - key: AuthApiKeyLookupKey<'_>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SNAPSHOT_COLUMNS); - match key { - AuthApiKeyLookupKey::KeyHash(key_hash) => { - builder - .push(" WHERE api_keys.key_hash = ") - .push_bind(key_hash); - } - AuthApiKeyLookupKey::ApiKeyId(api_key_id) => { - builder.push(" WHERE api_keys.id = ").push_bind(api_key_id); - } - AuthApiKeyLookupKey::UserApiKeyIds { - user_id, - api_key_id, - } => { - builder - .push(" WHERE api_keys.id = ") - .push_bind(api_key_id) - .push(" AND users.id = ") - .push_bind(user_id); - } - } - builder.push(" LIMIT 1"); - Ok(self.fetch_snapshot_rows(builder).await?.into_iter().next()) - } - - async fn list_api_key_snapshots_by_ids( - &self, - api_key_ids: &[String], - ) -> Result, DataLayerError> { - if api_key_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(SNAPSHOT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.id IN (", api_key_ids); - builder.push(" ORDER BY api_keys.id ASC"); - self.fetch_snapshot_rows(builder).await - } - - async fn list_export_api_keys_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.user_id IN (", user_ids); - builder - .push(" AND api_keys.is_standalone = 0 ORDER BY api_keys.user_id ASC, api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_api_keys_by_ids( - &self, - api_key_ids: &[String], - ) -> Result, DataLayerError> { - if api_key_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.id IN (", api_key_ids); - builder.push(" ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_api_keys_by_name_search( - &self, - name_search: &str, - ) -> Result, DataLayerError> { - let name_search = name_search.trim(); - if name_search.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder - .push(" WHERE LOWER(COALESCE(api_keys.name, '')) LIKE ") - .push_bind(format!("%{}%", name_search.to_ascii_lowercase())) - .push(" ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_standalone_api_keys_page( - &self, - query: &StandaloneApiKeyExportListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder.push(" WHERE api_keys.is_standalone = 1"); - if let Some(is_active) = query.is_active { - builder - .push(" AND api_keys.is_active = ") - .push_bind(is_active); - } - builder - .push(" ORDER BY api_keys.id ASC LIMIT ") - .push_bind(i64::try_from(query.limit).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid standalone api key export limit: {}", - query.limit - )) - })?) - .push(" OFFSET ") - .push_bind(i64::try_from(query.skip).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid standalone api key export skip: {}", - query.skip - )) - })?); - self.fetch_export_rows(builder).await - } - - async fn count_export_standalone_api_keys( - &self, - is_active: Option, - ) -> Result { - let mut builder = QueryBuilder::::new( - "SELECT COUNT(*) AS total FROM api_keys WHERE is_standalone = 1", - ); - if let Some(is_active) = is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - async fn summarize_export_api_keys_by_user_ids( - &self, - user_ids: &[String], - now_unix_secs: u64, - ) -> Result { - if user_ids.is_empty() { - return Ok(AuthApiKeyExportSummary::default()); - } - let now_unix_secs = i64_from_u64(now_unix_secs, "api_keys.summary_now")?; - - let mut builder = QueryBuilder::::new( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 AND (expires_at IS NULL OR expires_at >= -"#, - ); - builder.push_bind(now_unix_secs); - builder.push( - r#") THEN 1 ELSE 0 END) AS active -FROM api_keys -"#, - ); - push_in_clause(&mut builder, " WHERE user_id IN (", user_ids); - builder.push(" AND is_standalone = 0"); - summarize_row(builder.build().fetch_one(&self.pool).await.map_sql_err()?) - } - - async fn summarize_export_non_standalone_api_keys( - &self, - now_unix_secs: u64, - ) -> Result { - summarize_api_keys(&self.pool, false, now_unix_secs).await - } - - async fn summarize_export_standalone_api_keys( - &self, - now_unix_secs: u64, - ) -> Result { - summarize_api_keys(&self.pool, true, now_unix_secs).await - } - - async fn find_export_standalone_api_key_by_id( - &self, - api_key_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder - .push(" WHERE api_keys.is_standalone = 1 AND api_keys.id = ") - .push_bind(api_key_id) - .push(" LIMIT 1"); - Ok(self.fetch_export_rows(builder).await?.into_iter().next()) - } - - async fn list_export_standalone_api_keys( - &self, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder.push(" WHERE api_keys.is_standalone = 1 ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } -} - -#[async_trait] -impl AuthApiKeyWriteRepository for MysqlAuthApiKeyReadRepository { - async fn touch_last_used_at(&self, api_key_id: &str) -> Result { - let now = current_unix_secs() as i64; - let rows_affected = sqlx::query( - r#" -UPDATE api_keys -SET last_used_at = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(now) - .bind(now) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn create_user_api_key( - &self, - record: CreateUserApiKeyRecord, - ) -> Result, DataLayerError> { - self.create_api_key(CreateApiKeyInsertRecord { - user_id: record.user_id, - api_key_id: record.api_key_id, - key_hash: record.key_hash, - key_encrypted: record.key_encrypted, - name: record.name, - allowed_providers: record.allowed_providers, - allowed_api_formats: record.allowed_api_formats, - allowed_models: record.allowed_models, - ip_rules: record.ip_rules, - rate_limit: Some(record.rate_limit), - concurrent_limit: record.concurrent_limit, - force_capabilities: record.force_capabilities, - feature_settings: record.feature_settings, - is_active: record.is_active, - expires_at_unix_secs: record.expires_at_unix_secs, - auto_delete_on_expiry: record.auto_delete_on_expiry, - total_requests: record.total_requests, - total_tokens: record.total_tokens, - total_cost_usd: record.total_cost_usd, - is_standalone: false, - }) - .await - } - - async fn create_standalone_api_key( - &self, - record: CreateStandaloneApiKeyRecord, - ) -> Result, DataLayerError> { - self.create_api_key(CreateApiKeyInsertRecord { - user_id: record.user_id, - api_key_id: record.api_key_id, - key_hash: record.key_hash, - key_encrypted: record.key_encrypted, - name: record.name, - allowed_providers: record.allowed_providers, - allowed_api_formats: record.allowed_api_formats, - allowed_models: record.allowed_models, - ip_rules: record.ip_rules, - rate_limit: record.rate_limit, - concurrent_limit: record.concurrent_limit, - force_capabilities: record.force_capabilities, - feature_settings: None, - is_active: record.is_active, - expires_at_unix_secs: record.expires_at_unix_secs, - auto_delete_on_expiry: record.auto_delete_on_expiry, - total_requests: record.total_requests, - total_tokens: record.total_tokens, - total_cost_usd: record.total_cost_usd, - is_standalone: true, - }) - .await - } - - async fn update_user_api_key_basic( - &self, - record: UpdateUserApiKeyBasicRecord, - ) -> Result, DataLayerError> { - self.update_user_api_key_basic_scoped(record, false).await - } - - async fn compare_and_swap_api_key_ciphertext( - &self, - mutation: &CompareAndSwapAuthApiKeyCiphertext, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = ? -WHERE BINARY id = BINARY ? - AND BINARY user_id = BINARY ? - AND BINARY key_hash = BINARY ? - AND is_standalone = ? - AND BINARY key_encrypted = BINARY ? -"#, - ) - .bind(&mutation.key_encrypted) - .bind(&mutation.api_key_id) - .bind(&mutation.user_id) - .bind(&mutation.key_hash) - .bind(mutation.is_standalone) - .bind(&mutation.expected_key_encrypted) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn update_user_api_key_basic_if_unlocked( - &self, - record: UpdateUserApiKeyBasicRecord, - ) -> Result, DataLayerError> { - self.update_user_api_key_basic_scoped(record, true).await - } - - async fn update_standalone_api_key_basic( - &self, - record: UpdateStandaloneApiKeyBasicRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs() as i64; - sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = CASE WHEN ? THEN ? ELSE key_encrypted END, - name = CASE WHEN ? THEN ? ELSE name END, - force_capabilities = CASE WHEN ? THEN ? ELSE force_capabilities END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, - allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, - allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, - allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END, - ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, - expires_at = CASE WHEN ? THEN ? ELSE expires_at END, - auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END, - updated_at = ? -WHERE id = ? - AND is_standalone = 1 -"#, - ) - .bind(record.key_encrypted_present) - .bind(record.key_encrypted.as_deref()) - .bind(record.name_present) - .bind(record.name.as_deref()) - .bind(record.force_capabilities.is_some()) - .bind(optional_json_to_string( - &record.force_capabilities.clone().flatten(), - "api_keys.force_capabilities", - )?) - .bind(record.rate_limit_present) - .bind(record.rate_limit) - .bind(record.concurrent_limit_present) - .bind(record.concurrent_limit) - .bind(record.allowed_providers.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_providers, - "api_keys.allowed_providers", - )?) - .bind(record.allowed_api_formats.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_api_formats, - "api_keys.allowed_api_formats", - )?) - .bind(record.allowed_models.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_models, - "api_keys.allowed_models", - )?) - .bind(record.ip_rules.is_some()) - .bind(json_string_from_nested_string_list( - &record.ip_rules, - "api_keys.ip_rules", - )?) - .bind(record.expires_at_present) - .bind(optional_i64_from_u64( - record.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(record.auto_delete_on_expiry_present) - .bind(record.auto_delete_on_expiry) - .bind(now) - .bind(&record.api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(&record.api_key_id).await - } - - async fn restore_api_key_if_matches( - &self, - expected: &StoredAuthApiKeyExportRecord, - restored: &StoredAuthApiKeyExportRecord, - ) -> Result { - if restored.api_key_id != expected.api_key_id - || restored.user_id != expected.user_id - || restored.key_hash != expected.key_hash - || restored.is_standalone != expected.is_standalone - { - return Ok(false); - } - - let mut tx = self.pool.begin().await.map_sql_err()?; - let select_sql = format!("{EXPORT_COLUMNS} WHERE api_keys.id = ? LIMIT 1 FOR UPDATE"); - let row = sqlx::query(&select_sql) - .bind(&expected.api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_auth_api_key_export_row(&row)?; - if current != *expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = ?, - name = ?, - allowed_providers = ?, - allowed_api_formats = ?, - allowed_models = ?, - ip_rules = ?, - rate_limit = ?, - concurrent_limit = ?, - force_capabilities = ?, - feature_settings = ?, - is_active = ?, - expires_at = ?, - auto_delete_on_expiry = ?, - total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - last_used_at = ?, - updated_at = ? -WHERE id = ? - AND user_id = ? - AND key_hash = ? - AND is_standalone = ? -"#, - ) - .bind(restored.key_encrypted.as_deref()) - .bind(restored.name.as_deref()) - .bind(json_string_from_string_list( - restored.allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(json_string_from_string_list( - restored.allowed_api_formats.as_ref(), - "api_keys.allowed_api_formats", - )?) - .bind(json_string_from_string_list( - restored.allowed_models.as_ref(), - "api_keys.allowed_models", - )?) - .bind(json_string_from_string_list( - restored.ip_rules.as_ref(), - "api_keys.ip_rules", - )?) - .bind(restored.rate_limit) - .bind(restored.concurrent_limit) - .bind(optional_json_to_string( - &restored.force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(optional_json_to_string( - &restored.feature_settings, - "api_keys.feature_settings", - )?) - .bind(restored.is_active) - .bind(optional_i64_from_u64( - restored.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(restored.auto_delete_on_expiry) - .bind(i64_from_u64( - restored.total_requests, - "api_keys.total_requests", - )?) - .bind(i64_from_u64( - restored.total_tokens, - "api_keys.total_tokens", - )?) - .bind(restored.total_cost_usd) - .bind(optional_i64_from_u64( - restored.last_used_at_unix_secs, - "api_keys.last_used_at", - )?) - .bind(current_unix_secs() as i64) - .bind(&restored.api_key_id) - .bind(&restored.user_id) - .bind(&restored.key_hash) - .bind(restored.is_standalone) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn set_user_api_key_active( - &self, - user_id: &str, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, Some(user_id), is_active, false, false) - .await - } - - async fn set_user_api_key_active_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, Some(user_id), is_active, false, true) - .await - } - - async fn set_standalone_api_key_active( - &self, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, None, is_active, true, false) - .await - } - - async fn set_user_api_key_locked( - &self, - user_id: &str, - api_key_id: &str, - is_locked: bool, - ) -> Result { - let rows_affected = sqlx::query( - r#" -UPDATE api_keys -SET is_locked = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 -"#, - ) - .bind(is_locked) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn set_user_api_key_allowed_providers( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - ) -> Result, DataLayerError> { - self.set_user_api_key_allowed_providers_scoped( - user_id, - api_key_id, - allowed_providers, - false, - ) - .await - } - - async fn set_user_api_key_allowed_providers_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - ) -> Result, DataLayerError> { - self.set_user_api_key_allowed_providers_scoped(user_id, api_key_id, allowed_providers, true) - .await - } - - async fn set_user_api_key_force_capabilities( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_force_capabilities_scoped( - user_id, - api_key_id, - force_capabilities, - false, - ) - .await - } - - async fn set_user_api_key_force_capabilities_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_force_capabilities_scoped( - user_id, - api_key_id, - force_capabilities, - true, - ) - .await - } - - async fn set_user_api_key_feature_settings( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_feature_settings_scoped(user_id, api_key_id, feature_settings, false) - .await - } - - async fn set_user_api_key_feature_settings_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_feature_settings_scoped(user_id, api_key_id, feature_settings, true) - .await - } - - async fn set_api_key_usage_totals( - &self, - api_key_id: &str, - total_requests: u64, - total_tokens: u64, - total_cost_usd: f64, - ) -> Result, DataLayerError> { - if !total_cost_usd.is_finite() { - return Err(DataLayerError::InvalidInput( - "api_keys.total_cost_usd is not finite".to_string(), - )); - } - sqlx::query( - r#" -UPDATE api_keys -SET total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(i64_from_u64(total_requests, "api_keys.total_requests")?) - .bind(i64_from_u64(total_tokens, "api_keys.total_tokens")?) - .bind(total_cost_usd) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(api_key_id).await - } - - async fn delete_user_api_key( - &self, - user_id: &str, - api_key_id: &str, - ) -> Result { - self.delete_api_key(api_key_id, Some(user_id), false, false) - .await - } - - async fn delete_user_api_key_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - ) -> Result { - self.delete_api_key(api_key_id, Some(user_id), false, true) - .await - } - - async fn delete_standalone_api_key(&self, api_key_id: &str) -> Result { - self.delete_api_key(api_key_id, None, true, false).await - } - - async fn set_standalone_api_key_feature_settings( - &self, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - sqlx::query( - r#" -UPDATE api_keys -SET feature_settings = ?, updated_at = ? -WHERE id = ? - AND is_standalone = 1 -"#, - ) - .bind(optional_json_to_string( - &feature_settings, - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(api_key_id).await - } -} - -impl MysqlAuthApiKeyReadRepository { - async fn update_user_api_key_basic_scoped( - &self, - record: UpdateUserApiKeyBasicRecord, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = CASE WHEN ? THEN ? ELSE key_encrypted END, - name = CASE WHEN ? THEN ? ELSE name END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, - ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, - feature_settings = CASE WHEN ? THEN ? ELSE feature_settings END, - updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(record.key_encrypted_present) - .bind(record.key_encrypted.as_deref()) - .bind(record.name_present) - .bind(record.name.as_deref()) - .bind(record.rate_limit_present) - .bind(record.rate_limit) - .bind(record.concurrent_limit_present) - .bind(record.concurrent_limit) - .bind(record.ip_rules.is_some()) - .bind(json_string_from_nested_string_list( - &record.ip_rules, - "api_keys.ip_rules", - )?) - .bind(record.feature_settings.is_some()) - .bind(optional_json_to_string( - &record.feature_settings.clone().flatten(), - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(&record.api_key_id) - .bind(&record.user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(&record.api_key_id).await - } - - async fn set_active( - &self, - api_key_id: &str, - user_id: Option<&str>, - is_active: bool, - is_standalone: bool, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new("UPDATE api_keys SET is_active = "); - builder - .push_bind(is_active) - .push(", updated_at = ") - .push_bind(current_unix_secs() as i64) - .push(" WHERE id = ") - .push_bind(api_key_id) - .push(" AND is_standalone = ") - .push_bind(is_standalone); - if let Some(user_id) = user_id { - builder.push(" AND user_id = ").push_bind(user_id); - } - if require_unlocked { - builder.push(" AND is_locked = ").push_bind(false); - } - let result = builder.build().execute(&self.pool).await.map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_allowed_providers_scoped( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET allowed_providers = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(json_string_from_string_list( - allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_force_capabilities_scoped( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET force_capabilities = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(optional_json_to_string( - &force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_feature_settings_scoped( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET feature_settings = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(optional_json_to_string( - &feature_settings, - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn delete_api_key( - &self, - api_key_id: &str, - user_id: Option<&str>, - is_standalone: bool, - require_unlocked: bool, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let matching_api_key = if let Some(user_id) = user_id { - if require_unlocked { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND user_id = ? AND is_standalone = 0 AND is_locked = 0 FOR UPDATE", - ) - .bind(api_key_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND user_id = ? AND is_standalone = 0 FOR UPDATE", - ) - .bind(api_key_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } - } else { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND is_standalone = 1 FOR UPDATE", - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - }; - if matching_api_key.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - sqlx::query( - "UPDATE wallets SET status = 'disabled', updated_at = UNIX_TIMESTAMP() WHERE api_key_id = ? AND status <> 'disabled'", - ) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - for sql in MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL { - sqlx::query(sql) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - for sql in MYSQL_DELETE_API_KEY_DEPENDENTS_SQL { - sqlx::query(sql) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let result = sqlx::query("DELETE FROM api_keys WHERE id = ? AND is_standalone = ?") - .bind(api_key_id) - .bind(is_standalone) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } -} - -fn push_in_clause<'args>( - builder: &mut QueryBuilder<'args, MySql>, - prefix: &str, - values: &'args [String], -) { - builder.push(prefix); - { - let mut separated = builder.separated(", "); - for value in values { - separated.push_bind(value); - } - } - builder.push(")"); -} - -async fn summarize_api_keys( - pool: &MysqlPool, - is_standalone: bool, - now_unix_secs: u64, -) -> Result { - let now_unix_secs = i64_from_u64(now_unix_secs, "api_keys.summary_now")?; - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 AND (expires_at IS NULL OR expires_at >= ?) THEN 1 ELSE 0 END) AS active -FROM api_keys -WHERE is_standalone = ? -"#, - ) - .bind(now_unix_secs) - .bind(is_standalone) - .fetch_one(pool) - .await - .map_sql_err()?; - summarize_row(row) -} - -fn summarize_row(row: MySqlRow) -> Result { - Ok(AuthApiKeyExportSummary { - total: row.try_get::("total").map_sql_err()?.max(0) as u64, - active: row - .try_get::, _>("active") - .map_sql_err()? - .unwrap_or(0) - .max(0) as u64, - }) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| i64_from_u64(value, field_name)) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_string_list( - value: Option<&Vec>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable string list: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_nested_string_list( - value: &Option>>, - field_name: &str, -) -> Result, DataLayerError> { - match value { - Some(Some(values)) => json_string_from_string_list(Some(values), field_name), - Some(None) | None => Ok(None), - } -} - -fn map_auth_api_key_snapshot_row( - row: &MySqlRow, -) -> Result { - let snapshot = StoredAuthApiKeySnapshot::new( - row.try_get("user_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("user_role").map_sql_err()?, - row.try_get("user_auth_source").map_sql_err()?, - row.try_get("user_is_active").map_sql_err()?, - row.try_get("user_is_deleted").map_sql_err()?, - optional_json_from_string( - row.try_get("user_allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("user_allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("user_allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("api_key_is_active").map_sql_err()?, - row.try_get("api_key_is_locked").map_sql_err()?, - row.try_get("api_key_is_standalone").map_sql_err()?, - row.try_get("api_key_rate_limit").map_sql_err()?, - row.try_get("api_key_concurrent_limit").map_sql_err()?, - row.try_get("api_key_expires_at_unix_secs").map_sql_err()?, - optional_json_from_string( - row.try_get("api_key_allowed_providers").map_sql_err()?, - "api_keys.allowed_providers", - )?, - optional_json_from_string( - row.try_get("api_key_allowed_api_formats").map_sql_err()?, - "api_keys.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("api_key_allowed_models").map_sql_err()?, - "api_keys.allowed_models", - )?, - )? - .with_api_key_ip_rules(optional_json_from_string( - row.try_get("api_key_ip_rules").map_sql_err()?, - "api_keys.ip_rules", - )?)?; - Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?)) -} - -fn map_auth_api_key_export_row( - row: &MySqlRow, -) -> Result { - let feature_settings = optional_json_from_string( - row.try_get("feature_settings").map_sql_err()?, - "api_keys.feature_settings", - )?; - StoredAuthApiKeyExportRecord::new( - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("key_hash").map_sql_err()?, - row.try_get("key_encrypted").map_sql_err()?, - row.try_get("name").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "api_keys.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "api_keys.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "api_keys.allowed_models", - )?, - row.try_get("rate_limit").map_sql_err()?, - row.try_get("concurrent_limit").map_sql_err()?, - optional_json_from_string( - row.try_get("force_capabilities").map_sql_err()?, - "api_keys.force_capabilities", - )?, - row.try_get("is_active").map_sql_err()?, - row.try_get("expires_at_unix_secs").map_sql_err()?, - row.try_get("auto_delete_on_expiry").map_sql_err()?, - row.try_get("total_requests").map_sql_err()?, - row.try_get("total_tokens").map_sql_err()?, - row.try_get("total_cost_usd").map_sql_err()?, - row.try_get("is_standalone").map_sql_err()?, - ) - .and_then(|record| { - record.with_ip_rules(optional_json_from_string( - row.try_get("ip_rules").map_sql_err()?, - "api_keys.ip_rules", - )?) - }) - .map(|record| record.with_feature_settings(feature_settings)) - .and_then(|record| { - record.with_activity_timestamps( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - row.try_get("created_at_unix_secs").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - ) - }) -} - -#[cfg(test)] -mod tests { - use super::{ - MysqlAuthApiKeyReadRepository, MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL, - MYSQL_DELETE_API_KEY_DEPENDENTS_SQL, - }; - - #[test] - fn api_key_delete_sql_preserves_ids_and_removes_private_snapshots() { - for table in [ - "request_candidates", - "video_tasks", - "`usage`", - "stats_daily_api_key", - ] { - assert!(MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL.iter().any(|sql| { - sql.starts_with(&format!("UPDATE {table} ")) - && sql.contains("SET api_key_name = NULL") - && sql.ends_with("WHERE api_key_id = ?") - })); - } - assert!(MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL - .iter() - .any(|sql| sql - .starts_with("UPDATE audit_logs SET description = 'deleted API key event'"))); - assert!(MYSQL_ANONYMIZE_API_KEY_HISTORY_SQL.iter().any(|sql| sql - .starts_with("UPDATE payment_callbacks SET payload = NULL, error_message = NULL"))); - assert_eq!( - MYSQL_DELETE_API_KEY_DEPENDENTS_SQL, - &["DELETE FROM api_key_provider_mappings WHERE api_key_id = ?"] - ); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlAuthApiKeyReadRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/auth_modules.rs b/crates/aether-data/adapters/mysql/src/auth_modules.rs deleted file mode 100644 index 44ccc226c..000000000 --- a/crates/aether-data/adapters/mysql/src/auth_modules.rs +++ /dev/null @@ -1,439 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::auth_modules::*; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, WhereClause}; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const OAUTH_PROVIDER_COLUMNS: &str = r#" -SELECT - provider_type, - display_name, - client_id, - client_secret_encrypted, - redirect_uri -FROM oauth_providers -"#; - -const LDAP_CONFIG_COLUMNS: &str = r#" -SELECT - server_url, - bind_dn, - bind_password_encrypted, - base_dn, - user_search_filter, - username_attr, - email_attr, - display_name_attr, - is_enabled, - is_exclusive, - use_starttls, - connect_timeout -FROM ldap_configs -"#; - -const UPDATE_LDAP_CONFIG_PRESERVE_PASSWORD_SQL: &str = r#" -UPDATE ldap_configs -SET - server_url = ?, - bind_dn = ?, - base_dn = ?, - user_search_filter = ?, - username_attr = ?, - email_attr = ?, - display_name_attr = ?, - is_enabled = ?, - is_exclusive = ?, - use_starttls = ?, - connect_timeout = ?, - updated_at = GREATEST(updated_at + 1, ?) -WHERE singleton_key = 1 - AND server_url <=> ? - AND bind_dn <=> ? - AND BINARY bind_password_encrypted <=> BINARY ? - AND base_dn <=> ? - AND user_search_filter <=> ? - AND username_attr <=> ? - AND email_attr <=> ? - AND display_name_attr <=> ? - AND is_enabled <=> ? - AND is_exclusive <=> ? - AND use_starttls <=> ? - AND connect_timeout <=> ? -"#; - -const UPDATE_LDAP_CONFIG_REPLACE_PASSWORD_SQL: &str = r#" -UPDATE ldap_configs -SET - server_url = ?, - bind_dn = ?, - bind_password_encrypted = ?, - base_dn = ?, - user_search_filter = ?, - username_attr = ?, - email_attr = ?, - display_name_attr = ?, - is_enabled = ?, - is_exclusive = ?, - use_starttls = ?, - connect_timeout = ?, - updated_at = GREATEST(updated_at + 1, ?) -WHERE singleton_key = 1 - AND server_url <=> ? - AND bind_dn <=> ? - AND BINARY bind_password_encrypted <=> BINARY ? - AND base_dn <=> ? - AND user_search_filter <=> ? - AND username_attr <=> ? - AND email_attr <=> ? - AND display_name_attr <=> ? - AND is_enabled <=> ? - AND is_exclusive <=> ? - AND use_starttls <=> ? - AND connect_timeout <=> ? -"#; - -const INSERT_LDAP_CONFIG_SQL: &str = r#" -INSERT INTO ldap_configs ( - singleton_key, - server_url, - bind_dn, - bind_password_encrypted, - base_dn, - user_search_filter, - username_attr, - email_attr, - display_name_attr, - is_enabled, - is_exclusive, - use_starttls, - connect_timeout, - created_at, - updated_at -) VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#; - -#[derive(Debug, Clone)] -pub struct MysqlAuthModuleReadRepository { - pool: MysqlPool, -} - -impl MysqlAuthModuleReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[derive(Debug, Clone)] -pub struct MysqlAuthModuleRepository { - pool: MysqlPool, -} - -impl MysqlAuthModuleRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -async fn list_enabled_oauth_providers( - pool: &MysqlPool, -) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(OAUTH_PROVIDER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "is_enabled", true); - builder.push(" ORDER BY provider_type ASC"); - let rows = builder.build().fetch_all(pool).await.map_sql_err()?; - rows.iter().map(map_oauth_row).collect() -} - -async fn get_ldap_config( - pool: &MysqlPool, -) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(LDAP_CONFIG_COLUMNS); - builder.push(" WHERE singleton_key = 1"); - let row = builder.build().fetch_optional(pool).await.map_sql_err()?; - row.as_ref().map(map_ldap_row).transpose() -} - -#[async_trait] -impl AuthModuleReadRepository for MysqlAuthModuleReadRepository { - async fn list_enabled_oauth_providers( - &self, - ) -> Result, DataLayerError> { - list_enabled_oauth_providers(&self.pool).await - } - - async fn get_ldap_config(&self) -> Result, DataLayerError> { - get_ldap_config(&self.pool).await - } -} - -#[async_trait] -impl AuthModuleReadRepository for MysqlAuthModuleRepository { - async fn list_enabled_oauth_providers( - &self, - ) -> Result, DataLayerError> { - list_enabled_oauth_providers(&self.pool).await - } - - async fn get_ldap_config(&self) -> Result, DataLayerError> { - get_ldap_config(&self.pool).await - } -} - -#[async_trait] -impl AuthModuleWriteRepository for MysqlAuthModuleRepository { - async fn compare_and_swap_ldap_config( - &self, - expected: Option<&StoredLdapModuleConfig>, - replacement: &StoredLdapModuleConfig, - bind_password_update: &LdapBindPasswordUpdate, - ) -> Result { - let persisted = - ldap_config_after_password_update(expected, replacement, bind_password_update)?; - let now = now_unix_secs(); - let Some(expected) = expected else { - let insert = sqlx::query(INSERT_LDAP_CONFIG_SQL) - .bind(&persisted.server_url) - .bind(&persisted.bind_dn) - .bind(persisted.bind_password_encrypted.as_deref()) - .bind(&persisted.base_dn) - .bind(persisted.user_search_filter.as_deref()) - .bind(persisted.username_attr.as_deref()) - .bind(persisted.email_attr.as_deref()) - .bind(persisted.display_name_attr.as_deref()) - .bind(persisted.is_enabled) - .bind(persisted.is_exclusive) - .bind(persisted.use_starttls) - .bind(persisted.connect_timeout) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await; - return match insert { - Ok(result) if result.rows_affected() == 1 => { - Ok(CompareAndSwapLdapConfigResult::Applied(persisted)) - } - Ok(_) => Ok(CompareAndSwapLdapConfigResult::Conflict), - Err(error) - if error - .as_database_error() - .is_some_and(|error| error.is_unique_violation()) => - { - Ok(CompareAndSwapLdapConfigResult::Conflict) - } - Err(error) => Err(DataLayerError::sql(error)), - }; - }; - - let rows_affected = match bind_password_update { - LdapBindPasswordUpdate::Preserve => { - sqlx::query(UPDATE_LDAP_CONFIG_PRESERVE_PASSWORD_SQL) - .bind(&replacement.server_url) - .bind(&replacement.bind_dn) - .bind(&replacement.base_dn) - .bind(replacement.user_search_filter.as_deref()) - .bind(replacement.username_attr.as_deref()) - .bind(replacement.email_attr.as_deref()) - .bind(replacement.display_name_attr.as_deref()) - .bind(replacement.is_enabled) - .bind(replacement.is_exclusive) - .bind(replacement.use_starttls) - .bind(replacement.connect_timeout) - .bind(now as i64) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - } - LdapBindPasswordUpdate::Set(_) | LdapBindPasswordUpdate::Clear => { - sqlx::query(UPDATE_LDAP_CONFIG_REPLACE_PASSWORD_SQL) - .bind(&replacement.server_url) - .bind(&replacement.bind_dn) - .bind(persisted.bind_password_encrypted.as_deref()) - .bind(&replacement.base_dn) - .bind(replacement.user_search_filter.as_deref()) - .bind(replacement.username_attr.as_deref()) - .bind(replacement.email_attr.as_deref()) - .bind(replacement.display_name_attr.as_deref()) - .bind(replacement.is_enabled) - .bind(replacement.is_exclusive) - .bind(replacement.use_starttls) - .bind(replacement.connect_timeout) - .bind(now as i64) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - } - }; - if rows_affected == 1 { - Ok(CompareAndSwapLdapConfigResult::Applied(persisted)) - } else { - Ok(CompareAndSwapLdapConfigResult::Conflict) - } - } - - async fn delete_ldap_config_if_matches( - &self, - expected: &StoredLdapModuleConfig, - ) -> Result { - let rows_affected = sqlx::query( - r#" -DELETE FROM ldap_configs -WHERE singleton_key = 1 - AND server_url <=> ? - AND bind_dn <=> ? - AND BINARY bind_password_encrypted <=> BINARY ? - AND base_dn <=> ? - AND user_search_filter <=> ? - AND username_attr <=> ? - AND email_attr <=> ? - AND display_name_attr <=> ? - AND is_enabled <=> ? - AND is_exclusive <=> ? - AND use_starttls <=> ? - AND connect_timeout <=> ? -"#, - ) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - async fn compare_and_swap_ldap_bind_password( - &self, - expected: &str, - replacement: &str, - ) -> Result { - let rows_affected = sqlx::query( - r#" -UPDATE ldap_configs -SET bind_password_encrypted = ?, updated_at = GREATEST(updated_at + 1, ?) -WHERE singleton_key = 1 - AND BINARY bind_password_encrypted = BINARY ? -"#, - ) - .bind(replacement) - .bind(now_unix_secs() as i64) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } -} - -fn ldap_config_after_password_update( - expected: Option<&StoredLdapModuleConfig>, - replacement: &StoredLdapModuleConfig, - bind_password_update: &LdapBindPasswordUpdate, -) -> Result { - let bind_password_encrypted = match bind_password_update { - LdapBindPasswordUpdate::Preserve => expected - .ok_or_else(|| { - DataLayerError::InvalidConfiguration( - "LDAP bind password cannot be preserved while creating the singleton" - .to_string(), - ) - })? - .bind_password_encrypted - .clone(), - LdapBindPasswordUpdate::Set(ciphertext) => Some(ciphertext.clone()), - LdapBindPasswordUpdate::Clear => None, - }; - Ok(StoredLdapModuleConfig { - bind_password_encrypted, - ..replacement.clone() - }) -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn map_oauth_row(row: &MySqlRow) -> Result { - StoredOAuthProviderModuleConfig::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("client_id").map_sql_err()?, - row.try_get("client_secret_encrypted").map_sql_err()?, - row.try_get("redirect_uri").map_sql_err()?, - ) -} - -fn map_ldap_row(row: &MySqlRow) -> Result { - Ok(StoredLdapModuleConfig { - server_url: row.try_get("server_url").map_sql_err()?, - bind_dn: row.try_get("bind_dn").map_sql_err()?, - bind_password_encrypted: row.try_get("bind_password_encrypted").map_sql_err()?, - base_dn: row.try_get("base_dn").map_sql_err()?, - user_search_filter: row.try_get("user_search_filter").map_sql_err()?, - username_attr: row.try_get("username_attr").map_sql_err()?, - email_attr: row.try_get("email_attr").map_sql_err()?, - display_name_attr: row.try_get("display_name_attr").map_sql_err()?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - is_exclusive: row.try_get("is_exclusive").map_sql_err()?, - use_starttls: row.try_get("use_starttls").map_sql_err()?, - connect_timeout: row.try_get("connect_timeout").map_sql_err()?, - }) -} - -#[cfg(test)] -mod tests { - use super::{MysqlAuthModuleReadRepository, MysqlAuthModuleRepository}; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlAuthModuleReadRepository::new(pool.clone()); - let _writable_repository = MysqlAuthModuleRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/background_tasks.rs b/crates/aether-data/adapters/mysql/src/background_tasks.rs deleted file mode 100644 index e4b8479a4..000000000 --- a/crates/aether-data/adapters/mysql/src/background_tasks.rs +++ /dev/null @@ -1,450 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::background_tasks::*; - -use crate::error::SqlResultExt; -use crate::{DataLayerError, MysqlPool}; - -const RUN_COLUMNS: &str = r#" -SELECT - id, - task_key, - kind, - `trigger`, - status, - attempt, - max_attempts, - owner_instance, - progress_percent, - progress_message, - payload_json, - result_json, - error_message, - cancel_requested, - created_by, - created_at_unix_secs, - started_at_unix_secs, - finished_at_unix_secs, - updated_at_unix_secs -FROM background_task_runs -"#; - -const EVENT_COLUMNS: &str = r#" -SELECT - id, - run_id, - event_type, - message, - payload_json, - created_at_unix_secs -FROM background_task_events -"#; - -#[derive(Debug, Clone)] -pub struct MysqlBackgroundTaskRepository { - pool: MysqlPool, -} - -impl MysqlBackgroundTaskRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - fn apply_run_filter(builder: &mut QueryBuilder<'_, MySql>, query: &BackgroundTaskListQuery) { - let mut has_where = false; - if let Some(kind) = query.kind { - if !has_where { - builder.push(" WHERE "); - has_where = true; - } else { - builder.push(" AND "); - } - builder.push("kind = ").push_bind(kind.as_database()); - } - if let Some(status) = query.status { - if !has_where { - builder.push(" WHERE "); - has_where = true; - } else { - builder.push(" AND "); - } - builder.push("status = ").push_bind(status.as_database()); - } - if let Some(trigger) = query.trigger.as_deref() { - if !has_where { - builder.push(" WHERE "); - has_where = true; - } else { - builder.push(" AND "); - } - builder.push("`trigger` = ").push_bind(trigger.to_string()); - } - if let Some(task_key_substring) = query.task_key_substring.as_deref() { - if !has_where { - builder.push(" WHERE "); - } else { - builder.push(" AND "); - } - builder.push("LOWER(task_key) LIKE ").push_bind(format!( - "%{}%", - task_key_substring.trim().to_ascii_lowercase() - )); - } - } -} - -#[async_trait] -impl BackgroundTaskReadRepository for MysqlBackgroundTaskRepository { - async fn find_run( - &self, - run_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{RUN_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(run_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_run_row).transpose() - } - - async fn list_runs( - &self, - query: &BackgroundTaskListQuery, - ) -> Result { - let limit = query.limit.max(1); - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(id) AS total FROM background_task_runs"); - Self::apply_run_filter(&mut count_builder, query); - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let mut builder = QueryBuilder::::new(RUN_COLUMNS); - Self::apply_run_filter(&mut builder, query); - builder - .push(" ORDER BY created_at_unix_secs DESC, updated_at_unix_secs DESC") - .push(" LIMIT ") - .push_bind(i64_from_usize(limit, "run limit")?) - .push(" OFFSET ") - .push_bind(i64_from_usize(query.offset, "run offset")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let items = rows - .iter() - .map(map_run_row) - .collect::, _>>()?; - Ok(StoredBackgroundTaskRunPage { - items, - total: usize::try_from(total).unwrap_or_default(), - }) - } - - async fn list_events( - &self, - run_id: &str, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - let limit = limit.max(1); - let rows = sqlx::query(&format!( - "{EVENT_COLUMNS} WHERE run_id = ? ORDER BY created_at_unix_secs ASC, id ASC LIMIT ? OFFSET ?" - )) - .bind(run_id) - .bind(i64_from_usize(limit, "event limit")?) - .bind(i64_from_usize(offset, "event offset")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_event_row).collect() - } - - async fn summarize_runs(&self) -> Result { - let total = sqlx::query_scalar::<_, i64>("SELECT COUNT(id) FROM background_task_runs") - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let running_count = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(id) FROM background_task_runs WHERE status = 'running'", - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let status_rows = sqlx::query( - "SELECT status, COUNT(id) AS total FROM background_task_runs GROUP BY status", - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let kind_rows = - sqlx::query("SELECT kind, COUNT(id) AS total FROM background_task_runs GROUP BY kind") - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let mut by_status = std::collections::BTreeMap::new(); - for row in status_rows { - let key: String = row.try_get("status").map_sql_err()?; - let count: i64 = row.try_get("total").map_sql_err()?; - by_status.insert(key, u64::try_from(count).unwrap_or_default()); - } - let mut by_kind = std::collections::BTreeMap::new(); - for row in kind_rows { - let key: String = row.try_get("kind").map_sql_err()?; - let count: i64 = row.try_get("total").map_sql_err()?; - by_kind.insert(key, u64::try_from(count).unwrap_or_default()); - } - - Ok(BackgroundTaskSummary { - total: u64::try_from(total).unwrap_or_default(), - running_count: u64::try_from(running_count).unwrap_or_default(), - by_status, - by_kind, - }) - } -} - -#[async_trait] -impl BackgroundTaskWriteRepository for MysqlBackgroundTaskRepository { - async fn upsert_run( - &self, - mut run: UpsertBackgroundTaskRun, - ) -> Result { - run.sanitize_for_persistence(); - run.validate()?; - sqlx::query( - r#" -INSERT INTO background_task_runs ( - id, - task_key, - kind, - `trigger`, - status, - attempt, - max_attempts, - owner_instance, - progress_percent, - progress_message, - payload_json, - result_json, - error_message, - cancel_requested, - created_by, - created_at_unix_secs, - started_at_unix_secs, - finished_at_unix_secs, - updated_at_unix_secs -) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) -ON DUPLICATE KEY UPDATE - task_key = VALUES(task_key), - kind = VALUES(kind), - `trigger` = VALUES(`trigger`), - status = VALUES(status), - attempt = VALUES(attempt), - max_attempts = VALUES(max_attempts), - owner_instance = VALUES(owner_instance), - progress_percent = VALUES(progress_percent), - progress_message = VALUES(progress_message), - payload_json = VALUES(payload_json), - result_json = VALUES(result_json), - error_message = VALUES(error_message), - cancel_requested = VALUES(cancel_requested), - created_by = VALUES(created_by), - created_at_unix_secs = VALUES(created_at_unix_secs), - started_at_unix_secs = VALUES(started_at_unix_secs), - finished_at_unix_secs = VALUES(finished_at_unix_secs), - updated_at_unix_secs = VALUES(updated_at_unix_secs) -"#, - ) - .bind(&run.id) - .bind(&run.task_key) - .bind(run.kind.as_database()) - .bind(&run.trigger) - .bind(run.status.as_database()) - .bind(i64::from(run.attempt)) - .bind(i64::from(run.max_attempts)) - .bind(run.owner_instance.as_deref()) - .bind(i32::from(run.progress_percent)) - .bind(run.progress_message.as_deref()) - .bind(json_to_string(&run.payload_json, "payload_json")?) - .bind(json_to_string(&run.result_json, "result_json")?) - .bind(run.error_message.as_deref()) - .bind(run.cancel_requested) - .bind(run.created_by.as_deref()) - .bind(u64_to_i64( - run.created_at_unix_secs, - "created_at_unix_secs", - )?) - .bind(run.started_at_unix_secs.map(|value| value as i64)) - .bind(run.finished_at_unix_secs.map(|value| value as i64)) - .bind(u64_to_i64( - run.updated_at_unix_secs, - "updated_at_unix_secs", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.find_run(&run.id).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("background task run missing after upsert".to_string()) - }) - } - - async fn request_cancel( - &self, - run_id: &str, - updated_at_unix_secs: u64, - ) -> Result { - let affected = sqlx::query( - "UPDATE background_task_runs SET cancel_requested = TRUE, updated_at_unix_secs = ? WHERE id = ?", - ) - .bind(u64_to_i64(updated_at_unix_secs, "updated_at_unix_secs")?) - .bind(run_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(affected > 0) - } - - async fn upsert_event( - &self, - mut event: UpsertBackgroundTaskEvent, - ) -> Result { - event.sanitize_for_persistence(); - event.validate()?; - sqlx::query( - r#" -INSERT INTO background_task_events ( - id, run_id, event_type, message, payload_json, created_at_unix_secs -) VALUES (?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - run_id = VALUES(run_id), - event_type = VALUES(event_type), - message = VALUES(message), - payload_json = VALUES(payload_json), - created_at_unix_secs = VALUES(created_at_unix_secs) -"#, - ) - .bind(&event.id) - .bind(&event.run_id) - .bind(&event.event_type) - .bind(&event.message) - .bind(json_to_string(&event.payload_json, "payload_json")?) - .bind(u64_to_i64( - event.created_at_unix_secs, - "created_at_unix_secs", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - - let row = sqlx::query(&format!("{EVENT_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(&event.id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - map_event_row(&row) - } -} - -fn map_run_row(row: &MySqlRow) -> Result { - let kind: String = row.try_get("kind").map_sql_err()?; - let status: String = row.try_get("status").map_sql_err()?; - let attempt: i64 = row.try_get("attempt").map_sql_err()?; - let max_attempts: i64 = row.try_get("max_attempts").map_sql_err()?; - let progress_percent: i32 = row.try_get("progress_percent").map_sql_err()?; - let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?; - let started_at_unix_secs: Option = row.try_get("started_at_unix_secs").map_sql_err()?; - let finished_at_unix_secs: Option = row.try_get("finished_at_unix_secs").map_sql_err()?; - let updated_at_unix_secs: i64 = row.try_get("updated_at_unix_secs").map_sql_err()?; - - let mut run = StoredBackgroundTaskRun { - id: row.try_get("id").map_sql_err()?, - task_key: row.try_get("task_key").map_sql_err()?, - kind: BackgroundTaskKind::from_database(&kind)?, - trigger: row.try_get("trigger").map_sql_err()?, - status: BackgroundTaskStatus::from_database(&status)?, - attempt: u32::try_from(attempt).unwrap_or_default(), - max_attempts: u32::try_from(max_attempts).unwrap_or_default(), - owner_instance: row.try_get("owner_instance").map_sql_err()?, - progress_percent: u16::try_from(progress_percent).unwrap_or_default(), - progress_message: row.try_get("progress_message").map_sql_err()?, - payload_json: parse_optional_json( - row.try_get("payload_json").ok().flatten(), - "payload_json", - )?, - result_json: parse_optional_json(row.try_get("result_json").ok().flatten(), "result_json")?, - error_message: row.try_get("error_message").map_sql_err()?, - cancel_requested: row.try_get("cancel_requested").map_sql_err()?, - created_by: row.try_get("created_by").map_sql_err()?, - created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(), - started_at_unix_secs: started_at_unix_secs.and_then(|value| u64::try_from(value).ok()), - finished_at_unix_secs: finished_at_unix_secs.and_then(|value| u64::try_from(value).ok()), - updated_at_unix_secs: u64::try_from(updated_at_unix_secs).unwrap_or_default(), - }; - run.sanitize_persisted_data(); - Ok(run) -} - -fn map_event_row(row: &MySqlRow) -> Result { - let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?; - let mut event = StoredBackgroundTaskEvent { - id: row.try_get("id").map_sql_err()?, - run_id: row.try_get("run_id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - message: row.try_get("message").map_sql_err()?, - payload_json: parse_optional_json( - row.try_get("payload_json").ok().flatten(), - "payload_json", - )?, - created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(), - }; - event.sanitize_persisted_data(); - Ok(event) -} - -fn i64_from_usize(value: usize, label: &str) -> Result { - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}")) - }) -} - -fn u64_to_i64(value: u64, label: &str) -> Result { - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}")) - }) -} - -fn json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "background task {field_name} is unserializable: {err}" - )) - }) - }) - .transpose() -} - -fn parse_optional_json( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "background task {field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} diff --git a/crates/aether-data/adapters/mysql/src/billing.rs b/crates/aether-data/adapters/mysql/src/billing.rs deleted file mode 100644 index fa3705087..000000000 --- a/crates/aether-data/adapters/mysql/src/billing.rs +++ /dev/null @@ -1,1716 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; - -use aether_data_contracts::repository::billing::{ - AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, - AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput, - BillingPlanRecord, BillingPlanWriteInput, BillingReadRepository, - PaymentGatewayConfigCasWriteInput, PaymentGatewayConfigRecord, PaymentGatewayConfigWriteInput, - PaymentGatewaySecretCasUpdate, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord, - UserPlanEntitlementRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const MODEL_CONTEXT_COLUMNS: &str = r#" -SELECT - p.id AS provider_id, - p.billing_type AS provider_billing_type, - pak.id AS provider_api_key_id, - pak.rate_multipliers AS provider_api_key_rate_multipliers, - pak.cache_ttl_minutes AS provider_api_key_cache_ttl_minutes, - gm.id AS global_model_id, - gm.name AS global_model_name, - gm.config AS global_model_config, - gm.default_price_per_request AS default_price_per_request, - gm.default_tiered_pricing AS default_tiered_pricing, - m.id AS model_id, - m.provider_model_name AS model_provider_model_name, - m.config AS model_config, - m.price_per_request AS model_price_per_request, - m.tiered_pricing AS model_tiered_pricing, - m.provider_model_mappings AS provider_model_mappings, - m.is_available AS model_is_available, - m.created_at AS model_created_at -FROM providers p -"#; - -#[derive(Debug, Clone)] -pub struct MysqlBillingReadRepository { - pool: MysqlPool, -} - -impl MysqlBillingReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl BillingReadRepository for MysqlBillingReadRepository { - async fn find_model_context( - &self, - provider_id: &str, - provider_api_key_id: Option<&str>, - global_model_name: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{MODEL_CONTEXT_COLUMNS} -INNER JOIN global_models gm - ON gm.is_active = 1 -LEFT JOIN models m - ON m.global_model_id = gm.id - AND m.provider_id = p.id - AND m.is_active = 1 -LEFT JOIN provider_api_keys pak - ON pak.id = ? - AND pak.provider_id = p.id -WHERE p.id = ? - AND ( - gm.name = ? - OR m.provider_model_name = ? - OR m.provider_model_mappings IS NOT NULL - ) -"# - )) - .bind(provider_api_key_id) - .bind(provider_id) - .bind(global_model_name) - .bind(global_model_name) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - rows.iter() - .filter_map(|row| match_rank(row, global_model_name).transpose()) - .collect::, _>>()? - .into_iter() - .min_by_key(|candidate| { - ( - candidate.rank, - !candidate.is_available, - candidate.pricing_rank, - candidate.created_at, - ) - }) - .map(|candidate| candidate.context) - .transpose() - } - - async fn find_model_context_by_model_id( - &self, - provider_id: &str, - provider_api_key_id: Option<&str>, - model_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - r#" -{MODEL_CONTEXT_COLUMNS} -INNER JOIN models m - ON m.id = ? - AND m.provider_id = p.id - AND m.is_active = 1 -INNER JOIN global_models gm - ON gm.id = m.global_model_id - AND gm.is_active = 1 -LEFT JOIN provider_api_keys pak - ON pak.id = ? - AND pak.provider_id = p.id -WHERE p.id = ? -LIMIT 1 -"# - )) - .bind(model_id) - .bind(provider_api_key_id) - .bind(provider_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_row).transpose() - } - - async fn admin_billing_enabled_default_value_exists( - &self, - api_format: &str, - task_type: &str, - dimension_name: &str, - existing_id: Option<&str>, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM dimension_collectors -WHERE api_format = ? - AND task_type = ? - AND dimension_name = ? - AND is_enabled = 1 - AND default_value IS NOT NULL - AND (? IS NULL OR id <> ?) - "#, - ) - .bind(api_format) - .bind(task_type) - .bind(dimension_name) - .bind(existing_id) - .bind(existing_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(Some(read_count_mysql(&row)? > 0)) - } - - async fn create_admin_billing_rule( - &self, - input: &AdminBillingRuleWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -INSERT INTO billing_rules ( - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&id) - .bind(&input.name) - .bind(&input.task_type) - .bind(input.global_model_id.as_deref()) - .bind(input.model_id.as_deref()) - .bind(&input.expression) - .bind(json_to_string(&input.variables)?) - .bind(json_to_string(&input.dimension_mappings)?) - .bind(input.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - if let Err(err) = result { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))); - } - match find_admin_billing_rule_mysql(&self.pool, &id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing rule missing".to_string(), - )), - } - } - - async fn list_admin_billing_rules( - &self, - task_type: Option<&str>, - is_enabled: Option, - page: u32, - page_size: u32, - ) -> Result, u64)>, DataLayerError> { - let total_row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM billing_rules -WHERE (? IS NULL OR task_type = ?) - AND (? IS NULL OR is_enabled = ?) - "#, - ) - .bind(task_type) - .bind(task_type) - .bind(is_enabled) - .bind(is_enabled) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total = read_count_mysql(&total_row)?; - let offset = u64::from(page.saturating_sub(1) * page_size); - let rows = sqlx::query( - r#" -SELECT - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM billing_rules -WHERE (? IS NULL OR task_type = ?) - AND (? IS NULL OR is_enabled = ?) -ORDER BY updated_at DESC, id DESC -LIMIT ? OFFSET ? - "#, - ) - .bind(task_type) - .bind(task_type) - .bind(is_enabled) - .bind(is_enabled) - .bind(i64::from(page_size)) - .bind( - i64::try_from(offset) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_admin_billing_rule_mysql) - .collect::, _>>()?; - Ok(Some((items, total))) - } - - async fn find_admin_billing_rule( - &self, - rule_id: &str, - ) -> Result, DataLayerError> { - find_admin_billing_rule_mysql(&self.pool, rule_id).await - } - - async fn update_admin_billing_rule( - &self, - rule_id: &str, - input: &AdminBillingRuleWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE billing_rules -SET name = ?, - task_type = ?, - global_model_id = ?, - model_id = ?, - expression = ?, - variables = ?, - dimension_mappings = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&input.name) - .bind(&input.task_type) - .bind(input.global_model_id.as_deref()) - .bind(input.model_id.as_deref()) - .bind(&input.expression) - .bind(json_to_string(&input.variables)?) - .bind(json_to_string(&input.dimension_mappings)?) - .bind(input.is_enabled) - .bind(current_unix_secs_i64()) - .bind(rule_id) - .execute(&self.pool) - .await; - let affected = match result { - Ok(result) => result.rows_affected(), - Err(err) => { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))) - } - }; - if affected == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match find_admin_billing_rule_mysql(&self.pool, rule_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn create_admin_billing_collector( - &self, - input: &AdminBillingCollectorWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -INSERT INTO dimension_collectors ( - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&id) - .bind(&input.api_format) - .bind(&input.task_type) - .bind(&input.dimension_name) - .bind(&input.source_type) - .bind(input.source_path.as_deref()) - .bind(&input.value_type) - .bind(input.transform_expression.as_deref()) - .bind(input.default_value.as_deref()) - .bind(input.priority) - .bind(input.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - if let Err(err) = result { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))); - } - match find_admin_billing_collector_mysql(&self.pool, &id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing collector missing".to_string(), - )), - } - } - - async fn list_admin_billing_collectors( - &self, - api_format: Option<&str>, - task_type: Option<&str>, - dimension_name: Option<&str>, - is_enabled: Option, - page: u32, - page_size: u32, - ) -> Result, u64)>, DataLayerError> { - let total_row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM dimension_collectors -WHERE (? IS NULL OR api_format = ?) - AND (? IS NULL OR task_type = ?) - AND (? IS NULL OR dimension_name = ?) - AND (? IS NULL OR is_enabled = ?) - "#, - ) - .bind(api_format) - .bind(api_format) - .bind(task_type) - .bind(task_type) - .bind(dimension_name) - .bind(dimension_name) - .bind(is_enabled) - .bind(is_enabled) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total = read_count_mysql(&total_row)?; - let offset = u64::from(page.saturating_sub(1) * page_size); - let rows = sqlx::query( - r#" -SELECT - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, - created_at AS created_at_unix_ms, updated_at AS updated_at_unix_secs -FROM dimension_collectors -WHERE (? IS NULL OR api_format = ?) - AND (? IS NULL OR task_type = ?) - AND (? IS NULL OR dimension_name = ?) - AND (? IS NULL OR is_enabled = ?) -ORDER BY updated_at DESC, priority DESC, id ASC -LIMIT ? OFFSET ? - "#, - ) - .bind(api_format) - .bind(api_format) - .bind(task_type) - .bind(task_type) - .bind(dimension_name) - .bind(dimension_name) - .bind(is_enabled) - .bind(is_enabled) - .bind(i64::from(page_size)) - .bind( - i64::try_from(offset) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_admin_billing_collector_mysql) - .collect::, _>>()?; - Ok(Some((items, total))) - } - - async fn find_admin_billing_collector( - &self, - collector_id: &str, - ) -> Result, DataLayerError> { - find_admin_billing_collector_mysql(&self.pool, collector_id).await - } - - async fn update_admin_billing_collector( - &self, - collector_id: &str, - input: &AdminBillingCollectorWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE dimension_collectors -SET api_format = ?, - task_type = ?, - dimension_name = ?, - source_type = ?, - source_path = ?, - value_type = ?, - transform_expression = ?, - default_value = ?, - priority = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&input.api_format) - .bind(&input.task_type) - .bind(&input.dimension_name) - .bind(&input.source_type) - .bind(input.source_path.as_deref()) - .bind(&input.value_type) - .bind(input.transform_expression.as_deref()) - .bind(input.default_value.as_deref()) - .bind(input.priority) - .bind(input.is_enabled) - .bind(current_unix_secs_i64()) - .bind(collector_id) - .execute(&self.pool) - .await; - let affected = match result { - Ok(result) => result.rows_affected(), - Err(err) => { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))) - } - }; - if affected == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match find_admin_billing_collector_mysql(&self.pool, collector_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn apply_admin_billing_preset( - &self, - preset: &str, - mode: &str, - collectors: &[AdminBillingCollectorWriteInput], - ) -> Result, DataLayerError> { - let mut created = 0_u64; - let mut updated = 0_u64; - let mut skipped = 0_u64; - let mut errors = Vec::new(); - - for collector in collectors { - let existing_id = match sqlx::query_scalar::<_, String>( - r#" -SELECT id -FROM dimension_collectors -WHERE api_format = ? - AND task_type = ? - AND dimension_name = ? - AND priority = ? - AND is_enabled = 1 -LIMIT 1 - "#, - ) - .bind(&collector.api_format) - .bind(&collector.task_type) - .bind(&collector.dimension_name) - .bind(collector.priority) - .fetch_optional(&self.pool) - .await - { - Ok(value) => value, - Err(err) => { - errors.push(format!( - "Failed to query collector: api_format={} task_type={} dim={}: {}", - collector.api_format, collector.task_type, collector.dimension_name, err - )); - continue; - } - }; - - if let Some(existing_id) = existing_id { - if mode == "overwrite" { - match sqlx::query( - r#" -UPDATE dimension_collectors -SET source_type = ?, - source_path = ?, - value_type = ?, - transform_expression = ?, - default_value = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&collector.source_type) - .bind(collector.source_path.as_deref()) - .bind(&collector.value_type) - .bind(collector.transform_expression.as_deref()) - .bind(collector.default_value.as_deref()) - .bind(collector.is_enabled) - .bind(current_unix_secs_i64()) - .bind(&existing_id) - .execute(&self.pool) - .await - { - Ok(_) => updated += 1, - Err(err) => errors.push(format!( - "Failed to update collector {}: {}", - existing_id, err - )), - } - } else { - skipped += 1; - } - continue; - } - - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - match sqlx::query( - r#" -INSERT INTO dimension_collectors ( - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(id) - .bind(&collector.api_format) - .bind(&collector.task_type) - .bind(&collector.dimension_name) - .bind(&collector.source_type) - .bind(collector.source_path.as_deref()) - .bind(&collector.value_type) - .bind(collector.transform_expression.as_deref()) - .bind(collector.default_value.as_deref()) - .bind(collector.priority) - .bind(collector.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - { - Ok(_) => created += 1, - Err(err) => errors.push(format!( - "Failed to create collector: api_format={} task_type={} dim={}: {}", - collector.api_format, collector.task_type, collector.dimension_name, err - )), - } - } - - Ok(AdminBillingMutationOutcome::Applied( - AdminBillingPresetApplyResult { - preset: preset.to_string(), - mode: mode.to_string(), - created, - updated, - skipped, - errors, - }, - )) - } - - async fn find_payment_gateway_config( - &self, - provider: &str, - ) -> Result, DataLayerError> { - 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.trim().to_ascii_lowercase()) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref() - .map(map_payment_gateway_config_mysql) - .transpose() - } - - async fn compare_and_swap_payment_gateway_secret( - &self, - update: &PaymentGatewaySecretCasUpdate, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE payment_gateway_configs -SET merchant_key_encrypted = ? -WHERE provider = ? - AND BINARY merchant_key_encrypted = BINARY ? - "#, - ) - .bind(&update.merchant_key_encrypted) - .bind(update.provider.trim().to_ascii_lowercase()) - .bind(&update.expected_merchant_key_encrypted) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn compare_and_swap_payment_gateway_config( - &self, - mutation: &PaymentGatewayConfigCasWriteInput, - ) -> Result, DataLayerError> { - let input = &mutation.input; - let provider = input.provider.trim().to_ascii_lowercase(); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - - if mutation.expected_existing { - let current = sqlx::query( - r#" -SELECT merchant_key_encrypted -FROM payment_gateway_configs -WHERE provider = ? -LIMIT 1 -FOR UPDATE - "#, - ) - .bind(&provider) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let current_secret = match current.as_ref() { - Some(row) => row - .try_get::, _>("merchant_key_encrypted") - .map_sql_err()?, - None => { - tx.rollback().await.map_sql_err()?; - return Ok(AdminBillingMutationOutcome::NotFound); - } - }; - if current_secret != mutation.expected_merchant_key_encrypted { - tx.rollback().await.map_sql_err()?; - return Ok(AdminBillingMutationOutcome::NotFound); - } - - sqlx::query( - r#" -UPDATE payment_gateway_configs -SET - enabled = ?, - endpoint_url = ?, - callback_base_url = ?, - merchant_id = ?, - merchant_key_encrypted = CASE - WHEN ? THEN merchant_key_encrypted - ELSE ? - END, - pay_currency = ?, - usd_exchange_rate = ?, - min_recharge_usd = ?, - channels_json = ?, - updated_at = ? -WHERE provider = ? - "#, - ) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .bind(input.preserve_existing_secret) - .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(&provider) - .execute(&mut *tx) - .await - .map_sql_err()?; - } else { - let inserted = sqlx::query( - r#" -INSERT INTO payment_gateway_configs ( - provider, enabled, endpoint_url, callback_base_url, merchant_id, - merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, - channels_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&provider) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .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(&mut *tx) - .await; - if let Err(err) = inserted { - let unique = matches!( - &err, - sqlx::Error::Database(database_error) if database_error.is_unique_violation() - ); - tx.rollback().await.map_sql_err()?; - if unique { - return Ok(AdminBillingMutationOutcome::NotFound); - } - return Err(DataLayerError::sql(err)); - } - } - - 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 upsert_payment_gateway_config( - &self, - input: &PaymentGatewayConfigWriteInput, - ) -> Result, DataLayerError> { - let provider = input.provider.trim().to_ascii_lowercase(); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query( - r#" -INSERT INTO payment_gateway_configs ( - provider, enabled, endpoint_url, callback_base_url, merchant_id, - merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, - channels_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - enabled = VALUES(enabled), - endpoint_url = VALUES(endpoint_url), - callback_base_url = VALUES(callback_base_url), - merchant_id = VALUES(merchant_id), - 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), - channels_json = VALUES(channels_json), - updated_at = VALUES(updated_at) - "#, - ) - .bind(&provider) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .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) - .bind(input.preserve_existing_secret) - .execute(&mut *tx) - .await - .map_sql_err()?; - 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( - &self, - include_disabled: bool, - ) -> Result>, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM billing_plans -WHERE (? = TRUE OR enabled = TRUE) -ORDER BY sort_order ASC, price_amount ASC, id ASC - "#, - ) - .bind(include_disabled) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(Some( - rows.iter() - .map(map_billing_plan_mysql) - .collect::, _>>()?, - )) - } - - async fn find_billing_plan( - &self, - plan_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM billing_plans -WHERE id = ? -LIMIT 1 - "#, - ) - .bind(plan_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_billing_plan_mysql).transpose() - } - - async fn create_billing_plan( - &self, - input: &BillingPlanWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - sqlx::query(BILLING_PLAN_INSERT_MYSQL) - .bind(&id) - .bind(&input.title) - .bind(input.description.as_deref()) - .bind(input.price_amount) - .bind(&input.price_currency) - .bind(&input.duration_unit) - .bind(input.duration_value) - .bind(input.enabled) - .bind(input.sort_order) - .bind(input.max_active_per_user) - .bind(&input.purchase_limit_scope) - .bind(json_to_string(&input.entitlements_json)?) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - match self.find_billing_plan(&id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing plan missing".to_string(), - )), - } - } - - async fn update_billing_plan( - &self, - plan_id: &str, - input: &BillingPlanWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query(BILLING_PLAN_UPDATE_MYSQL) - .bind(&input.title) - .bind(input.description.as_deref()) - .bind(input.price_amount) - .bind(&input.price_currency) - .bind(&input.duration_unit) - .bind(input.duration_value) - .bind(input.enabled) - .bind(input.sort_order) - .bind(input.max_active_per_user) - .bind(&input.purchase_limit_scope) - .bind(json_to_string(&input.entitlements_json)?) - .bind(current_unix_secs_i64()) - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match self.find_billing_plan(plan_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn set_billing_plan_enabled( - &self, - plan_id: &str, - enabled: bool, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE billing_plans SET enabled = ?, updated_at = ? WHERE id = ?") - .bind(enabled) - .bind(current_unix_secs_i64()) - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match self.find_billing_plan(plan_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn delete_billing_plan( - &self, - plan_id: &str, - ) -> Result, DataLayerError> { - let exists = - sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM billing_plans WHERE id = ?") - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - if exists == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - - let order_count = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM payment_orders -WHERE product_id = ? - AND order_kind = 'plan_purchase' - "#, - ) - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let entitlement_count = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE plan_id = ?", - ) - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - if order_count > 0 || entitlement_count > 0 { - return Ok(AdminBillingMutationOutcome::Invalid( - "套餐已有订单或权益,不能删除,请停用该套餐".to_string(), - )); - } - - let result = sqlx::query("DELETE FROM billing_plans WHERE id = ?") - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - Ok(AdminBillingMutationOutcome::NotFound) - } else { - Ok(AdminBillingMutationOutcome::Applied(())) - } - } - - async fn list_user_plan_entitlements( - &self, - user_id: &str, - ) -> Result>, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, user_id, plan_id, payment_order_id, status, - starts_at AS starts_at_unix_secs, expires_at AS expires_at_unix_secs, - entitlements_snapshot, created_at AS created_at_unix_secs, - updated_at AS updated_at_unix_secs -FROM user_plan_entitlements -WHERE user_id = ? - AND status = 'active' - AND expires_at > ? -ORDER BY expires_at ASC, created_at ASC - "#, - ) - .bind(user_id) - .bind(current_unix_secs_i64()) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(Some( - rows.iter() - .map(map_user_plan_entitlement_mysql) - .collect::, _>>()?, - )) - } - - async fn revoke_user_plan_entitlement( - &self, - user_id: &str, - entitlement_id: &str, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -UPDATE user_plan_entitlements -SET status = 'revoked', - expires_at = LEAST(expires_at, ?), - updated_at = ? -WHERE id = ? - AND user_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(now) - .bind(now) - .bind(entitlement_id) - .bind(user_id) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - Ok(AdminBillingMutationOutcome::NotFound) - } else { - Ok(AdminBillingMutationOutcome::Applied(())) - } - } - - async fn find_user_daily_quota_availability( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let now_unix_secs = current_unix_secs_i64(); - let rows = sqlx::query( - r#" -SELECT - user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json -FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id -WHERE user_plan_entitlements.user_id = ? - AND user_plan_entitlements.status = 'active' - AND user_plan_entitlements.starts_at <= ? - AND user_plan_entitlements.expires_at > ? -ORDER BY user_plan_entitlements.expires_at ASC, - user_plan_entitlements.created_at ASC, - user_plan_entitlements.id ASC - "#, - ) - .bind(user_id) - .bind(now_unix_secs) - .bind(now_unix_secs) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let now = chrono::Utc::now(); - let mut grants = Vec::new(); - for row in rows { - let entitlement_id: String = row.try_get("id").map_sql_err()?; - let entitlements = parse_json(row.try_get("entitlements_snapshot").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])); - let plan_entitlements = - parse_json(row.try_get("plan_entitlements_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])); - grants.extend(daily_quota_grants_from_entitlement( - &entitlement_id, - &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), - now, - )?); - } - - let mut total_quota_usd = 0.0; - let mut used_usd = 0.0; - let mut remaining_usd = 0.0; - let mut allow_wallet_overage = true; - for grant in &grants { - allow_wallet_overage &= grant.allow_wallet_overage; - let used = sqlx::query_scalar::<_, f64>( - r#" -SELECT COALESCE(SUM(amount_usd), 0) -FROM entitlement_usage_ledgers -WHERE user_entitlement_id = ? - AND usage_date = ? - "#, - ) - .bind(&grant.entitlement_id) - .bind(&grant.usage_date) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - total_quota_usd += grant.daily_quota_usd; - used_usd += used.min(grant.daily_quota_usd).max(0.0); - remaining_usd += (grant.daily_quota_usd - used).max(0.0); - } - let has_active_daily_quota = !grants.is_empty(); - Ok(Some(UserDailyQuotaAvailabilityRecord { - has_active_daily_quota, - total_quota_usd, - used_usd, - remaining_usd, - allow_wallet_overage: has_active_daily_quota && allow_wallet_overage, - })) - } -} - -const BILLING_PLAN_INSERT_MYSQL: &str = r#" -INSERT INTO billing_plans ( - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#; - -const BILLING_PLAN_UPDATE_MYSQL: &str = r#" -UPDATE billing_plans -SET title = ?, - description = ?, - price_amount = ?, - price_currency = ?, - duration_unit = ?, - duration_value = ?, - enabled = ?, - sort_order = ?, - max_active_per_user = ?, - purchase_limit_scope = ?, - entitlements_json = ?, - updated_at = ? -WHERE id = ? -"#; - -struct RankedContext { - rank: u8, - is_available: bool, - pricing_rank: u8, - created_at: i64, - context: Result, -} - -fn match_rank( - row: &MySqlRow, - requested_model: &str, -) -> Result, DataLayerError> { - let provider_model_name: Option = - row.try_get("model_provider_model_name").map_sql_err()?; - let global_model_name: String = row.try_get("global_model_name").map_sql_err()?; - let mappings: Option = row.try_get("provider_model_mappings").ok().flatten(); - - let rank = if provider_model_name.as_deref() == Some(requested_model) { - 0 - } else if mappings - .as_deref() - .is_some_and(|mappings| provider_model_mappings_match(mappings, requested_model)) - { - 1 - } else if global_model_name == requested_model { - 2 - } else { - return Ok(None); - }; - - let has_model_price = row - .try_get::, _>("model_price_per_request") - .map_sql_err()? - .is_some() - || row - .try_get::, _>("model_tiered_pricing") - .ok() - .flatten() - .is_some(); - let has_default_price = row - .try_get::, _>("default_price_per_request") - .map_sql_err()? - .is_some() - || row - .try_get::, _>("default_tiered_pricing") - .ok() - .flatten() - .is_some(); - let pricing_rank = if has_model_price { - 0 - } else if has_default_price { - 1 - } else { - 2 - }; - - Ok(Some(RankedContext { - rank, - is_available: row - .try_get::, _>("model_is_available") - .map_sql_err()? - .unwrap_or(false), - pricing_rank, - created_at: row - .try_get::, _>("model_created_at") - .map_sql_err()? - .unwrap_or(i64::MAX), - context: map_row(row), - })) -} - -fn provider_model_mappings_match(raw: &str, requested_model: &str) -> bool { - let Ok(value) = serde_json::from_str::(raw) else { - return raw == requested_model; - }; - json_mapping_matches(&value, requested_model) -} - -fn json_mapping_matches(value: &serde_json::Value, requested_model: &str) -> bool { - match value { - serde_json::Value::String(value) => value == requested_model, - serde_json::Value::Array(values) => values - .iter() - .any(|value| json_mapping_matches(value, requested_model)), - serde_json::Value::Object(map) => map - .get("name") - .is_some_and(|value| json_mapping_matches(value, requested_model)), - _ => false, - } -} - -fn map_row(row: &MySqlRow) -> Result { - StoredBillingModelContext::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_billing_type").map_sql_err()?, - row.try_get("provider_api_key_id").map_sql_err()?, - parse_json( - row.try_get("provider_api_key_rate_multipliers") - .ok() - .flatten(), - )?, - row.try_get::, _>("provider_api_key_cache_ttl_minutes") - .map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - row.try_get("global_model_name").map_sql_err()?, - parse_json(row.try_get("global_model_config").ok().flatten())?, - row.try_get("default_price_per_request").map_sql_err()?, - parse_json(row.try_get("default_tiered_pricing").ok().flatten())?, - row.try_get("model_id").map_sql_err()?, - row.try_get("model_provider_model_name").map_sql_err()?, - parse_json(row.try_get("model_config").ok().flatten())?, - row.try_get("model_price_per_request").map_sql_err()?, - parse_json(row.try_get("model_tiered_pricing").ok().flatten())?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON field is invalid: {err}")) - }) - }) - .transpose() -} - -fn current_unix_secs_i64() -> i64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs() as i64 -} - -fn json_to_string(value: &serde_json::Value) -> Result { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON encode failed: {err}")) - }) -} - -#[derive(Debug)] -struct DailyQuotaGrant { - entitlement_id: String, - daily_quota_usd: f64, - usage_date: String, - allow_wallet_overage: bool, -} - -fn daily_quota_usage_date( - reset_timezone: Option<&str>, - now: chrono::DateTime, -) -> Result { - let timezone = reset_timezone - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("Asia/Shanghai") - .parse::() - .map_err(|err| DataLayerError::InvalidInput(format!("invalid reset_timezone: {err}")))?; - Ok(now.with_timezone(&timezone).date_naive().to_string()) -} - -fn daily_quota_grants_from_entitlement( - entitlement_id: &str, - entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, - now: chrono::DateTime, -) -> Result, DataLayerError> { - let mut grants = Vec::new(); - let Some(items) = entitlements.as_array() else { - return Ok(grants); - }; - for item in items { - if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") { - continue; - } - let daily_quota_usd = item - .get("daily_quota_usd") - .and_then(serde_json::Value::as_f64) - .unwrap_or(0.0); - if !daily_quota_usd.is_finite() || daily_quota_usd <= 0.0 { - continue; - } - grants.push(DailyQuotaGrant { - entitlement_id: entitlement_id.to_string(), - daily_quota_usd, - usage_date: daily_quota_usage_date( - item.get("reset_timezone") - .and_then(serde_json::Value::as_str), - now, - )?, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), - }); - } - Ok(grants) -} - -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - -fn read_count_mysql(row: &MySqlRow) -> Result { - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) -} - -fn map_payment_gateway_config_mysql( - row: &MySqlRow, -) -> Result { - Ok(PaymentGatewayConfigRecord { - provider: row.try_get("provider").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - endpoint_url: row.try_get("endpoint_url").map_sql_err()?, - callback_base_url: row.try_get("callback_base_url").map_sql_err()?, - merchant_id: row.try_get("merchant_id").map_sql_err()?, - merchant_key_encrypted: row.try_get("merchant_key_encrypted").map_sql_err()?, - pay_currency: row.try_get("pay_currency").map_sql_err()?, - usd_exchange_rate: row.try_get("usd_exchange_rate").map_sql_err()?, - min_recharge_usd: row.try_get("min_recharge_usd").map_sql_err()?, - channels_json: parse_json(row.try_get("channels_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn map_billing_plan_mysql(row: &MySqlRow) -> Result { - Ok(BillingPlanRecord { - id: row.try_get("id").map_sql_err()?, - title: row.try_get("title").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - price_amount: row.try_get("price_amount").map_sql_err()?, - price_currency: row.try_get("price_currency").map_sql_err()?, - duration_unit: row.try_get("duration_unit").map_sql_err()?, - duration_value: row.try_get("duration_value").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - sort_order: row.try_get("sort_order").map_sql_err()?, - max_active_per_user: row.try_get("max_active_per_user").map_sql_err()?, - purchase_limit_scope: row - .try_get::, _>("purchase_limit_scope") - .map_sql_err()? - .unwrap_or_else(|| "active_period".to_string()), - entitlements_json: parse_json(row.try_get("entitlements_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn map_user_plan_entitlement_mysql( - row: &MySqlRow, -) -> Result { - Ok(UserPlanEntitlementRecord { - id: row.try_get("id").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - plan_id: row.try_get("plan_id").map_sql_err()?, - payment_order_id: row.try_get("payment_order_id").map_sql_err()?, - status: row.try_get("status").map_sql_err()?, - starts_at_unix_secs: row - .try_get::("starts_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - expires_at_unix_secs: row - .try_get::("expires_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - entitlements_snapshot: parse_json(row.try_get("entitlements_snapshot").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -async fn find_admin_billing_rule_mysql( - pool: &MysqlPool, - rule_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM billing_rules -WHERE id = ? - "#, - ) - .bind(rule_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_admin_billing_rule_mysql).transpose() -} - -fn map_admin_billing_rule_mysql(row: &MySqlRow) -> Result { - Ok(AdminBillingRuleRecord { - id: row.try_get("id").map_sql_err()?, - name: row.try_get("name").map_sql_err()?, - task_type: row.try_get("task_type").map_sql_err()?, - global_model_id: row.try_get("global_model_id").map_sql_err()?, - model_id: row.try_get("model_id").map_sql_err()?, - expression: row.try_get("expression").map_sql_err()?, - variables: parse_required_json(row.try_get("variables").map_sql_err()?)?, - dimension_mappings: parse_required_json(row.try_get("dimension_mappings").map_sql_err()?)?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - created_at_unix_ms: row - .try_get::("created_at_unix_ms") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -async fn find_admin_billing_collector_mysql( - pool: &MysqlPool, - collector_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, - created_at AS created_at_unix_ms, updated_at AS updated_at_unix_secs -FROM dimension_collectors -WHERE id = ? - "#, - ) - .bind(collector_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref() - .map(map_admin_billing_collector_mysql) - .transpose() -} - -fn map_admin_billing_collector_mysql( - row: &MySqlRow, -) -> Result { - Ok(AdminBillingCollectorRecord { - id: row.try_get("id").map_sql_err()?, - api_format: row.try_get("api_format").map_sql_err()?, - task_type: row.try_get("task_type").map_sql_err()?, - dimension_name: row.try_get("dimension_name").map_sql_err()?, - source_type: row.try_get("source_type").map_sql_err()?, - source_path: row.try_get("source_path").map_sql_err()?, - value_type: row.try_get("value_type").map_sql_err()?, - transform_expression: row.try_get("transform_expression").map_sql_err()?, - default_value: row.try_get("default_value").map_sql_err()?, - priority: row.try_get("priority").map_sql_err()?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - created_at_unix_ms: row - .try_get::("created_at_unix_ms") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn parse_required_json(raw: String) -> Result { - serde_json::from_str(&raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON field is invalid: {err}")) - }) -} - -#[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() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - 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/candidate_selection.rs b/crates/aether-data/adapters/mysql/src/candidate_selection.rs deleted file mode 100644 index 5bdfe04a0..000000000 --- a/crates/aether-data/adapters/mysql/src/candidate_selection.rs +++ /dev/null @@ -1,1405 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; - -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::candidate_selection::{ - provider_model_mapping_api_format_covers, MinimalCandidateSelectionReadRepository, - StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow, - StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, - StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, - StoredRequestedModelCandidateRowsQuery, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const CANDIDATE_SELECTION_COLUMNS: &str = r#" -SELECT - p.id AS provider_id, - p.name AS provider_name, - p.provider_type AS provider_type, - p.provider_priority AS provider_priority, - p.is_active AS provider_is_active, - p.config AS provider_config, - pe.id AS endpoint_id, - COALESCE(pe.api_format, '') AS endpoint_api_format, - pe.api_family AS endpoint_api_family, - pe.endpoint_kind AS endpoint_kind, - pe.is_active AS endpoint_is_active, - pak.id AS key_id, - pak.name AS key_name, - pak.auth_type AS key_auth_type, - pak.auth_config AS key_auth_config, - pak.is_active AS key_is_active, - pak.api_formats AS key_api_formats, - pak.allowed_models AS key_allowed_models, - pak.capabilities AS key_capabilities, - pak.internal_priority AS key_internal_priority, - pak.global_priority_by_format AS key_global_priority_by_format, - m.id AS model_id, - m.global_model_id AS global_model_id, - gm.name AS global_model_name, - gm.config AS global_model_config, - m.provider_model_name AS model_provider_model_name, - m.provider_model_mappings AS model_provider_model_mappings, - m.supports_streaming AS model_supports_streaming, - m.is_active AS model_is_active, - m.is_available AS model_is_available -FROM providers p -INNER JOIN provider_endpoints pe ON pe.provider_id = p.id -INNER JOIN provider_api_keys pak ON pak.provider_id = p.id -INNER JOIN models m ON m.provider_id = p.id -INNER JOIN global_models gm ON gm.id = m.global_model_id -WHERE p.is_active = 1 - AND pe.is_active = 1 - AND pak.is_active = 1 - AND m.is_active = 1 - AND m.is_available = 1 - AND gm.is_active = 1 -"#; - -const REQUESTED_MODEL_RAW_PAGE_SIZE: u32 = 256; -const REQUESTED_MODEL_RAW_SCAN_LIMIT: u32 = 2048; - -#[derive(Debug, Clone)] -pub struct MysqlMinimalCandidateSelectionReadRepository { - pool: MysqlPool, -} - -#[derive(Debug, Clone)] -struct CandidateSelectionRow { - row: StoredMinimalCandidateSelectionRow, - provider_pool_enabled: bool, - key_auth_config: Option, -} - -#[derive(Debug)] -struct ExactPageAccumulator { - rows: Vec, - offset: usize, - limit: usize, - target_len: usize, -} - -impl ExactPageAccumulator { - fn new(offset: u32, limit: u32) -> Self { - let offset = usize::try_from(offset).unwrap_or(usize::MAX); - let limit = usize::try_from(limit).unwrap_or(usize::MAX); - Self { - rows: Vec::new(), - offset, - limit, - target_len: offset.saturating_add(limit), - } - } - - fn is_full(&self) -> bool { - self.rows.len() >= self.target_len - } - - fn push_matching(&mut self, rows: I, mut predicate: F) - where - I: IntoIterator, - F: FnMut(&T) -> bool, - { - let remaining = self.target_len.saturating_sub(self.rows.len()); - self.rows.extend( - rows.into_iter() - .filter(|row| predicate(row)) - .take(remaining), - ); - } - - fn into_page(self) -> Vec { - self.rows - .into_iter() - .skip(self.offset) - .take(self.limit) - .collect() - } -} - -impl MysqlMinimalCandidateSelectionReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn load_rows_for_api_format( - &self, - api_format: &str, - ) -> Result, DataLayerError> { - let canonical_api_format = normalize_api_format(api_format); - let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = sql_match_aliases(&storage_aliases); - - let mut builder = QueryBuilder::::new(CANDIDATE_SELECTION_COLUMNS); - builder.push(" AND LOWER(pe.api_format) IN ("); - { - let mut separated = builder.separated(", "); - for alias in &match_aliases { - separated.push_bind(alias); - } - } - builder.push(")"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut items = rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - items.retain(|item| { - api_format_matches(&item.row.endpoint_api_format, &canonical_api_format) - && item.row.key_supports_api_format(&canonical_api_format) - && key_auth_channel_matches(item, &canonical_api_format) - }); - Ok(items) - } - - async fn selected_rows_for_api_format( - &self, - api_format: &str, - ) -> Result, DataLayerError> { - let rows = self.load_rows_for_api_format(api_format).await?; - Ok(sort_rows(select_pool_rows(rows), true)) - } - - async fn load_selected_rows_for_api_format_page( - &self, - api_format: &str, - limit: u32, - offset: u32, - ) -> Result, DataLayerError> { - let mut builder = api_format_page_query(api_format, limit, offset); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_candidate_selection_row).collect() - } - - async fn selected_rows_for_api_format_page( - &self, - query: &StoredApiFormatCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - let target_len = query.offset.saturating_add(query.limit); - let mut raw_offset = 0_u32; - let mut selected = Vec::new(); - while selected.len() < target_len as usize && raw_offset < REQUESTED_MODEL_RAW_SCAN_LIMIT { - let raw_limit = - target_len.min(REQUESTED_MODEL_RAW_SCAN_LIMIT.saturating_sub(raw_offset)); - let rows = self - .load_selected_rows_for_api_format_page(&query.api_format, raw_limit, raw_offset) - .await?; - let raw_len = rows.len() as u32; - selected.extend(rows.into_iter().filter(|item| { - api_format_matches(&item.row.endpoint_api_format, &query.api_format) - && item.row.key_supports_api_format(&query.api_format) - && key_auth_channel_matches(item, &query.api_format) - })); - if raw_len < raw_limit || raw_len == 0 { - break; - } - let next_offset = raw_offset.saturating_add(raw_len); - if next_offset == raw_offset { - break; - } - raw_offset = next_offset; - } - let rows = selected.into_iter().map(|item| item.row).collect(); - Ok(sort_rows(dedupe_candidate_selection_rows(rows), true) - .into_iter() - .skip(query.offset as usize) - .take(query.limit as usize) - .collect()) - } -} - -fn api_format_page_query( - api_format: &str, - limit: u32, - offset: u32, -) -> QueryBuilder<'static, MySql> { - let mut builder = QueryBuilder::::new("WITH candidate_rows AS ("); - builder.push(CANDIDATE_SELECTION_COLUMNS); - push_candidate_api_format_filters(&mut builder, api_format); - push_selected_pool_rows(&mut builder); - builder.push( - r#" -ORDER BY - global_model_name ASC, - provider_priority ASC, - key_internal_priority ASC, - provider_id ASC, - endpoint_id ASC, - key_id ASC, - model_id ASC -LIMIT "#, - ); - builder.push_bind(i64::from(limit)); - builder.push(" OFFSET "); - builder.push_bind(i64::from(offset)); - builder -} - -fn requested_model_page_query( - api_format: &str, - requested_model_name: &str, - limit: u32, - offset: u32, -) -> QueryBuilder<'static, MySql> { - let mut builder = QueryBuilder::::new("WITH candidate_rows AS ("); - builder.push(CANDIDATE_SELECTION_COLUMNS); - push_candidate_api_format_filters(&mut builder, api_format); - push_requested_model_sql_filter(&mut builder, requested_model_name); - push_selected_pool_rows(&mut builder); - builder.push( - r#" -ORDER BY - global_model_name ASC, - provider_priority ASC, - key_internal_priority ASC, - provider_id ASC, - endpoint_id ASC, - key_id ASC, - model_id ASC -LIMIT "#, - ); - builder.push_bind(i64::from(limit)); - builder.push(" OFFSET "); - builder.push_bind(i64::from(offset)); - builder -} - -fn pool_key_group_query(query: &StoredPoolKeyCandidateRowsQuery) -> QueryBuilder<'static, MySql> { - let mut builder = QueryBuilder::::new(CANDIDATE_SELECTION_COLUMNS); - push_candidate_api_format_filters(&mut builder, &query.api_format); - push_pool_key_group_filters( - &mut builder, - &query.provider_id, - &query.endpoint_id, - &query.model_id, - ); - push_pool_key_order(&mut builder, &query.order); - builder.push(" LIMIT "); - builder.push_bind(i64::from(query.limit)); - builder.push(" OFFSET "); - builder.push_bind(i64::from(query.offset)); - builder -} - -fn pool_key_group_by_key_ids_query( - query: &StoredPoolKeyCandidateRowsByKeyIdsQuery, -) -> QueryBuilder<'static, MySql> { - let mut builder = QueryBuilder::::new(CANDIDATE_SELECTION_COLUMNS); - push_candidate_api_format_filters(&mut builder, &query.api_format); - push_pool_key_group_filters( - &mut builder, - &query.provider_id, - &query.endpoint_id, - &query.model_id, - ); - builder.push(" AND pak.id IN ("); - { - let mut separated = builder.separated(", "); - for key_id in &query.key_ids { - separated.push_bind(key_id.clone()); - } - } - builder.push(") ORDER BY FIELD(pak.id, "); - { - let mut separated = builder.separated(", "); - for key_id in &query.key_ids { - separated.push_bind(key_id.clone()); - } - } - builder.push(") ASC, pak.id ASC"); - builder -} - -fn push_candidate_api_format_filters(builder: &mut QueryBuilder<'_, MySql>, api_format: &str) { - let canonical_api_format = normalize_api_format(api_format); - let storage_aliases = sql_match_aliases(&api_format_aliases(&canonical_api_format)); - let permission_aliases = - sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); - builder.push(" AND LOWER(pe.api_format) IN ("); - { - let mut separated = builder.separated(", "); - for alias in storage_aliases { - separated.push_bind(alias); - } - } - builder.push(") AND (pak.api_formats IS NULL OR TRIM(pak.api_formats) = ''"); - for alias in permission_aliases { - builder.push(" OR JSON_SEARCH(LOWER(pak.api_formats), 'one', "); - builder.push_bind(alias); - builder.push(") IS NOT NULL"); - } - builder.push(")"); - push_key_auth_channel_filter(builder, &canonical_api_format); -} - -fn push_requested_model_sql_filter( - builder: &mut QueryBuilder<'_, MySql>, - requested_model_name: &str, -) { - builder.push(" AND (gm.name = "); - builder.push_bind(requested_model_name.to_string()); - builder.push(" OR m.provider_model_name = "); - builder.push_bind(requested_model_name.to_string()); - builder.push(" OR (m.provider_model_mappings IS NOT NULL AND LOCATE("); - builder.push_bind(requested_model_name.to_string()); - builder.push(", m.provider_model_mappings) > 0))"); -} - -fn push_selected_pool_rows(builder: &mut QueryBuilder<'_, MySql>) { - builder.push( - r#" -), -ranked_rows AS ( - SELECT - candidate_rows.*, - CASE - WHEN JSON_EXTRACT(provider_config, '$.pool_advanced') IS NOT NULL - AND JSON_TYPE(JSON_EXTRACT(provider_config, '$.pool_advanced')) <> 'NULL' - THEN 1 ELSE 0 - END AS provider_pool_enabled, - ROW_NUMBER() OVER ( - PARTITION BY provider_id, endpoint_id, model_id - ORDER BY key_internal_priority ASC, key_id ASC - ) AS pool_rank - FROM candidate_rows -), -selected_rows AS ( - SELECT * - FROM ranked_rows - WHERE provider_pool_enabled = 0 OR pool_rank = 1 -) -SELECT * -FROM selected_rows"#, - ); -} - -fn push_pool_key_group_filters( - builder: &mut QueryBuilder<'_, MySql>, - provider_id: &str, - endpoint_id: &str, - model_id: &str, -) { - builder.push(" AND p.id = "); - builder.push_bind(provider_id.to_string()); - builder.push(" AND pe.id = "); - builder.push_bind(endpoint_id.to_string()); - builder.push(" AND m.id = "); - builder.push_bind(model_id.to_string()); -} - -fn push_pool_key_order(builder: &mut QueryBuilder<'_, MySql>, order: &StoredPoolKeyCandidateOrder) { - match order { - StoredPoolKeyCandidateOrder::InternalPriority => { - builder.push(" ORDER BY pak.internal_priority ASC, pak.id ASC"); - } - StoredPoolKeyCandidateOrder::Lru => { - builder.push( - " ORDER BY pak.last_used_at IS NOT NULL ASC, pak.last_used_at ASC, pak.internal_priority ASC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::CacheAffinity => { - builder.push( - " ORDER BY pak.last_used_at IS NULL ASC, pak.last_used_at DESC, pak.internal_priority ASC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::SingleAccount => { - builder.push( - " ORDER BY pak.internal_priority ASC, pak.last_used_at IS NULL ASC, pak.last_used_at DESC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::LoadBalance { seed } => { - builder.push(" ORDER BY MD5(CONCAT("); - builder.push_bind(seed.clone()); - builder.push(", ':', pak.id)) ASC, pak.id ASC"); - } - } -} - -#[async_trait] -impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionReadRepository { - async fn list_for_exact_api_format( - &self, - api_format: &str, - ) -> Result, DataLayerError> { - self.selected_rows_for_api_format(api_format).await - } - - async fn list_for_exact_api_format_page( - &self, - query: &StoredApiFormatCandidateRowsQuery, - ) -> Result, DataLayerError> { - self.selected_rows_for_api_format_page(query).await - } - - async fn list_for_exact_api_format_and_global_model( - &self, - api_format: &str, - global_model_name: &str, - ) -> Result, DataLayerError> { - Ok(sort_rows( - self.selected_rows_for_api_format(api_format) - .await? - .into_iter() - .filter(|row| row.global_model_name == global_model_name) - .collect(), - false, - )) - } - - async fn list_for_exact_api_format_and_requested_model( - &self, - api_format: &str, - requested_model_name: &str, - ) -> Result, DataLayerError> { - let rows = self - .selected_rows_for_api_format(api_format) - .await? - .into_iter() - .filter(|row| row_matches_requested_model(row, requested_model_name, api_format)) - .collect::>(); - Ok(sort_rows(rows, true)) - } - - async fn list_for_exact_api_format_and_requested_model_page( - &self, - query: &StoredRequestedModelCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - // The SQL model predicate is a coarse superset, so fill the exact page across raw windows. - let mut exact_page = ExactPageAccumulator::new(query.offset, query.limit); - let mut raw_offset = 0_u32; - while !exact_page.is_full() && raw_offset < REQUESTED_MODEL_RAW_SCAN_LIMIT { - let raw_limit = REQUESTED_MODEL_RAW_PAGE_SIZE - .min(REQUESTED_MODEL_RAW_SCAN_LIMIT.saturating_sub(raw_offset)); - let mut builder = requested_model_page_query( - &query.api_format, - &query.requested_model_name, - raw_limit, - raw_offset, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let raw_len = u32::try_from(rows.len()).unwrap_or(u32::MAX); - let items = rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - exact_page.push_matching(items, |item| { - row_matches_requested_model( - &item.row, - &query.requested_model_name, - &query.api_format, - ) - }); - raw_offset = raw_offset.saturating_add(raw_len); - if raw_len < raw_limit || raw_len == 0 { - break; - } - } - let rows = exact_page - .into_page() - .into_iter() - .map(|item| item.row) - .collect(); - Ok(sort_rows(dedupe_candidate_selection_rows(rows), true)) - } - - async fn list_pool_key_rows_for_group( - &self, - query: &StoredPoolKeyCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - let mut builder = pool_key_group_query(query); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let rows = rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - Ok(dedupe_candidate_selection_rows( - rows.into_iter().map(|item| item.row).collect(), - )) - } - - async fn list_pool_key_rows_for_group_key_ids( - &self, - query: &StoredPoolKeyCandidateRowsByKeyIdsQuery, - ) -> Result, DataLayerError> { - if query.key_ids.is_empty() { - return Ok(Vec::new()); - } - let key_order = query - .key_ids - .iter() - .enumerate() - .map(|(index, key_id)| (key_id.as_str(), index)) - .collect::>(); - let mut builder = pool_key_group_by_key_ids_query(query); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut rows = rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()? - .into_iter() - .map(|item| item.row) - .collect::>(); - rows = dedupe_candidate_selection_rows(rows); - rows.sort_by(|left, right| { - key_order - .get(left.key_id.as_str()) - .cmp(&key_order.get(right.key_id.as_str())) - .then(left.key_id.cmp(&right.key_id)) - }); - Ok(rows) - } -} - -fn select_pool_rows(rows: Vec) -> Vec { - let mut selected = Vec::new(); - let mut pool_rows = - BTreeMap::<(String, String, String), StoredMinimalCandidateSelectionRow>::new(); - for item in rows { - if !item.provider_pool_enabled { - selected.push(item.row); - continue; - } - let key = ( - item.row.provider_id.clone(), - item.row.endpoint_id.clone(), - item.row.model_id.clone(), - ); - match pool_rows.get(&key) { - Some(existing) - if (existing.key_internal_priority, existing.key_id.as_str()) - <= (item.row.key_internal_priority, item.row.key_id.as_str()) => {} - _ => { - pool_rows.insert(key, item.row); - } - } - } - selected.extend(pool_rows.into_values()); - dedupe_candidate_selection_rows(selected) -} - -fn sort_rows( - mut rows: Vec, - include_global_model: bool, -) -> Vec { - rows.sort_by(|left, right| { - if include_global_model { - let ordering = left.global_model_name.cmp(&right.global_model_name); - if !ordering.is_eq() { - return ordering; - } - } - left.provider_priority - .cmp(&right.provider_priority) - .then(left.key_internal_priority.cmp(&right.key_internal_priority)) - .then(left.provider_id.cmp(&right.provider_id)) - .then(left.endpoint_id.cmp(&right.endpoint_id)) - .then(left.key_id.cmp(&right.key_id)) - .then(left.model_id.cmp(&right.model_id)) - }); - rows -} - -fn row_matches_requested_model( - row: &StoredMinimalCandidateSelectionRow, - requested_model_name: &str, - api_format: &str, -) -> bool { - (row_has_available_provider_model(row, api_format) - && row.global_model_name == requested_model_name) - || (row_default_provider_model_name_available(row, api_format) - && row.model_provider_model_name == requested_model_name) - || row - .model_provider_model_mappings - .as_ref() - .is_some_and(|mappings| { - mappings.iter().any(|mapping| { - mapping_scope_matches(mapping, row, api_format) - && mapping.name == requested_model_name - }) - }) -} - -fn row_has_available_provider_model( - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - row_mapping_matches_scope(row, api_format) - || row_default_provider_model_name_available(row, api_format) -} - -fn row_default_provider_model_name_available( - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - let Some(mappings) = row.model_provider_model_mappings.as_ref() else { - return true; - }; - let mut has_explicit_default_mapping = false; - for mapping in mappings { - if mapping.name != row.model_provider_model_name { - continue; - } - has_explicit_default_mapping = true; - if mapping_scope_matches(mapping, row, api_format) { - return true; - } - } - !has_explicit_default_mapping -} - -fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool { - row.model_provider_model_mappings - .as_ref() - .is_some_and(|mappings| { - mappings - .iter() - .any(|mapping| mapping_scope_matches(mapping, row, api_format)) - }) -} - -fn mapping_scope_matches( - mapping: &StoredProviderModelMapping, - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - mapping.api_formats.as_ref().is_none_or(|formats| { - formats.iter().any(|value| { - provider_model_mapping_api_format_covers(&row.provider_type, value, api_format) - }) - }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { - endpoint_ids - .iter() - .any(|endpoint_id| endpoint_id == &row.endpoint_id) - }) -} - -fn push_key_auth_channel_filter(builder: &mut QueryBuilder<'_, MySql>, api_format: &str) { - builder.push(" AND ((LOWER(TRIM(p.provider_type)) = 'codex'"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) = 'oauth' AND "); - builder.push_bind(api_format.to_string()); - builder.push( - " IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live'))", - ); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'chatgpt_web'"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) IN ('oauth', 'bearer') AND "); - builder.push_bind(api_format.to_string()); - builder.push(" = 'openai:image')"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'claude_code'"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) = 'oauth' AND "); - builder.push_bind(api_format.to_string()); - builder.push(" = 'claude:messages')"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'kiro' AND "); - builder.push_bind(api_format.to_string()); - builder.push(" = 'claude:messages' AND (LOWER(TRIM(pak.auth_type)) = 'oauth'"); - builder.push(" OR (LOWER(TRIM(pak.auth_type)) = 'bearer'"); - builder.push(" AND pak.auth_config IS NOT NULL AND TRIM(pak.auth_config) <> '')))"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) = 'oauth' AND "); - builder.push_bind(api_format.to_string()); - builder.push(" = 'gemini:generate_content')"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'grok'"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) = 'oauth' AND "); - builder.push_bind(api_format.to_string()); - builder.push(" IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image'))"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'windsurf'"); - builder.push(" AND LOWER(TRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer') AND "); - builder.push_bind(api_format.to_string()); - builder.push(" = 'openai:chat')"); - - builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'vertex_ai'"); - builder.push( - " AND LOWER(TRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') AND ", - ); - builder.push_bind(api_format.to_string()); - builder.push(" IN ('gemini:generate_content', 'gemini:embedding'))"); - - builder.push( - " OR (LOWER(TRIM(p.provider_type)) NOT IN ('chatgpt_web', 'claude_code', 'codex', 'gemini_cli', 'grok', 'vertex_ai', 'antigravity', 'kiro', 'windsurf') AND LOWER(TRIM(pak.auth_type)) <> 'oauth'))", - ); -} - -fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool { - let provider_type = row.row.provider_type.trim().to_ascii_lowercase(); - let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase(); - let api_format = normalize_api_format(api_format); - match provider_type.as_str() { - "codex" => { - auth_type == "oauth" - && matches!( - api_format.as_str(), - "openai:responses" - | "openai:responses:compact" - | "openai:search" - | "openai:image" - | "codex:live" - ) - } - "chatgpt_web" => { - matches!(auth_type.as_str(), "oauth" | "bearer") && api_format == "openai:image" - } - "claude_code" => auth_type == "oauth" && api_format == "claude:messages", - "kiro" => { - api_format == "claude:messages" - && (auth_type == "oauth" - || (auth_type == "bearer" - && row - .key_auth_config - .as_deref() - .is_some_and(|value| !value.trim().is_empty()))) - } - "gemini_cli" | "antigravity" => { - auth_type == "oauth" && api_format == "gemini:generate_content" - } - "grok" => { - auth_type == "oauth" - && matches!( - api_format.as_str(), - "openai:chat" | "openai:responses" | "claude:messages" | "openai:image" - ) - } - "windsurf" => { - matches!(auth_type.as_str(), "oauth" | "api_key" | "bearer") - && api_format == "openai:chat" - } - "vertex_ai" => vertex_key_auth_channel_matches(&auth_type, &api_format), - _ => auth_type != "oauth", - } -} - -fn vertex_key_auth_channel_matches(auth_type: &str, api_format: &str) -> bool { - matches!(auth_type, "api_key" | "service_account" | "vertex_ai") - && matches!(api_format, "gemini:generate_content" | "gemini:embedding") -} - -fn dedupe_candidate_selection_rows( - rows: Vec, -) -> Vec { - let mut seen = BTreeSet::new(); - rows.into_iter() - .filter(|row| { - seen.insert(( - row.endpoint_id.clone(), - row.key_id.clone(), - row.model_id.clone(), - )) - }) - .collect() -} - -fn map_candidate_selection_row(row: &MySqlRow) -> Result { - let provider_config = parse_json(row.try_get("provider_config").ok().flatten())?; - let global_model_config = parse_json(row.try_get("global_model_config").ok().flatten())?; - let provider_pool_enabled = json_object_field_present(&provider_config, "pool_advanced"); - let global_model_mappings = global_model_config - .as_ref() - .and_then(|value| value.get("model_mappings").cloned()); - let global_model_supports_streaming = global_model_config - .as_ref() - .and_then(|value| value.get("streaming")) - .and_then(json_bool); - Ok(CandidateSelectionRow { - row: StoredMinimalCandidateSelectionRow { - provider_id: row.try_get("provider_id").map_sql_err()?, - provider_name: row.try_get("provider_name").map_sql_err()?, - provider_type: row.try_get("provider_type").map_sql_err()?, - provider_priority: row.try_get("provider_priority").map_sql_err()?, - provider_is_active: row.try_get("provider_is_active").map_sql_err()?, - endpoint_id: row.try_get("endpoint_id").map_sql_err()?, - endpoint_api_format: row.try_get("endpoint_api_format").map_sql_err()?, - endpoint_api_family: row.try_get("endpoint_api_family").map_sql_err()?, - endpoint_kind: row.try_get("endpoint_kind").map_sql_err()?, - endpoint_is_active: row.try_get("endpoint_is_active").map_sql_err()?, - key_id: row.try_get("key_id").map_sql_err()?, - key_name: row.try_get("key_name").map_sql_err()?, - key_auth_type: row.try_get("key_auth_type").map_sql_err()?, - key_is_active: row.try_get("key_is_active").map_sql_err()?, - key_api_formats: parse_stored_key_policy_string_list( - row.try_get("key_api_formats").map_sql_err()?, - "provider_api_keys.api_formats", - )?, - key_allowed_models: parse_stored_key_policy_string_list( - row.try_get("key_allowed_models").map_sql_err()?, - "provider_api_keys.allowed_models", - )?, - key_capabilities: parse_json(row.try_get("key_capabilities").ok().flatten())?, - key_internal_priority: row.try_get("key_internal_priority").map_sql_err()?, - key_global_priority_by_format: parse_json( - row.try_get("key_global_priority_by_format").ok().flatten(), - )?, - model_id: row.try_get("model_id").map_sql_err()?, - global_model_id: row.try_get("global_model_id").map_sql_err()?, - global_model_name: row.try_get("global_model_name").map_sql_err()?, - global_model_mappings: parse_string_list( - global_model_mappings, - "global_models.config.model_mappings", - )?, - global_model_supports_streaming, - model_provider_model_name: row.try_get("model_provider_model_name").map_sql_err()?, - model_provider_model_mappings: parse_provider_model_mappings(parse_json( - row.try_get("model_provider_model_mappings").ok().flatten(), - )?)?, - model_supports_streaming: row.try_get("model_supports_streaming").map_sql_err()?, - model_is_active: row.try_get("model_is_active").map_sql_err()?, - model_is_available: row.try_get("model_is_available").map_sql_err()?, - }, - provider_pool_enabled, - key_auth_config: row.try_get("key_auth_config").map_sql_err()?, - }) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "candidate selection JSON field is invalid: {err}" - )) - }) - }) - .transpose() -} - -fn json_object_field_present(value: &Option, field: &str) -> bool { - value - .as_ref() - .and_then(|value| value.get(field)) - .is_some_and(|value| !value.is_null()) -} - -fn json_bool(value: &serde_json::Value) -> Option { - value.as_bool().or_else(|| { - value - .as_str() - .and_then(|value| value.trim().parse::().ok()) - }) -} - -fn parse_string_list( - value: Option, - field_name: &str, -) -> Result>, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - parse_string_list_value(&value, field_name) -} - -fn parse_stored_key_policy_string_list( - raw: Option, - field_name: &str, -) -> Result>, DataLayerError> { - let Some(raw) = raw else { - return Ok(None); - }; - let value = serde_json::from_str::(&raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}")) - })?; - parse_key_policy_string_list_value(&value, field_name) -} - -fn parse_key_policy_string_list_value( - value: &serde_json::Value, - field_name: &str, -) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains JSON null; use SQL NULL for an unset policy" - ))), - serde_json::Value::Array(array) => { - parse_key_policy_string_list_array(array, field_name).map(Some) - } - serde_json::Value::String(raw) => parse_embedded_key_policy_string_list(raw, field_name), - _ => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} is not a JSON array" - ))), - } -} - -fn parse_embedded_key_policy_string_list( - raw: &str, - field_name: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains an empty string" - ))); - } - if raw.eq_ignore_ascii_case("null") { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains stringified JSON null; use SQL NULL for an unset policy" - ))); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_key_policy_string_list_value(&decoded, field_name); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_key_policy_string_list_array( - array: &[serde_json::Value], - field_name: &str, -) -> Result, DataLayerError> { - let mut items = Vec::with_capacity(array.len()); - for item in array { - let Some(item) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains a non-string item" - ))); - }; - let item = item.trim(); - if item.is_empty() { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains an empty item" - ))); - } - items.push(item.to_string()); - } - Ok(items) -} - -fn parse_string_list_value( - value: &serde_json::Value, - field_name: &str, -) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some), - serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name), - _ => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} is not a JSON array" - ))), - } -} - -fn parse_embedded_string_list( - raw: &str, - field_name: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_string_list_value(&decoded, field_name); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_string_list_array( - array: &[serde_json::Value], - field_name: &str, -) -> Result, DataLayerError> { - let mut items = Vec::with_capacity(array.len()); - for item in array { - let Some(item) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains a non-string item" - ))); - }; - let item = item.trim(); - if !item.is_empty() { - items.push(item.to_string()); - } - } - Ok(items) -} - -fn parse_provider_model_mappings( - value: Option, -) -> Result>, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(array) => parse_provider_model_mappings_array(&array), - serde_json::Value::Object(object) => parse_provider_model_mapping_object_lenient(&object) - .map(|mapping| mapping.map(|value| vec![value])), - serde_json::Value::String(raw) => parse_embedded_provider_model_mappings(&raw), - _ => Err(DataLayerError::UnexpectedValue( - "models.provider_model_mappings is not a JSON array".to_string(), - )), - } -} - -fn parse_embedded_provider_model_mappings( - raw: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_provider_model_mappings(Some(decoded)); - } - - Ok(Some(vec![StoredProviderModelMapping { - name: raw.to_string(), - priority: 1, - api_formats: None, - endpoint_ids: None, - operations: None, - }])) -} - -fn parse_provider_model_mappings_array( - array: &[serde_json::Value], -) -> Result>, DataLayerError> { - let mut mappings = Vec::with_capacity(array.len()); - for raw in array { - match raw { - serde_json::Value::Object(object) => { - if let Some(mapping) = parse_provider_model_mapping_object_lenient(object)? { - mappings.push(mapping); - } - } - serde_json::Value::String(raw) if !raw.trim().is_empty() => { - mappings.push(StoredProviderModelMapping { - name: raw.trim().to_string(), - priority: 1, - api_formats: None, - endpoint_ids: None, - operations: None, - }); - } - _ => {} - } - } - - if mappings.is_empty() { - Ok(None) - } else { - Ok(Some(mappings)) - } -} - -fn parse_provider_model_mapping_object_lenient( - object: &serde_json::Map, -) -> Result, DataLayerError> { - let Some(name) = object - .get("name") - .and_then(serde_json::Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - else { - return Ok(None); - }; - - let priority = object - .get("priority") - .and_then(serde_json::Value::as_i64) - .unwrap_or(1) - .max(1); - let api_formats = parse_string_list( - object.get("api_formats").cloned(), - "models.provider_model_mappings.api_formats", - )? - .map(|formats| { - formats - .into_iter() - .map(|value| normalize_api_format(&value)) - .collect() - }); - let endpoint_ids = parse_string_list( - object.get("endpoint_ids").cloned(), - "models.provider_model_mappings.endpoint_ids", - )?; - let operations = parse_string_list( - object.get("operations").cloned(), - "models.provider_model_mappings.operations", - )? - .and_then(normalize_request_operations); - - Ok(Some(StoredProviderModelMapping { - name: name.to_string(), - priority: i32::try_from(priority).map_err(|_| { - DataLayerError::UnexpectedValue(format!( - "invalid models.provider_model_mappings.priority: {priority}" - )) - })?, - api_formats, - endpoint_ids, - operations, - })) -} - -fn normalize_request_operations(values: Vec) -> Option> { - let operations = values - .into_iter() - .map(|value| value.trim().to_ascii_lowercase()) - .filter(|value| !value.is_empty()) - .collect::>(); - (!operations.is_empty()).then_some(operations) -} - -fn api_format_aliases(api_format: &str) -> Vec { - aether_ai_formats::api_format_storage_aliases(api_format) -} - -fn api_format_permission_aliases(api_format: &str) -> Vec { - aether_ai_formats::api_format_permission_storage_aliases(api_format) -} - -fn normalize_api_format(api_format: &str) -> String { - aether_ai_formats::normalize_api_format_alias(api_format) -} - -fn api_format_matches(left: &str, right: &str) -> bool { - aether_ai_formats::api_format_alias_matches(left, right) -} - -fn sql_match_aliases(api_formats: &[String]) -> Vec { - api_formats - .iter() - .map(|value| value.trim().to_ascii_lowercase()) - .collect() -} - -#[cfg(test)] -mod tests { - use super::{ - api_format_page_query, parse_stored_key_policy_string_list, - pool_key_group_by_key_ids_query, pool_key_group_query, - provider_model_mapping_api_format_covers, push_key_auth_channel_filter, - requested_model_page_query, vertex_key_auth_channel_matches, ExactPageAccumulator, - MysqlMinimalCandidateSelectionReadRepository, REQUESTED_MODEL_RAW_SCAN_LIMIT, - }; - use aether_data_contracts::repository::candidate_selection::{ - StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, - StoredPoolKeyCandidateRowsQuery, - }; - - #[test] - fn api_format_page_query_uses_portable_sql_pagination_and_stable_order() { - let query = api_format_page_query("openai:chat", 256, 512); - let sql = query.sql(); - - assert!(sql.contains("ROW_NUMBER() OVER (")); - assert!(sql.contains("JSON_SEARCH(LOWER(pak.api_formats), 'one', ?")); - assert!(sql.contains("WHERE provider_pool_enabled = 0 OR pool_rank = 1")); - assert!(sql.contains( - "ORDER BY\n global_model_name ASC,\n provider_priority ASC,\n key_internal_priority ASC," - )); - assert!(sql.contains("LIMIT ? OFFSET ?")); - } - - #[test] - fn malformed_key_policy_never_degrades_to_unrestricted() { - for raw in ["null", "\"null\"", "\"\"", "[\"openai:chat\",null]"] { - assert!(parse_stored_key_policy_string_list( - Some(raw.to_string()), - "provider_api_keys.api_formats", - ) - .is_err()); - } - assert_eq!( - parse_stored_key_policy_string_list( - Some("[\"openai:chat\"]".to_string()), - "provider_api_keys.api_formats", - ) - .expect("valid key policy should parse"), - Some(vec!["openai:chat".to_string()]) - ); - } - - #[test] - fn requested_model_page_query_filters_and_pages_before_fetch() { - let query = requested_model_page_query("openai:chat", "gpt-5", 256, 256); - let sql = query.sql(); - - assert!(sql.contains("LOWER(pe.api_format) IN (")); - assert!(sql.contains("JSON_SEARCH(LOWER(pak.api_formats), 'one', ?")); - assert!(sql.contains("LOWER(TRIM(p.provider_type)) = 'codex'")); - assert!(sql.contains("AND (gm.name = ? OR m.provider_model_name = ?")); - assert!(sql.contains("LOCATE(?, m.provider_model_mappings) > 0")); - assert!(sql.contains("ROW_NUMBER() OVER (")); - assert!(sql.contains("LIMIT ? OFFSET ?")); - } - - #[test] - fn codex_auth_sql_allows_live_for_oauth_keys() { - let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); - push_key_auth_channel_filter(&mut builder, "codex:live"); - let sql = builder.sql(); - let codex_clause = sql - .split_once("LOWER(TRIM(p.provider_type)) = 'codex'") - .and_then(|(_, suffix)| { - suffix.split_once("LOWER(TRIM(p.provider_type)) = 'chatgpt_web'") - }) - .map(|(clause, _)| clause) - .expect("Codex auth clause should exist"); - - assert!(codex_clause.contains("LOWER(TRIM(pak.auth_type)) = 'oauth'")); - assert!(codex_clause.contains("'codex:live'")); - } - - #[test] - fn mysql_mapping_scope_keeps_legacy_responses_compatibility_codex_only() { - assert!(provider_model_mapping_api_format_covers( - "codex", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "openai", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "custom", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "codex", - "openai:chat", - "codex:live" - )); - } - - #[test] - fn exact_page_accumulator_continues_after_coarse_false_positives() { - let mut accumulator = ExactPageAccumulator::new(1, 2); - accumulator.push_matching(vec![("coarse-1", false), ("coarse-2", false)], |row| row.1); - assert!(!accumulator.is_full()); - - accumulator.push_matching( - vec![ - ("exact-1", true), - ("coarse-3", false), - ("exact-2", true), - ("exact-3", true), - ], - |row| row.1, - ); - - assert!(accumulator.is_full()); - assert_eq!( - accumulator.into_page(), - vec![("exact-2", true), ("exact-3", true)] - ); - assert_eq!(REQUESTED_MODEL_RAW_SCAN_LIMIT, 2048); - } - - #[test] - fn pool_group_query_pushes_group_filters_order_and_page_into_sql() { - let query = pool_key_group_query(&StoredPoolKeyCandidateRowsQuery { - api_format: "openai:chat".to_string(), - provider_id: "provider-1".to_string(), - endpoint_id: "endpoint-1".to_string(), - model_id: "model-1".to_string(), - selected_provider_model_name: "gpt-5".to_string(), - order: StoredPoolKeyCandidateOrder::Lru, - offset: 64, - limit: 64, - }); - let sql = query.sql(); - - assert!(sql.contains("LOWER(pe.api_format) IN (")); - assert!(sql.contains("JSON_SEARCH(LOWER(pak.api_formats), 'one', ?")); - assert!(sql.contains("LOWER(TRIM(p.provider_type)) = 'codex'")); - assert!(sql.contains("AND p.id = ? AND pe.id = ? AND m.id = ?")); - assert!(sql.contains( - "ORDER BY pak.last_used_at IS NOT NULL ASC, pak.last_used_at ASC, pak.internal_priority ASC, pak.id ASC" - )); - assert!(sql.contains("LIMIT ? OFFSET ?")); - } - - #[test] - fn load_balance_pool_group_query_uses_seeded_sql_order_before_paging() { - let query = pool_key_group_query(&StoredPoolKeyCandidateRowsQuery { - api_format: "openai:chat".to_string(), - provider_id: "provider-1".to_string(), - endpoint_id: "endpoint-1".to_string(), - model_id: "model-1".to_string(), - selected_provider_model_name: "gpt-5".to_string(), - order: StoredPoolKeyCandidateOrder::LoadBalance { - seed: "request-1".to_string(), - }, - offset: 128, - limit: 64, - }); - let sql = query.sql(); - - assert!(sql.contains("AND p.id = ? AND pe.id = ? AND m.id = ?")); - assert!( - sql.contains("ORDER BY MD5(CONCAT(?, ':', pak.id)) ASC, pak.id ASC LIMIT ? OFFSET ?") - ); - assert!(!sql.contains("ROW_NUMBER() OVER (")); - } - - #[test] - fn pool_group_by_key_ids_query_filters_ids_and_preserves_requested_order() { - let query = pool_key_group_by_key_ids_query(&StoredPoolKeyCandidateRowsByKeyIdsQuery { - api_format: "openai:chat".to_string(), - provider_id: "provider-1".to_string(), - endpoint_id: "endpoint-1".to_string(), - model_id: "model-1".to_string(), - selected_provider_model_name: "gpt-5".to_string(), - key_ids: vec!["key-2".to_string(), "key-1".to_string()], - }); - let sql = query.sql(); - - assert!(sql.contains("LOWER(pe.api_format) IN (")); - assert!(sql.contains("JSON_SEARCH(LOWER(pak.api_formats), 'one', ?")); - assert!(sql.contains("LOWER(TRIM(p.provider_type)) = 'codex'")); - assert!(sql.contains("AND p.id = ? AND pe.id = ? AND m.id = ?")); - assert!(sql.contains("AND pak.id IN (?, ?)")); - assert!(sql.contains("ORDER BY FIELD(pak.id, ?, ?) ASC, pak.id ASC")); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlMinimalCandidateSelectionReadRepository::new(pool); - } - - #[test] - fn vertex_auth_matrix_rejects_retired_claude_format() { - for auth_type in ["api_key", "service_account", "vertex_ai"] { - assert!(!vertex_key_auth_channel_matches( - auth_type, - "claude:messages" - )); - assert!(vertex_key_auth_channel_matches( - auth_type, - "gemini:generate_content" - )); - assert!(vertex_key_auth_channel_matches( - auth_type, - "gemini:embedding" - )); - } - } -} diff --git a/crates/aether-data/adapters/mysql/src/candidates.rs b/crates/aether-data/adapters/mysql/src/candidates.rs deleted file mode 100644 index 37944e3c6..000000000 --- a/crates/aether-data/adapters/mysql/src/candidates.rs +++ /dev/null @@ -1,1227 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; - -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, MySqlConnection, QueryBuilder, Row}; - -use aether_data_contracts::repository::candidates::{ - request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket, - RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, - StoredRequestCandidate, UpsertRequestCandidateRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const CANDIDATE_COLUMNS: &str = r#" -SELECT - id, - request_id, - user_id, - api_key_id, - username, - api_key_name, - candidate_index, - retry_index, - provider_id, - endpoint_id, - key_id, - status, - skip_reason, - is_cached, - status_code, - error_type, - error_message, - latency_ms, - concurrent_requests, - extra_data, - required_capabilities, - created_at AS created_at_unix_ms, - started_at AS started_at_unix_ms, - finished_at AS finished_at_unix_ms -FROM request_candidates -"#; - -#[derive(Debug, Clone)] -pub struct MysqlRequestCandidateRepository { - pool: MysqlPool, -} - -impl MysqlRequestCandidateRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl RequestCandidateReadRepository for MysqlRequestCandidateRepository { - async fn list_by_request_id( - &self, - request_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? ORDER BY candidate_index ASC, retry_index ASC, created_at ASC" - )) - .bind(request_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_attempted_by_request_id( - &self, - request_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? \ - AND (status IN ('streaming', 'success', 'failed', 'cancelled') \ - OR (status = 'pending' AND started_at IS NOT NULL)) \ - ORDER BY candidate_index ASC, retry_index ASC, created_at ASC" - )) - .bind(request_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_recent( - &self, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} ORDER BY created_at DESC LIMIT ?" - )) - .bind(limit_i64(limit, "recent request candidate limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_by_provider_id( - &self, - provider_id: &str, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE provider_id = ? ORDER BY created_at DESC LIMIT ?" - )) - .bind(provider_id) - .bind(limit_i64(limit, "provider request candidate limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_finalized_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() || limit == 0 { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(CANDIDATE_COLUMNS); - push_endpoint_in_clause(&mut builder, endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(unix_secs_to_ms_i64(since_unix_secs)?) - .push(" AND status IN ('success', 'failed', 'skipped')") - .push(" ORDER BY created_at DESC LIMIT ") - .push_bind(limit_i64(limit, "finalized request candidate limit")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn count_finalized_statuses_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - "SELECT endpoint_id, status, COUNT(id) AS count FROM request_candidates", - ); - push_endpoint_in_clause(&mut builder, endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(unix_secs_to_ms_i64(since_unix_secs)?) - .push(" AND status IN ('success', 'failed', 'skipped')") - .push(" GROUP BY endpoint_id, status"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(PublicHealthStatusCount { - endpoint_id: row.try_get("endpoint_id").map_sql_err()?, - status: RequestCandidateStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - count: u64::try_from(row.try_get::("count").map_sql_err()?).map_err( - |_| { - DataLayerError::UnexpectedValue( - "public health status count out of range".to_string(), - ) - }, - )?, - }) - }) - .collect() - } - - async fn aggregate_finalized_timeline_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - until_unix_secs: u64, - segments: u32, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs { - return Ok(Vec::new()); - } - let since_ms = unix_secs_to_ms_i64(since_unix_secs)?; - let until_ms = unix_secs_to_ms_i64(until_unix_secs)?; - let mut builder = QueryBuilder::::new(CANDIDATE_COLUMNS); - push_endpoint_in_clause(&mut builder, endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(since_ms) - .push(" AND created_at <= ") - .push_bind(until_ms) - .push(" AND status IN ('success', 'failed', 'skipped')"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - aggregate_timeline( - rows.iter() - .map(map_candidate_row) - .collect::, _>>()?, - since_unix_secs, - until_unix_secs, - segments, - ) - } -} - -#[async_trait] -impl RequestCandidateWriteRepository for MysqlRequestCandidateRepository { - async fn upsert( - &self, - mut candidate: UpsertRequestCandidateRecord, - ) -> Result { - candidate.sanitize_for_persistence(); - candidate.validate()?; - 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, - mut candidates: Vec, - ) -> Result { - if candidates.is_empty() { - return Ok(0); - } - for candidate in &mut candidates { - candidate.sanitize_for_persistence(); - 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( - &self, - created_before_unix_secs: u64, - limit: usize, - ) -> Result { - if limit == 0 { - return Ok(0); - } - let rows_affected = sqlx::query( - r#" -DELETE FROM request_candidates -WHERE id IN ( - SELECT id - FROM ( - SELECT id - FROM request_candidates - WHERE created_at < ? - ORDER BY created_at ASC, id ASC - LIMIT ? - ) AS old_request_candidates -) -"#, - ) - .bind(unix_secs_to_ms_i64(created_before_unix_secs)?) - .bind(limit_i64(limit, "request candidate delete limit")?) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -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( - connection: &mut MySqlConnection, - candidate: &StoredRequestCandidate, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -INSERT INTO request_candidates ( - id, request_id, user_id, api_key_id, username, api_key_name, - candidate_index, retry_index, provider_id, endpoint_id, key_id, status, - skip_reason, is_cached, status_code, error_type, error_message, latency_ms, - concurrent_requests, extra_data, required_capabilities, created_at, started_at, finished_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - user_id = VALUES(user_id), - api_key_id = VALUES(api_key_id), - username = VALUES(username), - api_key_name = VALUES(api_key_name), - provider_id = VALUES(provider_id), - endpoint_id = VALUES(endpoint_id), - key_id = VALUES(key_id), - status = CASE - WHEN status IN ('success', 'failed', 'cancelled', 'skipped') - AND VALUES(status) IN ('available', 'unused', 'pending', 'streaming') - THEN status - WHEN status = 'pending' AND VALUES(status) IN ('available', 'unused') - THEN status - WHEN status = 'streaming' AND VALUES(status) IN ('available', 'unused', 'pending') - THEN status - ELSE VALUES(status) - END, - skip_reason = VALUES(skip_reason), - is_cached = VALUES(is_cached), - status_code = CASE - WHEN status IN ('success', 'failed', 'cancelled', 'skipped') - AND VALUES(status) IN ('available', 'unused', 'pending', 'streaming') - THEN status_code - WHEN status = 'pending' AND VALUES(status) IN ('available', 'unused') - THEN status_code - WHEN status = 'streaming' AND VALUES(status) IN ('available', 'unused', 'pending') - THEN status_code - ELSE COALESCE(VALUES(status_code), status_code) - END, - error_type = VALUES(error_type), - error_message = NULL, - latency_ms = CASE - WHEN status IN ('success', 'failed', 'cancelled', 'skipped') - AND VALUES(status) IN ('available', 'unused', 'pending', 'streaming') - THEN latency_ms - WHEN status = 'pending' AND VALUES(status) IN ('available', 'unused') - THEN latency_ms - WHEN status = 'streaming' AND VALUES(status) IN ('available', 'unused', 'pending') - THEN latency_ms - ELSE COALESCE(VALUES(latency_ms), latency_ms) - END, - concurrent_requests = VALUES(concurrent_requests), - extra_data = VALUES(extra_data), - required_capabilities = VALUES(required_capabilities), - created_at = VALUES(created_at), - started_at = VALUES(started_at), - finished_at = CASE - WHEN status IN ('success', 'failed', 'cancelled', 'skipped') - AND VALUES(status) IN ('available', 'unused', 'pending', 'streaming') - THEN finished_at - WHEN status = 'pending' AND VALUES(status) IN ('available', 'unused') - THEN finished_at - WHEN status = 'streaming' AND VALUES(status) IN ('available', 'unused', 'pending') - THEN finished_at - ELSE COALESCE(VALUES(finished_at), finished_at) - END -"#, - ) - .bind(&candidate.id) - .bind(&candidate.request_id) - .bind(&candidate.user_id) - .bind(&candidate.api_key_id) - .bind(&candidate.username) - .bind(&candidate.api_key_name) - .bind(to_i32(candidate.candidate_index)?) - .bind(to_i32(candidate.retry_index)?) - .bind(&candidate.provider_id) - .bind(&candidate.endpoint_id) - .bind(&candidate.key_id) - .bind(status_to_database(candidate.status)) - .bind(&candidate.skip_reason) - .bind(candidate.is_cached) - .bind(candidate.status_code.map(i32::from)) - .bind(&candidate.error_type) - .bind(&candidate.error_message) - .bind(candidate.latency_ms.map(to_i32_u64).transpose()?) - .bind(candidate.concurrent_requests.map(to_i32).transpose()?) - .bind(json_to_string(&candidate.extra_data)?) - .bind(json_to_string(&candidate.required_capabilities)?) - .bind(u64_to_i64( - candidate.created_at_unix_ms, - "request candidate created_at", - )?) - .bind(optional_u64_to_i64( - candidate.started_at_unix_ms, - "request candidate started_at", - )?) - .bind(optional_u64_to_i64( - candidate.finished_at_unix_ms, - "request candidate finished_at", - )?) - .execute(connection) - .await - .map_sql_err()?; - Ok(()) -} - -fn push_endpoint_in_clause<'args>( - builder: &mut QueryBuilder<'args, MySql>, - endpoint_ids: &'args [String], -) { - builder.push(" WHERE endpoint_id IN ("); - { - let mut separated = builder.separated(", "); - for endpoint_id in endpoint_ids { - separated.push_bind(endpoint_id); - } - } - builder.push(")"); -} - -fn merge_candidate( - mut candidate: UpsertRequestCandidateRecord, - existing: Option, -) -> Result { - candidate.sanitize_for_persistence(); - let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| { - request_candidate_lifecycle_would_regress(value.status, candidate.status) - }); - let merged_status = if preserve_existing_lifecycle { - existing - .as_ref() - .map(|value| value.status) - .unwrap_or(candidate.status) - } else { - candidate.status - }; - let created_at_unix_ms = existing - .as_ref() - .map(|value| value.created_at_unix_ms) - .filter(|value| *value > 1000) - .or_else(|| candidate.created_at_unix_ms.filter(|value| *value > 1000)) - .or(candidate.started_at_unix_ms) - .or(candidate.finished_at_unix_ms) - .unwrap_or_else(current_unix_ms); - let id = existing - .as_ref() - .map(|value| value.id.clone()) - .unwrap_or(candidate.id); - let extra_data = merge_json_objects( - existing.as_ref().and_then(|value| value.extra_data.clone()), - candidate.extra_data, - ); - StoredRequestCandidate::new( - id, - candidate.request_id, - existing - .as_ref() - .and_then(|value| value.user_id.clone()) - .or(candidate.user_id), - existing - .as_ref() - .and_then(|value| value.api_key_id.clone()) - .or(candidate.api_key_id), - existing - .as_ref() - .and_then(|value| value.username.clone()) - .or(candidate.username), - existing - .as_ref() - .and_then(|value| value.api_key_name.clone()) - .or(candidate.api_key_name), - to_i32(candidate.candidate_index)?, - to_i32(candidate.retry_index)?, - existing - .as_ref() - .and_then(|value| value.provider_id.clone()) - .or(candidate.provider_id), - existing - .as_ref() - .and_then(|value| value.endpoint_id.clone()) - .or(candidate.endpoint_id), - existing - .as_ref() - .and_then(|value| value.key_id.clone()) - .or(candidate.key_id), - merged_status, - candidate.skip_reason.or_else(|| { - existing - .as_ref() - .and_then(|value| value.skip_reason.clone()) - }), - candidate - .is_cached - .unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)), - if preserve_existing_lifecycle { - existing - .as_ref() - .and_then(|value| value.status_code.map(i32::from)) - } else { - candidate.status_code.map(i32::from).or_else(|| { - existing - .as_ref() - .and_then(|value| value.status_code.map(i32::from)) - }) - }, - if preserve_existing_lifecycle { - existing.as_ref().and_then(|value| value.error_type.clone()) - } else { - candidate - .error_type - .or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())) - }, - None, - if preserve_existing_lifecycle { - match existing.as_ref().and_then(|value| value.latency_ms) { - Some(value) => Some(to_i32_u64(value)?), - None => None, - } - } else { - candidate.latency_ms.map(to_i32_u64).transpose()?.or( - match existing.as_ref().and_then(|value| value.latency_ms) { - Some(value) => Some(to_i32_u64(value)?), - None => None, - }, - ) - }, - candidate.concurrent_requests.map(to_i32).transpose()?.or( - match existing - .as_ref() - .and_then(|value| value.concurrent_requests) - { - Some(value) => Some(to_i32(value)?), - None => None, - }, - ), - extra_data, - candidate.required_capabilities.or_else(|| { - existing - .as_ref() - .and_then(|value| value.required_capabilities.clone()) - }), - u64_to_i64(created_at_unix_ms, "request candidate created_at")?, - existing - .as_ref() - .and_then(|value| value.started_at_unix_ms) - .or(candidate.started_at_unix_ms) - .map(|value| u64_to_i64(value, "request candidate started_at")) - .transpose()?, - if preserve_existing_lifecycle { - existing - .as_ref() - .and_then(|value| value.finished_at_unix_ms) - } else { - candidate.finished_at_unix_ms.or_else(|| { - existing - .as_ref() - .and_then(|value| value.finished_at_unix_ms) - }) - } - .map(|value| u64_to_i64(value, "request candidate finished_at")) - .transpose()?, - ) -} - -fn aggregate_timeline( - candidates: Vec, - since_unix_secs: u64, - until_unix_secs: u64, - segments: u32, -) -> Result, DataLayerError> { - let endpoint_ids = candidates - .iter() - .filter_map(|candidate| candidate.endpoint_id.clone()) - .collect::>(); - let span_ms = until_unix_secs - .saturating_sub(since_unix_secs) - .saturating_mul(1000) - .max(1); - let since_ms = since_unix_secs.saturating_mul(1000); - let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new(); - for candidate in candidates { - let Some(endpoint_id) = candidate.endpoint_id.clone() else { - continue; - }; - let offset = candidate.created_at_unix_ms.saturating_sub(since_ms); - let segment_idx = ((offset.saturating_mul(u64::from(segments))) / span_ms) - .min(u64::from(segments.saturating_sub(1))) as u32; - let bucket = buckets.entry((endpoint_id.clone(), segment_idx)).or_insert( - PublicHealthTimelineBucket { - endpoint_id, - segment_idx, - total_count: 0, - success_count: 0, - failed_count: 0, - min_created_at_unix_ms: Some(candidate.created_at_unix_ms), - max_created_at_unix_ms: Some(candidate.created_at_unix_ms), - }, - ); - bucket.total_count += 1; - if candidate.status == RequestCandidateStatus::Success { - bucket.success_count += 1; - } - if candidate.status == RequestCandidateStatus::Failed { - bucket.failed_count += 1; - } - bucket.min_created_at_unix_ms = bucket - .min_created_at_unix_ms - .map(|value| value.min(candidate.created_at_unix_ms)); - bucket.max_created_at_unix_ms = bucket - .max_created_at_unix_ms - .map(|value| value.max(candidate.created_at_unix_ms)); - } - for endpoint_id in endpoint_ids { - for segment_idx in 0..segments { - buckets.entry((endpoint_id.clone(), segment_idx)).or_insert( - PublicHealthTimelineBucket { - endpoint_id: endpoint_id.clone(), - segment_idx, - total_count: 0, - success_count: 0, - failed_count: 0, - min_created_at_unix_ms: None, - max_created_at_unix_ms: None, - }, - ); - } - } - Ok(buckets.into_values().collect()) -} - -fn map_candidate_row(row: &MySqlRow) -> Result { - StoredRequestCandidate::new( - row.try_get("id").map_sql_err()?, - row.try_get("request_id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("candidate_index").map_sql_err()?, - row.try_get("retry_index").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("endpoint_id").map_sql_err()?, - row.try_get("key_id").map_sql_err()?, - RequestCandidateStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - row.try_get("skip_reason").map_sql_err()?, - row.try_get("is_cached").map_sql_err()?, - row.try_get("status_code").map_sql_err()?, - row.try_get("error_type").map_sql_err()?, - row.try_get("error_message").map_sql_err()?, - row.try_get("latency_ms").map_sql_err()?, - row.try_get("concurrent_requests").map_sql_err()?, - parse_json(row.try_get("extra_data").ok().flatten())?, - parse_json(row.try_get("required_capabilities").ok().flatten())?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("started_at_unix_ms").map_sql_err()?, - row.try_get("finished_at_unix_ms").map_sql_err()?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "request_candidates JSON field is invalid: {err}" - )) - }) - }) - .transpose() -} - -fn json_to_string(value: &Option) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "request_candidates JSON field is unserializable: {err}" - )) - }) - }) - .transpose() -} - -fn merge_json_objects( - existing: Option, - overlay: Option, -) -> Option { - match (existing, overlay) { - ( - Some(serde_json::Value::Object(mut existing_object)), - Some(serde_json::Value::Object(overlay_object)), - ) => { - existing_object.extend(overlay_object); - Some(serde_json::Value::Object(existing_object)) - } - (_existing, Some(overlay)) => Some(overlay), - (existing, None) => existing, - } -} - -fn status_to_database(status: RequestCandidateStatus) -> &'static str { - match status { - RequestCandidateStatus::Available => "available", - RequestCandidateStatus::Unused => "unused", - RequestCandidateStatus::Pending => "pending", - RequestCandidateStatus::Streaming => "streaming", - RequestCandidateStatus::Success => "success", - RequestCandidateStatus::Failed => "failed", - RequestCandidateStatus::Cancelled => "cancelled", - RequestCandidateStatus::Skipped => "skipped", - } -} - -fn current_unix_ms() -> u64 { - chrono::Utc::now().timestamp_millis().max(0) as u64 -} - -fn unix_secs_to_ms_i64(value: u64) -> Result { - let value = value.checked_mul(1000).ok_or_else(|| { - DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string()) - })?; - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string()) - }) -} - -fn limit_i64(value: usize, name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {name}: {value}"))) -} - -fn to_i32(value: u32) -> Result { - i32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}")) - }) -} - -fn to_i32_u64(value: u64) -> Result { - i32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}")) - }) -} - -fn u64_to_i64(value: u64, name: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u64_to_i64(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u64_to_i64(value, name)).transpose() -} - -#[cfg(test)] -mod tests { - use super::MysqlRequestCandidateRepository; - use crate::run_migrations; - use aether_data_contracts::repository::candidates::{ - RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, - StoredRequestCandidate, UpsertRequestCandidateRecord, - }; - use serde_json::json; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlRequestCandidateRepository::new(pool); - } - - #[tokio::test] - 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()) - else { - eprintln!( - "skipping mysql candidate lifecycle test because AETHER_TEST_MYSQL_URL is unset" - ); - return; - }; - let pool = sqlx::mysql::MySqlPoolOptions::new() - .max_connections(12) - .connect(&database_url) - .await - .expect("mysql test pool should connect"); - run_migrations(&pool) - .await - .expect("mysql migrations should run"); - let repository = MysqlRequestCandidateRepository::new(pool.clone()); - let request_id = format!("candidate-concurrency-{}", uuid::Uuid::new_v4()); - let mut initial = sample_upsert( - &request_id, - "initial", - RequestCandidateStatus::Pending, - Some(json!({"gateway_execution_runtime": true})), - 3_000_000, - ); - initial.is_cached = Some(false); - repository - .upsert(initial) - .await - .expect("initial candidate should insert"); - sqlx::query( - "UPDATE request_candidates SET skip_reason = ?, error_type = ?, error_message = ?, extra_data = ?, required_capabilities = ? WHERE request_id = ?", - ) - .bind("legacy skip reason with tenant-secret") - .bind("legacy_error_type_with_token") - .bind("Bearer legacy-secret") - .bind(r#"{"gateway_execution_runtime":true,"request_body":{"password":"secret"}}"#) - .bind(r#"{"streaming":true,"internal_capability":"secret"}"#) - .bind(&request_id) - .execute(&pool) - .await - .expect("legacy diagnostics should be injected for the conflict test"); - - 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 extra_data = match writer { - 0 => json!({"stream_completed": true}), - 1 => json!({"cache_1h": true}), - 2 => json!({"first_byte_time_ms": 2}), - 3 => json!({"pool_key_index": 3}), - 4 => json!({"priority_slot": 4}), - 5 => json!({"ranking_index": 5}), - 6 => json!({"phase": "provider_request"}), - 7 => json!({"provider_api_format": "openai:responses"}), - _ => unreachable!("writer index is bounded by WRITERS"), - }; - let mut candidate = sample_upsert( - &request_id, - format!("writer-{writer}").as_str(), - status, - Some(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"); - 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)); - assert_eq!( - candidate.extra_data, - Some(json!({ - "cache_1h": true, - "first_byte_time_ms": 2, - "gateway_execution_runtime": true, - "phase": "provider_request", - "pool_key_index": 3, - "priority_slot": 4, - "provider_api_format": "openai:responses", - "ranking_index": 5, - "stream_completed": true - })) - ); - let raw = sqlx::query( - "SELECT skip_reason, error_type, error_message, extra_data, required_capabilities FROM request_candidates WHERE request_id = ?", - ) - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("raw candidate diagnostics should load"); - assert!( - sqlx::Row::try_get::, _>(&raw, "error_message") - .expect("error_message should decode") - .is_none() - ); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "skip_reason") - .expect("skip_reason should decode") - .as_deref(), - Some("unclassified_skip") - ); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "error_type") - .expect("error_type should decode") - .as_deref(), - Some("unclassified_error") - ); - let raw_extra = sqlx::Row::try_get::, _>(&raw, "extra_data") - .expect("extra_data should decode") - .and_then(|value| serde_json::from_str::(&value).ok()); - assert_eq!(raw_extra, candidate.extra_data); - let raw_capabilities = - sqlx::Row::try_get::, _>(&raw, "required_capabilities") - .expect("required_capabilities should decode") - .and_then(|value| serde_json::from_str::(&value).ok()); - assert_eq!(raw_capabilities, Some(json!({"streaming": true}))); - - 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!({"gateway_execution_runtime": true})), - 4_000_000, - ); - pending.is_cached = Some(false); - let mut streaming = sample_upsert( - &batch_request_id, - "batch-second", - RequestCandidateStatus::Streaming, - Some(json!({"stream_completed": true})), - 4_000_100, - ); - streaming.is_cached = None; - let mut success = sample_upsert( - &batch_request_id, - "batch-third", - RequestCandidateStatus::Success, - Some(json!({"cache_1h": true})), - 4_000_200, - ); - success.is_cached = Some(true); - let mut late_pending = sample_upsert( - &batch_request_id, - "batch-fourth", - RequestCandidateStatus::Pending, - Some(json!({"first_byte_time_ms": 42})), - 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!({ - "cache_1h": true, - "first_byte_time_ms": 42, - "gateway_execution_runtime": true, - "stream_completed": 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"); - } - - #[test] - fn merge_candidate_preserves_first_identity_and_terminal_fact() { - let existing = StoredRequestCandidate::new( - "candidate-1".to_string(), - "request-1".to_string(), - Some("user-1".to_string()), - Some("key-1".to_string()), - None, - None, - 0, - 0, - Some("provider-1".to_string()), - Some("endpoint-1".to_string()), - Some("provider-key-1".to_string()), - RequestCandidateStatus::Success, - None, - false, - Some(200), - None, - None, - Some(123), - None, - Some(serde_json::json!({"stream_completed": true})), - None, - 1_000, - Some(1_001), - Some(1_123), - ) - .expect("existing candidate should build"); - - let merged = super::merge_candidate( - UpsertRequestCandidateRecord { - id: "candidate-late".to_string(), - request_id: "request-1".to_string(), - user_id: Some("attacker-user".to_string()), - api_key_id: Some("attacker-api-key".to_string()), - username: Some("mallory".to_string()), - api_key_name: Some("attacker-key".to_string()), - candidate_index: 0, - retry_index: 0, - provider_id: Some("attacker-provider".to_string()), - endpoint_id: Some("attacker-endpoint".to_string()), - key_id: Some("attacker-provider-key".to_string()), - status: RequestCandidateStatus::Failed, - skip_reason: None, - is_cached: Some(false), - status_code: Some(200), - error_type: None, - error_message: Some("Bearer secret-token".to_string()), - latency_ms: Some(9_999), - concurrent_requests: None, - extra_data: Some(serde_json::json!({"gateway_execution_runtime": true})), - required_capabilities: None, - created_at_unix_ms: Some(1_050), - started_at_unix_ms: Some(1_051), - finished_at_unix_ms: None, - }, - Some(existing), - ) - .expect("candidate should merge"); - - assert_eq!(merged.id, "candidate-1"); - assert_eq!(merged.status, RequestCandidateStatus::Success); - assert_eq!(merged.user_id.as_deref(), Some("user-1")); - assert_eq!(merged.api_key_id.as_deref(), Some("key-1")); - assert_eq!(merged.provider_id.as_deref(), Some("provider-1")); - assert_eq!(merged.endpoint_id.as_deref(), Some("endpoint-1")); - assert_eq!(merged.key_id.as_deref(), Some("provider-key-1")); - assert!(merged.error_message.is_none()); - assert_eq!(merged.latency_ms, Some(123)); - assert_eq!(merged.finished_at_unix_ms, Some(1_123)); - assert_eq!( - merged.extra_data, - Some(serde_json::json!({ - "gateway_execution_runtime": true, - "stream_completed": true - })) - ); - } - - fn sample_upsert( - request_id: &str, - id: &str, - status: RequestCandidateStatus, - 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, - 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/error.rs b/crates/aether-data/adapters/mysql/src/error.rs deleted file mode 100644 index 0d97e23e3..000000000 --- a/crates/aether-data/adapters/mysql/src/error.rs +++ /dev/null @@ -1,11 +0,0 @@ -use crate::DataLayerError; - -pub(crate) trait SqlResultExt { - fn map_sql_err(self) -> Result; -} - -impl SqlResultExt for Result { - fn map_sql_err(self) -> Result { - self.map_err(DataLayerError::sql) - } -} diff --git a/crates/aether-data/adapters/mysql/src/gemini_file_mappings.rs b/crates/aether-data/adapters/mysql/src/gemini_file_mappings.rs deleted file mode 100644 index 13baa4492..000000000 --- a/crates/aether-data/adapters/mysql/src/gemini_file_mappings.rs +++ /dev/null @@ -1,581 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::gemini_file_mappings::{ - GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository, - GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping, - StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_ci_contains_any, push_limit_offset, SqlDialect, WhereClause}; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const OWNER_GUARDED_UPSERT_SQL: &str = r#" -INSERT INTO gemini_file_mappings ( - id, file_name, key_id, user_id, display_name, mime_type, source_hash, - created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - display_name = IF(BINARY file_name = BINARY VALUES(file_name) AND BINARY key_id = BINARY VALUES(key_id) AND BINARY user_id <=> BINARY VALUES(user_id), VALUES(display_name), display_name), - mime_type = IF(BINARY file_name = BINARY VALUES(file_name) AND BINARY key_id = BINARY VALUES(key_id) AND BINARY user_id <=> BINARY VALUES(user_id), VALUES(mime_type), mime_type), - source_hash = IF(BINARY file_name = BINARY VALUES(file_name) AND BINARY key_id = BINARY VALUES(key_id) AND BINARY user_id <=> BINARY VALUES(user_id), VALUES(source_hash), source_hash), - expires_at = IF(BINARY file_name = BINARY VALUES(file_name) AND BINARY key_id = BINARY VALUES(key_id) AND BINARY user_id <=> BINARY VALUES(user_id), VALUES(expires_at), expires_at) -"#; - -#[derive(Debug, Clone)] -pub struct MysqlGeminiFileMappingRepository { - pool: MysqlPool, -} - -impl MysqlGeminiFileMappingRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn reload_by_file_name( - &self, - file_name: &str, - ) -> Result { - self.find_by_file_name(file_name).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("gemini file mapping missing after write".to_string()) - }) - } -} - -#[async_trait] -impl GeminiFileMappingReadRepository for MysqlGeminiFileMappingRepository { - async fn find_by_file_name( - &self, - file_name: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE file_name = ? -LIMIT 1 -"#, - ) - .bind(file_name) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn find_active_by_file_name_for_user( - &self, - file_name: &str, - user_id: &str, - now_unix_secs: u64, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE BINARY file_name = BINARY ? - AND BINARY user_id = BINARY ? - AND expires_at > ? -LIMIT 1 -"#, - ) - .bind(file_name) - .bind(user_id) - .bind(i64_from_u64( - now_unix_secs, - "gemini_file_mappings.owner_read_now", - )?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn find_active_by_file_name_for_owner( - &self, - file_name: &str, - key_id: &str, - user_id: &str, - now_unix_secs: u64, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE BINARY file_name = BINARY ? - AND BINARY key_id = BINARY ? - AND BINARY user_id = BINARY ? - AND expires_at > ? -LIMIT 1 -"#, - ) - .bind(file_name) - .bind(key_id) - .bind(user_id) - .bind(i64_from_u64( - now_unix_secs, - "gemini_file_mappings.owner_read_now", - )?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn list_mappings( - &self, - query: &GeminiFileMappingListQuery, - ) -> Result { - let total = build_list_count_query(query) - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let rows = build_list_rows_query(query) - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows.iter().map(map_row).collect::, _>>()?; - Ok(StoredGeminiFileMappingListPage { - items, - total: usize::try_from(total).unwrap_or_default(), - }) - } - - async fn summarize_mappings( - &self, - now_unix_secs: u64, - ) -> Result { - let totals = sqlx::query( - r#" -SELECT - COUNT(*) AS total_mappings, - SUM(CASE WHEN expires_at > ? THEN 1 ELSE 0 END) AS active_mappings -FROM gemini_file_mappings -"#, - ) - .bind(now_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total_mappings = - usize::try_from(totals.try_get::("total_mappings").map_sql_err()?) - .unwrap_or_default(); - let active_mappings = usize::try_from( - totals - .try_get::, _>("active_mappings") - .map_sql_err()? - .unwrap_or(0), - ) - .unwrap_or_default(); - let by_mime_type_rows = sqlx::query( - r#" -SELECT - COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type, - COUNT(*) AS count -FROM gemini_file_mappings -WHERE expires_at > ? -GROUP BY COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') -ORDER BY mime_type ASC -"#, - ) - .bind(now_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let by_mime_type = by_mime_type_rows - .iter() - .map(|row| { - Ok(GeminiFileMappingMimeTypeCount { - mime_type: row.try_get("mime_type").map_sql_err()?, - count: usize::try_from(row.try_get::("count").map_sql_err()?) - .unwrap_or_default(), - }) - }) - .collect::, DataLayerError>>()?; - Ok(GeminiFileMappingStats { - total_mappings, - active_mappings, - expired_mappings: total_mappings.saturating_sub(active_mappings), - by_mime_type, - }) - } -} - -#[async_trait] -impl GeminiFileMappingWriteRepository for MysqlGeminiFileMappingRepository { - async fn upsert( - &self, - record: UpsertGeminiFileMappingRecord, - ) -> Result { - record.validate()?; - sqlx::query( - r#" -INSERT INTO gemini_file_mappings ( - id, file_name, key_id, user_id, display_name, mime_type, source_hash, - created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - key_id = VALUES(key_id), - user_id = VALUES(user_id), - display_name = VALUES(display_name), - mime_type = VALUES(mime_type), - source_hash = VALUES(source_hash), - expires_at = VALUES(expires_at) -"#, - ) - .bind(&record.id) - .bind(&record.file_name) - .bind(&record.key_id) - .bind(&record.user_id) - .bind(&record.display_name) - .bind(&record.mime_type) - .bind(&record.source_hash) - .bind(current_unix_secs() as i64) - .bind(i64_from_u64( - record.expires_at_unix_secs, - "gemini_file_mappings.expires_at", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_file_name(&record.file_name).await - } - - async fn upsert_if_owner_matches( - &self, - record: UpsertGeminiFileMappingRecord, - ) -> Result, DataLayerError> { - record.validate()?; - let mut transaction = self.pool.begin().await.map_sql_err()?; - sqlx::query(OWNER_GUARDED_UPSERT_SQL) - .bind(&record.id) - .bind(&record.file_name) - .bind(&record.key_id) - .bind(&record.user_id) - .bind(&record.display_name) - .bind(&record.mime_type) - .bind(&record.source_hash) - .bind(current_unix_secs() as i64) - .bind(i64_from_u64( - record.expires_at_unix_secs, - "gemini_file_mappings.expires_at", - )?) - .execute(&mut *transaction) - .await - .map_sql_err()?; - - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE file_name = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&record.file_name) - .fetch_one(&mut *transaction) - .await - .map_sql_err()?; - let stored = map_row(&row)?; - let owner_matches = stored.file_name == record.file_name - && stored.key_id == record.key_id - && stored.user_id == record.user_id; - transaction.commit().await.map_sql_err()?; - - Ok(owner_matches.then_some(stored)) - } - - async fn delete_by_file_name(&self, file_name: &str) -> Result { - let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE file_name = ?") - .bind(file_name) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_file_name_for_user( - &self, - file_name: &str, - user_id: &str, - ) -> Result { - let rows_affected = - sqlx::query( - "DELETE FROM gemini_file_mappings WHERE BINARY file_name = BINARY ? AND BINARY user_id = BINARY ?", - ) - .bind(file_name) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_file_name_for_owner( - &self, - file_name: &str, - key_id: &str, - user_id: &str, - ) -> Result { - let rows_affected = sqlx::query( - "DELETE FROM gemini_file_mappings WHERE BINARY file_name = BINARY ? AND BINARY key_id = BINARY ? AND BINARY user_id = BINARY ?", - ) - .bind(file_name) - .bind(key_id) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_id( - &self, - mapping_id: &str, - ) -> Result, DataLayerError> { - let existing = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(mapping_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - let Some(existing) = existing else { - return Ok(None); - }; - sqlx::query("DELETE FROM gemini_file_mappings WHERE id = ?") - .bind(mapping_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(Some(map_row(&existing)?)) - } - - async fn delete_expired_before(&self, now_unix_secs: u64) -> Result { - let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE expires_at <= ?") - .bind(now_unix_secs as i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -fn build_list_count_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, MySql> { - let mut builder = - QueryBuilder::::new("SELECT COUNT(*) AS total FROM gemini_file_mappings"); - let mut where_clause = WhereClause::new(); - apply_list_filters(&mut builder, &mut where_clause, query); - builder -} - -fn build_list_rows_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, MySql> { - let mut builder = QueryBuilder::::new( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -"#, - ); - let mut where_clause = WhereClause::new(); - apply_list_filters(&mut builder, &mut where_clause, query); - builder.push(" ORDER BY created_at DESC, file_name ASC"); - push_limit_offset( - &mut builder, - i64::try_from(query.limit).unwrap_or(i64::MAX), - i64::try_from(query.offset).unwrap_or(i64::MAX), - ); - builder -} - -fn apply_list_filters( - builder: &mut QueryBuilder<'_, MySql>, - where_clause: &mut WhereClause, - query: &GeminiFileMappingListQuery, -) { - if let Some(user_id) = query.user_id.as_deref() { - where_clause.push_next(builder); - builder.push("BINARY user_id = BINARY "); - builder.push_bind(user_id.to_string()); - } - if !query.include_expired { - where_clause.push_next(builder); - builder.push("expires_at > "); - builder.push_bind(query.now_unix_secs as i64); - } - if let Some(search) = query - .search - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - push_ci_contains_any( - builder, - where_clause, - SqlDialect::MySql, - &["file_name", "COALESCE(display_name, '')"], - search, - ); - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn map_row(row: &MySqlRow) -> Result { - Ok(StoredGeminiFileMapping { - id: row.try_get("id").map_sql_err()?, - file_name: row.try_get("file_name").map_sql_err()?, - key_id: row.try_get("key_id").map_sql_err()?, - user_id: row.try_get("user_id").ok().flatten(), - display_name: row.try_get("display_name").ok().flatten(), - mime_type: row.try_get("mime_type").ok().flatten(), - source_hash: row.try_get("source_hash").ok().flatten(), - created_at_unix_ms: u64::try_from( - row.try_get::("created_at_unix_ms").map_sql_err()?, - ) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "gemini_file_mappings.created_at is invalid".to_string(), - ) - })?, - expires_at_unix_secs: u64::try_from( - row.try_get::("expires_at_unix_secs") - .map_sql_err()?, - ) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "gemini_file_mappings.expires_at is invalid".to_string(), - ) - })?, - }) -} - -#[cfg(test)] -mod tests { - use super::{ - build_list_count_query, build_list_rows_query, MysqlGeminiFileMappingRepository, - OWNER_GUARDED_UPSERT_SQL, - }; - use aether_data_contracts::repository::gemini_file_mappings::GeminiFileMappingListQuery; - use sqlx::Execute; - - #[test] - fn list_query_uses_shared_mysql_filter_and_pagination_rendering() { - let query = GeminiFileMappingListQuery { - user_id: Some("user-1".to_string()), - include_expired: false, - search: Some(" Report ".to_string()), - offset: 5, - limit: 10, - now_unix_secs: 123, - }; - - let mut count = build_list_count_query(&query); - let count_sql = count.build().sql().to_string(); - assert!(count_sql.contains( - " WHERE BINARY user_id = BINARY ? AND expires_at > ? AND (LOWER(file_name) LIKE ?" - )); - assert!(!count_sql.contains("WHERE 1=1")); - - let mut rows = build_list_rows_query(&query); - let rows_sql = rows.build().sql().to_string(); - assert!(rows_sql.contains("LOWER(COALESCE(display_name, '')) LIKE ?")); - assert!(rows_sql.contains(" ORDER BY created_at DESC, file_name ASC LIMIT ? OFFSET ?")); - } - - #[test] - fn owner_guarded_upsert_uses_exact_file_key_and_user_identity() { - let identity = "BINARY file_name = BINARY VALUES(file_name) AND BINARY key_id = BINARY VALUES(key_id) AND BINARY user_id <=> BINARY VALUES(user_id)"; - assert_eq!(OWNER_GUARDED_UPSERT_SQL.matches(identity).count(), 4); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlGeminiFileMappingRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/global_models.rs b/crates/aether-data/adapters/mysql/src/global_models.rs deleted file mode 100644 index 863418f06..000000000 --- a/crates/aether-data/adapters/mysql/src/global_models.rs +++ /dev/null @@ -1,1394 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::global_models::{ - metadata_supports_embedding, AdminGlobalModelListQuery, AdminProviderModelListQuery, - 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, -} - -impl MysqlGlobalModelReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - pub async fn create_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - sqlx::query( - r#" -INSERT INTO models ( - id, - provider_id, - global_model_id, - provider_model_name, - provider_model_mappings, - price_per_request, - tiered_pricing, - supports_vision, - supports_function_calling, - supports_streaming, - supports_extended_thinking, - supports_image_generation, - is_active, - is_available, - config, - created_at, - updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.provider_id) - .bind(&record.global_model_id) - .bind(&record.provider_model_name) - .bind(optional_json_to_string( - &record.provider_model_mappings, - "models.provider_model_mappings", - )?) - .bind(record.price_per_request) - .bind(optional_json_to_string( - &record.tiered_pricing, - "models.tiered_pricing", - )?) - .bind(record.supports_vision) - .bind(record.supports_function_calling) - .bind(record.supports_streaming) - .bind(record.supports_extended_thinking) - .bind(record.supports_image_generation) - .bind(record.is_active) - .bind(record.is_available) - .bind(optional_json_to_string(&record.config, "models.config")?) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.get_admin_provider_model(&record.provider_id, &record.id) - .await - } - - pub async fn update_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let updated = sqlx::query( - r#" -UPDATE models -SET - global_model_id = ?, - provider_model_name = ?, - provider_model_mappings = ?, - price_per_request = ?, - tiered_pricing = ?, - supports_vision = ?, - supports_function_calling = ?, - supports_streaming = ?, - supports_extended_thinking = ?, - supports_image_generation = ?, - is_active = ?, - is_available = ?, - config = ?, - updated_at = ? -WHERE id = ? - AND provider_id = ? -"#, - ) - .bind(&record.global_model_id) - .bind(&record.provider_model_name) - .bind(optional_json_to_string( - &record.provider_model_mappings, - "models.provider_model_mappings", - )?) - .bind(record.price_per_request) - .bind(optional_json_to_string( - &record.tiered_pricing, - "models.tiered_pricing", - )?) - .bind(record.supports_vision) - .bind(record.supports_function_calling) - .bind(record.supports_streaming) - .bind(record.supports_extended_thinking) - .bind(record.supports_image_generation) - .bind(record.is_active) - .bind(record.is_available) - .bind(optional_json_to_string(&record.config, "models.config")?) - .bind(now as i64) - .bind(&record.id) - .bind(&record.provider_id) - .execute(&self.pool) - .await - .map_sql_err()?; - - if updated.rows_affected() == 0 { - return Ok(None); - } - - self.get_admin_provider_model(&record.provider_id, &record.id) - .await - } - - pub async fn delete_admin_provider_model( - &self, - provider_id: &str, - model_id: &str, - ) -> Result { - let deleted = sqlx::query( - r#" -DELETE FROM models -WHERE provider_id = ? - AND id = ? -"#, - ) - .bind(provider_id) - .bind(model_id) - .execute(&self.pool) - .await - .map_sql_err()?; - - Ok(deleted.rows_affected() > 0) - } - - pub async fn create_admin_global_model( - &self, - record: &CreateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let usage_count = - optional_admin_global_model_usage_count_i64(record.usage_count)?.unwrap_or_default(); - sqlx::query( - r#" -INSERT INTO global_models ( - id, - name, - display_name, - is_active, - default_price_per_request, - default_tiered_pricing, - supported_capabilities, - usage_count, - config, - created_at, - updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.name) - .bind(&record.display_name) - .bind(record.is_active) - .bind(record.default_price_per_request) - .bind(optional_json_to_string( - &record.default_tiered_pricing, - "global_models.default_tiered_pricing", - )?) - .bind(optional_json_to_string( - &record.supported_capabilities, - "global_models.supported_capabilities", - )?) - .bind(usage_count) - .bind(optional_json_to_string( - &record.config, - "global_models.config", - )?) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.get_admin_global_model_by_id(&record.id).await - } - - pub async fn update_admin_global_model( - &self, - record: &UpdateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let usage_count = optional_admin_global_model_usage_count_i64(record.usage_count)?; - let updated = sqlx::query( - r#" -UPDATE global_models -SET - display_name = ?, - is_active = ?, - default_price_per_request = ?, - default_tiered_pricing = ?, - supported_capabilities = ?, - config = ?, - usage_count = COALESCE(?, usage_count), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&record.display_name) - .bind(record.is_active) - .bind(record.default_price_per_request) - .bind(optional_json_to_string( - &record.default_tiered_pricing, - "global_models.default_tiered_pricing", - )?) - .bind(optional_json_to_string( - &record.supported_capabilities, - "global_models.supported_capabilities", - )?) - .bind(optional_json_to_string( - &record.config, - "global_models.config", - )?) - .bind(usage_count) - .bind(now as i64) - .bind(&record.id) - .execute(&self.pool) - .await - .map_sql_err()?; - - if updated.rows_affected() == 0 { - return Ok(None); - } - - self.get_admin_global_model_by_id(&record.id).await - } - - pub async fn delete_admin_global_model( - &self, - global_model_id: &str, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - - sqlx::query( - r#" -DELETE FROM models -WHERE global_model_id = ? -"#, - ) - .bind(global_model_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let deleted = sqlx::query( - r#" -DELETE FROM global_models -WHERE id = ? -"#, - ) - .bind(global_model_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - - Ok(deleted.rows_affected() > 0) - } -} - -#[async_trait] -impl GlobalModelReadRepository for MysqlGlobalModelReadRepository { - async fn list_public_models( - &self, - query: &PublicGlobalModelQuery, - ) -> Result { - 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> { - 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> { - 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> { - 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 { - 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> { - 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> { - 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( - &self, - provider_id: &str, - model_id: &str, - ) -> Result, DataLayerError> { - 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> { - 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> { - 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> { - 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> { - 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> { - 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() - } -} - -#[async_trait] -impl GlobalModelWriteRepository for MysqlGlobalModelReadRepository { - async fn create_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - Self::create_admin_provider_model(self, record).await - } - - async fn update_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - Self::update_admin_provider_model(self, record).await - } - - async fn delete_admin_provider_model( - &self, - provider_id: &str, - model_id: &str, - ) -> Result { - Self::delete_admin_provider_model(self, provider_id, model_id).await - } - - async fn create_admin_global_model( - &self, - record: &CreateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - Self::create_admin_global_model(self, record).await - } - - async fn update_admin_global_model( - &self, - record: &UpdateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - Self::update_admin_global_model(self, record).await - } - - async fn delete_admin_global_model( - &self, - global_model_id: &str, - ) -> Result { - Self::delete_admin_global_model(self, global_model_id).await - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_u64(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn first_tier_price(value: Option<&serde_json::Value>, key: &str) -> Option { - value - .and_then(|value| value.get("tiers")) - .and_then(serde_json::Value::as_array) - .and_then(|tiers| tiers.first()) - .and_then(|tier| tier.get(key)) - .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()?, - row.try_get("name").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("default_price_per_request").map_sql_err()?, - optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("supported_capabilities").map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "global_models.config")?, - row.try_get::("usage_count").map_sql_err()?.max(0) as u64, - ) -} - -fn map_admin_global_model_row(row: &MySqlRow) -> 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()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("default_price_per_request").map_sql_err()?, - optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("supported_capabilities").map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "global_models.config")?, - provider_count, - active_provider_count, - usage_count, - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "global_models.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "global_models.updated_at", - )?, - ) -} - -fn map_admin_provider_model_row( - row: &MySqlRow, -) -> Result { - StoredAdminProviderModel::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - row.try_get("provider_model_name").map_sql_err()?, - optional_json_from_string( - row.try_get("provider_model_mappings").map_sql_err()?, - "models.provider_model_mappings", - )?, - row.try_get("price_per_request").map_sql_err()?, - optional_json_from_string( - row.try_get("tiered_pricing").map_sql_err()?, - "models.tiered_pricing", - )?, - row.try_get("supports_vision").map_sql_err()?, - row.try_get("supports_function_calling").map_sql_err()?, - row.try_get("supports_streaming").map_sql_err()?, - row.try_get("supports_extended_thinking").map_sql_err()?, - row.try_get("supports_image_generation").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_available").map_sql_err()?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "models.config")?, - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "models.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "models.updated_at", - )?, - row.try_get("global_model_name").map_sql_err()?, - row.try_get("global_model_display_name").map_sql_err()?, - row.try_get("global_model_default_price_per_request") - .map_sql_err()?, - optional_json_from_string( - row.try_get("global_model_default_tiered_pricing") - .map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("global_model_supported_capabilities") - .map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string( - row.try_get("global_model_config").map_sql_err()?, - "global_models.config", - )?, - ) -} - -fn map_public_catalog_model_row( - row: &MySqlRow, -) -> Result { - let global_model_config = optional_json_from_string( - row.try_get("global_model_config").map_sql_err()?, - "global_models.config", - )?; - let global_model_supported_capabilities = optional_json_from_string( - row.try_get("global_model_supported_capabilities") - .map_sql_err()?, - "global_models.supported_capabilities", - )?; - let model_config = - optional_json_from_string(row.try_get("model_config").map_sql_err()?, "models.config")?; - let tiered_pricing = optional_json_from_string( - row.try_get("tiered_pricing").map_sql_err()?, - "models.tiered_pricing", - )?; - let default_tiered_pricing = optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?; - let pricing = tiered_pricing.as_ref().or(default_tiered_pricing.as_ref()); - let global_model_is_active = row - .try_get::, _>("global_model_is_active") - .map_sql_err()? - .unwrap_or(true); - let model_is_active: bool = row.try_get("is_active").map_sql_err()?; - let provider_is_active: bool = row.try_get("provider_is_active").map_sql_err()?; - - StoredPublicCatalogModel::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_name").map_sql_err()?, - row.try_get("provider_model_name").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - global_model_config - .as_ref() - .and_then(|value| value.get("description")) - .and_then(serde_json::Value::as_str) - .map(ToString::to_string), - global_model_config - .as_ref() - .and_then(|value| value.get("icon_url")) - .and_then(serde_json::Value::as_str) - .map(ToString::to_string), - 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()?, - row.try_get("supports_function_calling").map_sql_err()?, - row.try_get("supports_streaming").map_sql_err()?, - metadata_supports_embedding( - global_model_supported_capabilities.as_ref(), - global_model_config.as_ref(), - model_config.as_ref(), - ), - model_is_active && provider_is_active && global_model_is_active, - ) -} - -fn map_provider_model_stats_row( - row: &MySqlRow, -) -> Result { - StoredProviderModelStats::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("total_models").map_sql_err()?, - row.try_get::, _>("active_models") - .map_sql_err()? - .unwrap_or(0), - ) -} - -fn map_active_global_model_row( - row: &MySqlRow, -) -> Result { - StoredProviderActiveGlobalModel::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - ) -} - -fn optional_admin_global_model_usage_count_i64( - value: Option, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput( - "global_models.usage_count exceeds i64 range".to_string(), - ) - }) - }) - .transpose() -} - -#[cfg(test)] -mod tests { - 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() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - 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 deleted file mode 100644 index 5f5a21ac8..000000000 --- a/crates/aether-data/adapters/mysql/src/lib.rs +++ /dev/null @@ -1,53 +0,0 @@ -//! MySQL pool adapter primitives. - -mod announcements; -mod audit; -mod auth; -mod auth_modules; -mod background_tasks; -mod billing; -mod candidate_selection; -mod candidates; -mod error; -mod gemini_file_mappings; -mod global_models; -mod management_tokens; -mod migrations; -mod oauth_providers; -mod pool; -mod pool_scores; -mod provider_catalog; -mod proxy_nodes; -mod quota; -mod routing_profiles; -mod settlement; -mod usage; -mod users; -mod video_tasks; -mod wallet; - -pub use aether_data_contracts::{DataLayerError, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; -pub use announcements::MysqlAnnouncementRepository; -pub use audit::MysqlAuditLogReadRepository; -pub use auth::MysqlAuthApiKeyReadRepository; -pub use auth_modules::{MysqlAuthModuleReadRepository, MysqlAuthModuleRepository}; -pub use background_tasks::MysqlBackgroundTaskRepository; -pub use billing::MysqlBillingReadRepository; -pub use candidate_selection::MysqlMinimalCandidateSelectionReadRepository; -pub use candidates::MysqlRequestCandidateRepository; -pub use gemini_file_mappings::MysqlGeminiFileMappingRepository; -pub use global_models::MysqlGlobalModelReadRepository; -pub use management_tokens::MysqlManagementTokenRepository; -pub use migrations::{pending_migrations, prepare_database_for_startup, run_migrations, MIGRATOR}; -pub use oauth_providers::MysqlOAuthProviderRepository; -pub use pool::{MysqlPool, MysqlPoolConfig, MysqlPoolFactory}; -pub use pool_scores::MysqlPoolMemberScoreRepository; -pub use provider_catalog::MysqlProviderCatalogReadRepository; -pub use proxy_nodes::MysqlProxyNodeReadRepository; -pub use quota::MysqlProviderQuotaRepository; -pub use routing_profiles::MysqlRoutingGroupRepository; -pub use settlement::MysqlSettlementRepository; -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/management_tokens.rs b/crates/aether-data/adapters/mysql/src/management_tokens.rs deleted file mode 100644 index aff5831e2..000000000 --- a/crates/aether-data/adapters/mysql/src/management_tokens.rs +++ /dev/null @@ -1,786 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::management_tokens::{ - ActivateManagementTokenIfMatches, CreateManagementTokenRecord, ManagementTokenListQuery, - ManagementTokenReadRepository, ManagementTokenWriteRepository, RegenerateManagementTokenSecret, - StoredManagementToken, StoredManagementTokenListPage, StoredManagementTokenUserSummary, - StoredManagementTokenWithUser, UpdateManagementTokenRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, push_limit, push_limit_offset, push_optional_eq, WhereClause}; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -#[derive(Debug, Clone)] -pub struct MysqlManagementTokenRepository { - pool: MysqlPool, -} - -impl MysqlManagementTokenRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn get_token( - &self, - token_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", token_id.to_string()); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_row).transpose() - } - - async fn get_token_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", token_id.to_string()); - push_optional_eq( - &mut builder, - &mut where_clause, - "user_id", - expected_user_id.map(ToOwned::to_owned), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_row).transpose() - } - - async fn update_management_token_scoped( - &self, - record: &UpdateManagementTokenRecord, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - record.validate()?; - let allowed_ips = json_to_string(record.allowed_ips.as_ref())?; - let permissions = json_to_string(record.permissions.as_ref())?; - let now = now_unix_secs(); - - sqlx::query(UPDATE_MANAGEMENT_TOKEN_SQL) - .bind(record.name.as_deref()) - .bind(record.clear_description) - .bind(record.description.as_deref()) - .bind(record.clear_allowed_ips) - .bind(allowed_ips) - .bind(permissions) - .bind(record.clear_expires_at) - .bind( - record - .expires_at_unix_secs - .and_then(|value| i64::try_from(value).ok()), - ) - .bind(record.is_active) - .bind(now as i64) - .bind(&record.token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_err(|err| map_mysql_write_error(err, record.name.as_deref()))?; - self.get_token_scoped(&record.token_id, expected_user_id) - .await - } - - async fn delete_management_token_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - ) -> Result { - let result = sqlx::query( - "DELETE FROM management_tokens WHERE id = ? AND (? IS NULL OR user_id = ?)", - ) - .bind(token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn set_management_token_active_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - is_active: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - "UPDATE management_tokens SET is_active = ?, updated_at = ? WHERE id = ? AND (? IS NULL OR user_id = ?)", - ) - .bind(is_active) - .bind(now_unix_secs() as i64) - .bind(token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token_scoped(token_id, expected_user_id).await - } - - async fn regenerate_management_token_secret_scoped( - &self, - mutation: &RegenerateManagementTokenSecret, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - mutation.validate()?; - let result = sqlx::query( - r#" -UPDATE management_tokens -SET token_hash = ?, token_prefix = ?, updated_at = ? -WHERE id = ? AND (? IS NULL OR user_id = ?) -"#, - ) - .bind(&mutation.token_hash) - .bind(mutation.token_prefix.as_deref()) - .bind(now_unix_secs() as i64) - .bind(&mutation.token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token_scoped(&mutation.token_id, expected_user_id) - .await - } -} - -const UPDATE_MANAGEMENT_TOKEN_SQL: &str = r#" -UPDATE management_tokens -SET name = COALESCE(?, name), - description = CASE WHEN ? THEN NULL ELSE COALESCE(?, description) END, - allowed_ips = CASE WHEN ? THEN NULL ELSE COALESCE(?, allowed_ips) END, - permissions = COALESCE(?, permissions), - expires_at = CASE WHEN ? THEN NULL ELSE COALESCE(?, expires_at) END, - is_active = COALESCE(?, is_active), - updated_at = ? -WHERE id = ? AND (? IS NULL OR user_id = ?) -"#; - -const TOKEN_COLUMNS: &str = r#" -SELECT - id, - user_id, - name, - description, - token_prefix, - allowed_ips, - permissions, - expires_at AS expires_at_unix_secs, - last_used_at AS last_used_at_unix_secs, - last_used_ip, - COALESCE(usage_count, 0) AS usage_count, - is_active, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM management_tokens -"#; - -const LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL: &str = r#" -SELECT id -FROM users -WHERE id = ? - AND is_active = TRUE - AND is_deleted = FALSE - AND LOWER(role) = 'admin' - AND security_version = ? -FOR UPDATE -"#; - -const LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL: &str = r#" -SELECT - id, - user_id, - token_hash, - name, - description, - token_prefix, - allowed_ips, - permissions, - expires_at AS expires_at_unix_secs, - last_used_at AS last_used_at_unix_secs, - last_used_ip, - COALESCE(usage_count, 0) AS usage_count, - is_active, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM management_tokens -WHERE id = ? -FOR UPDATE -"#; - -const TOKEN_WITH_USER_COLUMNS: &str = r#" -SELECT - mt.id, - mt.user_id, - mt.name, - mt.description, - mt.token_prefix, - mt.allowed_ips, - mt.permissions, - mt.expires_at AS expires_at_unix_secs, - mt.last_used_at AS last_used_at_unix_secs, - mt.last_used_ip, - COALESCE(mt.usage_count, 0) AS usage_count, - mt.is_active, - mt.created_at AS created_at_unix_ms, - mt.updated_at AS updated_at_unix_secs, - u.id AS user_row_id, - u.email AS user_email, - u.username AS user_username, - u.role AS user_role -FROM management_tokens mt -JOIN users u ON u.id = mt.user_id -"#; - -#[async_trait] -impl ManagementTokenReadRepository for MysqlManagementTokenRepository { - async fn list_management_tokens( - &self, - query: &ManagementTokenListQuery, - ) -> Result { - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(mt.id) AS total FROM management_tokens mt"); - let mut count_where = WhereClause::new(); - apply_management_token_filters(&mut count_builder, &mut count_where, query); - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let mut list_builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut list_where = WhereClause::new(); - apply_management_token_filters(&mut list_builder, &mut list_where, query); - list_builder.push(" ORDER BY mt.created_at DESC, mt.id DESC"); - push_limit_offset( - &mut list_builder, - i64::try_from(query.limit).unwrap_or(i64::MAX), - i64::try_from(query.offset).unwrap_or(i64::MAX), - ); - let rows = list_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - Ok(StoredManagementTokenListPage { - items: rows - .iter() - .map(map_token_with_user_row) - .collect::, _>>()?, - total: usize::try_from(total.max(0)).unwrap_or(usize::MAX), - }) - } - - async fn get_management_token_with_user( - &self, - token_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "mt.id", - token_id.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_with_user_row).transpose() - } - - async fn get_management_token_with_user_by_hash( - &self, - token_hash: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "mt.token_hash", - token_hash.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_with_user_row).transpose() - } -} - -fn apply_management_token_filters( - builder: &mut QueryBuilder<'_, MySql>, - where_clause: &mut WhereClause, - query: &ManagementTokenListQuery, -) { - push_optional_eq(builder, where_clause, "mt.user_id", query.user_id.clone()); - push_optional_eq(builder, where_clause, "mt.is_active", query.is_active); -} - -#[async_trait] -impl ManagementTokenWriteRepository for MysqlManagementTokenRepository { - async fn create_management_token( - &self, - record: &CreateManagementTokenRecord, - ) -> Result { - record.validate()?; - let now = now_unix_secs(); - sqlx::query( - r#" -INSERT INTO management_tokens ( - id, user_id, token_hash, token_prefix, name, description, allowed_ips, - permissions, expires_at, is_active, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.user_id) - .bind(&record.token_hash) - .bind(record.token_prefix.as_deref()) - .bind(&record.name) - .bind(record.description.as_deref()) - .bind(json_to_string(record.allowed_ips.as_ref())?) - .bind(json_to_string(record.permissions.as_ref())?) - .bind( - record - .expires_at_unix_secs - .and_then(|value| i64::try_from(value).ok()), - ) - .bind(record.is_active) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_err(|err| map_mysql_write_error(err, Some(record.name.as_str())))?; - - self.get_token(&record.id).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("created management token missing".to_string()) - }) - } - - async fn update_management_token( - &self, - record: &UpdateManagementTokenRecord, - ) -> Result, DataLayerError> { - self.update_management_token_scoped(record, None).await - } - - async fn update_management_token_for_user( - &self, - record: &UpdateManagementTokenRecord, - user_id: &str, - ) -> Result, DataLayerError> { - self.update_management_token_scoped(record, Some(user_id)) - .await - } - - async fn delete_management_token(&self, token_id: &str) -> Result { - self.delete_management_token_scoped(token_id, None).await - } - - async fn delete_management_token_for_user( - &self, - token_id: &str, - user_id: &str, - ) -> Result { - self.delete_management_token_scoped(token_id, Some(user_id)) - .await - } - - async fn set_management_token_active( - &self, - token_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_management_token_active_scoped(token_id, None, is_active) - .await - } - - async fn set_management_token_active_for_user( - &self, - token_id: &str, - user_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_management_token_active_scoped(token_id, Some(user_id), is_active) - .await - } - - async fn activate_management_token_if_matches( - &self, - mutation: &ActivateManagementTokenIfMatches, - ) -> Result { - mutation.validate()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let eligible_user = - sqlx::query_scalar::<_, String>(LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL) - .bind(&mutation.expected_token.user_id) - .bind(mutation.expected_user_security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if eligible_user.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let locked = sqlx::query(LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL) - .bind(&mutation.expected_token.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let snapshot_matches = match locked.as_ref() { - Some(row) => { - let token_hash: String = row.try_get("token_hash").map_sql_err()?; - let token = map_token_row(row)?; - mutation.matches_locked_token_snapshot(&token, &token_hash) - } - None => false, - }; - if !snapshot_matches { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE management_tokens -SET is_active = TRUE, updated_at = ? -WHERE id = ? - AND BINARY token_hash = BINARY ? - AND is_active = FALSE - AND (expires_at IS NULL OR expires_at > ?) -"#, - ) - .bind(now_unix_secs() as i64) - .bind(&mutation.expected_token.id) - .bind(&mutation.token_hash) - .bind(i64::try_from(mutation.now_unix_secs).unwrap_or(i64::MAX)) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_inactive_management_token_if_matches( - &self, - mutation: &ActivateManagementTokenIfMatches, - ) -> Result { - mutation.validate()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let locked = sqlx::query(LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL) - .bind(&mutation.expected_token.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let snapshot_matches = match locked.as_ref() { - Some(row) => { - let token_hash: String = row.try_get("token_hash").map_sql_err()?; - let token = map_token_row(row)?; - mutation.matches_locked_token_snapshot(&token, &token_hash) - } - None => false, - }; - if !snapshot_matches { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "DELETE FROM management_tokens WHERE id = ? AND BINARY token_hash = BINARY ? AND is_active = FALSE", - ) - .bind(&mutation.expected_token.id) - .bind(&mutation.token_hash) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn regenerate_management_token_secret( - &self, - mutation: &RegenerateManagementTokenSecret, - ) -> Result, DataLayerError> { - self.regenerate_management_token_secret_scoped(mutation, None) - .await - } - - async fn regenerate_management_token_secret_for_user( - &self, - mutation: &RegenerateManagementTokenSecret, - user_id: &str, - ) -> Result, DataLayerError> { - self.regenerate_management_token_secret_scoped(mutation, Some(user_id)) - .await - } - - async fn record_management_token_usage( - &self, - token_id: &str, - last_used_ip: Option<&str>, - ) -> Result, DataLayerError> { - let now = now_unix_secs(); - let result = sqlx::query( - r#" -UPDATE management_tokens -SET last_used_at = ?, - last_used_ip = ?, - usage_count = COALESCE(usage_count, 0) + 1, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(now as i64) - .bind(last_used_ip) - .bind(now as i64) - .bind(token_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token(token_id).await - } -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn non_negative_u64(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!( - "management_tokens.{field_name} must not be negative" - )) - }) -} - -fn optional_unix_secs(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| non_negative_u64(value, field_name)) - .transpose() -} - -fn json_to_string(value: Option<&serde_json::Value>) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "invalid management token JSON field: {err}" - )) - }) - }) - .transpose() -} - -fn json_from_string(value: Option) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "invalid management token JSON field: {err}" - )) - }) - }) - .transpose() -} - -fn map_mysql_write_error(err: sqlx::Error, requested_name: Option<&str>) -> DataLayerError { - let message = err.to_string(); - if message.contains("uq_management_tokens_user_name") - || message.contains("management_tokens.user_id, management_tokens.name") - { - return DataLayerError::InvalidInput( - requested_name - .map(|name| format!("已存在名为 '{}' 的 Token", name)) - .unwrap_or_else(|| "Management Token 名称已存在".to_string()), - ); - } - DataLayerError::sql(err) -} - -fn map_token_row(row: &MySqlRow) -> Result { - Ok(StoredManagementToken::new( - row.try_get("id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - )? - .with_display_fields( - row.try_get("description").map_sql_err()?, - row.try_get("token_prefix").map_sql_err()?, - json_from_string(row.try_get("allowed_ips").map_sql_err()?)?, - ) - .with_permissions(json_from_string(row.try_get("permissions").map_sql_err()?)?) - .with_runtime_fields( - optional_unix_secs( - row.try_get("expires_at_unix_secs").map_sql_err()?, - "expires_at", - )?, - optional_unix_secs( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - "last_used_at", - )?, - row.try_get("last_used_ip").map_sql_err()?, - non_negative_u64( - row.try_get::("usage_count").map_sql_err()?, - "usage_count", - )?, - row.try_get("is_active").map_sql_err()?, - ) - .with_timestamps( - optional_unix_secs( - row.try_get("created_at_unix_ms").map_sql_err()?, - "created_at", - )?, - optional_unix_secs( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "updated_at", - )?, - )) -} - -fn map_user_summary_row( - row: &MySqlRow, -) -> Result { - StoredManagementTokenUserSummary::new( - row.try_get("user_row_id").map_sql_err()?, - row.try_get("user_email").map_sql_err()?, - row.try_get("user_username").map_sql_err()?, - row.try_get("user_role").map_sql_err()?, - ) -} - -fn map_token_with_user_row( - row: &MySqlRow, -) -> Result { - Ok(StoredManagementTokenWithUser::new( - map_token_row(row)?, - map_user_summary_row(row)?, - )) -} - -#[cfg(test)] -mod tests { - use super::{ - non_negative_u64, optional_unix_secs, MysqlManagementTokenRepository, - LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL, LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL, - UPDATE_MANAGEMENT_TOKEN_SQL, - }; - use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlManagementTokenRepository::new(pool); - } - - #[test] - fn mysql_install_activation_locks_admin_identity_and_token_snapshot() { - for predicate in [ - "is_active = TRUE", - "is_deleted = FALSE", - "LOWER(role) = 'admin'", - "security_version = ?", - "FOR UPDATE", - ] { - assert!(LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL.contains(predicate)); - } - for column in [ - "token_hash", - "name", - "description", - "token_prefix", - "allowed_ips", - "permissions", - "expires_at_unix_secs", - "last_used_at_unix_secs", - "last_used_ip", - "usage_count", - "is_active", - "created_at_unix_ms", - "updated_at_unix_secs", - "FOR UPDATE", - ] { - assert!(LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL.contains(column)); - } - } - - #[test] - fn mysql_management_token_mapping_rejects_negative_integer_state() { - assert!(optional_unix_secs(Some(-1), "expires_at").is_err()); - assert_eq!( - optional_unix_secs(None, "expires_at").expect("SQL NULL should remain optional"), - None - ); - assert!(non_negative_u64(-1, "usage_count").is_err()); - } - - #[test] - fn mysql_management_token_updates_patch_only_explicit_fields() { - for clause in [ - "name = COALESCE(?, name)", - "allowed_ips = CASE WHEN ? THEN NULL ELSE COALESCE(?, allowed_ips) END", - "permissions = COALESCE(?, permissions)", - "expires_at = CASE WHEN ? THEN NULL ELSE COALESCE(?, expires_at) END", - "is_active = COALESCE(?, is_active)", - ] { - assert!(UPDATE_MANAGEMENT_TOKEN_SQL.contains(clause)); - } - } - - #[test] - fn mysql_management_token_pool_config_remains_driver_specific() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "mysql://user:pass@localhost:3306/aether".to_string(), - pool: SqlPoolConfig::default(), - }; - - assert_eq!(config.driver, DatabaseDriver::Mysql); - } -} diff --git a/crates/aether-data/adapters/mysql/src/migrations.rs b/crates/aether-data/adapters/mysql/src/migrations.rs deleted file mode 100644 index da8046f5d..000000000 --- a/crates/aether-data/adapters/mysql/src/migrations.rs +++ /dev/null @@ -1,274 +0,0 @@ -use sqlx::{ - migrate::{AppliedMigration, Migrate, MigrateError, Migrator}, - MySqlPool, -}; - -use aether_data_contracts::PendingMigrationInfo; - -pub static MIGRATOR: Migrator = sqlx::migrate!("./migrations"); - -pub async fn run_migrations(pool: &MySqlPool) -> Result<(), MigrateError> { - MIGRATOR.run(pool).await -} - -pub async fn pending_migrations( - pool: &MySqlPool, -) -> Result, MigrateError> { - 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) => { - 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)) -} - -pub async fn prepare_database_for_startup( - pool: &MySqlPool, -) -> Result, MigrateError> { - pending_migrations(pool).await -} - -fn is_missing_sqlx_migrations_table_error(err: &MigrateError) -> bool { - let message = err.to_string().to_ascii_lowercase(); - message.contains("_sqlx_migrations") - && (message.contains("no such table") - || message.contains("doesn't exist") - || message.contains("does not exist") - || message.contains("unknown table")) -} - -fn pending_migrations_from_applied( - applied_migrations: &[sqlx::migrate::AppliedMigration], -) -> Vec { - let applied_versions = applied_migrations - .iter() - .map(|migration| migration.version) - .collect::>(); - MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .filter(|migration| !applied_versions.contains(&migration.version)) - .map(|migration| PendingMigrationInfo { - version: migration.version, - description: migration.description.to_string(), - }) - .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 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() { - let versions = MIGRATOR - .iter() - .map(|migration| migration.version) - .collect::>(); - assert!(!versions.is_empty()); - assert!(versions.windows(2).all(|pair| pair[0] < pair[1])); - } - - #[test] - fn embeds_scoped_codex_live_permission_migration() { - let migration = MIGRATOR - .iter() - .find(|migration| migration.version == 20260821000000) - .expect("Codex Live permission migration should be embedded"); - let sql = migration.sql.as_ref(); - - for required_fragment in [ - "UPDATE users", - "UPDATE user_groups", - "UPDATE api_keys", - "UPDATE provider_api_keys", - "provider.provider_type", - "openai:responses", - "codex:live", - ] { - assert!( - sql.contains(required_fragment), - "Codex Live permission migration is missing {required_fragment}" - ); - } - } - - #[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/oauth_providers.rs b/crates/aether-data/adapters/mysql/src/oauth_providers.rs deleted file mode 100644 index 44ba5949d..000000000 --- a/crates/aether-data/adapters/mysql/src/oauth_providers.rs +++ /dev/null @@ -1,491 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; - -use aether_data_contracts::repository::oauth_providers::{ - OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig, - UpsertOAuthProviderConfigOutcome, UpsertOAuthProviderConfigRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -#[derive(Debug, Clone)] -pub struct MysqlOAuthProviderRepository { - pool: MysqlPool, -} - -impl MysqlOAuthProviderRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn get_provider( - &self, - provider_type: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(GET_OAUTH_PROVIDER_CONFIG_SQL) - .bind(provider_type) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_oauth_provider_row).transpose() - } -} - -const LIST_OAUTH_PROVIDER_CONFIGS_SQL: &str = r#" -SELECT - provider_type, - display_name, - client_id, - client_secret_encrypted, - authorization_url_override, - token_url_override, - userinfo_url_override, - scopes, - redirect_uri, - frontend_callback_url, - attribute_mapping, - extra_config, - icon_url, - is_enabled, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM oauth_providers -ORDER BY provider_type ASC -"#; - -const GET_OAUTH_PROVIDER_CONFIG_SQL: &str = r#" -SELECT - provider_type, - display_name, - client_id, - client_secret_encrypted, - authorization_url_override, - token_url_override, - userinfo_url_override, - scopes, - redirect_uri, - frontend_callback_url, - attribute_mapping, - extra_config, - icon_url, - is_enabled, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM oauth_providers -WHERE provider_type = ? -LIMIT 1 -"#; - -const COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL: &str = r#" -SELECT COUNT(DISTINCT users.id) AS locked_count -FROM users -JOIN user_oauth_links - ON users.id = user_oauth_links.user_id -WHERE users.is_active = 1 - AND users.is_deleted = 0 - AND user_oauth_links.provider_type = ? - AND ( - ( - users.auth_source = 'oauth' - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links other_links - JOIN oauth_providers other_provider - ON other_links.provider_type = other_provider.provider_type - WHERE other_links.user_id = users.id - AND other_links.provider_type <> ? - AND other_provider.is_enabled = 1 - ) - ) OR ( - ? = 1 - AND users.auth_source = 'local' - AND users.role <> 'admin' - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links other_links - JOIN oauth_providers other_provider - ON other_links.provider_type = other_provider.provider_type - WHERE other_links.user_id = users.id - AND other_links.provider_type <> ? - AND other_provider.is_enabled = 1 - ) - ) - ) -"#; - -const COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL: &str = r#" -UPDATE oauth_providers -SET client_secret_encrypted = ? -WHERE BINARY provider_type = BINARY ? - AND BINARY client_secret_encrypted = BINARY ? -"#; - -#[async_trait] -impl OAuthProviderReadRepository for MysqlOAuthProviderRepository { - async fn list_oauth_provider_configs( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query(LIST_OAUTH_PROVIDER_CONFIGS_SQL) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_oauth_provider_row).collect() - } - - async fn get_oauth_provider_config( - &self, - provider_type: &str, - ) -> Result, DataLayerError> { - self.get_provider(provider_type).await - } - - async fn count_locked_users_if_provider_disabled( - &self, - provider_type: &str, - ldap_exclusive: bool, - ) -> Result { - let row = sqlx::query(COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL) - .bind(provider_type) - .bind(provider_type) - .bind(ldap_exclusive) - .bind(provider_type) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let locked_count = row.try_get::("locked_count").map_sql_err()?; - usize::try_from(locked_count.max(0)).map_err(|_| { - DataLayerError::UnexpectedValue( - "oauth_providers.locked_user_count overflowed".to_string(), - ) - }) - } -} - -#[async_trait] -impl OAuthProviderWriteRepository for MysqlOAuthProviderRepository { - async fn upsert_oauth_provider_config_guarded( - &self, - record: &UpsertOAuthProviderConfigRecord, - ldap_exclusive: bool, - force_disable: bool, - _locked_users_snapshot: usize, - ) -> Result { - record.validate()?; - let now = now_unix_secs(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let existing_enabled: Option = if record.is_enabled || force_disable { - None - } else { - sqlx::query_scalar::<_, String>( - "SELECT provider_type FROM oauth_providers ORDER BY provider_type FOR UPDATE", - ) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - sqlx::query_scalar("SELECT is_enabled FROM oauth_providers WHERE provider_type = ?") - .bind(&record.provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - }; - if existing_enabled == Some(true) { - let row = sqlx::query(COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL) - .bind(&record.provider_type) - .bind(&record.provider_type) - .bind(ldap_exclusive) - .bind(&record.provider_type) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let affected_count = - usize::try_from(row.try_get::("locked_count").map_sql_err()?.max(0)) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "oauth_providers.locked_user_count overflowed".to_string(), - ) - })?; - if affected_count > 0 { - tx.rollback().await.map_sql_err()?; - return Ok( - UpsertOAuthProviderConfigOutcome::DisableRequiresConfirmation { - affected_count, - }, - ); - } - } - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, - display_name, - client_id, - client_secret_encrypted, - authorization_url_override, - token_url_override, - userinfo_url_override, - scopes, - redirect_uri, - frontend_callback_url, - attribute_mapping, - extra_config, - icon_url, - is_enabled, - created_at, - updated_at -) VALUES ( - ?, ?, ?, - CASE ? WHEN 'set' THEN ? WHEN 'clear' THEN NULL ELSE NULL END, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? -) -ON DUPLICATE KEY UPDATE - display_name = VALUES(display_name), - client_id = VALUES(client_id), - client_secret_encrypted = CASE ? - WHEN 'set' THEN ? - WHEN 'clear' THEN NULL - ELSE client_secret_encrypted - END, - authorization_url_override = VALUES(authorization_url_override), - token_url_override = VALUES(token_url_override), - userinfo_url_override = VALUES(userinfo_url_override), - scopes = VALUES(scopes), - redirect_uri = VALUES(redirect_uri), - frontend_callback_url = VALUES(frontend_callback_url), - attribute_mapping = VALUES(attribute_mapping), - extra_config = VALUES(extra_config), - icon_url = VALUES(icon_url), - is_enabled = VALUES(is_enabled), - updated_at = VALUES(updated_at) -"#, - ) - .bind(&record.provider_type) - .bind(&record.display_name) - .bind(&record.client_id) - .bind(record.client_secret_encrypted.mode_name()) - .bind(record.client_secret_encrypted.value()) - .bind(record.authorization_url_override.as_deref()) - .bind(record.token_url_override.as_deref()) - .bind(record.userinfo_url_override.as_deref()) - .bind(scopes_to_json_string(record.scopes.as_ref())?) - .bind(&record.redirect_uri) - .bind(&record.frontend_callback_url) - .bind(json_to_string(record.attribute_mapping.as_ref())?) - .bind(json_to_string(record.extra_config.as_ref())?) - .bind(record.icon_url.as_deref()) - .bind(record.is_enabled) - .bind(now as i64) - .bind(now as i64) - .bind(record.client_secret_encrypted.mode_name()) - .bind(record.client_secret_encrypted.value()) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = sqlx::query(GET_OAUTH_PROVIDER_CONFIG_SQL) - .bind(&record.provider_type) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let provider = map_oauth_provider_row(&row)?; - tx.commit().await.map_sql_err()?; - Ok(UpsertOAuthProviderConfigOutcome::Upserted(provider)) - } - - async fn compare_and_swap_oauth_provider_client_secret( - &self, - provider_type: &str, - expected: &str, - replacement: &str, - ) -> Result { - let result = sqlx::query(COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL) - .bind(replacement) - .bind(provider_type) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn delete_oauth_provider_config_if_unlinked( - &self, - provider_type: &str, - has_links_snapshot: bool, - ) -> Result { - if has_links_snapshot { - return Ok(false); - } - let mut tx = self.pool.begin().await.map_sql_err()?; - let provider_exists: Option = sqlx::query_scalar( - "SELECT provider_type FROM oauth_providers WHERE provider_type = ? FOR UPDATE", - ) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "DELETE FROM oauth_providers WHERE provider_type = ? AND NOT EXISTS (SELECT 1 FROM user_oauth_links WHERE user_oauth_links.provider_type = oauth_providers.provider_type)", - ) - .bind(provider_type) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(result.rows_affected() == 1) - } -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_unix_secs(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn json_to_string(value: Option<&serde_json::Value>) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("invalid OAuth provider JSON field: {err}")) - }) - }) - .transpose() -} - -fn json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn scopes_to_json_string(scopes: Option<&Vec>) -> Result, DataLayerError> { - let value = scopes.map(|items| { - serde_json::Value::Array( - items - .iter() - .cloned() - .map(serde_json::Value::String) - .collect(), - ) - }); - json_to_string(value.as_ref()) -} - -fn parse_scopes(value: Option) -> Result>, DataLayerError> { - let Some(value) = json_from_string(value, "oauth_providers.scopes")? else { - return Ok(None); - }; - parse_scopes_value(&value) -} - -fn parse_scopes_value(value: &serde_json::Value) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(items) => parse_scopes_array(items).map(Some), - serde_json::Value::String(raw) => parse_embedded_scopes(raw), - _ => Err(DataLayerError::UnexpectedValue( - "oauth_providers.scopes is not a JSON array".to_string(), - )), - } -} - -fn parse_embedded_scopes(raw: &str) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_scopes_value(&decoded); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_scopes_array(items: &[serde_json::Value]) -> Result, DataLayerError> { - let mut scopes = Vec::with_capacity(items.len()); - for item in items { - let Some(scope) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue( - "oauth_providers.scopes contains non-string value".to_string(), - )); - }; - let scope = scope.trim(); - if !scope.is_empty() { - scopes.push(scope.to_string()); - } - } - Ok(scopes) -} - -fn map_oauth_provider_row(row: &MySqlRow) -> Result { - Ok(StoredOAuthProviderConfig::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("client_id").map_sql_err()?, - row.try_get("redirect_uri").map_sql_err()?, - row.try_get("frontend_callback_url").map_sql_err()?, - )? - .with_config_fields( - row.try_get("client_secret_encrypted").map_sql_err()?, - row.try_get("authorization_url_override").map_sql_err()?, - row.try_get("token_url_override").map_sql_err()?, - row.try_get("userinfo_url_override").map_sql_err()?, - parse_scopes(row.try_get("scopes").map_sql_err()?)?, - json_from_string( - row.try_get("attribute_mapping").map_sql_err()?, - "oauth_providers.attribute_mapping", - )?, - json_from_string( - row.try_get("extra_config").map_sql_err()?, - "oauth_providers.extra_config", - )?, - row.try_get("icon_url").map_sql_err()?, - row.try_get("is_enabled").map_sql_err()?, - ) - .with_timestamps( - optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - optional_unix_secs(row.try_get("updated_at_unix_secs").map_sql_err()?), - )) -} - -#[cfg(test)] -mod tests { - use super::{MysqlOAuthProviderRepository, COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL}; - - #[test] - fn client_secret_cas_updates_only_the_secret_column() { - assert!(COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL - .contains("SET client_secret_encrypted = ?")); - assert!(COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL - .contains("BINARY client_secret_encrypted = BINARY ?")); - assert!(!COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL.contains("updated_at")); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlOAuthProviderRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/pool.rs b/crates/aether-data/adapters/mysql/src/pool.rs deleted file mode 100644 index 062140152..000000000 --- a/crates/aether-data/adapters/mysql/src/pool.rs +++ /dev/null @@ -1,184 +0,0 @@ -use std::str::FromStr; -use std::time::Duration; - -use crate::{DataLayerError, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; -use sqlx::mysql::{MySqlConnectOptions, MySqlPoolOptions, MySqlSslMode}; -use sqlx::MySqlPool as SqlxMysqlPool; - -pub type MysqlPool = SqlxMysqlPool; -pub type MysqlPoolConfig = SqlDatabaseConfig; - -#[derive(Debug, Clone)] -pub struct MysqlPoolFactory { - config: MysqlPoolConfig, -} - -impl MysqlPoolFactory { - pub fn new(config: MysqlPoolConfig) -> Result { - if config.driver != DatabaseDriver::Mysql { - return Err(DataLayerError::InvalidConfiguration(format!( - "mysql pool requires mysql driver, got {}", - config.driver - ))); - } - config.validate()?; - Ok(Self { config }) - } - - pub fn config(&self) -> &MysqlPoolConfig { - &self.config - } - - pub fn connect_options(&self) -> Result { - MySqlConnectOptions::from_str(self.config.url.trim()) - .map(|options| { - // Preserve explicit VERIFY_CA/VERIFY_IDENTITY from the URL. - // `require_ssl` is a minimum transport guarantee: upgrade - // weaker modes to Required, never downgrade verification. - let ssl_mode = if self.config.pool.require_ssl - && !matches!( - options.get_ssl_mode(), - MySqlSslMode::VerifyCa | MySqlSslMode::VerifyIdentity - ) { - MySqlSslMode::Required - } else { - options.get_ssl_mode() - }; - options - .ssl_mode(ssl_mode) - .statement_cache_capacity(self.config.pool.statement_cache_capacity) - }) - .map_err(|err| { - DataLayerError::InvalidConfiguration(format!("invalid mysql database url: {err}")) - }) - } - - pub fn connect_lazy(&self) -> Result { - let SqlPoolConfig { - min_connections, - max_connections, - acquire_timeout_ms, - idle_timeout_ms, - max_lifetime_ms, - .. - } = self.config.pool; - - Ok(MySqlPoolOptions::new() - .min_connections(min_connections) - .max_connections(max_connections) - .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()?)) - } -} - -#[cfg(test)] -mod tests { - use super::MysqlPoolFactory; - use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; - use sqlx::mysql::MySqlSslMode; - - fn ssl_mode(url: &str, require_ssl: bool) -> MySqlSslMode { - MysqlPoolFactory::new(SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: url.to_string(), - pool: SqlPoolConfig { - require_ssl, - ..SqlPoolConfig::default() - }, - }) - .expect("mysql config should build") - .connect_options() - .expect("mysql options should parse") - .get_ssl_mode() - } - - #[test] - fn preserves_explicit_mysql_verification_modes() { - assert!(matches!( - ssl_mode( - "mysql://user:pass@localhost/aether?ssl-mode=VERIFY_IDENTITY", - false - ), - MySqlSslMode::VerifyIdentity - )); - assert!(matches!( - ssl_mode( - "mysql://user:pass@localhost/aether?ssl-mode=VERIFY_CA", - true - ), - MySqlSslMode::VerifyCa - )); - } - - #[test] - fn require_ssl_only_upgrades_weak_mysql_modes() { - for mode in ["DISABLED", "PREFERRED", "REQUIRED"] { - let url = format!("mysql://user:pass@localhost/aether?ssl-mode={mode}"); - assert!(matches!(ssl_mode(&url, true), MySqlSslMode::Required)); - } - assert!(matches!( - ssl_mode("mysql://user:pass@localhost/aether", false), - MySqlSslMode::Preferred - )); - } - - #[tokio::test] - async fn factory_builds_lazy_pool_from_valid_config() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "mysql://user:pass@localhost:3306/aether".to_string(), - pool: SqlPoolConfig { - min_connections: 1, - max_connections: 4, - acquire_timeout_ms: 1_000, - idle_timeout_ms: 5_000, - max_lifetime_ms: 30_000, - statement_cache_capacity: 64, - require_ssl: false, - }, - }; - - 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/pool_scores.rs b/crates/aether-data/adapters/mysql/src/pool_scores.rs deleted file mode 100644 index 54aba8d20..000000000 --- a/crates/aether-data/adapters/mysql/src/pool_scores.rs +++ /dev/null @@ -1,693 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::pool_scores::*; - -use crate::error::SqlResultExt; -use crate::{DataLayerError, MysqlPool}; - -const SCORE_COLUMNS: &str = r#" -SELECT - id, - pool_kind, - pool_id, - member_kind, - member_id, - capability, - scope_kind, - scope_id, - score, - hard_state, - score_version, - score_reason, - last_ranked_at, - last_scheduled_at, - last_success_at, - last_failure_at, - failure_count, - last_probe_attempt_at, - last_probe_success_at, - last_probe_failure_at, - probe_failure_count, - probe_status, - updated_at -FROM pool_member_scores -"#; - -const UPSERT_PRESERVING_NULLABLE_TIMESTAMPS_SQL: &str = r#" -INSERT INTO pool_member_scores ( - id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id, - score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at, - last_success_at, last_failure_at, failure_count, last_probe_attempt_at, - last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at -) VALUES ( - ?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,? -) -ON DUPLICATE KEY UPDATE - pool_kind = VALUES(pool_kind), - pool_id = VALUES(pool_id), - member_kind = VALUES(member_kind), - member_id = VALUES(member_id), - capability = VALUES(capability), - scope_kind = VALUES(scope_kind), - scope_id = VALUES(scope_id), - score = VALUES(score), - hard_state = VALUES(hard_state), - score_version = VALUES(score_version), - score_reason = VALUES(score_reason), - last_ranked_at = VALUES(last_ranked_at), - last_scheduled_at = COALESCE(VALUES(last_scheduled_at), last_scheduled_at), - last_success_at = COALESCE(VALUES(last_success_at), last_success_at), - last_failure_at = COALESCE(VALUES(last_failure_at), last_failure_at), - failure_count = VALUES(failure_count), - last_probe_attempt_at = COALESCE(VALUES(last_probe_attempt_at), last_probe_attempt_at), - last_probe_success_at = COALESCE(VALUES(last_probe_success_at), last_probe_success_at), - last_probe_failure_at = COALESCE(VALUES(last_probe_failure_at), last_probe_failure_at), - probe_failure_count = VALUES(probe_failure_count), - probe_status = VALUES(probe_status), - updated_at = VALUES(updated_at) -"#; - -const UPSERT_OAUTH_RECOVERY_SQL: &str = r#" -INSERT INTO pool_member_scores ( - id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id, - score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at, - last_success_at, last_failure_at, failure_count, last_probe_attempt_at, - last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at -) VALUES ( - ?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,? -) -ON DUPLICATE KEY UPDATE - pool_kind = VALUES(pool_kind), - pool_id = VALUES(pool_id), - member_kind = VALUES(member_kind), - member_id = VALUES(member_id), - capability = VALUES(capability), - scope_kind = VALUES(scope_kind), - scope_id = VALUES(scope_id), - score = IF(updated_at <= VALUES(updated_at), VALUES(score), score), - hard_state = IF(updated_at <= VALUES(updated_at), VALUES(hard_state), hard_state), - score_version = IF(updated_at <= VALUES(updated_at), VALUES(score_version), score_version), - score_reason = IF(updated_at <= VALUES(updated_at), VALUES(score_reason), score_reason), - last_ranked_at = IF(updated_at <= VALUES(updated_at), VALUES(last_ranked_at), last_ranked_at), - failure_count = IF( - last_failure_at IS NULL OR last_failure_at <= VALUES(updated_at), - VALUES(failure_count), failure_count), - last_failure_at = IF( - last_failure_at IS NULL OR last_failure_at <= VALUES(updated_at), - VALUES(last_failure_at), last_failure_at), - probe_status = IF( - (last_probe_attempt_at IS NOT NULL AND last_probe_attempt_at > VALUES(updated_at)) - OR (last_probe_success_at IS NOT NULL AND last_probe_success_at > VALUES(updated_at)) - OR (last_probe_failure_at IS NOT NULL AND last_probe_failure_at > VALUES(updated_at)), - probe_status, VALUES(probe_status)), - probe_failure_count = IF( - last_probe_failure_at IS NULL OR last_probe_failure_at <= VALUES(updated_at), - VALUES(probe_failure_count), probe_failure_count), - last_probe_failure_at = IF( - last_probe_failure_at IS NULL OR last_probe_failure_at <= VALUES(updated_at), - VALUES(last_probe_failure_at), last_probe_failure_at), - updated_at = GREATEST(updated_at, VALUES(updated_at)) -"#; - -fn pool_member_score_upsert_sql(mode: PoolMemberScoreUpsertMode) -> &'static str { - match mode { - PoolMemberScoreUpsertMode::PreserveExistingNullableTimestamps => { - UPSERT_PRESERVING_NULLABLE_TIMESTAMPS_SQL - } - PoolMemberScoreUpsertMode::OAuthRecovery => UPSERT_OAUTH_RECOVERY_SQL, - } -} - -#[derive(Debug, Clone)] -pub struct MysqlPoolMemberScoreRepository { - pool: MysqlPool, -} - -impl MysqlPoolMemberScoreRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn find_scores_by_identity( - &self, - identity: &PoolMemberIdentity, - scope: Option<&PoolScoreScope>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - builder - .push(" WHERE pool_kind = ") - .push_bind(identity.pool_kind.clone()) - .push(" AND pool_id = ") - .push_bind(identity.pool_id.clone()) - .push(" AND member_kind = ") - .push_bind(identity.member_kind.clone()) - .push(" AND member_id = ") - .push_bind(identity.member_id.clone()); - if let Some(scope) = scope { - builder - .push(" AND capability = ") - .push_bind(scope.capability.clone()) - .push(" AND scope_kind = ") - .push_bind(scope.scope_kind.clone()); - if let Some(scope_id) = &scope.scope_id { - builder.push(" AND scope_id = ").push_bind(scope_id.clone()); - } else { - builder.push(" AND scope_id IS NULL"); - } - } - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } -} - -#[async_trait] -impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository { - async fn list_ranked_pool_members( - &self, - query: &ListRankedPoolMembersQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - builder - .push(" WHERE pool_kind = ") - .push_bind(query.pool_kind.clone()) - .push(" AND pool_id = ") - .push_bind(query.pool_id.clone()) - .push(" AND capability = ") - .push_bind(query.capability.clone()) - .push(" AND scope_kind = ") - .push_bind(query.scope_kind.clone()); - if let Some(scope_id) = &query.scope_id { - builder.push(" AND scope_id = ").push_bind(scope_id.clone()); - } else { - builder.push(" AND scope_id IS NULL"); - } - if !query.hard_states.is_empty() { - builder.push(" AND hard_state IN ("); - let mut separated = builder.separated(", "); - for state in &query.hard_states { - separated.push_bind(state.as_database()); - } - separated.push_unseparated(")"); - } - if let Some(statuses) = &query.probe_statuses { - if !statuses.is_empty() { - builder.push(" AND probe_status IN ("); - let mut separated = builder.separated(", "); - for status in statuses { - separated.push_bind(status.as_database()); - } - separated.push_unseparated(")"); - } - } - builder - .push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC") - .push(" LIMIT ") - .push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?) - .push(" OFFSET ") - .push_bind(i64_from_usize(query.offset, "pool score offset")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn list_pool_member_scores( - &self, - query: &ListPoolMemberScoresQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - builder - .push(" WHERE pool_kind = ") - .push_bind(query.pool_kind.clone()) - .push(" AND pool_id = ") - .push_bind(query.pool_id.clone()); - if let Some(capability) = &query.capability { - builder - .push(" AND capability = ") - .push_bind(capability.clone()); - } - if let Some(scope_kind) = &query.scope_kind { - builder - .push(" AND scope_kind = ") - .push_bind(scope_kind.clone()); - } - if let Some(scope_id) = &query.scope_id { - builder.push(" AND scope_id = ").push_bind(scope_id.clone()); - } - if !query.hard_states.is_empty() { - builder.push(" AND hard_state IN ("); - let mut separated = builder.separated(", "); - for state in &query.hard_states { - separated.push_bind(state.as_database()); - } - separated.push_unseparated(")"); - } - if let Some(statuses) = &query.probe_statuses { - if !statuses.is_empty() { - builder.push(" AND probe_status IN ("); - let mut separated = builder.separated(", "); - for status in statuses { - separated.push_bind(status.as_database()); - } - separated.push_unseparated(")"); - } - } - builder - .push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC") - .push(" LIMIT ") - .push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?) - .push(" OFFSET ") - .push_bind(i64_from_usize(query.offset, "pool score offset")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn list_pool_member_probe_candidates( - &self, - query: &ListPoolMemberProbeCandidatesQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - builder - .push(" WHERE pool_kind = ") - .push_bind(query.pool_kind.clone()) - .push(" AND pool_id = ") - .push_bind(query.pool_id.clone()); - if let Some(capability) = &query.capability { - builder - .push(" AND capability = ") - .push_bind(capability.clone()); - } - builder - .push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')") - .push(" AND (probe_status IN ('never','failed','stale')") - .push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ") - .push_bind(i64_from_u64( - query.stale_before_unix_secs, - "pool probe stale_before_unix_secs", - )?) - .push("))") - .push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ") - .push_bind(i64_from_u64( - query.stale_before_unix_secs, - "pool probe stale_before_unix_secs", - )?) - .push(")))") - .push( - r#" - ORDER BY - CASE - WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0 - WHEN hard_state = 'quota_exhausted' THEN 1 - WHEN hard_state = 'unknown' THEN 2 - WHEN probe_status = 'stale' THEN 3 - ELSE 4 - END ASC, - probe_failure_count DESC, - COALESCE(last_probe_success_at, 0) ASC, - COALESCE(last_scheduled_at, 0) DESC, - member_id ASC -"#, - ) - .push(" LIMIT ") - .push_bind(i64_from_usize( - query.limit.max(1), - "pool probe candidate limit", - )?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn get_pool_member_scores_by_ids( - &self, - query: &GetPoolMemberScoresByIdsQuery, - ) -> Result, DataLayerError> { - if query.ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - builder.push(" WHERE id IN ("); - let mut separated = builder.separated(", "); - for id in &query.ids { - separated.push_bind(id.clone()); - } - separated.push_unseparated(")"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } -} - -#[async_trait] -impl PoolMemberScoreWriteRepository for MysqlPoolMemberScoreRepository { - async fn upsert_pool_member_score_with_mode( - &self, - score: UpsertPoolMemberScore, - mode: PoolMemberScoreUpsertMode, - ) -> Result { - score.validate()?; - let stored = score.into_stored(); - let score_reason = serde_json::to_string(&stored.score_reason) - .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; - sqlx::query(pool_member_score_upsert_sql(mode)) - .bind(stored.id.as_str()) - .bind(stored.pool_kind.as_str()) - .bind(stored.pool_id.as_str()) - .bind(stored.member_kind.as_str()) - .bind(stored.member_id.as_str()) - .bind(stored.capability.as_str()) - .bind(stored.scope_kind.as_str()) - .bind(stored.scope_id.as_deref()) - .bind(stored.score) - .bind(stored.hard_state.as_database()) - .bind(i64_from_u64(stored.score_version, "pool score version")?) - .bind(score_reason) - .bind(i64_opt_from_u64( - stored.last_ranked_at, - "pool score last_ranked_at", - )?) - .bind(i64_opt_from_u64( - stored.last_scheduled_at, - "pool score last_scheduled_at", - )?) - .bind(i64_opt_from_u64( - stored.last_success_at, - "pool score last_success_at", - )?) - .bind(i64_opt_from_u64( - stored.last_failure_at, - "pool score last_failure_at", - )?) - .bind(i64_from_u64( - stored.failure_count, - "pool score failure_count", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_attempt_at, - "pool score last_probe_attempt_at", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_success_at, - "pool score last_probe_success_at", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_failure_at, - "pool score last_probe_failure_at", - )?) - .bind(i64_from_u64( - stored.probe_failure_count, - "pool score probe_failure_count", - )?) - .bind(stored.probe_status.as_database()) - .bind(i64_from_u64(stored.updated_at, "pool score updated_at")?) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(stored) - } - - async fn mark_pool_member_probe_in_progress( - &self, - attempt: PoolMemberProbeAttempt, - ) -> Result { - let rows = self - .find_scores_by_identity(&attempt.identity, attempt.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_probe_attempt_at = Some(attempt.attempted_at); - row.probe_status = PoolMemberProbeStatus::InProgress; - row.score_reason = - merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone()); - row.updated_at = attempt.attempted_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn record_pool_member_probe_result( - &self, - result: PoolMemberProbeResult, - ) -> Result { - let rows = self - .find_scores_by_identity(&result.identity, result.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_probe_attempt_at = Some(result.attempted_at); - row.probe_status = result.probe_status; - if result.succeeded { - row.last_probe_success_at = Some(result.attempted_at); - row.probe_failure_count = 0; - } else { - row.last_probe_failure_at = Some(result.attempted_at); - row.probe_failure_count = row.probe_failure_count.saturating_add(1); - } - if let Some(hard_state) = result.hard_state { - row.hard_state = hard_state; - } - row.score_reason = - merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone()); - row.updated_at = result.attempted_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn record_pool_member_schedule_feedback( - &self, - feedback: PoolMemberScheduleFeedback, - ) -> Result { - let rows = self - .find_scores_by_identity(&feedback.identity, feedback.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_scheduled_at = Some(feedback.scheduled_at); - match feedback.succeeded { - Some(true) => row.last_success_at = Some(feedback.scheduled_at), - Some(false) => { - row.last_failure_at = Some(feedback.scheduled_at); - row.failure_count = row.failure_count.saturating_add(1); - } - None => {} - } - if let Some(hard_state) = feedback.hard_state { - row.hard_state = hard_state; - } - row.score = score_with_delta(row.score, feedback.score_delta); - row.score_reason = - merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone()); - row.updated_at = feedback.scheduled_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn mark_pool_member_hard_state( - &self, - identity: &PoolMemberIdentity, - scope: Option<&PoolScoreScope>, - hard_state: PoolMemberHardState, - updated_at: u64, - ) -> Result { - let rows = self.find_scores_by_identity(identity, scope).await?; - let count = rows.len(); - for mut row in rows { - row.hard_state = hard_state; - row.updated_at = updated_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn delete_pool_member_scores_for_member( - &self, - identity: &PoolMemberIdentity, - ) -> Result { - let result = sqlx::query( - r#" -DELETE FROM pool_member_scores -WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ? -"#, - ) - .bind(identity.pool_kind.as_str()) - .bind(identity.pool_id.as_str()) - .bind(identity.member_kind.as_str()) - .bind(identity.member_id.as_str()) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } -} - -fn map_score_row(row: &MySqlRow) -> Result { - let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?; - Ok(StoredPoolMemberScore { - id: row.try_get("id").map_sql_err()?, - pool_kind: row.try_get("pool_kind").map_sql_err()?, - pool_id: row.try_get("pool_id").map_sql_err()?, - member_kind: row.try_get("member_kind").map_sql_err()?, - member_id: row.try_get("member_id").map_sql_err()?, - capability: row.try_get("capability").map_sql_err()?, - scope_kind: row.try_get("scope_kind").map_sql_err()?, - scope_id: row.try_get("scope_id").map_sql_err()?, - score: row.try_get("score").map_sql_err()?, - hard_state: PoolMemberHardState::from_database( - row.try_get::("hard_state") - .map_sql_err()? - .as_str(), - )?, - score_version: u64_from_i64( - row.try_get("score_version").map_sql_err()?, - "pool_member_scores.score_version", - )?, - score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null), - last_ranked_at: u64_opt_from_i64( - row.try_get("last_ranked_at").map_sql_err()?, - "pool_member_scores.last_ranked_at", - )?, - last_scheduled_at: u64_opt_from_i64( - row.try_get("last_scheduled_at").map_sql_err()?, - "pool_member_scores.last_scheduled_at", - )?, - last_success_at: u64_opt_from_i64( - row.try_get("last_success_at").map_sql_err()?, - "pool_member_scores.last_success_at", - )?, - last_failure_at: u64_opt_from_i64( - row.try_get("last_failure_at").map_sql_err()?, - "pool_member_scores.last_failure_at", - )?, - failure_count: u64_from_i64( - row.try_get("failure_count").map_sql_err()?, - "pool_member_scores.failure_count", - )?, - last_probe_attempt_at: u64_opt_from_i64( - row.try_get("last_probe_attempt_at").map_sql_err()?, - "pool_member_scores.last_probe_attempt_at", - )?, - last_probe_success_at: u64_opt_from_i64( - row.try_get("last_probe_success_at").map_sql_err()?, - "pool_member_scores.last_probe_success_at", - )?, - last_probe_failure_at: u64_opt_from_i64( - row.try_get("last_probe_failure_at").map_sql_err()?, - "pool_member_scores.last_probe_failure_at", - )?, - probe_failure_count: u64_from_i64( - row.try_get("probe_failure_count").map_sql_err()?, - "pool_member_scores.probe_failure_count", - )?, - probe_status: PoolMemberProbeStatus::from_database( - row.try_get::("probe_status") - .map_sql_err()? - .as_str(), - )?, - updated_at: u64_from_i64( - row.try_get("updated_at").map_sql_err()?, - "pool_member_scores.updated_at", - )?, - }) -} - -fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore { - UpsertPoolMemberScore { - id: score.id, - identity: PoolMemberIdentity { - pool_kind: score.pool_kind, - pool_id: score.pool_id, - member_kind: score.member_kind, - member_id: score.member_id, - }, - scope: PoolScoreScope { - capability: score.capability, - scope_kind: score.scope_kind, - scope_id: score.scope_id, - }, - score: score.score, - hard_state: score.hard_state, - score_version: score.score_version, - score_reason: score.score_reason, - last_ranked_at: score.last_ranked_at, - last_scheduled_at: score.last_scheduled_at, - last_success_at: score.last_success_at, - last_failure_at: score.last_failure_at, - failure_count: score.failure_count, - last_probe_attempt_at: score.last_probe_attempt_at, - last_probe_success_at: score.last_probe_success_at, - last_probe_failure_at: score.last_probe_failure_at, - probe_failure_count: score.probe_failure_count, - probe_status: score.probe_status, - updated_at: score.updated_at, - } -} - -fn i64_from_usize(value: usize, field: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn mysql_upsert_mode_selects_nullable_timestamp_semantics() { - let preserving = pool_member_score_upsert_sql( - PoolMemberScoreUpsertMode::PreserveExistingNullableTimestamps, - ); - let recovering = pool_member_score_upsert_sql(PoolMemberScoreUpsertMode::OAuthRecovery); - - for field in [ - "last_scheduled_at", - "last_success_at", - "last_failure_at", - "last_probe_attempt_at", - "last_probe_success_at", - "last_probe_failure_at", - ] { - assert!(preserving.contains(&format!("{field} = COALESCE(VALUES({field}), {field})"))); - } - for field in [ - "last_scheduled_at", - "last_success_at", - "last_probe_attempt_at", - "last_probe_success_at", - ] { - assert!(!recovering.contains(&format!("\n {field} ="))); - } - for field in [ - "score", - "hard_state", - "score_version", - "score_reason", - "last_ranked_at", - ] { - assert!(recovering.contains(&format!( - "{field} = IF(updated_at <= VALUES(updated_at), VALUES({field}), {field})" - ))); - } - assert!( - recovering.contains("last_failure_at IS NULL OR last_failure_at <= VALUES(updated_at)") - ); - assert!(recovering.contains("failure_count = IF(")); - assert!(recovering.contains( - "last_probe_failure_at IS NULL OR last_probe_failure_at <= VALUES(updated_at)" - )); - assert!(recovering.contains("probe_failure_count = IF(")); - assert!(recovering.contains("updated_at = GREATEST(updated_at, VALUES(updated_at))")); - assert!( - recovering.find("failure_count =").unwrap() - < recovering.find("last_failure_at =").unwrap() - ); - assert!( - recovering.find("probe_status =").unwrap() - < recovering.find("last_probe_failure_at =").unwrap() - ); - assert_eq!(preserving.matches('?').count(), 23); - assert_eq!(recovering.matches('?').count(), 23); - } -} diff --git a/crates/aether-data/adapters/mysql/src/provider_catalog.rs b/crates/aether-data/adapters/mysql/src/provider_catalog.rs deleted file mode 100644 index b6b42f1e1..000000000 --- a/crates/aether-data/adapters/mysql/src/provider_catalog.rs +++ /dev/null @@ -1,3780 +0,0 @@ -use std::collections::BTreeMap; - -use async_trait::async_trait; -use sqlx::{ - mysql::{MySqlArguments, MySqlRow}, - query::Query, - MySql, QueryBuilder, Row, -}; - -use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, - ProviderCatalogKeyCredentialsCasUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - ProviderCatalogProviderConfigCasUpdate, ProviderCatalogProxyCasUpdate, - 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; - -#[derive(Debug, Clone)] -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, - auth_type_by_format, allow_auth_channel_mismatch_formats, - COALESCE(api_key, encrypted_key) AS api_key, - auth_config, note, internal_priority, rate_multipliers, - global_priority_by_format, allowed_models, - expires_at AS expires_at_unix_secs, - cache_ttl_minutes, max_probe_interval_minutes, proxy, fingerprint, - rpm_limit, concurrent_limit, learned_rpm_limit, concurrent_429_count, - rpm_429_count, last_429_at AS last_429_at_unix_secs, last_429_type, - adjustment_history, utilization_samples, - last_probe_increase_at AS last_probe_increase_at_unix_secs, - last_rpm_peak, request_count, total_tokens, total_cost_usd, - success_count, error_count, total_response_time_ms, - last_used_at AS last_used_at_unix_secs, auto_fetch_models, - last_models_fetch_at AS last_models_fetch_at_unix_secs, - last_models_fetch_error, locked_models, model_include_patterns, - model_exclude_patterns, upstream_metadata, - oauth_invalid_at AS oauth_invalid_at_unix_secs, - oauth_invalid_reason, status_snapshot, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs, - health_by_format, circuit_breaker_by_format -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 } - } - - pub async fn list_providers_by_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - 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() - } - - 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() - } - - 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() - } - - pub async fn list_keys_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - 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()?; - 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( - &self, - provider: &StoredProviderCatalogProvider, - shift_existing_priorities_from: Option, - ) -> Result { - validate_provider(provider)?; - let now = current_unix_secs(); - let created_at = provider.created_at_unix_ms.unwrap_or(now) as i64; - let updated_at = provider.updated_at_unix_secs.unwrap_or(now) as i64; - let mut tx = self.pool.begin().await.map_sql_err()?; - - if let Some(target_priority) = shift_existing_priorities_from { - sqlx::query( - r#" -UPDATE providers -SET provider_priority = provider_priority + 1 -WHERE provider_priority >= ? -"#, - ) - .bind(target_priority) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query( - r#" -INSERT INTO providers ( - id, name, description, website, provider_type, billing_type, - monthly_quota_usd, monthly_used_usd, quota_reset_day, - quota_last_reset_at, quota_expires_at, 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, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&provider.id) - .bind(&provider.name) - .bind(&provider.description) - .bind(&provider.website) - .bind(&provider.provider_type) - .bind( - provider - .billing_type - .clone() - .unwrap_or_else(|| "pay_as_you_go".to_string()), - ) - .bind(provider.monthly_quota_usd) - .bind(provider.monthly_used_usd) - .bind(optional_i64_from_u64( - provider.quota_reset_day, - "providers.quota_reset_day", - )?) - .bind(optional_i64_from_u64( - provider.quota_last_reset_at_unix_secs, - "providers.quota_last_reset_at", - )?) - .bind(optional_i64_from_u64( - provider.quota_expires_at_unix_secs, - "providers.quota_expires_at", - )?) - .bind(provider.provider_priority) - .bind(provider.is_active) - .bind(provider.keep_priority_on_conversion) - .bind(provider.enable_format_conversion) - .bind(provider.concurrent_limit) - .bind(provider.max_retries) - .bind(optional_json_to_string(&provider.proxy, "providers.proxy")?) - .bind(provider.request_timeout_secs) - .bind(provider.stream_first_byte_timeout_secs) - .bind(optional_json_to_string( - &provider.config, - "providers.config", - )?) - .bind(created_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - self.reload_provider(&provider.id, "created").await - } - - pub async fn update_provider( - &self, - provider: &StoredProviderCatalogProvider, - ) -> Result { - validate_provider(provider)?; - let updated_at = provider - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET - name = ?, - description = ?, - website = ?, - provider_type = ?, - billing_type = ?, - monthly_quota_usd = ?, - monthly_used_usd = ?, - quota_reset_day = ?, - quota_last_reset_at = ?, - quota_expires_at = ?, - provider_priority = ?, - is_active = ?, - keep_priority_on_conversion = ?, - enable_format_conversion = ?, - concurrent_limit = ?, - max_retries = ?, - proxy = ?, - request_timeout = ?, - stream_first_byte_timeout = ?, - config = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&provider.name) - .bind(&provider.description) - .bind(&provider.website) - .bind(&provider.provider_type) - .bind( - provider - .billing_type - .clone() - .unwrap_or_else(|| "pay_as_you_go".to_string()), - ) - .bind(provider.monthly_quota_usd) - .bind(provider.monthly_used_usd) - .bind(optional_i64_from_u64( - provider.quota_reset_day, - "providers.quota_reset_day", - )?) - .bind(optional_i64_from_u64( - provider.quota_last_reset_at_unix_secs, - "providers.quota_last_reset_at", - )?) - .bind(optional_i64_from_u64( - provider.quota_expires_at_unix_secs, - "providers.quota_expires_at", - )?) - .bind(provider.provider_priority) - .bind(provider.is_active) - .bind(provider.keep_priority_on_conversion) - .bind(provider.enable_format_conversion) - .bind(provider.concurrent_limit) - .bind(provider.max_retries) - .bind(optional_json_to_string(&provider.proxy, "providers.proxy")?) - .bind(provider.request_timeout_secs) - .bind(provider.stream_first_byte_timeout_secs) - .bind(optional_json_to_string( - &provider.config, - "providers.config", - )?) - .bind(updated_at) - .bind(&provider.id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog provider {} not found", - provider.id - ))); - } - self.reload_provider(&provider.id, "updated").await - } - - pub async fn compare_and_swap_provider_config( - &self, - update: &ProviderCatalogProviderConfigCasUpdate, - ) -> Result { - validate_non_empty(&update.provider_id, "provider catalog provider_id")?; - let expected_config = - optional_json_to_string(&update.expected_config, "providers.expected_config")?; - let config = optional_json_to_string(&update.config, "providers.config")?; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET config = ?, updated_at = ? -WHERE id = ? - AND config <=> ? -"#, - ) - .bind(config) - .bind(current_unix_secs() as i64) - .bind(&update.provider_id) - .bind(expected_config) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - pub async fn compare_and_swap_provider_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog provider_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM providers WHERE id = ?", - "UPDATE providers SET proxy = ?, updated_at = ? WHERE id = ? AND BINARY proxy <=> BINARY ?", - update, - "providers.proxy", - ) - .await - } - - pub async fn delete_provider(&self, provider_id: &str) -> Result { - validate_non_empty(provider_id, "provider catalog provider_id")?; - let rows_affected = sqlx::query("DELETE FROM providers WHERE id = ?") - .bind(provider_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn cleanup_deleted_provider_refs( - &self, - provider_id: &str, - provider_deleted: bool, - endpoint_ids: &[String], - key_ids: &[String], - ) -> Result<(), DataLayerError> { - validate_non_empty(provider_id, "provider catalog provider_id")?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - if provider_deleted { - sqlx::query( - "UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id = ?", - ) - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE video_tasks SET provider_id = NULL WHERE provider_id = ?") - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM request_candidates WHERE provider_id = ?") - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - for endpoint_id in endpoint_ids { - sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = ?") - .bind(endpoint_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM request_candidates WHERE endpoint_id = ?") - .bind(endpoint_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - for key_id in key_ids { - sqlx::query("DELETE FROM gemini_file_mappings WHERE key_id = ?") - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE video_tasks SET key_id = NULL WHERE key_id = ?") - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - tx.commit().await.map_sql_err()?; - Ok(()) - } - - pub async fn create_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - validate_endpoint(endpoint)?; - let now = current_unix_secs(); - sqlx::query( - r#" -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, api_family, endpoint_kind, - is_active, health_score, header_rules, body_rules, max_retries, - custom_path, config, format_acceptance_config, proxy, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&endpoint.id) - .bind(&endpoint.provider_id) - .bind(&endpoint.api_format) - .bind(&endpoint.base_url) - .bind(&endpoint.api_format) - .bind(&endpoint.api_family) - .bind(&endpoint.endpoint_kind) - .bind(endpoint.is_active) - .bind(endpoint.health_score) - .bind(optional_json_to_string( - &endpoint.header_rules, - "provider_endpoints.header_rules", - )?) - .bind(optional_json_to_string( - &endpoint.body_rules, - "provider_endpoints.body_rules", - )?) - .bind(endpoint.max_retries) - .bind(&endpoint.custom_path) - .bind(optional_json_to_string( - &endpoint.config, - "provider_endpoints.config", - )?) - .bind(optional_json_to_string( - &endpoint.format_acceptance_config, - "provider_endpoints.format_acceptance_config", - )?) - .bind(optional_json_to_string( - &endpoint.proxy, - "provider_endpoints.proxy", - )?) - .bind(endpoint.created_at_unix_ms.unwrap_or(now) as i64) - .bind(endpoint.updated_at_unix_secs.unwrap_or(now) as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.reload_endpoint(&endpoint.id, "created").await - } - - pub async fn update_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - validate_endpoint(endpoint)?; - let updated_at = endpoint - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64; - let rows_affected = sqlx::query( - r#" -UPDATE provider_endpoints -SET - provider_id = ?, - name = ?, - base_url = ?, - api_format = ?, - api_family = ?, - endpoint_kind = ?, - is_active = ?, - health_score = ?, - header_rules = ?, - body_rules = ?, - max_retries = ?, - custom_path = ?, - config = ?, - format_acceptance_config = ?, - proxy = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&endpoint.provider_id) - .bind(&endpoint.api_format) - .bind(&endpoint.base_url) - .bind(&endpoint.api_format) - .bind(&endpoint.api_family) - .bind(&endpoint.endpoint_kind) - .bind(endpoint.is_active) - .bind(endpoint.health_score) - .bind(optional_json_to_string( - &endpoint.header_rules, - "provider_endpoints.header_rules", - )?) - .bind(optional_json_to_string( - &endpoint.body_rules, - "provider_endpoints.body_rules", - )?) - .bind(endpoint.max_retries) - .bind(&endpoint.custom_path) - .bind(optional_json_to_string( - &endpoint.config, - "provider_endpoints.config", - )?) - .bind(optional_json_to_string( - &endpoint.format_acceptance_config, - "provider_endpoints.format_acceptance_config", - )?) - .bind(optional_json_to_string( - &endpoint.proxy, - "provider_endpoints.proxy", - )?) - .bind(updated_at) - .bind(&endpoint.id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog endpoint {} not found", - endpoint.id - ))); - } - self.reload_endpoint(&endpoint.id, "updated").await - } - - pub async fn compare_and_swap_endpoint_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog endpoint_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM provider_endpoints WHERE id = ?", - "UPDATE provider_endpoints SET proxy = ?, updated_at = ? WHERE id = ? AND BINARY proxy <=> BINARY ?", - update, - "provider_endpoints.proxy", - ) - .await - } - - pub async fn delete_endpoint(&self, endpoint_id: &str) -> Result { - validate_non_empty(endpoint_id, "provider catalog endpoint_id")?; - let rows_affected = sqlx::query("DELETE FROM provider_endpoints WHERE id = ?") - .bind(endpoint_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn create_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - validate_key(key)?; - let now = current_unix_secs(); - sqlx::query(key_insert_sql()) - .bind(&key.id) - .bind(&key.provider_id) - .bind(&key.name) - .bind(&key.encrypted_api_key) - .bind(&key.auth_type) - .bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .bind(key.is_active) - .bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .bind(&key.encrypted_auth_config) - .bind(&key.note) - .bind(key.internal_priority) - .bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .bind(key.cache_ttl_minutes) - .bind(key.max_probe_interval_minutes) - .bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .bind(optional_i64_from_u32(key.rpm_limit)) - .bind(key.concurrent_limit) - .bind(optional_i64_from_u32(key.learned_rpm_limit)) - .bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0)) - .bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0)) - .bind(optional_i64_from_u64( - key.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .bind(&key.last_429_type) - .bind(optional_json_to_string( - &key.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .bind(optional_json_to_string( - &key.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .bind(optional_i64_from_u64( - key.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .bind(optional_i64_from_u32(key.last_rpm_peak)) - .bind(optional_i64_from_u32(key.request_count).unwrap_or(0)) - .bind(i64::try_from(key.total_tokens).map_err(|_| { - DataLayerError::InvalidInput(format!( - "provider catalog key.total_tokens exceeds i64: {}", - key.total_tokens - )) - })?) - .bind(key.total_cost_usd) - .bind(optional_i64_from_u32(key.success_count).unwrap_or(0)) - .bind(optional_i64_from_u32(key.error_count).unwrap_or(0)) - .bind( - optional_i64_from_u64( - key.total_response_time_ms, - "provider_api_keys.total_response_time_ms", - )? - .unwrap_or(0), - ) - .bind(optional_i64_from_u64( - key.last_used_at_unix_secs, - "provider_api_keys.last_used_at", - )?) - .bind(key.auto_fetch_models) - .bind(optional_i64_from_u64( - key.last_models_fetch_at_unix_secs, - "provider_api_keys.last_models_fetch_at", - )?) - .bind(&key.last_models_fetch_error) - .bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .bind(optional_json_to_string( - &key.upstream_metadata, - "provider_api_keys.upstream_metadata", - )?) - .bind(optional_i64_from_u64( - key.oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .bind(&key.oauth_invalid_reason) - .bind(optional_json_to_string( - &key.status_snapshot, - "provider_api_keys.status_snapshot", - )?) - .bind(optional_json_to_string( - &key.health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &key.circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(key.created_at_unix_ms.unwrap_or(now) as i64) - .bind(key.updated_at_unix_secs.unwrap_or(now) as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.reload_key(&key.id, "created").await - } - - pub async fn update_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - validate_key(key)?; - let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let rows_affected = key_update_query(key, updated_at)? - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog key {} not found", - key.id - ))); - } - self.reload_key(&key.id, "updated").await - } - - pub async fn compare_and_swap_key_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog key_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM provider_api_keys WHERE id = ?", - "UPDATE provider_api_keys SET proxy = ?, updated_at = ? WHERE id = ? AND BINARY proxy <=> BINARY ?", - update, - "provider_api_keys.proxy", - ) - .await - } - - pub async fn compare_and_swap_key_credentials( - &self, - update: &ProviderCatalogKeyCredentialsCasUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.expected_provider_id, - "provider catalog expected provider_id", - )?; - let rows_affected = sqlx::query( - "UPDATE provider_api_keys SET api_key = ?, encrypted_key = NULL, auth_config = ? WHERE id = ? AND BINARY provider_id = BINARY ? AND BINARY COALESCE(api_key, encrypted_key) <=> BINARY ? AND BINARY auth_config <=> BINARY ?", - ) - .bind(update.encrypted_api_key.as_deref()) - .bind(update.encrypted_auth_config.as_deref()) - .bind(&update.key_id) - .bind(&update.expected_provider_id) - .bind(update.expected_encrypted_api_key.as_deref()) - .bind(update.expected_encrypted_auth_config.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - pub async fn compare_and_update_key_admin_state( - &self, - update: &ProviderCatalogKeyAdminCasUpdate, - ) -> Result { - validate_admin_key_cas_update(update)?; - let key = &update.key; - let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let rotation_json = update - .codex_rotation - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys Codex rotation is not serializable: {err}" - )) - })?; - let mut builder = QueryBuilder::::new("UPDATE provider_api_keys SET provider_id = "); - push_admin_key_assignments(&mut builder, key, updated_at)?; - if update.reset_oauth_runtime { - builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0"); - } - if let Some(rotation_json) = rotation_json.as_deref() { - builder - .push(", upstream_metadata = JSON_SET(COALESCE(upstream_metadata, JSON_OBJECT()), '$.codex', CAST(") - .push_bind(rotation_json) - .push(" AS JSON))"); - } - if update.codex_rotation.is_some() || update.reset_oauth_runtime { - builder.push(", status_snapshot = "); - match (update.codex_rotation.is_some(), update.reset_oauth_runtime) { - (true, true) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.quota', CAST('null' AS JSON), '$.oauth', CAST('null' AS JSON))"), - (true, false) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.quota', CAST('null' AS JSON))"), - (false, true) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.oauth', CAST('null' AS JSON))"), - (false, false) => unreachable!(), - }; - } - builder - .push(" WHERE id = ") - .push_bind(&key.id) - .push(" AND BINARY api_key <=> BINARY ") - .push_bind(update.expected_credential.encrypted_api_key.as_deref()) - .push(" AND BINARY auth_config <=> BINARY ") - .push_bind(update.expected_encrypted_auth_config.as_deref()) - .push(" AND BINARY auth_type = BINARY ") - .push_bind(&update.expected_credential.auth_type) - .push(" AND BINARY provider_id = BINARY ") - .push_bind(&update.expected_credential.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE BINARY providers.id = BINARY provider_api_keys.provider_id AND BINARY providers.provider_type = BINARY ", - ) - .push_bind(&update.expected_credential.provider_type) - .push(")"); - if update.codex_rotation.is_some() { - builder - .push(" AND JSON_TYPE(COALESCE(upstream_metadata, JSON_OBJECT())) = 'OBJECT'") - .push(" AND NOT (BINARY api_key <=> BINARY ") - .push_bind(key.encrypted_api_key.as_deref()) - .push(" AND BINARY auth_config <=> BINARY ") - .push_bind(key.encrypted_auth_config.as_deref()) - .push(" AND BINARY auth_type = BINARY ") - .push_bind(&key.auth_type) - .push(" AND BINARY provider_id = BINARY ") - .push_bind(&key.provider_id) - .push(")"); - } - if update.codex_rotation.is_some() || update.reset_oauth_runtime { - builder.push(" AND JSON_TYPE(COALESCE(status_snapshot, JSON_OBJECT())) = 'OBJECT'"); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_keys( - &self, - keys: &[StoredProviderCatalogKey], - ) -> Result, DataLayerError> { - if keys.is_empty() { - return Ok(Vec::new()); - } - for key in keys { - validate_key(key)?; - } - - let updated_at = current_unix_secs() as i64; - let mut transaction = self.pool.begin().await.map_sql_err()?; - for key in keys { - let key_updated_at = key.updated_at_unix_secs.unwrap_or(updated_at as u64) as i64; - let rows_affected = key_update_query(key, key_updated_at)? - .execute(&mut *transaction) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog key {} not found", - key.id - ))); - } - } - transaction.commit().await.map_sql_err()?; - let key_ids = keys.iter().map(|key| key.id.clone()).collect::>(); - let mut reloaded = self - .list_keys_by_ids(&key_ids) - .await? - .into_iter() - .map(|key| (key.id.clone(), key)) - .collect::>(); - keys.iter() - .map(|key| { - reloaded.remove(&key.id).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "updated provider catalog key {} could not be reloaded", - key.id - )) - }) - }) - .collect() - } - - pub async fn delete_key(&self, key_id: &str) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query("DELETE FROM provider_api_keys WHERE id = ?") - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_delete_key_oauth_credential( - &self, - delete: &ProviderCatalogKeyOAuthCredentialCasDelete, - ) -> Result { - validate_non_empty(&delete.key_id, "provider catalog key_id")?; - let expected = &delete.expected_credential; - if expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - || delete - .expected_upstream_metadata_namespace - .as_ref() - .is_some_and(|expected| expected.namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth credential CAS delete contains empty fields".to_string(), - )); - } - let mut builder = QueryBuilder::::new("DELETE FROM provider_api_keys WHERE id = "); - builder - .push_bind(&delete.key_id) - .push(" AND BINARY auth_config <=> BINARY ") - .push_bind(delete.expected_encrypted_auth_config.as_deref()) - .push(" AND BINARY api_key <=> BINARY ") - .push_bind(expected.encrypted_api_key.as_deref()) - .push(" AND BINARY auth_type = BINARY ") - .push_bind(&expected.auth_type) - .push(" AND BINARY provider_id = BINARY ") - .push_bind(&expected.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE \ - BINARY providers.id = BINARY provider_api_keys.provider_id \ - AND BINARY providers.provider_type = BINARY ", - ) - .push_bind(&expected.provider_type) - .push(")"); - if let Some(expected) = delete.expected_upstream_metadata_namespace.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(&expected.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let expected_value = expected - .expected_value - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - builder - .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") - .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") <=> CAST(") - .push_bind(expected_value) - .push(" AS JSON)"); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_upstream_metadata( - &self, - key_id: &str, - upstream_metadata: Option<&serde_json::Value>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_json_ref_to_string( - upstream_metadata, - "provider_api_keys.upstream_metadata", - )?) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn upsert_key_upstream_metadata_namespace( - &self, - key_id: &str, - namespace: &str, - value: &serde_json::Value, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - validate_non_empty(namespace, "provider catalog upstream metadata namespace")?; - let value_json = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })?; - let namespace_path = format!( - "$.{}", - serde_json::to_string(namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = JSON_SET( - COALESCE(NULLIF(upstream_metadata, ''), '{}'), - ?, CAST(? AS JSON) - ), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(namespace_path) - .bind(value_json) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_model_fetch_state( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: Option, - last_models_fetch_error: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_json_ref_to_string( - allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - last_models_fetch_at_unix_secs, - "provider_api_keys.last_models_fetch_at", - )?) - .bind(last_models_fetch_error) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_model_fetch_success( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: u64, - upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let allowed_models = - optional_json_ref_to_string(allowed_models, "provider_api_keys.allowed_models")?; - let namespace_updates = upstream_metadata_updates - .iter() - .map(|update| { - validate_non_empty( - &update.namespace, - "provider catalog upstream metadata namespace", - )?; - let path = format!( - "$.{}", - serde_json::to_string(&update.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(&update.value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })?; - Ok((path, value)) - }) - .collect::, DataLayerError>>()?; - let updated_at = updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let mut tx = self.pool.begin().await.map_sql_err()?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = NULL, updated_at = ? -WHERE id = ? -"#, - ) - .bind(allowed_models) - .bind(optional_i64_from_u64( - Some(last_models_fetch_at_unix_secs), - "provider_api_keys.last_models_fetch_at", - )?) - .bind(updated_at) - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - for (path, value) in namespace_updates { - sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = JSON_SET( - COALESCE(NULLIF(upstream_metadata, ''), '{}'), - ?, CAST(? AS JSON) - ) -WHERE id = ? -"#, - ) - .bind(path) - .bind(value) - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - pub async fn clear_key_oauth_invalid_marker( - &self, - key_id: &str, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET oauth_invalid_at = NULL, oauth_invalid_reason = NULL, updated_at = ? -WHERE id = ? -"#, - ) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_oauth_runtime_state( - &self, - key_id: &str, - oauth_invalid_at_unix_secs: Option, - oauth_invalid_reason: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET oauth_invalid_at = ?, oauth_invalid_reason = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_i64_from_u64( - oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .bind(oauth_invalid_reason) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_oauth_runtime_state( - &self, - update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.encrypted_auth_config, - "provider catalog OAuth auth_config", - )?; - if update - .encrypted_api_key_update - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth api_key update must not be empty".to_string(), - )); - } - if update.expected_credential.as_ref().is_some_and(|expected| { - expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - }) { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth credential fence must not contain empty fields".to_string(), - )); - } - if update - .expected_upstream_metadata_namespace - .as_ref() - .is_some_and(|expected| expected.namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace must not be empty".to_string(), - )); - } - if update - .upstream_metadata_namespace_to_remove - .as_deref() - .is_some_and(|namespace| namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace to remove must not be empty" - .to_string(), - )); - } - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - )); - } - if update - .upstream_metadata_patch - .as_ref() - .is_some_and(|patch| !patch.is_object()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog upstream metadata patch must be an object".to_string(), - )); - } - if update - .upstream_metadata_namespace_to_remove - .as_ref() - .is_some_and(|namespace| { - update - .upstream_metadata_patch - .as_ref() - .and_then(serde_json::Value::as_object) - .is_some_and(|patch| patch.contains_key(namespace)) - }) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace cannot be patched and removed in the same update" - .to_string(), - )); - } - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET oauth_invalid_at = "); - builder - .push_bind(optional_i64_from_u64( - update.oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .push(", oauth_invalid_reason = ") - .push_bind(update.oauth_invalid_reason.as_deref()) - .push(", auth_config = ") - .push_bind(&update.encrypted_auth_config); - if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_deref() { - builder.push(", api_key = ").push_bind(encrypted_api_key); - } - if let Some(expires_at_unix_secs) = update.expires_at_unix_secs_update { - builder - .push(", expires_at = ") - .push_bind(optional_i64_from_u64( - expires_at_unix_secs, - "provider_api_keys.expires_at", - )?); - } - if update.upstream_metadata_patch.is_some() - || update.upstream_metadata_namespace_to_remove.is_some() - { - builder.push(", upstream_metadata = "); - if update.upstream_metadata_namespace_to_remove.is_some() { - builder.push("JSON_REMOVE("); - } - if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { - push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; - } else { - builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - } - if let Some(namespace) = update.upstream_metadata_namespace_to_remove.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - builder.push(", ").push_bind(namespace_path).push(")"); - } - } - builder.push(", status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - if update.reset_error_count { - builder.push(", error_count = 0"); - } - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND BINARY auth_config <=> BINARY ") - .push_bind(update.expected_encrypted_auth_config.as_deref()); - if let Some(expected) = update.expected_credential.as_ref() { - builder - .push(" AND BINARY api_key <=> BINARY ") - .push_bind(expected.encrypted_api_key.as_deref()) - .push(" AND BINARY auth_type = BINARY ") - .push_bind(&expected.auth_type) - .push(" AND BINARY provider_id = BINARY ") - .push_bind(&expected.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE \ - BINARY providers.id = BINARY provider_api_keys.provider_id \ - AND BINARY providers.provider_type = BINARY ", - ) - .push_bind(&expected.provider_type) - .push(")"); - } - if update.expected_upstream_metadata_namespace.is_some() - || update.upstream_metadata_patch.is_some() - || update.upstream_metadata_namespace_to_remove.is_some() - { - builder - .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'"); - } - if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(&expected.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let expected_value = expected - .expected_value - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - builder - .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") <=> CAST(") - .push_bind(expected_value) - .push(" AS JSON)"); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_health_state( - &self, - key_id: &str, - is_active: bool, - health_by_format: Option<&serde_json::Value>, - circuit_breaker_by_format: Option<&serde_json::Value>, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET is_active = ?, health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(is_active) - .bind(optional_json_ref_to_string( - health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_ref_to_string( - circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn reset_key_error_count(&self, key_id: &str) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET error_count = 0, updated_at = ? -WHERE id = ? -"#, - ) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_adaptive_state( - &self, - update: &ProviderCatalogKeyAdaptiveStateUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - let status_snapshot_patch = adaptive_status_snapshot_patch(&update.status_snapshot_patch)?; - let expected = update.expected.canonicalized(); - let next = update.next.canonicalized(); - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET learned_rpm_limit = "); - builder - .push_bind(optional_i64_from_u32(next.learned_rpm_limit)) - .push(", rpm_429_count = ") - .push_bind(optional_i64_from_u32(next.rpm_429_count)) - .push(", last_429_at = ") - .push_bind(optional_i64_from_u64( - next.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .push(", last_429_type = ") - .push_bind(&next.last_429_type) - .push(", adjustment_history = ") - .push_bind(optional_json_to_string( - &next.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .push(", utilization_samples = ") - .push_bind(optional_json_to_string( - &next.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .push(", last_probe_increase_at = ") - .push_bind(optional_i64_from_u64( - next.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .push(", last_rpm_peak = ") - .push_bind(optional_i64_from_u32(next.last_rpm_peak)) - .push(", concurrent_429_count = ") - .push_bind(optional_i64_from_u32(next.concurrent_429_count)) - .push(", status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND learned_rpm_limit <=> ") - .push_bind(optional_i64_from_u32(expected.learned_rpm_limit)) - .push(" AND rpm_429_count <=> ") - .push_bind(optional_i64_from_u32(expected.rpm_429_count)) - .push(" AND last_429_at <=> ") - .push_bind(optional_i64_from_u64( - expected.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .push(" AND last_429_type <=> ") - .push_bind(&expected.last_429_type) - .push(" AND JSON_EXTRACT(adjustment_history, '$') <=> CAST(") - .push_bind(optional_json_to_string( - &expected.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .push(" AS JSON)") - .push(" AND JSON_EXTRACT(utilization_samples, '$') <=> CAST(") - .push_bind(optional_json_to_string( - &expected.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .push(" AS JSON)") - .push(" AND last_probe_increase_at <=> ") - .push_bind(optional_i64_from_u64( - expected.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .push(" AND last_rpm_peak <=> ") - .push_bind(optional_i64_from_u32(expected.last_rpm_peak)) - .push(" AND concurrent_429_count <=> ") - .push_bind(optional_i64_from_u32(expected.concurrent_429_count)); - if let Some(expected_encrypted_auth_config) = - update.expected_encrypted_auth_config.as_deref() - { - builder - .push(" AND BINARY auth_config <=> BINARY ") - .push_bind(expected_encrypted_auth_config); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_runtime_metadata( - &self, - update: &ProviderCatalogKeyRuntimeMetadataUpdate, - ) -> Result { - validate_runtime_metadata_update(update)?; - let namespace_path = format!( - "$.{}", - serde_json::to_string(&update.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let metadata_value = - serde_json::to_string(&update.upstream_metadata_value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata value is not serializable: {err}" - )) - })?; - let expected_metadata_value = update - .expected_upstream_metadata_value - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - let mut builder = QueryBuilder::::new( - "UPDATE provider_api_keys SET upstream_metadata = JSON_SET(\ - COALESCE(NULLIF(upstream_metadata, ''), '{}'), ", - ); - builder - .push_bind(namespace_path.clone()) - .push(", CAST(") - .push_bind(metadata_value) - .push(" AS JSON)), status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") - .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") <=> CAST(") - .push_bind(expected_metadata_value) - .push(" AS JSON)"); - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_status_snapshot( - &self, - update: &ProviderCatalogKeyStatusSnapshotUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - )); - } - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id); - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_health_state( - &self, - update: &ProviderCatalogKeyHealthStateUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ? -WHERE id = ? - AND JSON_EXTRACT(health_by_format, '$') <=> CAST(? AS JSON) - AND JSON_EXTRACT(circuit_breaker_by_format, '$') <=> CAST(? AS JSON) - AND (? IS NULL OR BINARY auth_config <=> BINARY ?) -"#, - ) - .bind(optional_json_to_string( - &update.health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &update.circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(current_unix_secs() as i64) - .bind(&update.key_id) - .bind(optional_json_to_string( - &update.expected_health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &update.expected_circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(update.expected_encrypted_auth_config.as_deref()) - .bind(update.expected_encrypted_auth_config.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn reload_provider( - &self, - provider_id: &str, - operation: &str, - ) -> Result { - self.list_providers_by_ids(&[provider_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog provider {provider_id} could not be reloaded" - )) - }) - } - - async fn reload_endpoint( - &self, - endpoint_id: &str, - operation: &str, - ) -> Result { - self.list_endpoints_by_ids(&[endpoint_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog endpoint {endpoint_id} could not be reloaded" - )) - }) - } - - async fn reload_key( - &self, - key_id: &str, - operation: &str, - ) -> Result { - self.list_keys_by_ids(&[key_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog key {key_id} could not be reloaded" - )) - }) - } -} - -#[async_trait] -impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository { - async fn list_providers( - &self, - active_only: bool, - ) -> Result, DataLayerError> { - Self::list_providers(self, active_only).await - } - - async fn list_providers_by_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_providers_by_ids(self, provider_ids).await - } - - async fn list_endpoints_by_ids( - &self, - endpoint_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_endpoints_by_ids(self, endpoint_ids).await - } - - async fn list_endpoints_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_endpoints_by_provider_ids(self, provider_ids).await - } - - async fn list_keys_by_ids( - &self, - key_ids: &[String], - ) -> Result, DataLayerError> { - 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(self, provider_ids).await - } - - async fn list_key_summaries_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - 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> { - Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await - } - - async fn list_keys_page( - &self, - query: &ProviderCatalogKeyListQuery, - ) -> Result { - Self::list_keys_page(self, query).await - } - - async fn list_key_stats_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_key_stats_by_provider_ids(self, provider_ids).await - } -} - -#[async_trait] -impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository { - async fn create_provider( - &self, - provider: &StoredProviderCatalogProvider, - shift_existing_priorities_from: Option, - ) -> Result { - Self::create_provider(self, provider, shift_existing_priorities_from).await - } - - async fn update_provider( - &self, - provider: &StoredProviderCatalogProvider, - ) -> Result { - Self::update_provider(self, provider).await - } - - async fn compare_and_swap_provider_config( - &self, - update: &ProviderCatalogProviderConfigCasUpdate, - ) -> Result { - Self::compare_and_swap_provider_config(self, update).await - } - - async fn compare_and_swap_provider_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_provider_proxy(self, update).await - } - - async fn delete_provider(&self, provider_id: &str) -> Result { - Self::delete_provider(self, provider_id).await - } - - async fn cleanup_deleted_provider_refs( - &self, - provider_id: &str, - provider_deleted: bool, - endpoint_ids: &[String], - key_ids: &[String], - ) -> Result<(), DataLayerError> { - Self::cleanup_deleted_provider_refs( - self, - provider_id, - provider_deleted, - endpoint_ids, - key_ids, - ) - .await - } - - async fn create_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - Self::create_endpoint(self, endpoint).await - } - - async fn update_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - Self::update_endpoint(self, endpoint).await - } - - async fn compare_and_swap_endpoint_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_endpoint_proxy(self, update).await - } - - async fn delete_endpoint(&self, endpoint_id: &str) -> Result { - Self::delete_endpoint(self, endpoint_id).await - } - - async fn create_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - Self::create_key(self, key).await - } - - async fn update_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - Self::update_key(self, key).await - } - - async fn compare_and_swap_key_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_key_proxy(self, update).await - } - - async fn compare_and_swap_key_credentials( - &self, - update: &ProviderCatalogKeyCredentialsCasUpdate, - ) -> Result { - Self::compare_and_swap_key_credentials(self, update).await - } - - async fn compare_and_update_key_admin_state( - &self, - update: &ProviderCatalogKeyAdminCasUpdate, - ) -> Result { - Self::compare_and_update_key_admin_state(self, update).await - } - - async fn update_keys( - &self, - keys: &[StoredProviderCatalogKey], - ) -> Result, DataLayerError> { - Self::update_keys(self, keys).await - } - - async fn update_key_upstream_metadata( - &self, - key_id: &str, - upstream_metadata: Option<&serde_json::Value>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_upstream_metadata(self, key_id, upstream_metadata, updated_at_unix_secs) - .await - } - - async fn upsert_key_upstream_metadata_namespace( - &self, - key_id: &str, - namespace: &str, - value: &serde_json::Value, - updated_at_unix_secs: Option, - ) -> Result { - Self::upsert_key_upstream_metadata_namespace( - self, - key_id, - namespace, - value, - updated_at_unix_secs, - ) - .await - } - - async fn update_key_model_fetch_state( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: Option, - last_models_fetch_error: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_model_fetch_state( - self, - key_id, - allowed_models, - last_models_fetch_at_unix_secs, - last_models_fetch_error, - updated_at_unix_secs, - ) - .await - } - - async fn update_key_model_fetch_success( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: u64, - upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_model_fetch_success( - self, - key_id, - allowed_models, - last_models_fetch_at_unix_secs, - upstream_metadata_updates, - updated_at_unix_secs, - ) - .await - } - - async fn delete_key(&self, key_id: &str) -> Result { - Self::delete_key(self, key_id).await - } - - async fn compare_and_delete_key_oauth_credential( - &self, - delete: &ProviderCatalogKeyOAuthCredentialCasDelete, - ) -> Result { - Self::compare_and_delete_key_oauth_credential(self, delete).await - } - - async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result { - Self::clear_key_oauth_invalid_marker(self, key_id).await - } - - async fn update_key_oauth_runtime_state( - &self, - key_id: &str, - oauth_invalid_at_unix_secs: Option, - oauth_invalid_reason: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_oauth_runtime_state( - self, - key_id, - oauth_invalid_at_unix_secs, - oauth_invalid_reason, - updated_at_unix_secs, - ) - .await - } - - async fn compare_and_update_key_oauth_runtime_state( - &self, - update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ) -> Result { - Self::compare_and_update_key_oauth_runtime_state(self, update).await - } - - async fn update_key_health_state( - &self, - key_id: &str, - is_active: bool, - health_by_format: Option<&serde_json::Value>, - circuit_breaker_by_format: Option<&serde_json::Value>, - ) -> Result { - Self::update_key_health_state( - self, - key_id, - is_active, - health_by_format, - circuit_breaker_by_format, - ) - .await - } - - async fn reset_key_error_count(&self, key_id: &str) -> Result { - Self::reset_key_error_count(self, key_id).await - } - - async fn compare_and_update_key_adaptive_state( - &self, - update: &ProviderCatalogKeyAdaptiveStateUpdate, - ) -> Result { - Self::compare_and_update_key_adaptive_state(self, update).await - } - - async fn update_key_runtime_metadata( - &self, - update: &ProviderCatalogKeyRuntimeMetadataUpdate, - ) -> Result { - Self::update_key_runtime_metadata(self, update).await - } - - async fn update_key_status_snapshot( - &self, - update: &ProviderCatalogKeyStatusSnapshotUpdate, - ) -> Result { - Self::update_key_status_snapshot(self, update).await - } - - async fn compare_and_update_key_health_state( - &self, - update: &ProviderCatalogKeyHealthStateUpdate, - ) -> Result { - Self::compare_and_update_key_health_state(self, update).await - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn validate_non_empty(value: &str, field_name: &str) -> Result<(), DataLayerError> { - if value.trim().is_empty() { - return Err(DataLayerError::InvalidInput(format!( - "{field_name} is empty" - ))); - } - Ok(()) -} - -fn adaptive_status_snapshot_patch( - patch: &serde_json::Value, -) -> Result { - const OWNED_FIELDS: [&str; 6] = [ - "observation_count", - "header_observation_count", - "latest_upstream_limit", - "learning_confidence", - "enforcement_active", - "known_boundary", - ]; - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog adaptive status snapshot patch must be an object".to_string(), - ) - })?; - Ok(serde_json::Value::Object( - OWNED_FIELDS - .into_iter() - .filter_map(|field| { - object - .get(field) - .cloned() - .map(|value| (field.to_string(), value)) - }) - .collect(), - )) -} - -fn validate_runtime_metadata_update( - update: &ProviderCatalogKeyRuntimeMetadataUpdate, -) -> Result<(), DataLayerError> { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.namespace, - "provider catalog runtime metadata namespace", - )?; - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog runtime status snapshot patch must be an object".to_string(), - )); - } - Ok(()) -} - -fn push_status_snapshot_shallow_patch<'args>( - builder: &mut QueryBuilder<'args, MySql>, - patch: &serde_json::Value, -) -> Result<(), DataLayerError> { - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - ) - })?; - if object.is_empty() { - builder.push("COALESCE(NULLIF(status_snapshot, ''), '{}')"); - return Ok(()); - } - - builder.push("JSON_SET(COALESCE(NULLIF(status_snapshot, ''), '{}')"); - for (field, value) in object { - let path = format!( - "$.{}", - serde_json::to_string(field).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.status_snapshot field is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.status_snapshot value is not serializable: {err}" - )) - })?; - builder.push(", ").push_bind(path).push(", CAST("); - builder.push_bind(value).push(" AS JSON)"); - } - builder.push(")"); - Ok(()) -} - -fn push_upstream_metadata_shallow_patch<'args>( - builder: &mut QueryBuilder<'args, MySql>, - patch: &serde_json::Value, -) -> Result<(), DataLayerError> { - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog upstream metadata patch must be an object".to_string(), - ) - })?; - if object.is_empty() { - builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - return Ok(()); - } - - builder.push("JSON_SET(COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - for (field, value) in object { - let path = format!( - "$.{}", - serde_json::to_string(field).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata field is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata value is not serializable: {err}" - )) - })?; - builder.push(", ").push_bind(path).push(", CAST("); - builder.push_bind(value).push(" AS JSON)"); - } - builder.push(")"); - Ok(()) -} - -fn validate_provider(provider: &StoredProviderCatalogProvider) -> Result<(), DataLayerError> { - validate_non_empty(&provider.id, "provider catalog provider.id")?; - validate_non_empty(&provider.name, "provider catalog provider.name")?; - validate_non_empty( - &provider.provider_type, - "provider catalog provider.provider_type", - )?; - if provider - .billing_type - .as_deref() - .map(str::trim) - .is_some_and(str::is_empty) - { - return Err(DataLayerError::InvalidInput( - "provider catalog provider.billing_type is empty".to_string(), - )); - } - Ok(()) -} - -fn validate_endpoint(endpoint: &StoredProviderCatalogEndpoint) -> Result<(), DataLayerError> { - validate_non_empty(&endpoint.id, "provider catalog endpoint.id")?; - validate_non_empty( - &endpoint.provider_id, - "provider catalog endpoint.provider_id", - )?; - validate_non_empty(&endpoint.api_format, "provider catalog endpoint.api_format")?; - validate_non_empty(&endpoint.base_url, "provider catalog endpoint.base_url")?; - Ok(()) -} - -fn validate_key(key: &StoredProviderCatalogKey) -> Result<(), DataLayerError> { - validate_non_empty(&key.id, "provider catalog key.id")?; - validate_non_empty(&key.provider_id, "provider catalog key.provider_id")?; - validate_non_empty(&key.name, "provider catalog key.name")?; - validate_non_empty(&key.auth_type, "provider catalog key.auth_type")?; - Ok(()) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}")) - }) - }) - .transpose() -} - -fn optional_i64_from_u32(value: Option) -> Option { - value.map(i64::from) -} - -fn optional_json_ref_to_string( - value: Option<&serde_json::Value>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - optional_json_ref_to_string(value.as_ref(), field_name) -} - -async fn compare_and_swap_proxy_json( - pool: &MysqlPool, - select_sql: &'static str, - update_sql: &'static str, - update: &ProviderCatalogProxyCasUpdate, - field_name: &'static str, -) -> Result { - // Legacy catalog rows may contain semantically identical JSON with Python-style - // whitespace. Comparing a re-serialized serde_json::Value directly to a TEXT column - // would make lazy credential migration conflict forever. Compare the parsed value first, - // then fence the write against the exact raw bytes that were observed. - // Outer None means the row does not exist; inner None is an existing SQL NULL proxy. - let observed_raw: Option> = sqlx::query_scalar::<_, Option>(select_sql) - .bind(&update.record_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - let Some(observed_raw) = observed_raw else { - return Ok(false); - }; - let observed = optional_json_from_string(observed_raw.clone(), field_name)?; - if observed != update.expected_proxy { - return Ok(false); - } - - let replacement = optional_json_to_string(&update.proxy, field_name)?; - let rows_affected = sqlx::query(update_sql) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(&update.record_id) - .bind(observed_raw) - .execute(pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) -} - -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); - 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" - } - } -} - -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 { - r#" -INSERT INTO provider_api_keys ( - id, provider_id, name, api_key, auth_type, capabilities, is_active, - api_formats, auth_type_by_format, allow_auth_channel_mismatch_formats, - auth_config, note, internal_priority, rate_multipliers, - global_priority_by_format, allowed_models, expires_at, cache_ttl_minutes, - max_probe_interval_minutes, proxy, fingerprint, rpm_limit, concurrent_limit, - learned_rpm_limit, concurrent_429_count, rpm_429_count, last_429_at, - last_429_type, adjustment_history, utilization_samples, - last_probe_increase_at, last_rpm_peak, request_count, total_tokens, - total_cost_usd, success_count, error_count, total_response_time_ms, - last_used_at, auto_fetch_models, last_models_fetch_at, - last_models_fetch_error, locked_models, model_include_patterns, - model_exclude_patterns, upstream_metadata, oauth_invalid_at, - oauth_invalid_reason, status_snapshot, health_by_format, - circuit_breaker_by_format, created_at, updated_at -) -VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ? -) -"# -} - -fn key_update_sql() -> &'static str { - r#" -UPDATE provider_api_keys -SET - provider_id = ?, - name = ?, - api_key = ?, - auth_type = ?, - capabilities = ?, - is_active = ?, - api_formats = ?, - auth_type_by_format = ?, - allow_auth_channel_mismatch_formats = ?, - auth_config = ?, - note = ?, - internal_priority = ?, - rate_multipliers = ?, - global_priority_by_format = ?, - allowed_models = ?, - expires_at = ?, - cache_ttl_minutes = ?, - max_probe_interval_minutes = ?, - proxy = ?, - fingerprint = ?, - rpm_limit = ?, - concurrent_limit = ?, - auto_fetch_models = ?, - locked_models = ?, - model_include_patterns = ?, - model_exclude_patterns = ?, - updated_at = ? -WHERE id = ? - AND BINARY provider_id = BINARY ? - AND BINARY auth_type = BINARY ? - AND BINARY api_key <=> BINARY ? - AND BINARY auth_config <=> BINARY ? -"# -} - -fn key_update_query( - key: &StoredProviderCatalogKey, - updated_at: i64, -) -> Result, DataLayerError> { - Ok(sqlx::query(key_update_sql()) - .bind(&key.provider_id) - .bind(&key.name) - .bind(&key.encrypted_api_key) - .bind(&key.auth_type) - .bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .bind(key.is_active) - .bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .bind(&key.encrypted_auth_config) - .bind(&key.note) - .bind(key.internal_priority) - .bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .bind(key.cache_ttl_minutes) - .bind(key.max_probe_interval_minutes) - .bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .bind(optional_i64_from_u32(key.rpm_limit)) - .bind(key.concurrent_limit) - .bind(key.auto_fetch_models) - .bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .bind(updated_at) - .bind(&key.id) - .bind(&key.provider_id) - .bind(&key.auth_type) - .bind(&key.encrypted_api_key) - .bind(&key.encrypted_auth_config)) -} - -fn push_admin_key_assignments<'args>( - builder: &mut QueryBuilder<'args, MySql>, - key: &'args StoredProviderCatalogKey, - updated_at: i64, -) -> Result<(), DataLayerError> { - builder - .push_bind(&key.provider_id) - .push(", name = ") - .push_bind(&key.name) - .push(", api_key = ") - .push_bind(&key.encrypted_api_key) - .push(", auth_type = ") - .push_bind(&key.auth_type) - .push(", capabilities = ") - .push_bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .push(", is_active = ") - .push_bind(key.is_active) - .push(", api_formats = ") - .push_bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .push(", auth_type_by_format = ") - .push_bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .push(", allow_auth_channel_mismatch_formats = ") - .push_bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .push(", auth_config = ") - .push_bind(&key.encrypted_auth_config) - .push(", note = ") - .push_bind(&key.note) - .push(", internal_priority = ") - .push_bind(key.internal_priority) - .push(", rate_multipliers = ") - .push_bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .push(", global_priority_by_format = ") - .push_bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .push(", allowed_models = ") - .push_bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .push(", expires_at = ") - .push_bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .push(", cache_ttl_minutes = ") - .push_bind(key.cache_ttl_minutes) - .push(", max_probe_interval_minutes = ") - .push_bind(key.max_probe_interval_minutes) - .push(", proxy = ") - .push_bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .push(", fingerprint = ") - .push_bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .push(", rpm_limit = ") - .push_bind(optional_i64_from_u32(key.rpm_limit)) - .push(", concurrent_limit = ") - .push_bind(key.concurrent_limit) - .push(", auto_fetch_models = ") - .push_bind(key.auto_fetch_models) - .push(", locked_models = ") - .push_bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .push(", model_include_patterns = ") - .push_bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .push(", model_exclude_patterns = ") - .push_bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .push(", updated_at = ") - .push_bind(updated_at); - Ok(()) -} - -fn validate_admin_key_cas_update( - update: &ProviderCatalogKeyAdminCasUpdate, -) -> Result<(), DataLayerError> { - validate_key(&update.key)?; - let expected = &update.expected_credential; - if expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - || expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || update - .expected_encrypted_auth_config - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog admin credential fence contains empty fields".to_string(), - )); - } - let Some(rotation) = update.codex_rotation.as_ref() else { - return Ok(()); - }; - let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex") - && rotation.as_object().is_some_and(|object| { - object.len() == 1 - && object - .get("credential_generation") - .and_then(serde_json::Value::as_str) - .is_some_and(|generation| !generation.trim().is_empty()) - }); - if !valid_rotation { - return Err(DataLayerError::InvalidInput( - "provider catalog Codex rotation must contain only credential_generation".to_string(), - )); - } - Ok(()) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_u64(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn optional_u32(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn map_provider_row(row: &MySqlRow) -> Result { - Ok(StoredProviderCatalogProvider::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("website").map_sql_err()?, - row.try_get("provider_type").map_sql_err()?, - )? - .with_description(row.try_get("description").map_sql_err()?) - .with_billing_fields( - row.try_get("billing_type").map_sql_err()?, - row.try_get("monthly_quota_usd").map_sql_err()?, - row.try_get("monthly_used_usd").map_sql_err()?, - optional_u64( - row.try_get("quota_reset_day").map_sql_err()?, - "providers.quota_reset_day", - )?, - optional_u64( - row.try_get("quota_last_reset_at_unix_secs").map_sql_err()?, - "providers.quota_last_reset_at", - )?, - optional_u64( - row.try_get("quota_expires_at_unix_secs").map_sql_err()?, - "providers.quota_expires_at", - )?, - ) - .with_routing_fields(row.try_get("provider_priority").map_sql_err()?) - .with_transport_fields( - row.try_get("is_active").map_sql_err()?, - row.try_get("keep_priority_on_conversion").map_sql_err()?, - row.try_get("enable_format_conversion").map_sql_err()?, - row.try_get("concurrent_limit").map_sql_err()?, - row.try_get("max_retries").map_sql_err()?, - optional_json_from_string(row.try_get("proxy").map_sql_err()?, "providers.proxy")?, - row.try_get("request_timeout").map_sql_err()?, - row.try_get("stream_first_byte_timeout").map_sql_err()?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "providers.config")?, - ) - .with_timestamps( - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "providers.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "providers.updated_at", - )?, - )) -} - -fn map_endpoint_row(row: &MySqlRow) -> Result { - StoredProviderCatalogEndpoint::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("api_format").map_sql_err()?, - row.try_get("api_family").map_sql_err()?, - row.try_get("endpoint_kind").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - )? - .with_timestamps( - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "provider_endpoints.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "provider_endpoints.updated_at", - )?, - ) - .with_health_score( - row.try_get::, _>("health_score") - .map_sql_err()? - .unwrap_or(1.0), - ) - .with_transport_fields( - row.try_get("base_url").map_sql_err()?, - optional_json_from_string( - row.try_get("header_rules").map_sql_err()?, - "provider_endpoints.header_rules", - )?, - optional_json_from_string( - row.try_get("body_rules").map_sql_err()?, - "provider_endpoints.body_rules", - )?, - row.try_get("max_retries").map_sql_err()?, - row.try_get("custom_path").map_sql_err()?, - optional_json_from_string( - row.try_get("config").map_sql_err()?, - "provider_endpoints.config", - )?, - optional_json_from_string( - row.try_get("format_acceptance_config").map_sql_err()?, - "provider_endpoints.format_acceptance_config", - )?, - optional_json_from_string( - row.try_get("proxy").map_sql_err()?, - "provider_endpoints.proxy", - )?, - ) -} - -fn map_key_stats_row(row: &MySqlRow) -> 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") - .map_sql_err()? - .unwrap_or(0.0); - if !total_cost_usd.is_finite() { - return Err(DataLayerError::UnexpectedValue( - "invalid provider_api_keys.total_cost_usd".to_string(), - )); - } - - StoredProviderCatalogKey::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("auth_type").map_sql_err()?, - optional_json_from_string( - row.try_get("capabilities").map_sql_err()?, - "provider_api_keys.capabilities", - )?, - row.try_get("is_active").map_sql_err()?, - )? - .with_transport_fields( - optional_json_from_string( - row.try_get("api_formats").map_sql_err()?, - "provider_api_keys.api_formats", - )?, - row.try_get::, _>("api_key").map_sql_err()?, - row.try_get("auth_config").map_sql_err()?, - optional_json_from_string( - row.try_get("rate_multipliers").map_sql_err()?, - "provider_api_keys.rate_multipliers", - )?, - optional_json_from_string( - row.try_get("global_priority_by_format").map_sql_err()?, - "provider_api_keys.global_priority_by_format", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "provider_api_keys.allowed_models", - )?, - optional_u64( - row.try_get("expires_at_unix_secs").map_sql_err()?, - "provider_api_keys.expires_at", - )?, - optional_json_from_string( - row.try_get("proxy").map_sql_err()?, - "provider_api_keys.proxy", - )?, - optional_json_from_string( - row.try_get("fingerprint").map_sql_err()?, - "provider_api_keys.fingerprint", - )?, - ) - .map(|key| { - let mut key = key - .with_rate_limit_fields( - optional_u32( - row.try_get("rpm_limit").map_sql_err()?, - "provider_api_keys.rpm_limit", - )?, - row.try_get("concurrent_limit").map_sql_err()?, - optional_u32( - row.try_get("learned_rpm_limit").map_sql_err()?, - "provider_api_keys.learned_rpm_limit", - )?, - optional_u32( - row.try_get("concurrent_429_count").map_sql_err()?, - "provider_api_keys.concurrent_429_count", - )?, - optional_u32( - row.try_get("rpm_429_count").map_sql_err()?, - "provider_api_keys.rpm_429_count", - )?, - optional_u64( - row.try_get("last_429_at_unix_secs").map_sql_err()?, - "provider_api_keys.last_429_at", - )?, - optional_json_from_string( - row.try_get("adjustment_history").map_sql_err()?, - "provider_api_keys.adjustment_history", - )?, - optional_u32( - row.try_get("request_count").map_sql_err()?, - "provider_api_keys.request_count", - )?, - optional_u32( - row.try_get("success_count").map_sql_err()?, - "provider_api_keys.success_count", - )?, - ) - .with_usage_fields( - optional_u32( - row.try_get("error_count").map_sql_err()?, - "provider_api_keys.error_count", - )?, - optional_u64( - row.try_get("total_response_time_ms").map_sql_err()?, - "provider_api_keys.total_response_time_ms", - )?, - ) - .with_usage_totals( - optional_u64( - row.try_get("total_tokens").map_sql_err()?, - "provider_api_keys.total_tokens", - )? - .unwrap_or(0), - total_cost_usd, - ) - .with_health_fields( - optional_json_from_string( - row.try_get("health_by_format").map_sql_err()?, - "provider_api_keys.health_by_format", - )?, - optional_json_from_string( - row.try_get("circuit_breaker_by_format").map_sql_err()?, - "provider_api_keys.circuit_breaker_by_format", - )?, - ); - key.note = row.try_get("note").map_sql_err()?; - let auth_type_by_format = optional_json_from_string( - row.try_get("auth_type_by_format").map_sql_err()?, - "provider_api_keys.auth_type_by_format", - )?; - let allow_auth_channel_mismatch_formats = optional_json_from_string( - row.try_get("allow_auth_channel_mismatch_formats") - .map_sql_err()?, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?; - key.internal_priority = row.try_get("internal_priority").unwrap_or(50); - key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").unwrap_or(5); - key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").unwrap_or(32); - key.last_429_type = row.try_get("last_429_type").map_sql_err()?; - key.utilization_samples = optional_json_from_string( - row.try_get("utilization_samples").map_sql_err()?, - "provider_api_keys.utilization_samples", - )?; - key.last_probe_increase_at_unix_secs = optional_u64( - row.try_get("last_probe_increase_at_unix_secs") - .map_sql_err()?, - "provider_api_keys.last_probe_increase_at", - )?; - key.last_rpm_peak = optional_u32( - row.try_get("last_rpm_peak").map_sql_err()?, - "provider_api_keys.last_rpm_peak", - )?; - key.last_used_at_unix_secs = optional_u64( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - "provider_api_keys.last_used_at", - )?; - key.auto_fetch_models = row.try_get("auto_fetch_models").map_sql_err()?; - key.last_models_fetch_at_unix_secs = optional_u64( - row.try_get("last_models_fetch_at_unix_secs") - .map_sql_err()?, - "provider_api_keys.last_models_fetch_at", - )?; - key.last_models_fetch_error = row.try_get("last_models_fetch_error").map_sql_err()?; - key.locked_models = optional_json_from_string( - row.try_get("locked_models").map_sql_err()?, - "provider_api_keys.locked_models", - )?; - key.model_include_patterns = optional_json_from_string( - row.try_get("model_include_patterns").map_sql_err()?, - "provider_api_keys.model_include_patterns", - )?; - key.model_exclude_patterns = optional_json_from_string( - row.try_get("model_exclude_patterns").map_sql_err()?, - "provider_api_keys.model_exclude_patterns", - )?; - key.upstream_metadata = optional_json_from_string( - row.try_get("upstream_metadata").map_sql_err()?, - "provider_api_keys.upstream_metadata", - )?; - key.oauth_invalid_at_unix_secs = optional_u64( - row.try_get("oauth_invalid_at_unix_secs").map_sql_err()?, - "provider_api_keys.oauth_invalid_at", - )?; - key.oauth_invalid_reason = row.try_get("oauth_invalid_reason").map_sql_err()?; - key.status_snapshot = optional_json_from_string( - row.try_get("status_snapshot").map_sql_err()?, - "provider_api_keys.status_snapshot", - )?; - key.created_at_unix_ms = optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "provider_api_keys.created_at", - )?; - key.updated_at_unix_secs = optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "provider_api_keys.updated_at", - )?; - key.with_auth_channel_policy_fields( - auth_type_by_format, - allow_auth_channel_mismatch_formats, - ) - })? -} - -#[cfg(test)] -mod tests { - use super::MysqlProviderCatalogReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, StoredProviderCatalogEndpoint, - StoredProviderCatalogKey, StoredProviderCatalogProvider, - }; - use serde_json::json; - use sqlx::Execute; - - #[test] - fn ordinary_key_update_does_not_own_adaptive_runtime_fields() { - let sql = super::key_update_sql().to_ascii_lowercase(); - for runtime_assignment in [ - "learned_rpm_limit =", - "rpm_429_count =", - "last_429_at =", - "last_429_type =", - "adjustment_history =", - "utilization_samples =", - "last_probe_increase_at =", - "last_rpm_peak =", - ] { - assert!(!sql.contains(runtime_assignment)); - } - assert!(sql.contains("binary api_key <=> binary ?")); - assert!(sql.contains("binary auth_config <=> binary ?")); - } - - #[test] - fn credential_cas_migrates_legacy_encrypted_key_with_binary_fence() { - let source = include_str!("provider_catalog.rs"); - assert!(source.contains( - "SET api_key = ?, encrypted_key = NULL, auth_config = ? WHERE id = ? AND BINARY provider_id = BINARY ? AND BINARY COALESCE(api_key, encrypted_key) <=> BINARY ? AND BINARY auth_config <=> BINARY ?" - )); - } - - #[test] - fn admin_credential_cas_has_atomic_rotation_guards() { - let source = include_str!("provider_catalog.rs"); - for predicate in [ - "BINARY api_key <=> BINARY ", - "BINARY auth_config <=> BINARY ", - "JSON_TYPE(COALESCE(upstream_metadata, JSON_OBJECT())) = 'OBJECT'", - "JSON_TYPE(COALESCE(status_snapshot, JSON_OBJECT())) = 'OBJECT'", - "JSON_SET(COALESCE(upstream_metadata, JSON_OBJECT()), '$.codex'", - "oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0", - ] { - assert!( - source.contains(predicate), - "missing admin CAS guard: {predicate}" - ); - } - } - - #[test] - fn runtime_metadata_cas_requires_an_object_metadata_root() { - let source = include_str!("provider_catalog.rs"); - assert!( - source.contains("JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") - ); - } - #[tokio::test] - async fn empty_id_lists_do_not_connect_to_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let repository = MysqlProviderCatalogReadRepository::new(pool); - assert!(repository - .list_providers_by_ids(&[]) - .await - .expect("empty provider ids should not connect") - .is_empty()); - assert!(repository - .list_endpoints_by_ids(&[]) - .await - .expect("empty endpoint ids should not connect") - .is_empty()); - assert!(repository - .list_endpoints_by_provider_ids(&[]) - .await - .expect("empty endpoint provider ids should not connect") - .is_empty()); - assert!(repository - .list_keys_by_ids(&[]) - .await - .expect("empty key ids should not connect") - .is_empty()); - assert!(repository - .list_keys_by_provider_ids(&[]) - .await - .expect("empty key provider ids should not connect") - .is_empty()); - assert!(repository - .list_key_summaries_by_provider_ids(&[]) - .await - .expect("empty key summary provider ids should not connect") - .is_empty()); - assert!(repository - .list_key_maintenance_summaries_by_provider_ids(&[]) - .await - .expect("empty maintenance provider ids should not connect") - .is_empty()); - assert!(repository - .list_key_stats_by_provider_ids(&[]) - .await - .expect("empty stats provider ids should not connect") - .is_empty()); - } - - #[test] - fn list_query_targets_ids_and_index_aligned_ordering() { - let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()]; - let mut builder = super::build_in_query( - super::KEY_SELECT_SQL, - "provider_id", - &provider_ids, - " ORDER BY provider_id ASC, name ASC, id ASC", - ); - let query = builder.build(); - let sql = query.sql(); - - assert!(sql.contains("WHERE provider_id IN (")); - assert!(sql.contains("ORDER BY provider_id ASC, name ASC, id ASC")); - } - - #[test] - fn key_page_builders_apply_matching_filters_and_requested_order() { - let query = ProviderCatalogKeyListQuery { - provider_id: "provider-a".to_string(), - search: Some("Primary".to_string()), - is_active: Some(true), - offset: 20, - limit: 10, - order: ProviderCatalogKeyListOrder::LastUsedAtDesc, - }; - let mut count_builder = super::build_key_page_count_query(&query); - let count_sql = count_builder.build().sql().to_string(); - let mut list_builder = super::build_key_page_list_query(&query, 10, 20); - let list_sql = list_builder.build().sql().to_string(); - - for sql in [&count_sql, &list_sql] { - assert!(sql.contains("provider_id = ?")); - assert!(sql.contains("LOWER(name) LIKE ?")); - assert!(sql.contains("LOWER(id) LIKE ?")); - assert!(sql.contains("is_active = ?")); - } - assert!(list_sql - .contains("ORDER BY last_used_at IS NULL ASC, last_used_at DESC, name ASC, id ASC")); - assert!(list_sql.contains("LIMIT ? OFFSET ?")); - } - - #[test] - fn key_page_orders_match_cross_driver_contract() { - for (order, expected) in [ - ( - 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", - ), - ] { - assert_eq!(super::key_page_order_by(&order), expected); - } - } - - #[tokio::test] - async fn mysql_provider_catalog_repository_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 provider catalog smoke 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 endpoint_id = format!("endpoint-{suffix}"); - let key_id = format!("key-{suffix}"); - let repository = MysqlProviderCatalogReadRepository::new(pool); - - let provider = StoredProviderCatalogProvider::new( - provider_id.clone(), - format!("Provider {suffix}"), - Some(format!("https://{suffix}.example.com")), - "custom".to_string(), - ) - .expect("provider should build") - .with_description(Some("mysql smoke provider".to_string())) - .with_billing_fields( - Some("pay_as_you_go".to_string()), - Some(100.0), - Some(7.5), - Some(1), - Some(1_710_000_000), - None, - ) - .with_routing_fields(20) - .with_transport_fields( - true, - true, - true, - Some(4), - Some(2), - Some(json!({"http":"proxy"})), - Some(30.0), - Some(2.5), - Some(json!({"region":"us"})), - ); - let created_provider = repository - .create_provider(&provider, None) - .await - .expect("provider should create"); - assert_eq!(created_provider.provider_priority, 20); - - let endpoint = StoredProviderCatalogEndpoint::new( - endpoint_id.clone(), - provider_id.clone(), - "openai:chat".to_string(), - Some("openai".to_string()), - Some("chat".to_string()), - true, - ) - .expect("endpoint should build") - .with_health_score(0.88) - .with_transport_fields( - "https://api.example.com/v1".to_string(), - Some(json!({"Authorization":"Bearer"})), - Some(json!({"model":"gpt"})), - Some(3), - Some("/chat/completions".to_string()), - Some(json!({"timeout":30})), - Some(json!({"accept":["openai:chat"]})), - Some(json!({"https":"proxy"})), - ) - .expect("endpoint transport should build"); - let created_endpoint = repository - .create_endpoint(&endpoint) - .await - .expect("endpoint should create"); - assert_eq!(created_endpoint.health_score, 0.88); - - let key = StoredProviderCatalogKey::new( - key_id.clone(), - provider_id.clone(), - "Default Key".to_string(), - "api_key".to_string(), - Some(json!({"cache_1h":true})), - true, - ) - .expect("key should build") - .with_transport_fields( - Some(json!(["openai:chat"])), - Some("enc-key".to_string()), - Some("enc-auth".to_string()), - Some(json!({"openai:chat":1.0})), - Some(json!({"openai:chat":10})), - Some(json!(["gpt-4.1"])), - Some(1_730_000_000), - Some(json!({"http":"proxy"})), - Some(json!({"fp":"abc"})), - ) - .expect("key transport should build") - .with_rate_limit_fields( - Some(120), - Some(3), - Some(110), - Some(1), - Some(2), - Some(1_720_000_000), - Some(json!([{"rpm":110}])), - Some(10), - Some(9), - ) - .with_usage_fields(Some(1), Some(250)) - .with_usage_totals(1234, 1.5) - .with_health_fields( - Some(json!({"openai:chat":{"score":1}})), - Some(json!({"openai:chat":{"open":false}})), - ); - let created_key = repository - .create_key(&key) - .await - .expect("key should create"); - assert_eq!(created_key.concurrent_limit, Some(3)); - - let providers = repository - .list_providers_by_ids(std::slice::from_ref(&provider_id)) - .await - .expect("providers should list by ids"); - assert_eq!(providers.len(), 1); - assert_eq!(providers[0].id, provider_id); - assert!(repository - .list_providers(true) - .await - .expect("active providers should list") - .iter() - .any(|provider| provider.id == provider_id)); - - let endpoints = repository - .list_endpoints_by_ids(std::slice::from_ref(&endpoint_id)) - .await - .expect("endpoints should list by ids"); - assert_eq!(endpoints.len(), 1); - assert_eq!(endpoints[0].id, endpoint_id); - assert_eq!( - repository - .list_endpoints_by_provider_ids(std::slice::from_ref(&provider_id)) - .await - .expect("endpoints should list by provider") - .len(), - 1 - ); - - let keys = repository - .list_keys_by_provider_ids(std::slice::from_ref(&provider_id)) - .await - .expect("keys should list"); - assert_eq!(keys.len(), 1); - assert_eq!(keys[0].total_tokens, 1234); - - let summaries = repository - .list_key_summaries_by_provider_ids(std::slice::from_ref(&provider_id)) - .await - .expect("key summaries should list"); - assert_eq!(summaries.len(), 1); - assert_eq!(summaries[0].encrypted_api_key.as_deref(), Some("summary")); - assert_eq!(summaries[0].total_tokens, 0); - - let page = repository - .list_keys_page(&ProviderCatalogKeyListQuery { - provider_id: provider_id.clone(), - search: Some("default".to_string()), - is_active: Some(true), - offset: 0, - limit: 10, - order: ProviderCatalogKeyListOrder::Name, - }) - .await - .expect("filtered key page should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].id, key_id); - - let stats = repository - .list_key_stats_by_provider_ids(std::slice::from_ref(&provider_id)) - .await - .expect("key stats should list"); - assert_eq!(stats.len(), 1); - assert_eq!(stats[0].total_keys, 1); - assert_eq!(stats[0].active_keys, 1); - - assert!(repository - .update_key_upstream_metadata( - &key_id, - Some(&json!({"models":["gpt-4.1-mini"]})), - Some(1_740_000_000), - ) - .await - .expect("upstream metadata should update")); - assert!(repository - .update_key_health_state( - &key_id, - false, - Some(&json!({"openai:chat":{"score":0.9}})), - None, - ) - .await - .expect("health state should update")); - - let reloaded_key = repository - .list_keys_by_ids(std::slice::from_ref(&key_id)) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!( - reloaded_key.upstream_metadata, - Some(json!({"models":["gpt-4.1-mini"]})) - ); - assert!(!reloaded_key.is_active); - - let maintenance = repository - .list_key_maintenance_summaries_by_provider_ids(std::slice::from_ref(&provider_id)) - .await - .expect("maintenance summaries should list"); - assert_eq!(maintenance.len(), 1); - assert!(!maintenance[0].is_active); - assert_eq!( - maintenance[0].upstream_metadata, - Some(json!({"models":["gpt-4.1-mini"]})) - ); - - assert!(repository - .delete_key(&key_id) - .await - .expect("key should delete")); - assert!(repository - .delete_endpoint(&endpoint_id) - .await - .expect("endpoint should delete")); - assert!(repository - .delete_provider(&provider_id) - .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/proxy_nodes.rs b/crates/aether-data/adapters/mysql/src/proxy_nodes.rs deleted file mode 100644 index 53cd53fd3..000000000 --- a/crates/aether-data/adapters/mysql/src/proxy_nodes.rs +++ /dev/null @@ -1,2197 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; - -use aether_data_contracts::repository::proxy_nodes::{ - bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample, - merge_proxy_metadata_for_registration, normalize_heartbeat_proxy_metadata, - normalize_proxy_metadata, proxy_metadata_has_explicit_tunnel_security, - reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, - ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, - ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation, - ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, - ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent, - StoredProxyNodeMetricsBucket, TunnelErrorEventRecord, TunnelMetricsSample, - PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const PROXY_NODE_REGISTRATION_CAS_RETRIES: usize = 8; - -fn log_reported_tunnel_error_event( - node_id: &str, - event: &TunnelErrorEventRecord, - received_at_unix_secs: u64, -) { - tracing::warn!( - event_name = "proxy_tunnel_error_reported", - source = "heartbeat", - node_id = %node_id, - category = %event.category, - message = %event.message, - severity = ?event.severity, - component = ?event.component, - summary = ?event.summary, - operator_action = ?event.operator_action, - error_reported_at_unix_secs = event.timestamp_unix_secs, - error_reported_at_unix_ms = ?event.timestamp_unix_ms, - report_received_at_unix_secs = received_at_unix_secs, - "proxy reported tunnel error via heartbeat" - ); -} - -#[derive(Debug, Clone)] -pub struct MysqlProxyNodeReadRepository { - pool: MysqlPool, -} - -impl MysqlProxyNodeReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn write_node( - &self, - node: &StoredProxyNode, - update_existing: bool, - ) -> Result<(), DataLayerError> { - let now = current_unix_secs(); - let upsert_sql = r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, region, status, registered_by, last_heartbeat_at, - heartbeat_interval, active_connections, total_requests, avg_latency_ms, - is_manual, proxy_url, proxy_username, proxy_password, created_at, - updated_at, remote_config, config_version, hardware_info, - estimated_max_concurrency, tunnel_mode, tunnel_connected, tunnel_connected_at, - failed_requests, dns_failures, stream_errors, proxy_metadata -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - name = VALUES(name), - ip = VALUES(ip), - port = VALUES(port), - region = VALUES(region), - status = VALUES(status), - registered_by = VALUES(registered_by), - last_heartbeat_at = VALUES(last_heartbeat_at), - heartbeat_interval = VALUES(heartbeat_interval), - active_connections = VALUES(active_connections), - total_requests = VALUES(total_requests), - avg_latency_ms = VALUES(avg_latency_ms), - is_manual = VALUES(is_manual), - proxy_url = VALUES(proxy_url), - proxy_username = VALUES(proxy_username), - proxy_password = VALUES(proxy_password), - updated_at = VALUES(updated_at), - remote_config = VALUES(remote_config), - config_version = VALUES(config_version), - hardware_info = VALUES(hardware_info), - estimated_max_concurrency = VALUES(estimated_max_concurrency), - tunnel_mode = VALUES(tunnel_mode), - tunnel_connected = VALUES(tunnel_connected), - tunnel_connected_at = VALUES(tunnel_connected_at), - failed_requests = VALUES(failed_requests), - dns_failures = VALUES(dns_failures), - stream_errors = VALUES(stream_errors), - proxy_metadata = VALUES(proxy_metadata) -"#; - let sql = if update_existing { - upsert_sql - } else { - upsert_sql - .split_once("\nON DUPLICATE KEY UPDATE") - .map(|(insert_sql, _)| insert_sql) - .expect("proxy node upsert SQL should contain its conflict clause") - }; - sqlx::query(sql) - .bind(&node.id) - .bind(&node.tunnel_generation) - .bind(&node.name) - .bind(&node.ip) - .bind(node.port) - .bind(&node.region) - .bind(&node.status) - .bind(&node.registered_by) - .bind(optional_i64_from_u64( - node.last_heartbeat_at_unix_secs, - "proxy_nodes.last_heartbeat_at", - )?) - .bind(node.heartbeat_interval) - .bind(node.active_connections) - .bind(node.total_requests) - .bind(node.avg_latency_ms) - .bind(node.is_manual) - .bind(&node.proxy_url) - .bind(&node.proxy_username) - .bind(&node.proxy_password) - .bind(node.created_at_unix_ms.unwrap_or(now) as i64) - .bind(node.updated_at_unix_secs.unwrap_or(now) as i64) - .bind(optional_json_to_string( - &node.remote_config, - "proxy_nodes.remote_config", - )?) - .bind(node.config_version) - .bind(optional_json_to_string( - &node.hardware_info, - "proxy_nodes.hardware_info", - )?) - .bind(node.estimated_max_concurrency) - .bind(node.tunnel_mode) - .bind(node.tunnel_connected) - .bind(optional_i64_from_u64( - node.tunnel_connected_at_unix_secs, - "proxy_nodes.tunnel_connected_at", - )?) - .bind(node.failed_requests) - .bind(node.dns_failures) - .bind(node.stream_errors) - .bind(optional_json_to_string( - &node.proxy_metadata, - "proxy_nodes.proxy_metadata", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - async fn insert_node(&self, node: &StoredProxyNode) -> Result<(), DataLayerError> { - self.write_node(node, false).await - } - - async fn update_existing_registration_if_unchanged( - &self, - mutation: &ProxyNodeRegistrationMutation, - existing: &StoredProxyNode, - replacement_proxy_metadata: Option<&serde_json::Value>, - now: u64, - ) -> Result { - let hardware_info = - optional_json_to_string(&mutation.hardware_info, "proxy_nodes.hardware_info")?; - let replacement_proxy_metadata_json = optional_json_to_string( - &replacement_proxy_metadata.cloned(), - "proxy_nodes.proxy_metadata", - )?; - let expected_proxy_metadata = - optional_json_to_string(&existing.proxy_metadata, "proxy_nodes.proxy_metadata")?; - let result = sqlx::query(UPDATE_PROXY_NODE_REGISTRATION_SQL) - .bind(&mutation.name) - .bind(&mutation.ip) - .bind(mutation.port) - .bind(mutation.region.as_deref()) - .bind(mutation.registered_by.as_deref()) - .bind(now as i64) - .bind(mutation.heartbeat_interval) - .bind(mutation.active_connections) - .bind(mutation.total_requests) - .bind(mutation.avg_latency_ms) - .bind(hardware_info) - .bind(mutation.estimated_max_concurrency) - .bind(mutation.tunnel_mode) - .bind(replacement_proxy_metadata_json) - .bind(now as i64) - .bind(&existing.id) - .bind(&existing.tunnel_generation) - .bind(&existing.ip) - .bind(existing.port) - .bind(expected_proxy_metadata.as_deref()) - .bind(expected_proxy_metadata.as_deref()) - .bind(expected_proxy_metadata.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() != 0 { - return Ok(true); - } - - let Some(current) = self.find_proxy_node(&existing.id).await? else { - return Ok(false); - }; - Ok(proxy_node_registration_matches( - ¤t, - mutation, - existing, - replacement_proxy_metadata, - now, - )) - } - - async fn find_duplicate_proxy_node( - &self, - ip: &str, - port: i32, - excluding_node_id: Option<&str>, - ) -> Result, DataLayerError> { - let row = if let Some(excluding_node_id) = excluding_node_id { - sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE ip = ? AND port = ? AND id <> ? LIMIT 1" - )) - .bind(ip) - .bind(port) - .bind(excluding_node_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE ip = ? AND port = ? LIMIT 1" - )) - .bind(ip) - .bind(port) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - }; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn find_registered_proxy_node_by_endpoint( - &self, - ip: &str, - port: i32, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE BINARY ip = BINARY ? AND port = ? AND is_manual = 0 ORDER BY created_at ASC, id ASC LIMIT 1" - )) - .bind(ip) - .bind(port) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn insert_event( - &self, - node_id: &str, - expected_tunnel_generation: Option<&str>, - event_type: &str, - detail: Option<&str>, - event_metadata: Option<&serde_json::Value>, - created_at_unix_secs: Option, - ) -> Result<(), DataLayerError> { - sqlx::query( - r#" -INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at) -SELECT id, ?, ?, ?, ? -FROM proxy_nodes -WHERE id = ? AND (? IS NULL OR BINARY tunnel_generation = BINARY ?) -"#, - ) - .bind(event_type) - .bind(detail) - .bind(optional_json_to_string( - &event_metadata.cloned(), - "proxy_node_events.event_metadata", - )?) - .bind(created_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(node_id) - .bind(expected_tunnel_generation) - .bind(expected_tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - async fn upsert_metrics_bucket( - &self, - table: &str, - node_id: &str, - expected_tunnel_generation: Option<&str>, - bucket_start: u64, - sample: &TunnelMetricsSample, - ) -> Result<(), DataLayerError> { - sqlx::query(&format!( - r#" -INSERT INTO {table} ( - node_id, - bucket_start_unix_secs, - samples, - uptime_samples, - active_connections_sum, - active_connections_max, - heartbeat_rtt_ms_sum, - heartbeat_rtt_ms_max, - connect_errors_delta, - disconnects_delta, - error_events_delta, - ws_in_bytes_delta, - ws_out_bytes_delta, - ws_in_frames_delta, - ws_out_frames_delta -) -SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? -FROM proxy_nodes -WHERE id = ? AND (? IS NULL OR BINARY tunnel_generation = BINARY ?) -ON DUPLICATE KEY UPDATE - samples = samples + VALUES(samples), - uptime_samples = uptime_samples + VALUES(uptime_samples), - active_connections_sum = active_connections_sum + VALUES(active_connections_sum), - active_connections_max = GREATEST(active_connections_max, VALUES(active_connections_max)), - heartbeat_rtt_ms_sum = heartbeat_rtt_ms_sum + VALUES(heartbeat_rtt_ms_sum), - heartbeat_rtt_ms_max = GREATEST(heartbeat_rtt_ms_max, VALUES(heartbeat_rtt_ms_max)), - connect_errors_delta = connect_errors_delta + VALUES(connect_errors_delta), - disconnects_delta = disconnects_delta + VALUES(disconnects_delta), - error_events_delta = error_events_delta + VALUES(error_events_delta), - ws_in_bytes_delta = ws_in_bytes_delta + VALUES(ws_in_bytes_delta), - ws_out_bytes_delta = ws_out_bytes_delta + VALUES(ws_out_bytes_delta), - ws_in_frames_delta = ws_in_frames_delta + VALUES(ws_in_frames_delta), - ws_out_frames_delta = ws_out_frames_delta + VALUES(ws_out_frames_delta) -"# - )) - .bind(node_id) - .bind(i64::try_from(bucket_start).unwrap_or(i64::MAX)) - .bind(sample.samples) - .bind(sample.uptime_samples) - .bind(sample.active_connections_sum) - .bind(sample.active_connections_max) - .bind(sample.heartbeat_rtt_ms_sum) - .bind(sample.heartbeat_rtt_ms_max) - .bind(sample.connect_errors_delta) - .bind(sample.disconnects_delta) - .bind(sample.error_events_delta) - .bind(sample.ws_in_bytes_delta) - .bind(sample.ws_out_bytes_delta) - .bind(sample.ws_in_frames_delta) - .bind(sample.ws_out_frames_delta) - .bind(node_id) - .bind(expected_tunnel_generation) - .bind(expected_tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - fn normalize_remote_config( - mutation: &ProxyNodeRemoteConfigMutation, - existing: Option<&serde_json::Value>, - ) -> Option { - let mut config = match existing { - Some(serde_json::Value::Object(map)) => map.clone(), - _ => serde_json::Map::new(), - }; - - if let Some(node_name) = mutation.node_name.as_ref() { - config.insert( - "node_name".to_string(), - serde_json::Value::String(node_name.clone()), - ); - } - if let Some(allowed_ports) = mutation.allowed_ports.as_ref() { - config.insert( - "allowed_ports".to_string(), - serde_json::json!(allowed_ports), - ); - } - if let Some(log_level) = mutation.log_level.as_ref() { - config.insert( - "log_level".to_string(), - serde_json::Value::String(log_level.clone()), - ); - } - if let Some(heartbeat_interval) = mutation.heartbeat_interval { - config.insert( - "heartbeat_interval".to_string(), - serde_json::json!(heartbeat_interval), - ); - } - if let Some(scheduling_state) = mutation.scheduling_state.as_ref() { - match scheduling_state { - Some(state) => { - config.insert( - "scheduling_state".to_string(), - serde_json::Value::String(state.clone()), - ); - } - None => { - config.remove("scheduling_state"); - } - } - } - if let Some(upgrade_to) = mutation.upgrade_to.as_ref() { - match upgrade_to { - Some(version) => { - config.insert( - "upgrade_to".to_string(), - serde_json::Value::String(version.clone()), - ); - } - None => { - config.remove("upgrade_to"); - } - } - } - - (!config.is_empty()).then_some(serde_json::Value::Object(config)) - } -} - -const PROXY_NODE_COLUMNS: &str = r#" -SELECT - id, - tunnel_generation, - name, - ip, - port, - region, - is_manual, - proxy_url, - proxy_username, - proxy_password, - status, - registered_by, - last_heartbeat_at AS last_heartbeat_at_unix_secs, - heartbeat_interval, - active_connections, - total_requests, - avg_latency_ms, - failed_requests, - dns_failures, - stream_errors, - proxy_metadata, - hardware_info, - estimated_max_concurrency, - tunnel_mode, - tunnel_connected, - tunnel_connected_at AS tunnel_connected_at_unix_secs, - remote_config, - config_version, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM proxy_nodes -"#; - -const APPLY_HEARTBEAT_SQL: &str = r#" -UPDATE proxy_nodes -SET last_heartbeat_at = ?, - tunnel_connected_at = CASE - WHEN status <> 'online' OR tunnel_connected = 0 THEN ? - ELSE tunnel_connected_at - END, - updated_at = CASE - WHEN status <> 'online' OR tunnel_connected = 0 THEN ? - ELSE updated_at - END, - status = 'online', - tunnel_connected = 1, - heartbeat_interval = COALESCE(?, heartbeat_interval), - active_connections = COALESCE(?, active_connections), - avg_latency_ms = COALESCE(?, avg_latency_ms), - total_requests = total_requests + GREATEST(COALESCE(?, 0), 0), - failed_requests = failed_requests + GREATEST(COALESCE(?, 0), 0), - dns_failures = dns_failures + GREATEST(COALESCE(?, 0), 0), - stream_errors = stream_errors + GREATEST(COALESCE(?, 0), 0) -WHERE id = ? - AND tunnel_mode = 1 - AND BINARY tunnel_generation = BINARY ? -"#; - -const CAS_HEARTBEAT_PROXY_METADATA_SQL: &str = r#" -UPDATE proxy_nodes -SET proxy_metadata = ?, updated_at = ? -WHERE id = ? AND BINARY tunnel_generation = BINARY ? - AND ( - (proxy_metadata IS NULL AND ? IS NULL) - OR ( - proxy_metadata IS NOT NULL AND ? IS NOT NULL - AND JSON_VALID(proxy_metadata) = 1 - AND CAST(proxy_metadata AS JSON) = CAST(? AS JSON) - ) - ) -"#; - -const UPDATE_TUNNEL_STATUS_SQL: &str = r#" -UPDATE proxy_nodes -SET tunnel_connected = ?, - active_connections = CASE WHEN ? THEN active_connections ELSE 0 END, - tunnel_connected_at = ?, - status = CASE WHEN ? THEN 'online' ELSE 'offline' END, - updated_at = ? -WHERE id = ? - AND BINARY tunnel_generation = BINARY ? - AND (tunnel_connected_at IS NULL OR tunnel_connected_at <= ?) -"#; - -const UPDATE_MANUAL_PROXY_NODE_SQL: &str = r#" -UPDATE proxy_nodes -SET name = COALESCE(?, name), - ip = COALESCE(?, ip), - port = COALESCE(?, port), - region = COALESCE(?, region), - proxy_url = COALESCE(?, proxy_url), - proxy_username = COALESCE(?, proxy_username), - proxy_password = COALESCE(?, proxy_password), - updated_at = ? -WHERE id = ? AND is_manual = 1 - AND BINARY tunnel_generation = BINARY ? -"#; - -const UPDATE_PROXY_NODE_REGISTRATION_SQL: &str = r#" -UPDATE proxy_nodes -SET name = ?, ip = ?, port = ?, region = ?, registered_by = ?, - last_heartbeat_at = ?, heartbeat_interval = ?, - active_connections = COALESCE(?, active_connections), - total_requests = COALESCE(?, total_requests), - avg_latency_ms = COALESCE(?, avg_latency_ms), - hardware_info = COALESCE(?, hardware_info), - estimated_max_concurrency = COALESCE(?, estimated_max_concurrency), - tunnel_mode = ?, proxy_metadata = COALESCE(?, proxy_metadata), updated_at = ? -WHERE BINARY id = BINARY ? AND BINARY tunnel_generation = BINARY ? - AND is_manual = 0 AND BINARY ip = BINARY ? AND port = ? - AND ( - (proxy_metadata IS NULL AND ? IS NULL) - OR ( - proxy_metadata IS NOT NULL AND ? IS NOT NULL - AND JSON_VALID(proxy_metadata) = 1 - AND CAST(proxy_metadata AS JSON) = CAST(? AS JSON) - ) - ) -"#; - -const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#" -UPDATE proxy_nodes -SET name = COALESCE(?, name), remote_config = ?, - config_version = config_version + 1, updated_at = ? -WHERE id = ? AND BINARY tunnel_generation = BINARY ? AND config_version = ? - AND is_manual = 0 -"#; - -const RECORD_PROXY_NODE_TRAFFIC_SQL: &str = 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 = ? AND is_manual = 1 - AND BINARY tunnel_generation = BINARY ? -"#; - -const INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL: &str = r#" -UPDATE proxy_nodes -SET total_requests = total_requests + GREATEST(?, 0), - failed_requests = failed_requests + GREATEST(?, 0), - avg_latency_ms = COALESCE(?, avg_latency_ms), - updated_at = ? -WHERE id = ? AND is_manual = 1 - AND BINARY tunnel_generation = BINARY ? -"#; - -const UNREGISTER_PROXY_NODE_SQL: &str = r#" -UPDATE proxy_nodes -SET status = 'offline', tunnel_connected = 0, active_connections = 0, - tunnel_connected_at = ?, updated_at = ? -WHERE id = ? - AND BINARY tunnel_generation = BINARY ? -"#; - -// Run after the parent delete commits so delete never waits on an outbox row -// already claimed by the flusher (which acquires locks in the opposite order). -const RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL: &str = r#" -DELETE FROM usage_counter_deltas -WHERE kind = 'proxy_node' - AND target_id = ? - AND BINARY target_tunnel_generation = BINARY ? - AND processed_at IS NULL -"#; - -#[async_trait] -impl ProxyNodeReadRepository for MysqlProxyNodeReadRepository { - async fn list_proxy_nodes(&self) -> Result, DataLayerError> { - let rows = sqlx::query(&format!("{PROXY_NODE_COLUMNS} ORDER BY name ASC, id ASC")) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_node_row).collect() - } - - async fn find_proxy_node( - &self, - node_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn list_proxy_node_events( - &self, - node_id: &str, - limit: usize, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, - node_id, - event_type, - detail, - event_metadata, - created_at AS created_at_unix_ms -FROM proxy_node_events -WHERE node_id = ? -ORDER BY created_at DESC, id DESC -LIMIT ? -"#, - ) - .bind(node_id) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_node_event_row).collect() - } - - async fn list_proxy_node_events_filtered( - &self, - node_id: &str, - query: &ProxyNodeEventQuery, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, - node_id, - event_type, - detail, - event_metadata, - created_at AS created_at_unix_ms -FROM proxy_node_events -WHERE node_id = ? - AND (? IS NULL OR created_at >= ?) - AND (? IS NULL OR created_at <= ?) - AND (? IS NULL OR LOWER(event_type) = LOWER(?)) -ORDER BY created_at DESC, id DESC -LIMIT ? -"#, - ) - .bind(node_id) - .bind( - query - .from_unix_secs - .map(|v| i64::try_from(v).unwrap_or(i64::MAX)), - ) - .bind( - query - .from_unix_secs - .map(|v| i64::try_from(v).unwrap_or(i64::MAX)), - ) - .bind( - query - .to_unix_secs - .map(|v| i64::try_from(v).unwrap_or(i64::MAX)), - ) - .bind( - query - .to_unix_secs - .map(|v| i64::try_from(v).unwrap_or(i64::MAX)), - ) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .bind(i64::try_from(query.limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_node_event_row).collect() - } - - async fn list_proxy_node_metrics( - &self, - node_id: &str, - step: ProxyNodeMetricsStep, - from_unix_secs: u64, - to_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let table = match step { - ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m", - ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h", - }; - let rows = sqlx::query(&format!( - r#" -SELECT - node_id, - bucket_start_unix_secs, - samples, - uptime_samples, - active_connections_sum, - active_connections_max, - heartbeat_rtt_ms_sum, - heartbeat_rtt_ms_max, - connect_errors_delta, - disconnects_delta, - error_events_delta, - ws_in_bytes_delta, - ws_out_bytes_delta, - ws_in_frames_delta, - ws_out_frames_delta -FROM {table} -WHERE node_id = ? - AND bucket_start_unix_secs >= ? - AND bucket_start_unix_secs <= ? -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"# - )) - .bind(node_id) - .bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_node_metric_row).collect() - } - - async fn list_proxy_fleet_metrics( - &self, - step: ProxyNodeMetricsStep, - from_unix_secs: u64, - to_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let table = match step { - ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m", - ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h", - }; - let rows = sqlx::query(&format!( - r#" -SELECT - bucket_start_unix_secs, - SUM(samples) AS samples, - SUM(uptime_samples) AS uptime_samples, - SUM(active_connections_sum) AS active_connections_sum, - MAX(active_connections_max) AS active_connections_max, - SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum, - MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max, - SUM(connect_errors_delta) AS connect_errors_delta, - SUM(disconnects_delta) AS disconnects_delta, - SUM(error_events_delta) AS error_events_delta, - SUM(ws_in_bytes_delta) AS ws_in_bytes_delta, - SUM(ws_out_bytes_delta) AS ws_out_bytes_delta, - SUM(ws_in_frames_delta) AS ws_in_frames_delta, - SUM(ws_out_frames_delta) AS ws_out_frames_delta -FROM {table} -WHERE bucket_start_unix_secs >= ? - AND bucket_start_unix_secs <= ? -GROUP BY bucket_start_unix_secs -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"# - )) - .bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_fleet_metric_row).collect() - } -} - -#[async_trait] -impl ProxyNodeWriteRepository for MysqlProxyNodeReadRepository { - async fn reset_stale_tunnel_statuses(&self) -> Result { - let now = current_unix_secs() as i64; - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET tunnel_connected = 0, - status = 'offline', - active_connections = 0, - tunnel_connected_at = ?, - updated_at = ? -WHERE is_manual = 0 - AND tunnel_connected = 1 -"#, - ) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn compare_and_set_proxy_password( - &self, - node_id: &str, - expected: &str, - replacement: &str, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET proxy_password = ?, updated_at = ? -WHERE id = ? AND BINARY proxy_password = BINARY ? -"#, - ) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn compare_and_set_proxy_metadata( - &self, - node_id: &str, - expected: &serde_json::Value, - replacement: &serde_json::Value, - ) -> Result { - let expected = serde_json::to_string(expected).map_err(|err| { - DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}")) - })?; - let replacement = serde_json::to_string(replacement).map_err(|err| { - DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}")) - })?; - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET proxy_metadata = ?, updated_at = ? -WHERE id = ? - AND proxy_metadata IS NOT NULL - AND JSON_VALID(proxy_metadata) = 1 - AND CAST(proxy_metadata AS JSON) = CAST(? AS JSON) -"#, - ) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(&expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn create_manual_node( - &self, - mutation: &ProxyNodeManualCreateMutation, - ) -> Result { - if let Some(existing) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - return Err(duplicate_proxy_node_error(&existing)); - } - - let node_id = requested_proxy_node_id(mutation.node_id.as_deref())? - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - if let Some(existing) = self.find_proxy_node(&node_id).await? { - return Err(proxy_node_id_in_use_error(&existing)); - } - let now = Some(current_unix_secs()); - let node = StoredProxyNode::new( - node_id, - mutation.name.clone(), - mutation.ip.clone(), - mutation.port, - true, - "online".to_string(), - 0, - 0, - 0, - 0, - 0, - 0, - false, - false, - 0, - )? - .with_manual_proxy_fields( - Some(mutation.proxy_url.clone()), - mutation.proxy_username.clone(), - mutation.proxy_password.clone(), - ) - .with_runtime_fields( - mutation.region.clone(), - mutation.registered_by.clone(), - None, - None, - None, - None, - None, - None, - None, - now, - now, - ); - - if let Err(error) = self.insert_node(&node).await { - if let Some(duplicate) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - return Err(duplicate_proxy_node_error(&duplicate)); - } - if let Some(owner) = self.find_proxy_node(&node.id).await? { - return Err(proxy_node_id_in_use_error(&owner)); - } - return Err(error); - } - Ok(node) - } - - async fn update_manual_node( - &self, - mutation: &ProxyNodeManualUpdateMutation, - ) -> Result, DataLayerError> { - let Some(existing) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if !existing.is_manual { - return Err(DataLayerError::InvalidInput( - "只能编辑手动添加的代理节点".to_string(), - )); - } - - let next_ip = mutation.ip.as_deref().unwrap_or(existing.ip.as_str()); - let next_port = mutation.port.unwrap_or(existing.port); - if let Some(existing) = self - .find_duplicate_proxy_node(next_ip, next_port, Some(&mutation.node_id)) - .await? - { - return Err(duplicate_proxy_node_error(&existing)); - } - - let result = sqlx::query(UPDATE_MANUAL_PROXY_NODE_SQL) - .bind(mutation.name.as_deref()) - .bind(mutation.ip.as_deref()) - .bind(mutation.port) - .bind(mutation.region.as_deref()) - .bind(mutation.proxy_url.as_deref()) - .bind(mutation.proxy_username.as_deref()) - .bind(mutation.proxy_password.as_deref()) - .bind(current_unix_secs() as i64) - .bind(&mutation.node_id) - .bind(&existing.tunnel_generation) - .execute(&self.pool) - .await; - let result = match result { - Ok(result) => result, - Err(error) => { - if let Some(duplicate) = self - .find_duplicate_proxy_node(next_ip, next_port, Some(&mutation.node_id)) - .await? - { - return Err(duplicate_proxy_node_error(&duplicate)); - } - return Err(DataLayerError::sql(error)); - } - }; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_proxy_node(&mutation.node_id).await - } - - async fn register_node( - &self, - mutation: &ProxyNodeRegistrationMutation, - ) -> Result { - let requested_id = requested_proxy_node_id(mutation.node_id.as_deref())?; - let normalized_proxy_metadata = normalize_proxy_metadata( - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ); - let rotates_tunnel_security = - proxy_metadata_has_explicit_tunnel_security(normalized_proxy_metadata.as_ref()); - - let Some(initial_existing) = self - .find_registered_proxy_node_by_endpoint(&mutation.ip, mutation.port) - .await? - else { - let node_id = requested_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - if let Some(existing) = self.find_proxy_node(&node_id).await? { - return Err(proxy_node_id_in_use_error(&existing)); - } - let now = Some(current_unix_secs()); - let node = StoredProxyNode::new( - node_id, - mutation.name.clone(), - mutation.ip.clone(), - mutation.port, - false, - "offline".to_string(), - mutation.heartbeat_interval, - mutation.active_connections.unwrap_or(0), - mutation.total_requests.unwrap_or(0), - 0, - 0, - 0, - mutation.tunnel_mode, - false, - 0, - )? - .with_runtime_fields( - mutation.region.clone(), - mutation.registered_by.clone(), - now, - mutation.avg_latency_ms, - merge_proxy_metadata_for_registration(None, normalized_proxy_metadata.clone()), - mutation.hardware_info.clone(), - mutation.estimated_max_concurrency, - None, - None, - now, - now, - ); - if let Err(error) = self.insert_node(&node).await { - if let Some(winner) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - if winner.is_manual { - return Err(duplicate_proxy_node_error(&winner)); - } - if let Some(requested_id) = requested_id.as_deref() { - if requested_id != winner.id { - return Err(proxy_node_registration_identity_error( - requested_id, - &winner.id, - )); - } - } - return Ok(winner); - } - if let Some(owner) = self.find_proxy_node(&node.id).await? { - return Err(proxy_node_id_in_use_error(&owner)); - } - return Err(error); - } - return Ok(node); - }; - - if let Some(requested_id) = requested_id.as_deref() { - if requested_id != initial_existing.id { - return Err(proxy_node_registration_identity_error( - requested_id, - &initial_existing.id, - )); - } - } - - let pinned_id = initial_existing.id.clone(); - let pinned_generation = initial_existing.tunnel_generation.clone(); - let mut existing = initial_existing; - for attempt in 0..PROXY_NODE_REGISTRATION_CAS_RETRIES { - if attempt != 0 { - existing = self - .find_registered_proxy_node_by_endpoint(&mutation.ip, mutation.port) - .await? - .ok_or_else(proxy_node_registration_changed_error)?; - } - if existing.id != pinned_id || existing.tunnel_generation != pinned_generation { - return Err(proxy_node_registration_changed_error()); - } - - let replacement_proxy_metadata = merge_proxy_metadata_for_registration( - existing.proxy_metadata.as_ref(), - normalized_proxy_metadata.clone(), - ); - let now = current_unix_secs(); - if self - .update_existing_registration_if_unchanged( - mutation, - &existing, - replacement_proxy_metadata.as_ref(), - now, - ) - .await? - { - return self - .find_proxy_node(&pinned_id) - .await? - .filter(|current| current.tunnel_generation == pinned_generation) - .ok_or_else(proxy_node_registration_changed_error); - } - if rotates_tunnel_security { - return Err(DataLayerError::UnexpectedValue( - "proxy node changed during explicit tunnel security rotation".to_string(), - )); - } - } - - Err(DataLayerError::UnexpectedValue( - "proxy node registration changed during every CAS retry".to_string(), - )) - } - - async fn apply_heartbeat( - &self, - mutation: &ProxyNodeHeartbeatMutation, - ) -> Result, DataLayerError> { - let Some(existing) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != existing.tunnel_generation) - { - return Ok(None); - } - if !existing.tunnel_mode { - return Err(DataLayerError::InvalidInput( - "non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode" - .to_string(), - )); - } - - let tunnel_generation = existing.tunnel_generation.clone(); - let now_unix_secs = current_unix_secs(); - let now = i64::try_from(now_unix_secs).unwrap_or(i64::MAX); - let has_proxy_metadata_update = normalize_heartbeat_proxy_metadata( - None, - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ) - .is_some(); - - let result = sqlx::query(APPLY_HEARTBEAT_SQL) - .bind(now) - .bind(now) - .bind(now) - .bind(mutation.heartbeat_interval) - .bind(mutation.active_connections) - .bind(mutation.avg_latency_ms) - .bind(mutation.total_requests_delta) - .bind(mutation.failed_requests_delta) - .bind(mutation.dns_failures_delta) - .bind(mutation.stream_errors_delta) - .bind(&mutation.node_id) - .bind(&tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - - let mut updated = None; - let mut tunnel_metrics_sample = None; - if has_proxy_metadata_update { - for _ in 0..8 { - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != tunnel_generation { - return Ok(None); - } - let Some(replacement) = normalize_heartbeat_proxy_metadata( - current.proxy_metadata.as_ref(), - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ) else { - break; - }; - if current.proxy_metadata.as_ref() == Some(&replacement) { - tunnel_metrics_sample = build_tunnel_metrics_sample( - current.proxy_metadata.as_ref(), - Some(&replacement), - current.active_connections, - current.tunnel_connected, - ); - updated = Some(current); - break; - } - - let expected = - optional_json_to_string(¤t.proxy_metadata, "proxy_nodes.proxy_metadata")?; - let replacement_json = serde_json::to_string(&replacement).map_err(|error| { - DataLayerError::UnexpectedValue(format!( - "proxy_nodes.proxy_metadata contains unserializable JSON: {error}" - )) - })?; - let result = sqlx::query(CAS_HEARTBEAT_PROXY_METADATA_SQL) - .bind(replacement_json) - .bind(now) - .bind(&mutation.node_id) - .bind(&tunnel_generation) - .bind(expected.as_deref()) - .bind(expected.as_deref()) - .bind(expected.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - continue; - } - let Some(after_cas) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if after_cas.tunnel_generation != tunnel_generation { - return Ok(None); - } - tunnel_metrics_sample = build_tunnel_metrics_sample( - current.proxy_metadata.as_ref(), - after_cas.proxy_metadata.as_ref(), - after_cas.active_connections, - after_cas.tunnel_connected, - ); - updated = Some(after_cas); - break; - } - } - let updated = if let Some(updated) = updated { - updated - } else { - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != tunnel_generation { - return Ok(None); - } - current - }; - - if let Some(sample) = tunnel_metrics_sample.as_ref() { - self.upsert_metrics_bucket( - "proxy_node_metrics_1m", - &updated.id, - Some(tunnel_generation.as_str()), - bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute), - sample, - ) - .await?; - self.upsert_metrics_bucket( - "proxy_node_metrics_1h", - &updated.id, - Some(tunnel_generation.as_str()), - bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour), - sample, - ) - .await?; - - for error in &sample.recent_error_events { - log_reported_tunnel_error_event(&updated.id, error, now_unix_secs); - let detail = build_tunnel_error_event_detail(error); - let event_metadata = serde_json::json!({ - "source": "heartbeat", - "category": error.category, - "message": error.message, - "severity": error.severity.as_deref(), - "component": error.component.as_deref(), - "summary": error.summary.as_deref(), - "operator_action": error.operator_action.as_deref(), - "timestamp_unix_secs": error.timestamp_unix_secs, - "timestamp_unix_ms": error.timestamp_unix_ms, - }); - self.insert_event( - &updated.id, - Some(tunnel_generation.as_str()), - PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR, - Some(detail.as_str()), - Some(&event_metadata), - Some(if error.timestamp_unix_secs == 0 { - now_unix_secs - } else { - error.timestamp_unix_secs - }), - ) - .await?; - } - } - if reconcile_remote_config_after_heartbeat( - updated.remote_config.as_ref(), - mutation.proxy_version.as_deref(), - ) != updated.remote_config - { - return self - .update_remote_config(&ProxyNodeRemoteConfigMutation { - node_id: mutation.node_id.clone(), - expected_tunnel_generation: Some(tunnel_generation), - node_name: None, - allowed_ports: None, - log_level: None, - heartbeat_interval: None, - scheduling_state: None, - upgrade_to: Some(None), - }) - .await; - } - - Ok(Some(updated)) - } - - async fn record_traffic( - &self, - mutation: &ProxyNodeTrafficMutation, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(&mutation.node_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - let Some(expected_generation) = mutation.expected_tunnel_generation.as_deref() else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - if expected_generation != generation { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query(RECORD_PROXY_NODE_TRAFFIC_SQL) - .bind(mutation.total_requests_delta) - .bind(mutation.failed_requests_delta) - .bind(mutation.dns_failures_delta) - .bind(mutation.stream_errors_delta) - .bind(current_unix_secs() as i64) - .bind(&mutation.node_id) - .bind(expected_generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - let applied = result.rows_affected() > 0; - tx.commit().await.map_sql_err()?; - Ok(applied) - } - - async fn update_tunnel_status( - &self, - mutation: &ProxyNodeTunnelStatusMutation, - ) -> Result, DataLayerError> { - let Some(node) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != node.tunnel_generation) - { - return Ok(None); - } - - let event_time = mutation - .observed_at_unix_secs - .unwrap_or_else(current_unix_secs); - let event_type = if mutation.connected { - "connected" - } else { - "disconnected" - }; - let event_detail = mutation.detail.clone().unwrap_or_else(|| { - format!( - "[tunnel_node_status] conn_count={}", - i32::max(mutation.conn_count, 0) - ) - }); - - let event_time_i64 = i64::try_from(event_time).unwrap_or(i64::MAX); - let result = sqlx::query(UPDATE_TUNNEL_STATUS_SQL) - .bind(mutation.connected) - .bind(mutation.connected) - .bind(event_time_i64) - .bind(mutation.connected) - .bind(event_time_i64) - .bind(&mutation.node_id) - .bind(&node.tunnel_generation) - .bind(event_time_i64) - .execute(&self.pool) - .await - .map_sql_err()?; - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != node.tunnel_generation { - return Ok(None); - } - let stale = result.rows_affected() == 0 - && current - .tunnel_connected_at_unix_secs - .is_some_and(|last_transition| event_time < last_transition); - let persisted_detail = if stale { - format!("[stale_ignored] {event_detail}") - } else { - event_detail - }; - self.insert_event( - &mutation.node_id, - Some(node.tunnel_generation.as_str()), - event_type, - Some(&persisted_detail), - None, - Some(if stale { - current_unix_secs() - } else { - event_time - }), - ) - .await?; - Ok(Some(current)) - } - - async fn unregister_node( - &self, - node_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(node_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(None); - }; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - let now = current_unix_secs() as i64; - sqlx::query(UNREGISTER_PROXY_NODE_SQL) - .bind(now) - .bind(now) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err() - .and_then(|row| map_proxy_node_row(&row))?; - tx.commit().await.map_sql_err()?; - Ok(Some(updated)) - } - - async fn delete_node(&self, node_id: &str) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(node_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(None); - }; - let existing = map_proxy_node_row(&row)?; - let generation = existing.tunnel_generation.as_str(); - - // Child tables do not carry generation; retain the parent identity check - // so cleanup cannot target a replacement row if schema constraints differ. - sqlx::query( - "DELETE FROM proxy_node_events WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes p WHERE p.id = ? AND BINARY p.tunnel_generation = BINARY ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "DELETE FROM proxy_node_metrics_1m WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes p WHERE p.id = ? AND BINARY p.tunnel_generation = BINARY ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "DELETE FROM proxy_node_metrics_1h WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes p WHERE p.id = ? AND BINARY p.tunnel_generation = BINARY ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let deleted = sqlx::query( - "DELETE FROM proxy_nodes WHERE id = ? AND BINARY tunnel_generation = BINARY ?", - ) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - if deleted.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - tx.commit().await.map_sql_err()?; - if let Err(error) = sqlx::query(RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL) - .bind(node_id) - .bind(generation) - .execute(&self.pool) - .await - .map_sql_err() - { - tracing::warn!( - node_id = %node_id, - tunnel_generation = %generation, - error = ?error, - "failed to retire deleted proxy node counter rows" - ); - } - Ok(Some(existing)) - } - - async fn update_remote_config( - &self, - mutation: &ProxyNodeRemoteConfigMutation, - ) -> Result, DataLayerError> { - for _ in 0..8 { - let Some(node) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != node.tunnel_generation) - { - return Ok(None); - } - if node.is_manual { - return Err(DataLayerError::InvalidInput( - "手动节点不支持远程配置下发".to_string(), - )); - } - - let remote_config = - Self::normalize_remote_config(mutation, node.remote_config.as_ref()); - let remote_config = - optional_json_to_string(&remote_config, "proxy_nodes.remote_config")?; - let now = current_unix_secs() as i64; - let result = sqlx::query(UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL) - .bind(mutation.node_name.as_deref()) - .bind(remote_config) - .bind(now) - .bind(&mutation.node_id) - .bind(&node.tunnel_generation) - .bind(node.config_version) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - continue; - } - - let current = self.find_proxy_node(&mutation.node_id).await?; - return Ok( - current.filter(|current| current.tunnel_generation == node.tunnel_generation) - ); - } - - Err(DataLayerError::UnexpectedValue( - "proxy node remote config changed during every CAS retry".to_string(), - )) - } - - async fn increment_manual_node_requests( - &self, - node_id: &str, - total_delta: i64, - failed_delta: i64, - latency_ms: Option, - ) -> Result<(), DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(node_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(()); - }; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - sqlx::query(INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL) - .bind(total_delta) - .bind(failed_delta) - .bind(latency_ms.map(|value| value as f64)) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(()) - } - - async fn cleanup_proxy_node_metrics( - &self, - retain_1m_from_unix_secs: u64, - retain_1h_from_unix_secs: u64, - delete_limit: usize, - ) -> Result { - let delete_limit_i64 = i64::try_from(delete_limit.max(1)).unwrap_or(i64::MAX); - let deleted_1m = sqlx::query( - r#" -DELETE FROM proxy_node_metrics_1m -WHERE bucket_start_unix_secs < ? -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"#, - ) - .bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX)) - .bind(delete_limit_i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() as usize; - - let deleted_1h = sqlx::query( - r#" -DELETE FROM proxy_node_metrics_1h -WHERE bucket_start_unix_secs < ? -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"#, - ) - .bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX)) - .bind(delete_limit_i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() as usize; - - Ok(ProxyNodeMetricsCleanupSummary { - deleted_1m_rows: deleted_1m, - deleted_1h_rows: deleted_1h, - }) - } -} - -fn optional_unix_secs(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}")) - }) - }) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn duplicate_proxy_node_error(node: &StoredProxyNode) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "已存在相同地址的代理节点: {} ({}:{})", - node.name, node.ip, node.port - )) -} - -fn proxy_node_registration_matches( - current: &StoredProxyNode, - mutation: &ProxyNodeRegistrationMutation, - expected: &StoredProxyNode, - replacement_proxy_metadata: Option<&serde_json::Value>, - now: u64, -) -> bool { - current.id == expected.id - && current.tunnel_generation == expected.tunnel_generation - && !current.is_manual - && current.name == mutation.name - && current.ip == mutation.ip - && current.port == mutation.port - && current.region == mutation.region - && current.registered_by == mutation.registered_by - && current.last_heartbeat_at_unix_secs == Some(now) - && current.heartbeat_interval == mutation.heartbeat_interval - && mutation - .active_connections - .is_none_or(|value| current.active_connections == value) - && mutation - .total_requests - .is_none_or(|value| current.total_requests == value) - && mutation - .avg_latency_ms - .is_none_or(|value| current.avg_latency_ms == Some(value)) - && mutation - .hardware_info - .as_ref() - .is_none_or(|value| current.hardware_info.as_ref() == Some(value)) - && mutation - .estimated_max_concurrency - .is_none_or(|value| current.estimated_max_concurrency == Some(value)) - && current.tunnel_mode == mutation.tunnel_mode - && replacement_proxy_metadata - .is_none_or(|value| current.proxy_metadata.as_ref() == Some(value)) - && current.updated_at_unix_secs == Some(now) -} - -fn requested_proxy_node_id(value: Option<&str>) -> Result, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - if value.is_empty() || value.trim() != value { - return Err(DataLayerError::InvalidInput( - "proxy node id must be non-empty and unpadded".to_string(), - )); - } - Ok(Some(value.to_string())) -} - -fn proxy_node_registration_identity_error(requested_id: &str, existing_id: &str) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "proxy node registration identity changed: requested {requested_id}, existing {existing_id}" - )) -} - -fn proxy_node_registration_changed_error() -> DataLayerError { - DataLayerError::UnexpectedValue( - "registered proxy node identity changed during registration".to_string(), - ) -} - -fn proxy_node_id_in_use_error(node: &StoredProxyNode) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "proxy node id is already in use: {} ({}:{})", - node.id, node.ip, node.port - )) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn map_proxy_node_row(row: &MySqlRow) -> Result { - let tunnel_generation: String = row.try_get("tunnel_generation").map_sql_err()?; - if tunnel_generation.trim().is_empty() { - return Err(DataLayerError::UnexpectedValue( - "proxy_nodes.tunnel_generation must not be empty".to_string(), - )); - } - Ok(StoredProxyNode::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("ip").map_sql_err()?, - row.try_get("port").map_sql_err()?, - row.try_get("is_manual").map_sql_err()?, - row.try_get("status").map_sql_err()?, - row.try_get("heartbeat_interval").map_sql_err()?, - row.try_get("active_connections").map_sql_err()?, - row.try_get("total_requests").map_sql_err()?, - row.try_get("failed_requests").map_sql_err()?, - row.try_get("dns_failures").map_sql_err()?, - row.try_get("stream_errors").map_sql_err()?, - row.try_get("tunnel_mode").map_sql_err()?, - row.try_get("tunnel_connected").map_sql_err()?, - row.try_get("config_version").map_sql_err()?, - )? - .with_tunnel_generation(tunnel_generation) - .with_manual_proxy_fields( - row.try_get("proxy_url").map_sql_err()?, - row.try_get("proxy_username").map_sql_err()?, - row.try_get("proxy_password").map_sql_err()?, - ) - .with_runtime_fields( - row.try_get("region").map_sql_err()?, - row.try_get("registered_by").map_sql_err()?, - optional_unix_secs(row.try_get("last_heartbeat_at_unix_secs").map_sql_err()?), - row.try_get("avg_latency_ms").map_sql_err()?, - optional_json_from_string( - row.try_get("proxy_metadata").map_sql_err()?, - "proxy_nodes.proxy_metadata", - )?, - optional_json_from_string( - row.try_get("hardware_info").map_sql_err()?, - "proxy_nodes.hardware_info", - )?, - row.try_get("estimated_max_concurrency").map_sql_err()?, - optional_unix_secs(row.try_get("tunnel_connected_at_unix_secs").map_sql_err()?), - optional_json_from_string( - row.try_get("remote_config").map_sql_err()?, - "proxy_nodes.remote_config", - )?, - optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - optional_unix_secs(row.try_get("updated_at_unix_secs").map_sql_err()?), - )) -} - -fn map_proxy_node_event_row(row: &MySqlRow) -> Result { - Ok(StoredProxyNodeEvent { - id: row.try_get("id").map_sql_err()?, - node_id: row.try_get("node_id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - detail: row.try_get("detail").map_sql_err()?, - event_metadata: optional_json_from_string( - row.try_get("event_metadata").map_sql_err()?, - "proxy_node_events.event_metadata", - )?, - created_at_unix_ms: optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - }) -} - -fn map_proxy_node_metric_row( - row: &MySqlRow, -) -> Result { - Ok(StoredProxyNodeMetricsBucket { - node_id: row.try_get("node_id").map_sql_err()?, - bucket_start_unix_secs: optional_unix_secs( - row.try_get("bucket_start_unix_secs").map_sql_err()?, - ) - .unwrap_or_default(), - samples: row.try_get("samples").map_sql_err()?, - uptime_samples: row.try_get("uptime_samples").map_sql_err()?, - active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?, - active_connections_max: row.try_get("active_connections_max").map_sql_err()?, - heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?, - heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?, - connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?, - disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?, - error_events_delta: row.try_get("error_events_delta").map_sql_err()?, - ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?, - ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?, - ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?, - ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?, - }) -} - -fn map_proxy_fleet_metric_row( - row: &MySqlRow, -) -> Result { - Ok(StoredProxyFleetMetricsBucket { - bucket_start_unix_secs: optional_unix_secs( - row.try_get("bucket_start_unix_secs").map_sql_err()?, - ) - .unwrap_or_default(), - samples: row.try_get("samples").map_sql_err()?, - uptime_samples: row.try_get("uptime_samples").map_sql_err()?, - active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?, - active_connections_max: row.try_get("active_connections_max").map_sql_err()?, - heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?, - heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?, - connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?, - disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?, - error_events_delta: row.try_get("error_events_delta").map_sql_err()?, - ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?, - ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?, - ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?, - ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?, - }) -} - -#[cfg(test)] -mod tests { - use super::MysqlProxyNodeReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::proxy_nodes::{ - merge_proxy_metadata_for_registration, normalize_proxy_metadata, - ProxyNodeManualCreateMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation, - ProxyNodeWriteRepository, - }; - use serde_json::json; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlProxyNodeReadRepository::new(pool); - } - - #[test] - fn proxy_node_mutation_sql_is_atomic_and_field_scoped() { - assert!(super::APPLY_HEARTBEAT_SQL - .contains("total_requests = total_requests + GREATEST(COALESCE(?, 0), 0)")); - assert!(super::APPLY_HEARTBEAT_SQL - .contains("failed_requests = failed_requests + GREATEST(COALESCE(?, 0), 0)")); - assert!(!super::APPLY_HEARTBEAT_SQL.contains("remote_config =")); - assert!(!super::APPLY_HEARTBEAT_SQL.contains("config_version =")); - assert!(super::APPLY_HEARTBEAT_SQL.contains("BINARY tunnel_generation = BINARY ?")); - - assert!(super::UPDATE_TUNNEL_STATUS_SQL - .contains("tunnel_connected_at IS NULL OR tunnel_connected_at <= ?")); - assert!(super::UPDATE_TUNNEL_STATUS_SQL.contains("BINARY tunnel_generation = BINARY ?")); - assert!(super::RECORD_PROXY_NODE_TRAFFIC_SQL - .contains("total_requests = total_requests + GREATEST(?, 0)")); - assert!(super::RECORD_PROXY_NODE_TRAFFIC_SQL.contains("is_manual = 1")); - assert!(super::UPDATE_MANUAL_PROXY_NODE_SQL.contains("name = COALESCE(?, name)")); - assert!(!super::UPDATE_MANUAL_PROXY_NODE_SQL.contains("remote_config")); - assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL - .contains("BINARY id = BINARY ? AND BINARY tunnel_generation = BINARY ?")); - assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL - .contains("is_manual = 0 AND BINARY ip = BINARY ? AND port = ?")); - assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL - .contains("CAST(proxy_metadata AS JSON) = CAST(? AS JSON)")); - } - - #[tokio::test] - async fn proxy_metadata_cas_distinguishes_duplicate_array_elements_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 proxy metadata CAS test because AETHER_TEST_MYSQL_URL is unset" - ); - return; - }; - let pool = sqlx::mysql::MySqlPoolOptions::new() - .max_connections(2) - .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 repository = MysqlProxyNodeReadRepository::new(pool.clone()); - let node = repository - .create_manual_node(&ProxyNodeManualCreateMutation { - node_id: None, - name: format!("metadata-cas-{suffix}"), - ip: format!("metadata-cas-{suffix}"), - port: 1, - region: None, - proxy_url: "http://127.0.0.1:1".to_string(), - proxy_username: None, - proxy_password: None, - registered_by: None, - }) - .await - .expect("mysql proxy fixture should insert"); - let stored = json!({"nested": {"values": [1, 1]}}); - sqlx::query("UPDATE proxy_nodes SET proxy_metadata = ? WHERE id = ?") - .bind(serde_json::to_string(&stored).expect("stored metadata should serialize")) - .bind(&node.id) - .execute(&pool) - .await - .expect("mysql proxy metadata fixture should update"); - - let updated = repository - .compare_and_set_proxy_metadata( - &node.id, - &json!({"nested": {"values": [1]}}), - &json!({"replacement": true}), - ) - .await - .expect("mysql proxy metadata CAS should execute"); - let persisted = repository - .find_proxy_node(&node.id) - .await - .expect("mysql proxy fixture should read") - .expect("mysql proxy fixture should exist") - .proxy_metadata; - let cleanup = sqlx::query("DELETE FROM proxy_nodes WHERE id = ?") - .bind(&node.id) - .execute(&pool) - .await; - - assert!(!updated, "different JSON arrays must not compare equal"); - assert_eq!(persisted, Some(stored)); - cleanup.expect("mysql proxy fixture should clean up"); - } - - #[tokio::test] - async fn registration_cas_rejects_stale_security_snapshot() { - let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") - .ok() - .filter(|value| !value.trim().is_empty()) - else { - eprintln!( - "skipping mysql registration CAS test because AETHER_TEST_MYSQL_URL is unset" - ); - return; - }; - let pool = sqlx::mysql::MySqlPoolOptions::new() - .max_connections(2) - .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 node_id = format!("registration-security-{suffix}"); - let endpoint = format!("registration-security-{suffix}"); - let repository = MysqlProxyNodeReadRepository::new(pool.clone()); - let first = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some(node_id.clone()), - name: node_id.clone(), - ip: endpoint.clone(), - port: 7070, - region: None, - heartbeat_interval: 30, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({ - "tunnel_security": { - "mode": "non_tls_required", - "encryption_key_encrypted": "aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-old" - } - })), - proxy_version: Some("1.0.0".to_string()), - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("initial mysql registration should succeed"); - let stale = first.clone(); - - let rotated = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some(node_id.clone()), - name: node_id.clone(), - ip: endpoint.clone(), - port: 7070, - region: None, - heartbeat_interval: 30, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({ - "tunnel_security": { - "mode": "non_tls_required", - "encryption_key_encrypted": "aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new" - } - })), - proxy_version: Some("2.0.0".to_string()), - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("mysql security rotation should succeed"); - - let stale_refresh = ProxyNodeRegistrationMutation { - node_id: Some(node_id.clone()), - name: node_id.clone(), - ip: endpoint, - port: 7070, - region: None, - heartbeat_interval: 30, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({"runtime": "stale-writer"})), - proxy_version: Some("2.1.0".to_string()), - registered_by: None, - tunnel_mode: true, - }; - let stale_replacement = merge_proxy_metadata_for_registration( - stale.proxy_metadata.as_ref(), - normalize_proxy_metadata( - stale_refresh.proxy_metadata.as_ref(), - stale_refresh.proxy_version.as_deref(), - ), - ); - assert!(!repository - .update_existing_registration_if_unchanged( - &stale_refresh, - &stale, - stale_replacement.as_ref(), - super::current_unix_secs(), - ) - .await - .expect("stale mysql registration CAS should execute")); - - let committed = repository - .register_node(&ProxyNodeRegistrationMutation { - name: format!("{node_id}-committed"), - proxy_metadata: Some(json!({"runtime": "committed-after-rotation"})), - ..stale_refresh - }) - .await - .expect("mysql metadata refresh should merge current security state"); - let cleanup = sqlx::query("DELETE FROM proxy_nodes WHERE id = ?") - .bind(&node_id) - .execute(&pool) - .await; - - assert_eq!( - rotated - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new") - ); - assert_eq!( - committed - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new") - ); - assert_eq!( - committed - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.get("runtime")), - Some(&json!("committed-after-rotation")) - ); - cleanup.expect("mysql proxy fixture should clean up"); - } -} diff --git a/crates/aether-data/adapters/mysql/src/quota.rs b/crates/aether-data/adapters/mysql/src/quota.rs deleted file mode 100644 index 95fe21ac1..000000000 --- a/crates/aether-data/adapters/mysql/src/quota.rs +++ /dev/null @@ -1,174 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, Row}; - -use aether_data_contracts::repository::quota::{ - ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot, -}; -use aether_data_query::{DialectSql, SelectColumn, SelectQuery, SqlDialect}; - -use crate::error::SqlResultExt; -use crate::{DataLayerError, MysqlPool}; - -fn quota_snapshot_select() -> SelectQuery<'static> { - SelectQuery::new("providers").select_columns([ - SelectColumn::expr("id").alias("provider_id"), - SelectColumn::expr( - DialectSql::common("billing_type").with_postgres("CAST(billing_type AS TEXT)"), - ) - .alias("billing_type"), - SelectColumn::expr( - DialectSql::dialect( - "CAST(monthly_quota_usd AS DOUBLE PRECISION)", - "CAST(monthly_quota_usd AS REAL)", - ) - .with_mysql("CAST(monthly_quota_usd AS DOUBLE)"), - ) - .alias("monthly_quota_usd"), - SelectColumn::expr( - DialectSql::dialect( - "CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE PRECISION)", - "CAST(COALESCE(monthly_used_usd, 0) AS REAL)", - ) - .with_mysql("CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE)"), - ) - .alias("monthly_used_usd"), - SelectColumn::expr("quota_reset_day"), - SelectColumn::expr( - DialectSql::dialect( - "CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT)", - "quota_last_reset_at", - ) - .with_mysql("quota_last_reset_at"), - ) - .alias("quota_last_reset_at_unix_secs"), - SelectColumn::expr( - DialectSql::dialect( - "CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT)", - "quota_expires_at", - ) - .with_mysql("quota_expires_at"), - ) - .alias("quota_expires_at_unix_secs"), - SelectColumn::expr("is_active"), - ]) -} - -#[derive(Debug, Clone)] -pub struct MysqlProviderQuotaRepository { - pool: MysqlPool, -} - -impl MysqlProviderQuotaRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl ProviderQuotaReadRepository for MysqlProviderQuotaRepository { - async fn find_by_provider_id( - &self, - provider_id: &str, - ) -> Result, DataLayerError> { - let mut statement = quota_snapshot_select().statement::(SqlDialect::MySql); - statement.where_eq("id", provider_id.to_string()).limit(1); - let row = statement - .finish() - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_row).transpose() - } - - async fn find_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut statement = quota_snapshot_select().statement::(SqlDialect::MySql); - statement - .where_in("id", provider_ids) - .order_by_sql("id ASC"); - let rows = statement - .finish() - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_row).collect() - } -} - -#[async_trait] -impl ProviderQuotaWriteRepository for MysqlProviderQuotaRepository { - async fn reset_due(&self, now_unix_secs: u64) -> Result { - let now = i64::try_from(now_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("provider quota reset timestamp overflow".to_string()) - })?; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET monthly_used_usd = 0, - quota_last_reset_at = ?, - updated_at = ? -WHERE billing_type = 'monthly_quota' - AND is_active = 1 - AND ( - quota_last_reset_at IS NULL - OR (? - quota_last_reset_at) >= (quota_reset_day * 86400) - ) -"#, - ) - .bind(now) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -fn map_row(row: &MySqlRow) -> Result { - StoredProviderQuotaSnapshot::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("billing_type").map_sql_err()?, - row.try_get("monthly_quota_usd").map_sql_err()?, - row.try_get("monthly_used_usd").map_sql_err()?, - row.try_get("quota_reset_day").map_sql_err()?, - row.try_get("quota_last_reset_at_unix_secs").map_sql_err()?, - row.try_get("quota_expires_at_unix_secs").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - ) -} - -#[cfg(test)] -mod tests { - use super::{quota_snapshot_select, MysqlProviderQuotaRepository}; - use aether_data_query::SqlDialect; - - #[test] - fn quota_projection_renders_for_mysql() { - let sql = quota_snapshot_select().render(SqlDialect::MySql); - - assert!(sql.contains("id AS `provider_id`")); - assert!(sql.contains("CAST(monthly_quota_usd AS DOUBLE) AS `monthly_quota_usd`")); - assert!(sql.contains("quota_last_reset_at AS `quota_last_reset_at_unix_secs`")); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlProviderQuotaRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/routing_profiles.rs b/crates/aether-data/adapters/mysql/src/routing_profiles.rs deleted file mode 100644 index 6b602d91e..000000000 --- a/crates/aether-data/adapters/mysql/src/routing_profiles.rs +++ /dev/null @@ -1,518 +0,0 @@ -use async_trait::async_trait; -use serde_json::Value; -use sqlx::{mysql::MySqlRow, Acquire, Row}; - -use aether_data_contracts::repository::routing_profiles::*; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::pool::MysqlPool; - -const ROUTING_GROUP_SELECT: &str = r#" -SELECT - id, - name, - description, - enabled, - is_system_default, - sort_order, - config_json, - version, - created_at, - updated_at, - published_at -FROM routing_groups -"#; - -const ROUTING_GROUP_BINDING_SELECT: &str = r#" -SELECT - id, - group_id, - subject_type, - subject_id, - is_default, - allow_explicit_select, - created_at, - updated_at -FROM routing_group_bindings -"#; - -const ROUTING_GROUP_VERSION_SELECT: &str = r#" -SELECT - id, - group_id, - version, - config_json, - created_at, - created_by -FROM routing_group_versions -"#; - -#[derive(Debug, Clone)] -pub struct MysqlRoutingGroupRepository { - pool: MysqlPool, -} - -impl MysqlRoutingGroupRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl RoutingGroupReadRepository for MysqlRoutingGroupRepository { - async fn list_routing_groups(&self) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} ORDER BY enabled DESC, sort_order ASC, name ASC, id ASC" - )) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_group_row).collect() - } - - async fn find_routing_group( - &self, - lookup: RoutingGroupLookupKey<'_>, - ) -> Result, DataLayerError> { - let row = match lookup { - RoutingGroupLookupKey::Id(id) => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1" - )) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - RoutingGroupLookupKey::Name(name) => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE name = ? LIMIT 1" - )) - .bind(name) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE is_system_default = 1 AND enabled = 1 ORDER BY updated_at DESC, id ASC LIMIT 1" - )) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - }; - row.as_ref().map(map_group_row).transpose() - } - - async fn list_routing_group_bindings( - &self, - query: &RoutingGroupBindingQuery, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{ROUTING_GROUP_BINDING_SELECT} -WHERE (? IS NULL OR group_id = ?) - AND (? IS NULL OR subject_type = ?) - AND (? IS NULL OR subject_id = ?) -ORDER BY created_at ASC, id ASC -"# - )) - .bind(query.group_id.as_deref()) - .bind(query.group_id.as_deref()) - .bind(query.subject_type.map(binding_subject_to_database)) - .bind(query.subject_type.map(binding_subject_to_database)) - .bind(query.subject_id.as_deref()) - .bind(query.subject_id.as_deref()) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_binding_row).collect() - } - - async fn has_any_routing_group_binding(&self) -> Result { - let row = sqlx::query("SELECT 1 FROM routing_group_bindings LIMIT 1") - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(row.is_some()) - } - - async fn list_routing_group_versions( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = ? ORDER BY version DESC, created_at DESC, id ASC" - )) - .bind(group_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_version_row).collect() - } -} - -#[async_trait] -impl RoutingGroupWriteRepository for MysqlRoutingGroupRepository { - async fn create_routing_group( - &self, - record: CreateRoutingGroupRecord, - ) -> Result { - let group = StoredRoutingGroup::new(record)?; - let mut connection = self.pool.acquire().await.map_sql_err()?; - sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") - .execute(&mut *connection) - .await - .map_sql_err()?; - let mut tx = connection.begin().await.map_sql_err()?; - sqlx::query("SELECT id FROM routing_groups ORDER BY id FOR UPDATE") - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if group.is_system_default { - sqlx::query( - "UPDATE routing_groups SET is_system_default = 0 WHERE is_system_default = 1", - ) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -INSERT INTO routing_groups ( - id, name, description, enabled, is_system_default, sort_order, config_json, - version, created_at, updated_at, published_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&group.id) - .bind(&group.name) - .bind(&group.description) - .bind(group.enabled) - .bind(group.is_system_default) - .bind(group.sort_order) - .bind(json_to_string( - &group.config_json, - "routing_groups.config_json", - )?) - .bind(group.version) - .bind(group.created_at) - .bind(group.updated_at) - .bind(group.published_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(group) - } - - async fn update_routing_group( - &self, - id: &str, - patch: UpdateRoutingGroupRecord, - ) -> Result, DataLayerError> { - let mut connection = self.pool.acquire().await.map_sql_err()?; - sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") - .execute(&mut *connection) - .await - .map_sql_err()?; - let mut tx = connection.begin().await.map_sql_err()?; - sqlx::query("SELECT id FROM routing_groups ORDER BY id FOR UPDATE") - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let row = sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(mut group) = row.as_ref().map(map_group_row).transpose()? else { - return Ok(None); - }; - apply_group_patch(&mut group, patch)?; - if group.is_system_default { - sqlx::query( - "UPDATE routing_groups SET is_system_default = 0 WHERE is_system_default = 1 AND id <> ?", - ) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -UPDATE routing_groups -SET name = ?, - description = ?, - enabled = ?, - is_system_default = ?, - sort_order = ?, - config_json = ?, - version = ?, - updated_at = ?, - published_at = ? -WHERE id = ? -"#, - ) - .bind(&group.name) - .bind(&group.description) - .bind(group.enabled) - .bind(group.is_system_default) - .bind(group.sort_order) - .bind(json_to_string( - &group.config_json, - "routing_groups.config_json", - )?) - .bind(group.version) - .bind(group.updated_at) - .bind(group.published_at) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(Some(group)) - } - - async fn delete_routing_group(&self, id: &str) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM routing_group_versions WHERE group_id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - tx.commit().await.map_sql_err()?; - Ok(rows_affected > 0) - } - - async fn create_routing_group_binding( - &self, - record: CreateRoutingGroupBindingRecord, - ) -> Result { - let binding = StoredRoutingGroupBinding::new(record)?; - let mut connection = self.pool.acquire().await.map_sql_err()?; - sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") - .execute(&mut *connection) - .await - .map_sql_err()?; - let mut tx = connection.begin().await.map_sql_err()?; - sqlx::query("SELECT id FROM routing_group_bindings ORDER BY id FOR UPDATE") - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if binding.is_default { - sqlx::query( - r#" -UPDATE routing_group_bindings -SET is_default = 0 -WHERE is_default = 1 AND subject_type = ? AND subject_id = ? -"#, - ) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -INSERT INTO routing_group_bindings ( - id, group_id, subject_type, subject_id, is_default, - allow_explicit_select, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&binding.id) - .bind(&binding.group_id) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(binding.is_default) - .bind(binding.allow_explicit_select) - .bind(binding.created_at) - .bind(binding.updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(binding) - } - - async fn delete_routing_group_binding(&self, id: &str) -> Result { - Ok( - sqlx::query("DELETE FROM routing_group_bindings WHERE id = ?") - .bind(id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - > 0, - ) - } - - async fn update_routing_group_binding( - &self, - id: &str, - patch: UpdateRoutingGroupBindingRecord, - ) -> Result, DataLayerError> { - let mut connection = self.pool.acquire().await.map_sql_err()?; - sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") - .execute(&mut *connection) - .await - .map_sql_err()?; - let mut tx = connection.begin().await.map_sql_err()?; - sqlx::query("SELECT id FROM routing_group_bindings ORDER BY id FOR UPDATE") - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let row = sqlx::query(&format!( - "{ROUTING_GROUP_BINDING_SELECT} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(mut binding) = row.as_ref().map(map_binding_row).transpose()? else { - return Ok(None); - }; - apply_binding_patch(&mut binding, patch)?; - if binding.is_default { - sqlx::query( - r#" -UPDATE routing_group_bindings -SET is_default = 0 -WHERE is_default = 1 - AND subject_type = ? - AND subject_id = ? - AND id <> ? -"#, - ) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -UPDATE routing_group_bindings -SET group_id = ?, - subject_type = ?, - subject_id = ?, - is_default = ?, - allow_explicit_select = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&binding.group_id) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(binding.is_default) - .bind(binding.allow_explicit_select) - .bind(binding.updated_at) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(Some(binding)) - } - - async fn create_routing_group_version( - &self, - record: CreateRoutingGroupVersionRecord, - ) -> Result { - let version = StoredRoutingGroupVersion::new(record)?; - sqlx::query( - r#" -INSERT INTO routing_group_versions ( - id, group_id, version, config_json, created_at, created_by -) -VALUES (?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&version.id) - .bind(&version.group_id) - .bind(version.version) - .bind(json_to_string( - &version.config_json, - "routing_group_versions.config_json", - )?) - .bind(version.created_at) - .bind(&version.created_by) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(version) - } -} - -fn map_group_row(row: &MySqlRow) -> Result { - Ok(StoredRoutingGroup { - id: row.try_get("id").map_sql_err()?, - name: row.try_get("name").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - is_system_default: row.try_get("is_system_default").map_sql_err()?, - sort_order: row.try_get("sort_order").map_sql_err()?, - config_json: json_from_string( - row.try_get("config_json").map_sql_err()?, - "routing_groups.config_json", - )?, - version: row.try_get("version").map_sql_err()?, - created_at: row.try_get("created_at").map_sql_err()?, - updated_at: row.try_get("updated_at").map_sql_err()?, - published_at: row.try_get("published_at").map_sql_err()?, - }) -} - -fn map_binding_row(row: &MySqlRow) -> Result { - Ok(StoredRoutingGroupBinding { - id: row.try_get("id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - subject_type: binding_subject_from_database(row.try_get("subject_type").map_sql_err()?)?, - subject_id: row.try_get("subject_id").map_sql_err()?, - is_default: row.try_get("is_default").map_sql_err()?, - allow_explicit_select: row.try_get("allow_explicit_select").map_sql_err()?, - created_at: row.try_get("created_at").map_sql_err()?, - updated_at: row.try_get("updated_at").map_sql_err()?, - }) -} - -fn map_version_row(row: &MySqlRow) -> Result { - Ok(StoredRoutingGroupVersion { - id: row.try_get("id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - version: row.try_get("version").map_sql_err()?, - config_json: json_from_string( - row.try_get("config_json").map_sql_err()?, - "routing_group_versions.config_json", - )?, - created_at: row.try_get("created_at").map_sql_err()?, - created_by: row.try_get("created_by").map_sql_err()?, - }) -} - -fn json_to_string(value: &Value, field_name: &str) -> Result { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains unserializable JSON: {err}")) - }) -} - -fn json_from_string(value: String, field_name: &str) -> Result { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}")) - }) -} diff --git a/crates/aether-data/adapters/mysql/src/settlement.rs b/crates/aether-data/adapters/mysql/src/settlement.rs deleted file mode 100644 index adc141cd9..000000000 --- a/crates/aether-data/adapters/mysql/src/settlement.rs +++ /dev/null @@ -1,1604 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Acquire, Row}; - -use aether_data_contracts::repository::settlement::{ - finite_wallet_available_usd, plan_finite_wallet_debit, settlement_billable_cost_usd, - settlement_billing_status_for_usage_status, validate_wallet_settlement_values, - ReconcileUsagePolicyCostInput, ReleaseUsagePolicyRequestAdmissionInput, - ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput, - ReserveUsagePolicyRequestOutcome, SettlementWriteRepository, StoredUsagePolicyCostReservation, - StoredUsagePolicyRequestAdmission, StoredUsageSettlement, UsagePolicyCostReservationState, - UsagePolicyRequestAdmissionState, UsageSettlementInput, SETTLEMENT_EPSILON_USD, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#" -SELECT - usage_record.request_id, - COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id, - COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status) AS billing_status, - COALESCE( - usage_settlement_snapshots.wallet_balance_before, - usage_record.wallet_balance_before - ) AS wallet_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_balance_after, - usage_record.wallet_balance_after - ) AS wallet_balance_after, - COALESCE( - usage_settlement_snapshots.wallet_recharge_balance_before, - usage_record.wallet_recharge_balance_before - ) AS wallet_recharge_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_recharge_balance_after, - usage_record.wallet_recharge_balance_after - ) AS wallet_recharge_balance_after, - COALESCE( - usage_settlement_snapshots.wallet_gift_balance_before, - usage_record.wallet_gift_balance_before - ) AS wallet_gift_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_gift_balance_after, - usage_record.wallet_gift_balance_after - ) AS wallet_gift_balance_after, - usage_settlement_snapshots.provider_monthly_used_usd AS provider_monthly_used_usd, - usage_record.provider_id, - COALESCE(usage_settlement_snapshots.finalized_at, usage_record.finalized_at) AS finalized_at_unix_secs -FROM `usage` AS usage_record -LEFT JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = usage_record.request_id -WHERE usage_record.request_id = ? -FOR UPDATE -"#; - -const FINALIZE_USAGE_BILLING_SQL: &str = r#" -UPDATE `usage` -SET - billing_status = ?, - finalized_at = COALESCE(finalized_at, ?) -WHERE request_id = ? -"#; - -const UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL: &str = r#" -INSERT INTO usage_settlement_snapshots ( - request_id, - billing_status, - wallet_id, - wallet_balance_before, - wallet_balance_after, - wallet_recharge_balance_before, - wallet_recharge_balance_after, - wallet_gift_balance_before, - wallet_gift_balance_after, - provider_monthly_used_usd, - finalized_at, - created_at, - updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - billing_status = VALUES(billing_status), - wallet_id = COALESCE(VALUES(wallet_id), wallet_id), - wallet_balance_before = COALESCE(VALUES(wallet_balance_before), wallet_balance_before), - wallet_balance_after = COALESCE(VALUES(wallet_balance_after), wallet_balance_after), - wallet_recharge_balance_before = COALESCE( - VALUES(wallet_recharge_balance_before), - wallet_recharge_balance_before - ), - wallet_recharge_balance_after = COALESCE( - VALUES(wallet_recharge_balance_after), - wallet_recharge_balance_after - ), - wallet_gift_balance_before = COALESCE(VALUES(wallet_gift_balance_before), wallet_gift_balance_before), - wallet_gift_balance_after = COALESCE(VALUES(wallet_gift_balance_after), wallet_gift_balance_after), - provider_monthly_used_usd = COALESCE(VALUES(provider_monthly_used_usd), provider_monthly_used_usd), - finalized_at = COALESCE(VALUES(finalized_at), finalized_at), - updated_at = VALUES(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 MysqlSettlementRepository { - pool: MysqlPool, -} - -impl MysqlSettlementRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -fn usage_policy_cost_i64(value: u64, field: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds the integer range"))) -} - -fn usage_policy_cost_u64(value: i64, field: &str) -> Result { - u64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("{field} must not be negative"))) -} - -fn usage_policy_request_admission_from_mysql_row( - row: &MySqlRow, -) -> Result { - let state: String = row.try_get("state").map_sql_err()?; - Ok(StoredUsagePolicyRequestAdmission { - request_id: row.try_get("request_id").map_sql_err()?, - subject_id: row.try_get("subject_id").map_sql_err()?, - event_token: row.try_get("event_token").map_sql_err()?, - admitted_at_unix_secs: usage_policy_cost_u64( - row.try_get("admitted_at_unix_secs").map_sql_err()?, - "usage policy request admitted_at", - )?, - retain_until_unix_secs: usage_policy_cost_u64( - row.try_get("retain_until_unix_secs").map_sql_err()?, - "usage policy request retain_until", - )?, - state: UsagePolicyRequestAdmissionState::parse(&state).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "unknown usage policy request admission state {state}" - )) - })?, - released_at_unix_secs: row - .try_get::, _>("released_at_unix_secs") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy request released_at")) - .transpose()?, - }) -} - -const FIND_USAGE_POLICY_REQUEST_ADMISSION_MYSQL_SQL: &str = r#" -SELECT request_id, subject_id, event_token, - admitted_at AS admitted_at_unix_secs, - retain_until AS retain_until_unix_secs, - state, released_at AS released_at_unix_secs -FROM usage_request_admissions -WHERE event_token = ? -FOR UPDATE -"#; - -const USAGE_POLICY_REQUEST_TRANSACTION_ISOLATION_MYSQL_SQL: &str = - "SET TRANSACTION ISOLATION LEVEL READ COMMITTED"; - -fn usage_policy_cost_reservation_from_mysql_row( - row: &MySqlRow, -) -> Result { - let state: String = row.try_get("state").map_sql_err()?; - Ok(StoredUsagePolicyCostReservation { - request_id: row.try_get("request_id").map_sql_err()?, - subject_id: row.try_get("subject_id").map_sql_err()?, - reservation_token: row.try_get("reservation_token").map_sql_err()?, - admitted_at_unix_secs: usage_policy_cost_u64( - row.try_get("admitted_at_unix_secs").map_sql_err()?, - "usage policy admitted_at", - )?, - reserved_cost_units: usage_policy_cost_u64( - row.try_get("reserved_cost_units").map_sql_err()?, - "usage policy reserved_cost_units", - )?, - actual_cost_units: row - .try_get::, _>("actual_cost_units") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy actual_cost_units")) - .transpose()?, - state: UsagePolicyCostReservationState::parse(&state).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "unknown usage policy reservation state {state}" - )) - })?, - reservation_expires_at_unix_secs: usage_policy_cost_u64( - row.try_get("reservation_expires_at_unix_secs") - .map_sql_err()?, - "usage policy reservation_expires_at", - )?, - retain_until_unix_secs: usage_policy_cost_u64( - row.try_get("retain_until_unix_secs").map_sql_err()?, - "usage policy retain_until", - )?, - finalized_at_unix_secs: row - .try_get::, _>("finalized_at_unix_secs") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy finalized_at")) - .transpose()?, - }) -} - -const FIND_USAGE_POLICY_COST_RESERVATION_MYSQL_SQL: &str = r#" -SELECT - request_id, - subject_id, - reservation_token, - admitted_at AS admitted_at_unix_secs, - reserved_cost_units, - actual_cost_units, - state, - reservation_expires_at AS reservation_expires_at_unix_secs, - retain_until AS retain_until_unix_secs, - finalized_at AS finalized_at_unix_secs -FROM usage_cost_reservations -WHERE reservation_token = ? -FOR UPDATE -"#; - -async fn lock_usage_policy_subject_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - subject_id: &str, -) -> Result { - let exists = sqlx::query_scalar::<_, String>( - r#" -SELECT id -FROM users -WHERE id = ? -FOR UPDATE - "#, - ) - .bind(subject_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .is_some(); - Ok(exists) -} - -fn usage_policy_subject_missing() -> DataLayerError { - DataLayerError::InvalidInput("usage policy subject does not exist".to_string()) -} - -fn settlement_from_row(row: &MySqlRow) -> Result { - Ok(StoredUsageSettlement { - request_id: row.try_get("request_id").map_sql_err()?, - wallet_id: row.try_get("wallet_id").map_sql_err()?, - billing_status: row.try_get("billing_status").map_sql_err()?, - wallet_balance_before: row.try_get("wallet_balance_before").map_sql_err()?, - wallet_balance_after: row.try_get("wallet_balance_after").map_sql_err()?, - wallet_recharge_balance_before: row - .try_get("wallet_recharge_balance_before") - .map_sql_err()?, - wallet_recharge_balance_after: row - .try_get("wallet_recharge_balance_after") - .map_sql_err()?, - wallet_gift_balance_before: row.try_get("wallet_gift_balance_before").map_sql_err()?, - wallet_gift_balance_after: row.try_get("wallet_gift_balance_after").map_sql_err()?, - provider_monthly_used_usd: row.try_get("provider_monthly_used_usd").map_sql_err()?, - finalized_at_unix_secs: row - .try_get::, _>("finalized_at_unix_secs") - .map_sql_err()? - .map(|value| value as u64), - }) -} - -fn now_unix_secs() -> Result { - i64::try_from( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - ) - .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, - insufficient: bool, -} - -#[derive(Debug)] -struct DailyQuotaGrant { - entitlement_id: String, - daily_quota_usd: f64, - usage_date: String, - allow_wallet_overage: bool, -} - -fn daily_quota_usage_date( - reset_timezone: Option<&str>, - now: chrono::DateTime, -) -> Result { - let timezone = reset_timezone - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("Asia/Shanghai") - .parse::() - .map_err(|err| DataLayerError::InvalidInput(format!("invalid reset_timezone: {err}")))?; - Ok(now.with_timezone(&timezone).date_naive().to_string()) -} - -fn daily_quota_grants_from_entitlement( - entitlement_id: &str, - entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, - now: chrono::DateTime, -) -> Result, DataLayerError> { - let mut grants = Vec::new(); - let Some(items) = entitlements.as_array() else { - return Ok(grants); - }; - for item in items { - if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") { - continue; - } - let daily_quota_usd = item - .get("daily_quota_usd") - .and_then(serde_json::Value::as_f64) - .unwrap_or(0.0); - if !daily_quota_usd.is_finite() || daily_quota_usd <= 0.0 { - continue; - } - grants.push(DailyQuotaGrant { - entitlement_id: entitlement_id.to_string(), - daily_quota_usd, - usage_date: daily_quota_usage_date( - item.get("reset_timezone") - .and_then(serde_json::Value::as_str), - now, - )?, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), - }); - } - Ok(grants) -} - -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - -async fn consume_daily_quota_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - user_id: &str, - request_id: &str, - total_cost_usd: f64, - wallet_available_usd: Option, - wallet_can_overdraft: bool, - now_unix_secs: i64, -) -> Result { - if !total_cost_usd.is_finite() || total_cost_usd < 0.0 { - return Err(DataLayerError::InvalidInput( - "daily quota settlement cost must be finite and non-negative".to_string(), - )); - } - if total_cost_usd == 0.0 { - return Ok(DailyQuotaDebitResult::default()); - } - let rows = sqlx::query( - r#" -SELECT - user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json -FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id -WHERE user_plan_entitlements.user_id = ? - AND user_plan_entitlements.status = 'active' - AND user_plan_entitlements.starts_at <= ? - AND user_plan_entitlements.expires_at > ? -ORDER BY user_plan_entitlements.expires_at ASC, - user_plan_entitlements.created_at ASC, - user_plan_entitlements.id ASC -FOR UPDATE -"#, - ) - .bind(user_id) - .bind(now_unix_secs) - .bind(now_unix_secs) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - let now = chrono::Utc::now(); - let mut grants = Vec::new(); - for row in rows { - let entitlement_id: String = row.try_get("id").map_sql_err()?; - let entitlements_raw: String = row.try_get("entitlements_snapshot").map_sql_err()?; - let entitlements = - serde_json::from_str::(&entitlements_raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "user_plan_entitlements.entitlements_snapshot invalid json: {err}" - )) - })?; - let plan_entitlements_raw: String = row.try_get("plan_entitlements_json").map_sql_err()?; - let plan_entitlements = serde_json::from_str::(&plan_entitlements_raw) - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "billing_plans.entitlements_json invalid json: {err}" - )) - })?; - grants.extend(daily_quota_grants_from_entitlement( - &entitlement_id, - &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), - now, - )?); - } - if grants.is_empty() { - return Ok(DailyQuotaDebitResult::default()); - } - - let mut grants_with_remaining = Vec::new(); - let mut total_remaining = 0.0; - let mut allow_wallet_overage = true; - for grant in grants { - allow_wallet_overage &= grant.allow_wallet_overage; - let used = sqlx::query_scalar::<_, f64>( - r#" -SELECT COALESCE(SUM(amount_usd), 0) -FROM entitlement_usage_ledgers -WHERE user_entitlement_id = ? - AND usage_date = ? -"#, - ) - .bind(&grant.entitlement_id) - .bind(&grant.usage_date) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - if !used.is_finite() || used < 0.0 { - return Err(DataLayerError::UnexpectedValue( - "daily quota usage ledger total is invalid".to_string(), - )); - } - let remaining = (grant.daily_quota_usd - used).max(0.0); - total_remaining += remaining; - if !total_remaining.is_finite() { - return Err(DataLayerError::UnexpectedValue( - "daily quota remaining total overflowed".to_string(), - )); - } - grants_with_remaining.push((grant, remaining)); - } - let insufficient = (!allow_wallet_overage && total_remaining + 0.000_000_01 < total_cost_usd) - || (allow_wallet_overage - && !wallet_can_overdraft - && wallet_available_usd.is_some_and(|available| { - total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd - })); - - let mut remaining_cost = total_cost_usd; - let mut debited = 0.0; - for (grant, balance_before) in grants_with_remaining { - if remaining_cost <= 0.000_000_01 || balance_before <= 0.0 { - continue; - } - let amount = remaining_cost.min(balance_before); - let balance_after = balance_before - amount; - sqlx::query( - r#" -INSERT IGNORE INTO entitlement_usage_ledgers ( - id, user_entitlement_id, user_id, request_id, amount_usd, - balance_before, balance_after, usage_date, created_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&grant.entitlement_id) - .bind(user_id) - .bind(request_id) - .bind(amount) - .bind(balance_before) - .bind(balance_after) - .bind(&grant.usage_date) - .bind(now_unix_secs) - .execute(&mut **tx) - .await - .map_sql_err()?; - remaining_cost -= amount; - debited += amount; - } - Ok(DailyQuotaDebitResult { - debited_usd: debited, - insufficient, - }) -} - -#[async_trait] -impl SettlementWriteRepository for MysqlSettlementRepository { - async fn reserve_usage_policy_request( - &self, - input: ReserveUsagePolicyRequestInput, - ) -> Result { - input.validate()?; - let admitted_at = usage_policy_cost_i64( - input.admitted_at_unix_secs, - "usage policy request admitted_at", - )?; - let retain_until = usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy request retain_until", - )?; - let created_at = now_unix_secs()?; - // Different subjects lock different `users` rows. Under InnoDB's default REPEATABLE READ, - // two missing-token locking reads can retain compatible gap locks and then deadlock when - // both transactions try to insert the same unique event token. READ COMMITTED removes - // that gap-lock cycle while the subject row still serializes each subject's window count. - let mut connection = self.pool.acquire().await.map_sql_err()?; - sqlx::query(USAGE_POLICY_REQUEST_TRANSACTION_ISOLATION_MYSQL_SQL) - .execute(&mut *connection) - .await - .map_sql_err()?; - let mut tx = connection.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_mysql(&mut tx, &input.subject_id).await? { - return Err(usage_policy_subject_missing()); - } - - let existing_row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_MYSQL_SQL) - .bind(&input.event_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if let Some(row) = existing_row.as_ref() { - let existing = usage_policy_request_admission_from_mysql_row(row)?; - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Conflict); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy event_token must keep its original admitted_at".to_string(), - )); - } - sqlx::query( - "UPDATE usage_request_admissions SET retain_until = GREATEST(retain_until, ?) WHERE event_token = ?", - ) - .bind(retain_until) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - let outcome = match existing.state { - UsagePolicyRequestAdmissionState::Active => { - ReserveUsagePolicyRequestOutcome::Allowed - } - UsagePolicyRequestAdmissionState::Released => { - ReserveUsagePolicyRequestOutcome::AlreadyReleased - } - }; - tx.commit().await.map_sql_err()?; - return Ok(outcome); - } - - for (window_index, window) in input.windows.iter().enumerate() { - let used_requests = sqlx::query_scalar::<_, i64>( - r#" -SELECT CAST(COUNT(*) AS SIGNED) -FROM usage_request_admissions -WHERE subject_id = ? - AND state = 'active' - AND admitted_at >= ? - AND admitted_at < ? - "#, - ) - .bind(&input.subject_id) - .bind(usage_policy_cost_i64( - window.starts_at_unix_secs, - "usage policy request window start", - )?) - .bind(usage_policy_cost_i64( - window.ends_at_unix_secs, - "usage policy request window end", - )?) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let used_requests = - usage_policy_cost_u64(used_requests, "usage policy request used_requests")?; - if used_requests >= window.limit_requests { - let outcome = ReserveUsagePolicyRequestOutcome::Rejected { - window_index, - limit_requests: window.limit_requests, - used_requests, - }; - tx.commit().await.map_sql_err()?; - return Ok(outcome); - } - } - - sqlx::query( - r#" -INSERT INTO usage_request_admissions ( - request_id, subject_id, event_token, admitted_at, retain_until, - state, released_at, created_at -) VALUES (?, ?, ?, ?, ?, 'active', NULL, ?) -ON DUPLICATE KEY UPDATE event_token = VALUES(event_token) - "#, - ) - .bind(&input.request_id) - .bind(&input.subject_id) - .bind(&input.event_token) - .bind(admitted_at) - .bind(retain_until) - .bind(created_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_MYSQL_SQL) - .bind(&input.event_token) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let existing = usage_policy_request_admission_from_mysql_row(&row)?; - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Conflict); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy event_token must keep its original admitted_at".to_string(), - )); - } - sqlx::query( - "UPDATE usage_request_admissions SET retain_until = GREATEST(retain_until, ?) WHERE event_token = ?", - ) - .bind(retain_until) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - let outcome = match existing.state { - UsagePolicyRequestAdmissionState::Active => ReserveUsagePolicyRequestOutcome::Allowed, - UsagePolicyRequestAdmissionState::Released => { - ReserveUsagePolicyRequestOutcome::AlreadyReleased - } - }; - tx.commit().await.map_sql_err()?; - Ok(outcome) - } - - async fn release_usage_policy_request_admission( - &self, - input: ReleaseUsagePolicyRequestAdmissionInput, - ) -> Result, DataLayerError> { - input.validate()?; - let released_at = usage_policy_cost_i64( - input.released_at_unix_secs, - "usage policy request released_at", - )?; - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_mysql(&mut tx, &input.subject_id).await? { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - let row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_MYSQL_SQL) - .bind(&input.event_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - let mut admission = usage_policy_request_admission_from_mysql_row(&row)?; - if admission.request_id != input.request_id || admission.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - if input.released_at_unix_secs < admission.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy released_at must not precede admitted_at".to_string(), - )); - } - if admission.state == UsagePolicyRequestAdmissionState::Active { - sqlx::query( - "UPDATE usage_request_admissions SET state = 'released', released_at = ? WHERE event_token = ? AND state = 'active'", - ) - .bind(released_at) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - admission.state = UsagePolicyRequestAdmissionState::Released; - admission.released_at_unix_secs = Some(input.released_at_unix_secs); - } - tx.commit().await.map_sql_err()?; - Ok(Some(admission)) - } - - async fn cleanup_usage_policy_request_admissions( - &self, - now_unix_secs: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(0); - } - let now = usage_policy_cost_i64(now_unix_secs, "usage policy request cleanup timestamp")?; - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - let result = sqlx::query( - r#" -DELETE FROM usage_request_admissions -WHERE retain_until <= ? -ORDER BY retain_until, event_token -LIMIT ? - "#, - ) - .bind(now) - .bind(limit) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn reserve_usage_policy_cost( - &self, - input: ReserveUsagePolicyCostInput, - ) -> Result { - input.validate()?; - let admitted_at = - usage_policy_cost_i64(input.admitted_at_unix_secs, "usage policy admitted_at")?; - let reservation_expires_at = usage_policy_cost_i64( - input.reservation_expires_at_unix_secs, - "usage policy reservation_expires_at", - )?; - let updated_at = now_unix_secs()?; - - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_mysql(&mut tx, &input.subject_id).await? { - return Err(usage_policy_subject_missing()); - } - let existing_row = sqlx::query(FIND_USAGE_POLICY_COST_RESERVATION_MYSQL_SQL) - .bind(&input.reservation_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let existing = existing_row - .as_ref() - .map(usage_policy_cost_reservation_from_mysql_row) - .transpose()?; - if let Some(existing) = existing.as_ref() { - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyCostOutcome::Conflict); - } - if existing.state != UsagePolicyCostReservationState::Reserved { - let outcome = ReserveUsagePolicyCostOutcome::AlreadyTerminal { - state: existing.state, - }; - tx.commit().await.map_sql_err()?; - return Ok(outcome); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy reservation_token must keep its original admitted_at".to_string(), - )); - } - } - - let previous_reserved_cost_units = existing - .as_ref() - .map(|reservation| reservation.reserved_cost_units) - .unwrap_or(0); - let target_reserved_cost_units = - previous_reserved_cost_units.max(input.reserved_cost_units); - for (window_index, window) in input.windows.iter().enumerate() { - let window_start = - usage_policy_cost_i64(window.starts_at_unix_secs, "usage policy window start")?; - let window_end = - usage_policy_cost_i64(window.ends_at_unix_secs, "usage policy window end")?; - let used_cost_units = sqlx::query_scalar::<_, i64>( - r#" -SELECT CAST(COALESCE(SUM( - CASE - WHEN state = 'finalized' THEN COALESCE(actual_cost_units, 0) - WHEN state = 'reserved' AND reservation_expires_at > ? THEN reserved_cost_units - ELSE 0 - END -), 0) AS SIGNED) -FROM usage_cost_reservations -WHERE subject_id = ? - AND admitted_at >= ? - AND admitted_at < ? - AND reservation_token <> ? - "#, - ) - .bind(admitted_at) - .bind(&input.subject_id) - .bind(window_start) - .bind(window_end) - .bind(&input.reservation_token) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let used_cost_units = - usage_policy_cost_u64(used_cost_units, "usage policy used_cost_units")?; - if used_cost_units - .checked_add(target_reserved_cost_units) - .is_none_or(|total| total > window.limit_cost_units) - { - let outcome = ReserveUsagePolicyCostOutcome::Rejected { - window_index, - limit_cost_units: window.limit_cost_units, - used_cost_units, - }; - tx.commit().await.map_sql_err()?; - return Ok(outcome); - } - } - - let admitted_at = existing - .as_ref() - .map(|reservation| { - usage_policy_cost_i64( - reservation.admitted_at_unix_secs, - "usage policy admitted_at", - ) - }) - .transpose()? - .unwrap_or(admitted_at); - sqlx::query( - r#" -INSERT INTO usage_cost_reservations ( - request_id, subject_id, reservation_token, admitted_at, - reserved_cost_units, actual_cost_units, - state, reservation_expires_at, retain_until, finalized_at, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, NULL, 'reserved', ?, ?, NULL, ?, ?) -ON DUPLICATE KEY UPDATE - reserved_cost_units = GREATEST(reserved_cost_units, VALUES(reserved_cost_units)), - reservation_expires_at = GREATEST( - reservation_expires_at, - VALUES(reservation_expires_at) - ), - retain_until = GREATEST(retain_until, VALUES(retain_until)), - updated_at = VALUES(updated_at) - "#, - ) - .bind(&input.request_id) - .bind(&input.subject_id) - .bind(&input.reservation_token) - .bind(admitted_at) - .bind(usage_policy_cost_i64( - target_reserved_cost_units, - "usage policy reserved_cost_units", - )?) - .bind(reservation_expires_at) - .bind(usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy retain_until", - )?) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - Ok(ReserveUsagePolicyCostOutcome::Allowed { - reserved_cost_units: target_reserved_cost_units, - additional_reserved_cost_units: target_reserved_cost_units - .saturating_sub(previous_reserved_cost_units), - }) - } - - async fn reconcile_usage_policy_cost( - &self, - input: ReconcileUsagePolicyCostInput, - ) -> Result, DataLayerError> { - input.validate()?; - let actual_cost_units = - usage_policy_cost_i64(input.actual_cost_units, "usage policy actual_cost_units")?; - let finalized_at = - usage_policy_cost_i64(input.finalized_at_unix_secs, "usage policy finalized_at")?; - let updated_at = now_unix_secs()?; - - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_mysql(&mut tx, &input.subject_id).await? { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - let row = sqlx::query(FIND_USAGE_POLICY_COST_RESERVATION_MYSQL_SQL) - .bind(&input.reservation_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - let mut reservation = usage_policy_cost_reservation_from_mysql_row(&row)?; - if reservation.request_id != input.request_id || reservation.subject_id != input.subject_id - { - // The token selects the row; audit identity must still match before the reservation - // can be finalized. - tx.commit().await.map_sql_err()?; - return Ok(None); - } - if reservation.state == UsagePolicyCostReservationState::Reserved { - sqlx::query( - r#" -UPDATE usage_cost_reservations -SET state = ?, - actual_cost_units = ?, - finalized_at = ?, - updated_at = ? -WHERE reservation_token = ? - AND request_id = ? - AND subject_id = ? - AND state = 'reserved' - "#, - ) - .bind(input.terminal_state.as_str()) - .bind(actual_cost_units) - .bind(finalized_at) - .bind(updated_at) - .bind(&input.reservation_token) - .bind(&input.request_id) - .bind(&input.subject_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - reservation.state = input.terminal_state; - reservation.actual_cost_units = Some(input.actual_cost_units); - reservation.finalized_at_unix_secs = Some(input.finalized_at_unix_secs); - } - - tx.commit().await.map_sql_err()?; - Ok(Some(reservation)) - } - - async fn cleanup_usage_policy_cost_reservations( - &self, - now_unix_secs: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(0); - } - let now = usage_policy_cost_i64(now_unix_secs, "usage policy cleanup timestamp")?; - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - let result = sqlx::query( - r#" -DELETE FROM usage_cost_reservations -WHERE retain_until <= ? -ORDER BY retain_until, reservation_token -LIMIT ? - "#, - ) - .bind(now) - .bind(limit) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn settle_usage( - &self, - input: UsageSettlementInput, - ) -> Result, DataLayerError> { - input.validate()?; - let finalized_at = i64::try_from( - input - .finalized_at_unix_secs - .unwrap_or(now_unix_secs()? as u64), - ) - .map_err(|_| DataLayerError::InvalidInput("finalized_at overflow".to_string()))?; - let updated_at = now_unix_secs()?; - - let mut tx = self.pool.begin().await.map_sql_err()?; - let row = sqlx::query(FIND_USAGE_FOR_SETTLEMENT_SQL) - .bind(&input.request_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - - let Some(usage_row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let current_billing_status: String = usage_row.try_get("billing_status").map_sql_err()?; - if matches!( - current_billing_status.as_str(), - "settled" | "void" | "insufficient_quota" - ) { - let settlement = settlement_from_row(&usage_row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - let mut final_billing_status = - settlement_billing_status_for_usage_status(&input.status).to_string(); - let mut settlement = StoredUsageSettlement { - request_id: input.request_id.clone(), - wallet_id: None, - billing_status: final_billing_status.clone(), - wallet_balance_before: None, - wallet_balance_after: None, - wallet_recharge_balance_before: None, - wallet_recharge_balance_after: None, - wallet_gift_balance_before: None, - wallet_gift_balance_after: None, - provider_monthly_used_usd: None, - finalized_at_unix_secs: Some(finalized_at as u64), - }; - - if final_billing_status == "settled" { - let api_key_id = input - .api_key_id - .as_deref() - .filter(|value| !value.is_empty()); - let api_key_is_standalone = if input.api_key_is_standalone { - true - } else if let Some(api_key_id) = api_key_id { - sqlx::query_scalar::<_, bool>( - r#" -SELECT is_standalone -FROM api_keys -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .unwrap_or(false) - } else { - false - }; - - let wallet_row = if let Some(api_key_id) = api_key_id { - sqlx::query( - r#" -SELECT id, balance, gift_balance, total_consumed, limit_mode -FROM wallets -WHERE api_key_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - None - }; - - let wallet_row = if wallet_row.is_some() { - wallet_row - } else if !api_key_is_standalone { - if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) { - sqlx::query( - r#" -SELECT id, balance, gift_balance, total_consumed, limit_mode -FROM wallets -WHERE user_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - None - } - } else { - None - }; - - let wallet_can_overdraft = wallet_row.is_some(); - let wallet_available_usd = match wallet_row.as_ref() { - Some(row) => { - let recharge_balance: f64 = row.try_get("balance").map_sql_err()?; - let gift_balance: f64 = row.try_get("gift_balance").map_sql_err()?; - let total_consumed: f64 = row.try_get("total_consumed").map_sql_err()?; - validate_wallet_settlement_values( - recharge_balance, - gift_balance, - total_consumed, - 0.0, - )?; - let limit_mode: String = row.try_get("limit_mode").map_sql_err()?; - if limit_mode.eq_ignore_ascii_case("unlimited") { - None - } else { - Some(finite_wallet_available_usd(recharge_balance, gift_balance)) - } - } - None => Some(0.0), - }; - if let Some(row) = wallet_row.as_ref() { - let wallet_id: String = row.try_get("id").map_sql_err()?; - let before_recharge: f64 = row.try_get("balance").map_sql_err()?; - let before_gift: f64 = row.try_get("gift_balance").map_sql_err()?; - let before_total = before_recharge + before_gift; - settlement.wallet_id = Some(wallet_id); - settlement.wallet_balance_before = Some(before_total); - settlement.wallet_balance_after = Some(before_total); - settlement.wallet_recharge_balance_before = Some(before_recharge); - settlement.wallet_recharge_balance_after = Some(before_recharge); - settlement.wallet_gift_balance_before = Some(before_gift); - settlement.wallet_gift_balance_after = Some(before_gift); - } - - let billable_cost_usd = settlement_billable_cost_usd(&input); - let wallet_debit_cost_usd = if !api_key_is_standalone { - if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) { - let quota = consume_daily_quota_mysql( - &mut tx, - user_id, - &input.request_id, - billable_cost_usd, - wallet_available_usd, - wallet_can_overdraft, - updated_at, - ) - .await?; - if quota.insufficient { - final_billing_status = "insufficient_quota".to_string(); - settlement.billing_status = final_billing_status.clone(); - 0.0 - } else { - (billable_cost_usd - quota.debited_usd).max(0.0) - } - } else { - billable_cost_usd - } - } else { - billable_cost_usd - }; - if final_billing_status != "settled" { - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - if wallet_debit_cost_usd > SETTLEMENT_EPSILON_USD { - if let Some(wallet_row) = wallet_row { - let wallet_id: String = wallet_row.try_get("id").map_sql_err()?; - let before_recharge: f64 = wallet_row.try_get("balance").map_sql_err()?; - let before_gift: f64 = wallet_row.try_get("gift_balance").map_sql_err()?; - let total_consumed: f64 = wallet_row.try_get("total_consumed").map_sql_err()?; - let limit_mode: String = wallet_row.try_get("limit_mode").map_sql_err()?; - let before_total = before_recharge + before_gift; - let mut after_recharge = before_recharge; - let mut after_gift = before_gift; - if !limit_mode.eq_ignore_ascii_case("unlimited") { - let debit_plan = plan_finite_wallet_debit( - before_recharge, - before_gift, - wallet_debit_cost_usd, - ); - (after_recharge, after_gift) = - debit_plan.after_balances(before_recharge, before_gift); - } - let total_consumed_after = total_consumed + wallet_debit_cost_usd; - validate_wallet_settlement_values( - after_recharge, - after_gift, - total_consumed_after, - 0.0, - )?; - if final_billing_status == "settled" { - sqlx::query( - r#" -UPDATE wallets -SET - balance = ?, - gift_balance = ?, - total_consumed = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(total_consumed_after) - .bind(updated_at) - .bind(&wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - settlement.wallet_id = Some(wallet_id); - settlement.wallet_balance_before = Some(before_total); - settlement.wallet_balance_after = Some(after_recharge + after_gift); - settlement.wallet_recharge_balance_before = Some(before_recharge); - settlement.wallet_recharge_balance_after = Some(after_recharge); - settlement.wallet_gift_balance_before = Some(before_gift); - settlement.wallet_gift_balance_after = Some(after_gift); - } else { - final_billing_status = "insufficient_quota".to_string(); - settlement.billing_status = final_billing_status.clone(); - } - } - - if final_billing_status != "settled" { - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - if let Some(provider_id) = input - .provider_id - .as_deref() - .filter(|value| !value.is_empty()) - { - enqueue_provider_monthly_usage_delta_mysql( - &mut tx, - &input.request_id, - provider_id, - input.actual_total_cost_usd, - updated_at, - ) - .await?; - } - } - - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - Ok(Some(settlement)) - } -} - -#[cfg(test)] -mod tests { - use super::{MysqlSettlementRepository, USAGE_POLICY_REQUEST_TRANSACTION_ISOLATION_MYSQL_SQL}; - use crate::run_migrations; - use aether_data_contracts::repository::settlement::{ - ReserveUsagePolicyRequestInput, ReserveUsagePolicyRequestOutcome, - SettlementWriteRepository, UsagePolicyRequestWindow, UsageSettlementInput, - }; - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlSettlementRepository::new(pool); - } - - #[test] - fn request_admission_transactions_use_read_committed() { - assert_eq!( - USAGE_POLICY_REQUEST_TRANSACTION_ISOLATION_MYSQL_SQL, - "SET TRANSACTION ISOLATION LEVEL READ COMMITTED" - ); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn cross_subject_same_token_is_allowed_once_without_deadlock_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 request admission race test because AETHER_TEST_MYSQL_URL is unset" - ); - return; - }; - let pool = sqlx::mysql::MySqlPoolOptions::new() - .max_connections(2) - .connect(&database_url) - .await - .expect("mysql pool should connect"); - run_migrations(&pool) - .await - .expect("mysql migrations should run"); - cleanup_request_admission_rows(&pool).await; - sqlx::query( - r#" -INSERT INTO users (id, username, auth_source, created_at, updated_at) -VALUES - ('admission-race-user-1', 'admission-race-user-1', 'local', 1, 1), - ('admission-race-user-2', 'admission-race-user-2', 'local', 1, 1) - "#, - ) - .execute(&pool) - .await - .expect("race users should seed"); - - let repository = MysqlSettlementRepository::new(pool.clone()); - let reserve = |request_id: &str, subject_id: &str| ReserveUsagePolicyRequestInput { - request_id: request_id.to_string(), - subject_id: subject_id.to_string(), - event_token: "admission-race-token".to_string(), - admitted_at_unix_secs: 100, - retain_until_unix_secs: 1_000, - windows: vec![UsagePolicyRequestWindow { - starts_at_unix_secs: 0, - ends_at_unix_secs: 1_000, - limit_requests: 10, - }], - }; - let (first, second) = tokio::join!( - repository.reserve_usage_policy_request(reserve( - "admission-race-request-1", - "admission-race-user-1" - )), - repository.reserve_usage_policy_request(reserve( - "admission-race-request-2", - "admission-race-user-2" - )) - ); - let mut outcomes = vec![ - first.expect("first reserve should not deadlock"), - second.expect("second reserve should not deadlock"), - ]; - outcomes.sort_by_key(|outcome| match outcome { - ReserveUsagePolicyRequestOutcome::Allowed => 0, - ReserveUsagePolicyRequestOutcome::Conflict => 1, - _ => 2, - }); - assert_eq!( - outcomes, - vec![ - ReserveUsagePolicyRequestOutcome::Allowed, - ReserveUsagePolicyRequestOutcome::Conflict, - ] - ); - - cleanup_request_admission_rows(&pool).await; - } - - #[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"); - } - } - - async fn cleanup_request_admission_rows(pool: &sqlx::MySqlPool) { - sqlx::query( - "DELETE FROM usage_request_admissions WHERE event_token = 'admission-race-token'", - ) - .execute(pool) - .await - .expect("admission race row cleanup should succeed"); - sqlx::query( - "DELETE FROM users WHERE id IN ('admission-race-user-1', 'admission-race-user-2')", - ) - .execute(pool) - .await - .expect("admission race user cleanup should succeed"); - } -} diff --git a/crates/aether-data/adapters/mysql/src/usage.rs b/crates/aether-data/adapters/mysql/src/usage.rs deleted file mode 100644 index 2b851f1bf..000000000 --- a/crates/aether-data/adapters/mysql/src/usage.rs +++ /dev/null @@ -1,1873 +0,0 @@ -use std::collections::{BTreeMap, HashSet}; -use std::time::{SystemTime, UNIX_EPOCH}; - -use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::usage::{ - sanitize_usage_capture_controls_for_persistence, sanitize_usage_for_persistence, - sanitize_usage_request_metadata, usage_can_recover_terminal_failure, - usage_error_category_for_status_code, usage_lifecycle_update_allowed, - usage_request_metadata_client_family, PendingUsageCleanupSummary, StoredRequestUsageAudit, - StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardSummary, - StoredUsageUserTotals, UpsertUsageRecord, UsageCleanupExecutionMode, UsageCleanupPreviewCounts, - UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery, - UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageWriteRepository, -}; -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, - `usage`.request_id, - user_id, - api_key_id, - `usage`.username, - `usage`.api_key_name, - provider_name, - model, - target_model, - 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, - endpoint_kind, - endpoint_api_format, - provider_api_family, - provider_endpoint_kind, - 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, - COALESCE(usage_settlement_snapshots.billing_output_tokens, `usage`.output_tokens, 0) - AS output_tokens, - total_tokens, - 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, - 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, - 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#" -INSERT INTO `usage` ( - request_id, - 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, - input_tokens, - 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, - status_code, - error_message, - error_category, - response_time_ms, - first_byte_time_ms, - status, - billing_status, - request_metadata, - candidate_id, - candidate_index, - key_name, - planner_kind, - route_family, - route_kind, - execution_path, - local_execution_runtime_miss_reason, - finalized_at, - created_at_unix_ms, - updated_at_unix_secs -) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ? -) -ON DUPLICATE KEY UPDATE - user_id = VALUES(user_id), - api_key_id = VALUES(api_key_id), - provider_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_name ELSE VALUES(provider_name) END, - model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN model ELSE VALUES(model) END, - target_model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN target_model ELSE VALUES(target_model) END, - provider_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_id ELSE VALUES(provider_id) END, - provider_endpoint_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_id ELSE VALUES(provider_endpoint_id) END, - provider_api_key_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_key_id ELSE VALUES(provider_api_key_id) END, - request_type = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_type ELSE VALUES(request_type) END, - api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_format ELSE VALUES(api_format) END, - api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_family ELSE VALUES(api_family) END, - endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_kind ELSE VALUES(endpoint_kind) END, - endpoint_api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_api_format ELSE VALUES(endpoint_api_format) END, - provider_api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_family ELSE VALUES(provider_api_family) END, - provider_endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_kind ELSE VALUES(provider_endpoint_kind) END, - has_format_conversion = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN has_format_conversion ELSE VALUES(has_format_conversion) END, - is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN is_stream ELSE VALUES(is_stream) END, - upstream_is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN upstream_is_stream ELSE VALUES(upstream_is_stream) END, - input_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN input_tokens - ELSE VALUES(input_tokens) - END, - output_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_tokens - ELSE VALUES(output_tokens) - END, - total_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_tokens - ELSE VALUES(total_tokens) - END, - cache_creation_input_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_input_tokens - ELSE VALUES(cache_creation_input_tokens) - END, - cache_creation_ephemeral_5m_input_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_5m_input_tokens - ELSE VALUES(cache_creation_ephemeral_5m_input_tokens) - END, - cache_creation_ephemeral_1h_input_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_1h_input_tokens - ELSE VALUES(cache_creation_ephemeral_1h_input_tokens) - END, - cache_read_input_tokens = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_input_tokens - ELSE VALUES(cache_read_input_tokens) - END, - cache_creation_cost_usd = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_cost_usd - ELSE VALUES(cache_creation_cost_usd) - END, - cache_read_cost_usd = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_cost_usd - ELSE VALUES(cache_read_cost_usd) - END, - output_price_per_1m = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_price_per_1m - ELSE VALUES(output_price_per_1m) - END, - total_cost_usd = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_cost_usd - ELSE VALUES(total_cost_usd) - END, - actual_total_cost_usd = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN actual_total_cost_usd - ELSE VALUES(actual_total_cost_usd) - END, - status_code = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status_code - WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status_code - WHEN status = 'streaming' AND VALUES(status) = 'streaming' AND VALUES(status_code) IS NULL THEN status_code - ELSE VALUES(status_code) - END, - error_message = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_message - WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_message - ELSE VALUES(error_message) - END, - error_category = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_category - WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_category - ELSE VALUES(error_category) - END, - response_time_ms = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN response_time_ms - WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0 - THEN COALESCE(response_time_ms, VALUES(response_time_ms)) - ELSE VALUES(response_time_ms) - END, - first_byte_time_ms = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN first_byte_time_ms - WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0 - THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms)) - ELSE VALUES(first_byte_time_ms) - END, - billing_status = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN billing_status - ELSE VALUES(billing_status) - END, - request_metadata = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_metadata ELSE VALUES(request_metadata) END, - candidate_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_id ELSE VALUES(candidate_id) END, - candidate_index = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_index ELSE VALUES(candidate_index) END, - key_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN key_name ELSE VALUES(key_name) END, - planner_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN planner_kind ELSE VALUES(planner_kind) END, - route_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_family ELSE VALUES(route_family) END, - route_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_kind ELSE VALUES(route_kind) END, - execution_path = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN execution_path ELSE VALUES(execution_path) END, - local_execution_runtime_miss_reason = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN local_execution_runtime_miss_reason ELSE VALUES(local_execution_runtime_miss_reason) END, - finalized_at = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN finalized_at - ELSE VALUES(finalized_at) - END, - updated_at_unix_secs = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN updated_at_unix_secs - ELSE VALUES(updated_at_unix_secs) - END, - status = CASE - WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status - WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status - ELSE VALUES(status) - END -"#; - -const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#" -SELECT - `usage`.request_id, - COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) AS billing_status -FROM `usage` -LEFT JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = `usage`.request_id -WHERE `usage`.status IN ('pending', 'streaming') - AND `usage`.created_at_unix_ms < ? -ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC -LIMIT ? -"#; - -const SELECT_COMPLETED_REQUEST_CANDIDATES_SQL: &str = r#" -SELECT status, extra_data -FROM request_candidates -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, -} - -#[derive(Debug, Clone)] -pub struct MysqlUsageStorage { - pool: MysqlPool, -} - -impl MysqlUsageStorage { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - pub async fn read_usage_counter_health( - &self, - ) -> Result - { - counters::read_health(&self.pool).await - } - - 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( - &self, - created_from_unix_secs: u64, - created_until_unix_secs: u64, - user_id: Option<&str>, - ) -> Result, DataLayerError> { - 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({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"); - } - sql.push_str("GROUP BY date ORDER BY date ASC"); - - let mut query = sqlx::query(&sql) - .bind(to_i64(created_from_unix_secs, "usage.created_at_unix_ms")?) - .bind(to_i64(created_until_unix_secs, "usage.created_at_unix_ms")?); - if let Some(user_id) = user_id { - query = query.bind(user_id.to_string()); - } - let rows = query.fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_mysql_usage_daily_summary).collect() - } - - async fn summarize_usage_daily_heatmap_from_daily_aggregates( - &self, - created_from_unix_secs: u64, - created_until_unix_secs: u64, - user_id: Option<&str>, - ) -> Result, DataLayerError> { - let rows = if let Some(user_id) = user_id { - sqlx::query( - r#" -SELECT - DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, - total_requests AS requests, - input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, - total_cost AS total_cost_usd, - total_cost AS actual_total_cost_usd -FROM stats_user_daily -WHERE user_id = ? - AND `date` >= ? - AND `date` < ? - AND total_requests > 0 -ORDER BY `date` ASC -"#, - ) - .bind(user_id) - .bind(to_i64(created_from_unix_secs, "stats_user_daily.date")?) - .bind(to_i64(created_until_unix_secs, "stats_user_daily.date")?) - .fetch_all(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, - total_requests AS requests, - input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, - total_cost AS total_cost_usd, - actual_total_cost AS actual_total_cost_usd -FROM stats_daily -WHERE `date` >= ? - AND `date` < ? - AND total_requests > 0 -ORDER BY `date` ASC -"#, - ) - .bind(to_i64(created_from_unix_secs, "stats_daily.date")?) - .bind(to_i64(created_until_unix_secs, "stats_daily.date")?) - .fetch_all(&self.pool) - .await - .map_sql_err()? - }; - - rows.iter().map(map_mysql_usage_daily_summary).collect() - } - - pub async fn summarize_usage_daily_heatmap( - &self, - query: &UsageDailyHeatmapQuery, - ) -> Result, DataLayerError> { - let created_until_unix_secs = usage_current_unix_secs().saturating_add(1); - let user_id = query.user_id.as_deref(); - let mut summaries = BTreeMap::::new(); - - for item in self - .summarize_usage_daily_heatmap_from_daily_aggregates( - query.created_from_unix_secs, - created_until_unix_secs, - user_id, - ) - .await? - { - summaries.insert(item.date.clone(), item); - } - for item in self - .summarize_usage_daily_heatmap_raw_from_range( - query.created_from_unix_secs, - created_until_unix_secs, - user_id, - ) - .await? - { - summaries.entry(item.date.clone()).or_insert(item); - } - - Ok(summaries.into_values().collect()) - } - - pub async fn summarize_dashboard_usage_from_daily_aggregates( - &self, - query: &UsageDashboardSummaryQuery, - ) -> Result, DataLayerError> { - let row = if let Some(user_id) = query.user_id.as_deref() { - sqlx::query( - r#" -SELECT - CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests, - CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens, - CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens, - CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens, - CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, - CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens, - CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens, - CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context, - CAST(0.0 AS DOUBLE) AS cache_creation_cost_usd, - CAST(0.0 AS DOUBLE) AS cache_read_cost_usd, - CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, - CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd, - CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests, - CAST(0.0 AS DOUBLE) AS response_time_sum_ms, - CAST(0 AS SIGNED) AS response_time_samples -FROM stats_user_daily -WHERE user_id = ? - AND `date` >= ? - AND `date` < ? -"#, - ) - .bind(user_id) - .bind(to_i64( - query.created_from_unix_secs, - "stats_user_daily.date", - )?) - .bind(to_i64( - query.created_until_unix_secs, - "stats_user_daily.date", - )?) - .fetch_one(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests, - CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens, - CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens, - CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens, - CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, - CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens, - CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens, - CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context, - CAST(COALESCE(SUM(COALESCE(cache_creation_cost, 0)), 0) AS DOUBLE) AS cache_creation_cost_usd, - CAST(COALESCE(SUM(COALESCE(cache_read_cost, 0)), 0) AS DOUBLE) AS cache_read_cost_usd, - CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, - CAST(COALESCE(SUM(COALESCE(actual_total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd, - CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests, - CAST(0.0 AS DOUBLE) AS response_time_sum_ms, - CAST(0 AS SIGNED) AS response_time_samples -FROM stats_daily -WHERE `date` >= ? - AND `date` < ? -"#, - ) - .bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?) - .bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?) - .fetch_one(&self.pool) - .await - .map_sql_err()? - }; - - let total_requests = row_u64(&row, "total_requests")?; - if total_requests == 0 { - return Ok(None); - } - - Ok(Some(StoredUsageDashboardSummary { - total_requests, - input_tokens: row_u64(&row, "input_tokens")?, - effective_input_tokens: row_u64(&row, "effective_input_tokens")?, - output_tokens: row_u64(&row, "output_tokens")?, - total_tokens: row_u64(&row, "total_tokens")?, - cache_creation_tokens: row_u64(&row, "cache_creation_tokens")?, - cache_read_tokens: row_u64(&row, "cache_read_tokens")?, - total_input_context: row_u64(&row, "total_input_context")?, - cache_creation_cost_usd: row.try_get("cache_creation_cost_usd").map_sql_err()?, - cache_read_cost_usd: row.try_get("cache_read_cost_usd").map_sql_err()?, - total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, - actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?, - error_requests: row_u64(&row, "error_requests")?, - response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?, - response_time_samples: row_u64(&row, "response_time_samples")?, - })) - } - - pub async fn list_dashboard_daily_breakdown_from_daily_aggregates( - &self, - query: &UsageDashboardDailyBreakdownQuery, - ) -> Result, DataLayerError> { - let rows = if let Some(user_id) = query.user_id.as_deref() { - sqlx::query( - r#" -SELECT - DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, - 'aggregate' AS model, - 'aggregate' AS provider, - CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests, - CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, - CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, - CAST(0.0 AS DOUBLE) AS response_time_sum_ms, - CAST(0 AS SIGNED) AS response_time_samples -FROM stats_user_daily -WHERE user_id = ? - AND `date` >= ? - AND `date` < ? - AND total_requests > 0 -GROUP BY `date` -ORDER BY `date` ASC -"#, - ) - .bind(user_id) - .bind(to_i64( - query.created_from_unix_secs, - "stats_user_daily.date", - )?) - .bind(to_i64( - query.created_until_unix_secs, - "stats_user_daily.date", - )?) - .fetch_all(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date, - 'aggregate' AS model, - 'aggregate' AS provider, - CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests, - CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens, - CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd, - CAST(0.0 AS DOUBLE) AS response_time_sum_ms, - CAST(0 AS SIGNED) AS response_time_samples -FROM stats_daily -WHERE `date` >= ? - AND `date` < ? - AND total_requests > 0 -GROUP BY `date` -ORDER BY `date` ASC -"#, - ) - .bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?) - .bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?) - .fetch_all(&self.pool) - .await - .map_sql_err()? - }; - - rows.iter() - .map(|row| { - Ok(StoredUsageDashboardDailyBreakdownRow { - date: row.try_get("date").map_sql_err()?, - model: row.try_get("model").map_sql_err()?, - provider: row.try_get("provider").map_sql_err()?, - requests: row_u64(row, "requests")?, - total_tokens: row_u64(row, "total_tokens")?, - total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, - response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?, - response_time_samples: row_u64(row, "response_time_samples")?, - }) - }) - .collect() - } - - pub async fn summarize_usage_totals_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let unique_user_ids = user_ids - .iter() - .cloned() - .collect::>(); - let mut totals = BTreeMap::::new(); - let mut aggregate_cutoffs = BTreeMap::::new(); - - let mut aggregate_builder = QueryBuilder::::new( - r#" -SELECT - user_id, - CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS request_count, - CAST(COALESCE( - SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), - 0 - ) AS SIGNED) AS total_tokens, - MAX(`date`) AS latest_date -FROM stats_user_daily -WHERE user_id IN ( -"#, - ); - { - let mut separated = aggregate_builder.separated(", "); - for user_id in &unique_user_ids { - separated.push_bind(user_id.clone()); - } - } - aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC"); - - let aggregate_rows = aggregate_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - for row in aggregate_rows { - let user_id: String = row.try_get("user_id").map_sql_err()?; - let latest_date = row.try_get::("latest_date").map_sql_err()?.max(0) as u64; - aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400)); - totals.insert( - user_id.clone(), - StoredUsageUserTotals { - user_id, - request_count: row_u64(&row, "request_count")?, - total_tokens: row_u64(&row, "total_tokens")?, - }, - ); - } - - let mut raw_builder = QueryBuilder::::new(format!( - r#" -SELECT - `usage`.user_id, - CAST(COUNT(*) AS SIGNED) AS request_count, - 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 "); - } - let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default(); - raw_builder - .push("SELECT ") - .push_bind(user_id.clone()) - .push(" AS user_id, ") - .push_bind(to_i64(cutoff, "usage aggregate cutoff")?) - .push(" AS cutoff_unix_secs"); - } - 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') -GROUP BY `usage`.user_id -ORDER BY `usage`.user_id ASC -"#, - ); - - let raw_rows = raw_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - for row in raw_rows { - let user_id: String = row.try_get("user_id").map_sql_err()?; - let entry = totals - .entry(user_id.clone()) - .or_insert_with(|| StoredUsageUserTotals { - user_id, - request_count: 0, - total_tokens: 0, - }); - entry.request_count = entry - .request_count - .saturating_add(row_u64(&row, "request_count")?); - entry.total_tokens = entry - .total_tokens - .saturating_add(row_u64(&row, "total_tokens")?); - } - - Ok(totals.into_values().collect()) - } -} - -impl MysqlUsageWriteRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - 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), - } - } -} - -#[async_trait] -impl UsageWriteRepository for MysqlUsageWriteRepository { - async fn upsert( - &self, - usage: UpsertUsageRecord, - ) -> Result { - usage.validate()?; - // Auxiliary tables may receive only clear tombstones, never request or response content. - let capture_usage = usage.clone(); - let mut usage = sanitize_usage_for_persistence(usage); - usage.validate()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let existing = counters::lock_and_load_usage(&mut tx, &usage.request_id).await?; - if let Some(existing) = existing.as_ref() { - if !usage_lifecycle_update_allowed( - &existing.status, - &existing.billing_status, - existing.updated_at_unix_secs, - existing.finalized_at_unix_secs, - &usage.status, - &usage.billing_status, - usage.updated_at_unix_secs, - usage.finalized_at_unix_secs, - ) { - let existing = existing.clone(); - tx.rollback().await.map_sql_err()?; - return http_capture::hydrate_usage_body_refs(&self.pool, existing).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") - && !recovers_terminal_failure - { - let existing = existing.clone(); - tx.rollback().await.map_sql_err()?; - return http_capture::hydrate_usage_body_refs(&self.pool, existing).await; - } - } - - let mut capture_usage = sanitize_usage_capture_controls_for_persistence(capture_usage); - let prepared_capture = http_capture::prepare_usage_http_capture(&mut capture_usage)?; - 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 capture_usage, existing.as_ref()); - usage.request_metadata = - sanitize_usage_request_metadata(capture_usage.request_metadata.clone()); - } - let prepared_snapshots = capture_update_allowed - // The control projection preserves safe typed routing and allow-listed billing facts. - .then(|| snapshots::from_usage(&capture_usage)) - .transpose()?; - bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? - .execute(&mut *tx) - .await - .map_sql_err()?; - 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(|| { - DataLayerError::UnexpectedValue("usage upsert returned no row".to_string()) - }) - } - - 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 -SET total_requests = 0, - total_tokens = 0, - total_cost_usd = 0, - last_used_at = NULL -"#, - ) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let rows_affected = sqlx::query(&format!( - r#" -UPDATE api_keys -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 -"#, - 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 -SET request_count = 0, - success_count = 0, - error_count = 0, - total_tokens = 0, - total_cost_usd = 0, - total_response_time_ms = 0, - last_used_at = NULL -"#, - ) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let rows_affected = sqlx::query(&format!( - r#" -UPDATE provider_api_keys -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 -"#, - 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( - &self, - cutoff_unix_secs: u64, - now_unix_secs: u64, - _timeout_minutes: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(PendingUsageCleanupSummary::default()); - } - - 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(|_| { - DataLayerError::InvalidInput(format!( - "invalid stale pending usage batch size: {batch_size}" - )) - })?; - - 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_secs, "stale pending usage cutoff")?) - .bind(to_i64(batch_size_u64, "stale pending usage batch size")?) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - - if stale_rows.is_empty() { - tx.rollback().await.map_sql_err()?; - break; - } - - let stale_rows = stale_rows - .iter() - .map(|row| { - Ok(StalePendingUsageRow { - request_id: row.try_get("request_id").map_sql_err()?, - billing_status: row.try_get("billing_status").map_sql_err()?, - }) - }) - .collect::, DataLayerError>>()?; - let completed_request_ids = - completed_request_ids_mysql(&mut tx, stale_rows.iter().map(|row| &row.request_id)) - .await?; - - for row in stale_rows { - if completed_request_ids.contains(&row.request_id) { - sqlx::query( - r#" -UPDATE `usage` -SET status = 'completed', - status_code = 200, - error_message = NULL, - error_category = NULL -WHERE request_id = ? -"#, - ) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE request_candidates -SET status = 'success', - finished_at = ? -WHERE request_id = ? - AND status = 'streaming' -"#, - ) - .bind(to_i64(now_unix_ms, "request candidate finished_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - summary.recovered += 1; - continue; - } - - let candidate_info = - latest_failed_candidate_mysql(&mut tx, &row.request_id).await?; - let status_code = resolve_stale_pending_status_code(candidate_info.as_ref()); - let error_category = usage_error_category_for_status_code(status_code); - let status_code_i64 = i64::from(status_code); - if row.billing_status == "pending" { - sqlx::query( - r#" -UPDATE `usage` -SET status = 'failed', - status_code = ?, - error_message = NULL, - error_category = ?, - billing_status = 'void', - finalized_at = ?, - total_cost_usd = 0, - actual_total_cost_usd = 0 -WHERE request_id = ? -"#, - ) - .bind(status_code_i64) - .bind(error_category) - .bind(to_i64(now_unix_secs, "usage finalized_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - upsert_void_usage_settlement_snapshot_mysql( - &mut tx, - &row.request_id, - now_unix_secs, - ) - .await?; - } else { - sqlx::query( - r#" -UPDATE `usage` -SET status = 'failed', - status_code = ?, - error_message = NULL, - error_category = ? -WHERE request_id = ? -"#, - ) - .bind(status_code_i64) - .bind(error_category) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query( - r#" -UPDATE request_candidates -SET status = 'failed', - finished_at = ?, - error_type = 'internal', - error_message = NULL -WHERE request_id = ? - AND status IN ('pending', 'streaming') -"#, - ) - .bind(to_i64(now_unix_ms, "request candidate finished_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - summary.failed += 1; - } - - tx.commit().await.map_sql_err()?; - } - - 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 { - request_id: String, - billing_status: String, -} - -async fn completed_request_ids_mysql<'a>( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - request_ids: impl Iterator, -) -> Result, DataLayerError> { - let mut completed = HashSet::new(); - for request_id in request_ids { - let rows = sqlx::query(SELECT_COMPLETED_REQUEST_CANDIDATES_SQL) - .bind(request_id) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - let mut is_completed = false; - for row in &rows { - if candidate_row_is_completed(row)? { - is_completed = true; - break; - } - } - if is_completed { - completed.insert(request_id.clone()); - } - } - Ok(completed) -} - -fn candidate_row_is_completed(row: &MySqlRow) -> Result { - let status: String = row.try_get("status").map_sql_err()?; - if status == "streaming" { - return Ok(true); - } - if status != "success" { - return Ok(false); - } - let Some(extra_data) = row - .try_get::, _>("extra_data") - .map_sql_err()? - else { - return Ok(false); - }; - let Ok(value) = serde_json::from_str::(&extra_data) else { - return Ok(false); - }; - Ok(value - .get("stream_completed") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false)) -} - -async fn upsert_void_usage_settlement_snapshot_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - request_id: &str, - now_unix_secs: u64, -) -> Result<(), DataLayerError> { - let now = to_i64(now_unix_secs, "usage settlement snapshot timestamp")?; - sqlx::query( - r#" -INSERT INTO usage_settlement_snapshots ( - request_id, - billing_status, - finalized_at, - created_at, - updated_at -) 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) -"#, - ) - .bind(request_id) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -struct FailedCandidateCleanupInfo { - status_code: Option, -} - -fn resolve_stale_pending_status_code(candidate: Option<&FailedCandidateCleanupInfo>) -> u16 { - candidate - .and_then(|info| info.status_code) - .unwrap_or(if candidate.is_some() { 502 } else { 504 }) -} - -async fn latest_failed_candidate_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - request_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT status_code -FROM request_candidates -WHERE request_id = ? - AND status IN ('failed', 'cancelled') -ORDER BY - COALESCE(finished_at, started_at, created_at) DESC, - retry_index DESC, - candidate_index DESC -LIMIT 1 -"#, - ) - .bind(request_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()?; - - let Some(row) = row else { - return Ok(None); - }; - let status_code = row - .try_get::, _>("status_code") - .map_sql_err()? - .and_then(|value| u16::try_from(value).ok()); - Ok(Some(FailedCandidateCleanupInfo { status_code })) -} - -fn bind_upsert<'q>( - mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, - usage: &'q UpsertUsageRecord, -) -> Result, DataLayerError> { - let input_tokens = usage.input_tokens.unwrap_or_default(); - let output_tokens = usage.output_tokens.unwrap_or_default(); - let cache_creation_tokens = usage - .cache_creation_input_tokens - .or_else(|| { - Some( - usage - .cache_creation_ephemeral_5m_input_tokens - .unwrap_or_default() - + usage - .cache_creation_ephemeral_1h_input_tokens - .unwrap_or_default(), - ) - }) - .unwrap_or_default(); - let cache_read_tokens = usage.cache_read_input_tokens.unwrap_or_default(); - let total_tokens = usage - .total_tokens - .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); - let request_metadata = usage - .request_metadata - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; - - query = query - .bind(&usage.request_id) - .bind(&usage.request_id) - .bind(usage.user_id.as_deref()) - .bind(usage.api_key_id.as_deref()) - .bind(&usage.provider_name) - .bind(&usage.model) - .bind(usage.target_model.as_deref()) - .bind(usage.provider_id.as_deref()) - .bind(usage.provider_endpoint_id.as_deref()) - .bind(usage.provider_api_key_id.as_deref()) - .bind(usage.request_type.as_deref()) - .bind(usage.api_format.as_deref()) - .bind(usage.api_family.as_deref()) - .bind(usage.endpoint_kind.as_deref()) - .bind(usage.endpoint_api_format.as_deref()) - .bind(usage.provider_api_family.as_deref()) - .bind(usage.provider_endpoint_kind.as_deref()) - .bind(usage.has_format_conversion.unwrap_or(false)) - .bind(usage.is_stream.unwrap_or(false)) - .bind(usage_upstream_is_stream(usage)) - .bind(to_i64(input_tokens, "input_tokens")?) - .bind(to_i64(output_tokens, "output_tokens")?) - .bind(to_i64(total_tokens, "total_tokens")?) - .bind(to_i64( - cache_creation_tokens, - "cache_creation_input_tokens", - )?) - .bind(to_i64( - usage - .cache_creation_ephemeral_5m_input_tokens - .unwrap_or_default(), - "cache_creation_ephemeral_5m_input_tokens", - )?) - .bind(to_i64( - usage - .cache_creation_ephemeral_1h_input_tokens - .unwrap_or_default(), - "cache_creation_ephemeral_1h_input_tokens", - )?) - .bind(to_i64(cache_read_tokens, "cache_read_input_tokens")?) - .bind(usage.cache_creation_cost_usd.unwrap_or_default()) - .bind(usage.cache_read_cost_usd.unwrap_or_default()) - .bind(usage.output_price_per_1m) - .bind(usage.total_cost_usd.unwrap_or_default()) - .bind(usage.actual_total_cost_usd.unwrap_or_default()) - .bind(usage.status_code.map(i64::from)) - .bind(usage.error_message.as_deref()) - .bind(usage.error_category.as_deref()) - .bind(usage.response_time_ms.map(|value| value as i64)) - .bind(usage.first_byte_time_ms.map(|value| value as i64)) - .bind(&usage.status) - .bind(&usage.billing_status) - .bind(request_metadata) - .bind(usage.candidate_id.as_deref()) - .bind(usage.candidate_index.map(|value| value as i64)) - .bind(usage.key_name.as_deref()) - .bind(usage.planner_kind.as_deref()) - .bind(usage.route_family.as_deref()) - .bind(usage.route_kind.as_deref()) - .bind(usage.execution_path.as_deref()) - .bind(usage.local_execution_runtime_miss_reason.as_deref()) - .bind(usage.finalized_at_unix_secs.map(|value| value as i64)) - .bind(to_i64(created_at, "created_at_unix_ms")?) - .bind(to_i64(usage.updated_at_unix_secs, "updated_at_unix_secs")?); - Ok(query) -} - -fn map_usage_row( - row: &MySqlRow, - resolve_legacy_compressed: bool, -) -> Result { - let id = row - .try_get::, _>("id") - .map_sql_err()? - .unwrap_or_else(|| { - row.try_get::("request_id") - .unwrap_or_else(|_| "unknown".to_string()) - }); - let mut audit = StoredRequestUsageAudit::new( - id, - row.try_get("request_id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("provider_name").map_sql_err()?, - row.try_get("model").map_sql_err()?, - row.try_get("target_model").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_endpoint_id").map_sql_err()?, - row.try_get("provider_api_key_id").map_sql_err()?, - row.try_get("request_type").map_sql_err()?, - row.try_get("api_format").map_sql_err()?, - row.try_get("api_family").map_sql_err()?, - row.try_get("endpoint_kind").map_sql_err()?, - row.try_get("endpoint_api_format").map_sql_err()?, - row.try_get("provider_api_family").map_sql_err()?, - row.try_get("provider_endpoint_kind").map_sql_err()?, - row.try_get::("has_format_conversion") - .map_sql_err()?, - row.try_get::("is_stream").map_sql_err()?, - row_i32(row, "input_tokens")?, - row_i32(row, "output_tokens")?, - row_i32(row, "total_tokens")?, - row.try_get("total_cost_usd").map_sql_err()?, - row.try_get("actual_total_cost_usd").map_sql_err()?, - row_optional_i32(row, "status_code")?, - row.try_get("error_message").map_sql_err()?, - row.try_get("error_category").map_sql_err()?, - row_optional_i32(row, "response_time_ms")?, - row_optional_i32(row, "first_byte_time_ms")?, - row.try_get("status").map_sql_err()?, - row.try_get("billing_status").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - row.try_get("finalized_at_unix_secs").map_sql_err()?, - )?; - audit.cache_creation_input_tokens = row_u64(row, "cache_creation_input_tokens")?; - audit.cache_creation_ephemeral_5m_input_tokens = - row_u64(row, "cache_creation_ephemeral_5m_input_tokens")?; - audit.cache_creation_ephemeral_1h_input_tokens = - row_u64(row, "cache_creation_ephemeral_1h_input_tokens")?; - audit.cache_read_input_tokens = row_u64(row, "cache_read_input_tokens")?; - audit.cache_creation_cost_usd = row.try_get("cache_creation_cost_usd").map_sql_err()?; - audit.cache_read_cost_usd = row.try_get("cache_read_cost_usd").map_sql_err()?; - audit.output_price_per_1m = row.try_get("output_price_per_1m").map_sql_err()?; - audit.request_metadata = row - .try_get::, _>("request_metadata") - .map_sql_err()? - .map(|raw| serde_json::from_str(&raw)) - .transpose() - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; - audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref()) - .map(ToOwned::to_owned); - http_capture::hydrate_usage_row(row, &mut audit, resolve_legacy_compressed)?; - let upstream_is_stream = row - .try_get::, _>("upstream_is_stream") - .map_sql_err()?; - merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream); - snapshots::hydrate_row(row, &mut audit)?; - Ok(audit) -} - -fn to_i64(value: u64, field: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow"))) -} - -fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool { - 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_else(|| usage.is_stream.unwrap_or(false)) -} - -fn merge_usage_stream_metadata(metadata: &mut Option, upstream: Option) { - let Some(upstream) = upstream else { - return; - }; - let value = metadata.get_or_insert_with(|| serde_json::json!({})); - let Some(object) = value.as_object_mut() else { - return; - }; - object - .entry(UPSTREAM_IS_STREAM_KEY) - .or_insert(serde_json::Value::Bool(upstream)); -} - -fn row_i32(row: &MySqlRow, field: &str) -> Result { - let value: i64 = row.try_get(field).map_sql_err()?; - i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) -} - -fn row_optional_i32(row: &MySqlRow, field: &str) -> Result, DataLayerError> { - row.try_get::, _>(field) - .map_sql_err()? - .map(|value| { - i32::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) - }) - .transpose() -} - -fn row_u64(row: &MySqlRow, field: &str) -> Result { - let value: i64 = row.try_get(field).map_sql_err()?; - u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative"))) -} - -fn map_mysql_usage_daily_summary( - row: &MySqlRow, -) -> Result { - Ok(StoredUsageDailySummary { - date: row.try_get("date").map_sql_err()?, - requests: row_u64(row, "requests")?, - total_tokens: row_u64(row, "total_tokens")?, - total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?, - actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?, - }) -} - -fn usage_current_unix_secs() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs()) - .unwrap_or_default() -} - -#[cfg(test)] -mod tests; diff --git a/crates/aether-data/adapters/mysql/src/usage/cleanup.rs b/crates/aether-data/adapters/mysql/src/usage/cleanup.rs deleted file mode 100644 index 81aea0694..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/cleanup.rs +++ /dev/null @@ -1,865 +0,0 @@ -use aether_data_contracts::repository::usage::{ - UsageCleanupExecutionMode, UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, - UsageCleanupWindow, -}; -use chrono::{DateTime, Utc}; -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 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 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 - ) -) -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, 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, - } - } -} - -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 { - purge_legacy_body_refs(pool, window.detail_cutoff, detail_newer_than, batch_size).await? - } else { - 0 - }; - let body_externalized = if targets.detail_body { - cleanup_body_fields( - pool, - window.detail_cutoff, - detail_newer_than, - batch_size, - BodyCleanupKind::All, - ) - .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, - cost_reservations_deleted: 0, - request_admissions_deleted: 0, - }) -} - -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()?; - 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 purge_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 purge 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 purged = 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(metadata) = legacy_body_ref_purge_plan(metadata.as_deref())? else { - continue; - }; - 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(); - purge_detached_body_capture(&mut tx, &row.request_id).await?; - if updated > 0 { - purged += 1; - } - } - tx.commit().await.map_sql_err()?; - total = total.saturating_add(purged); - if row_count < batch_size || purged == 0 { - break; - } - } - Ok(total) -} - -fn legacy_body_ref_purge_plan( - 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 removed = false; - for key in [ - "request_body_ref", - "provider_request_body_ref", - "response_body_ref", - "client_response_body_ref", - ] { - if object.remove(key).is_some() { - removed = true; - } - } - 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(metadata)) -} - -async fn purge_detached_body_capture( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - request_id: &str, -) -> Result<(), DataLayerError> { - sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") - .bind(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(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - delete_empty_http_audit(tx, request_id).await -} - -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 serde_json::json; - - use super::{legacy_body_ref_purge_plan, DETAIL_BODY_PREDICATE}; - - #[test] - fn mysql_cleanup_legacy_body_ref_purge_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 metadata = legacy_body_ref_purge_plan(Some(&metadata)) - .expect("legacy plan should build") - .expect("legacy refs should be present"); - 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_detail_cleanup_includes_detached_capture() { - assert!(DETAIL_BODY_PREDICATE.contains("usage_body_blobs")); - assert!(DETAIL_BODY_PREDICATE.contains("usage_http_audits")); - assert!(!DETAIL_BODY_PREDICATE.contains("payload_gzip")); - } -} diff --git a/crates/aether-data/adapters/mysql/src/usage/counters.rs b/crates/aether-data/adapters/mysql/src/usage/counters.rs deleted file mode 100644 index be41956c1..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/counters.rs +++ /dev/null @@ -1,1210 +0,0 @@ -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, - target_tunnel_generation, - 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, - target_tunnel_generation: Option, - 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, - // Keep the node incarnation in the aggregation key. A node id can be - // reused after deletion, so a bare id would route old deltas to the new - // node. - proxy_nodes: BTreeMap<(String, String), ProxyNodeCounterDelta>, - 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 Some(tunnel_generation) = row - .target_tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - else { - // Legacy rows have no identity fence. Mark them - // processed without applying them to any node. - continue; - }; - let aggregate_key = (row.target_id.clone(), tunnel_generation.clone()); - let entry = aggregates.proxy_nodes.entry(aggregate_key).or_insert( - ProxyNodeCounterDelta { - node_id: row.target_id.clone(), - expected_tunnel_generation: Some(tunnel_generation), - 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, tunnel_generation), delta) in &aggregates.proxy_nodes { - apply_proxy_node(&mut tx, target_id, tunnel_generation, 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 Some(expected_tunnel_generation) = delta - .expected_tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .filter(|value| value.len() <= 64) - .map(ToOwned::to_owned) - else { - // A bare id is not an identity fence. Reject it instead of rebinding - // the delta to whichever incarnation currently owns that id. - 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()?; - // Keep the parent lookup lock-free because flush claims outbox rows before - // updating proxy_nodes. The generation is stored in the outbox row and is - // checked again by flush, so a concurrent id reuse can only discard this - // delta, never apply it to the replacement row. - let tunnel_generation: Option = sqlx::query_scalar( - "SELECT tunnel_generation FROM proxy_nodes WHERE id = ? AND BINARY tunnel_generation = BINARY ? LIMIT 1", - ) - .bind(&node_id) - .bind(&expected_tunnel_generation) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(_tunnel_generation) = tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - 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, - target_tunnel_generation: Some(&expected_tunnel_generation), - ..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, - target_tunnel_generation: Option<&'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, target_tunnel_generation, - 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.target_tunnel_generation) - .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()?, - target_tunnel_generation: row.try_get("target_tunnel_generation").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, - tunnel_generation: &str, - delta: &ProxyNodeCounterDelta, -) -> Result<(), DataLayerError> { - if target_id.trim().is_empty() || tunnel_generation.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 = ? AND BINARY tunnel_generation = BINARY ? -"#, - ) - .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) - .bind(tunnel_generation) - .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 deleted file mode 100644 index d5b120530..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/http_capture.rs +++ /dev/null @@ -1,759 +0,0 @@ -use aether_data_contracts::repository::usage::{ - canonical_usage_body_ref_for, parse_usage_body_ref, read_decompressed_usage_json, - usage_body_ref, StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, - UsageBodyField, -}; -use flate2::read::GzDecoder; -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 { - payload_gzip: Option>, -} - -#[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, -} - -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(request_body_value)?; - let provider_request_body = prepare_body(provider_request_body_value)?; - let response_body = prepare_body(response_body_value)?; - let client_response_body = prepare_body(client_response_body_value)?; - - 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(value: Option<&Value>) -> Result { - if value.is_some() { - return Err(DataLayerError::InvalidInput( - "usage body persistence is disabled".to_string(), - )); - } - Ok(PreparedBody { payload_gzip: None }) -} - -fn resolved_write_ref( - explicit_ref: Option<&str>, - request_id: &str, - field: UsageBodyField, - has_blob: bool, -) -> Option { - explicit_ref - .and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field)) - .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> { - let bodies = [ - &prepared.request_body, - &prepared.provider_request_body, - &prepared.response_body, - &prepared.client_response_body, - ]; - let contains_capture = prepared.request_headers.is_some() - || prepared.provider_request_headers.is_some() - || prepared.response_headers.is_some() - || prepared.client_response_headers.is_some() - || prepared.refs.any_present() - || bodies.iter().any(|body| body.payload_gzip.is_some()) - || prepared.capture_mode != "none"; - if contains_capture { - return Err(DataLayerError::InvalidInput( - "usage HTTP capture persistence is disabled".to_string(), - )); - } - - sqlx::query("DELETE FROM usage_http_audits WHERE request_id = ?") - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE `usage` -SET request_headers = NULL, - request_body = NULL, - provider_request_headers = NULL, - provider_request_body = NULL, - response_headers = NULL, - response_body = NULL, - client_response_headers = NULL, - client_response_body = NULL, - request_body_compressed = NULL, - provider_request_body_compressed = NULL, - response_body_compressed = NULL, - client_response_body_compressed = NULL -WHERE request_id = ? -"#, - ) - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - - 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() { - 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(()) -} - -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 - .and_then(|body_ref| canonical_usage_body_ref_for(&body_ref, request_id, field)) - .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) - .and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, 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 Some(body_ref) = canonical_usage_body_ref_for(body_ref, &usage.request_id, 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> { - let Some((request_id, field)) = parse_usage_body_ref(body_ref) else { - return Ok(None); - }; - let canonical_ref = usage_body_ref(&request_id, field); - if let Some(payload_gzip) = sqlx::query_scalar::<_, Vec>( - "SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? AND request_id = ? AND body_field = ? LIMIT 1", - ) - .bind(&canonical_ref) - .bind(&request_id) - .bind(field.as_storage_field()) - .fetch_optional(pool) - .await - .map_sql_err()? - { - return inflate_json(&payload_gzip).map(Some); - } - 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 decoded = read_decompressed_usage_json(GzDecoder::new(bytes))?; - 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 deleted file mode 100644 index 9aac4e1d3..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/read.rs +++ /dev/null @@ -1,852 +0,0 @@ -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 let Some(is_websocket) = query.is_websocket { - push_where(builder, has_where); - builder - .push("COALESCE(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode')), 'false') = ") - .push_bind(if is_websocket { "true" } else { "false" }); - } - 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, - is_websocket: query.is_websocket, - 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 ?")); - } - - #[test] - fn websocket_filter_is_applied_to_list_and_keyword_queries() { - let mut list_query = QueryBuilder::::new(USAGE_COLUMNS); - let mut has_where = false; - push_list_filters( - &mut list_query, - &UsageAuditListQuery { - is_websocket: Some(true), - ..UsageAuditListQuery::default() - }, - &mut has_where, - ) - .expect("WebSocket list query should build"); - assert!(list_query - .sql() - .contains("JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode'))")); - - let mut keyword_query = QueryBuilder::::new(USAGE_COLUMNS); - let mut has_where = false; - push_keyword_filters( - &mut keyword_query, - &UsageAuditKeywordSearchQuery { - is_websocket: Some(true), - keywords: vec!["live".to_string()], - ..UsageAuditKeywordSearchQuery::default() - }, - &mut has_where, - ) - .expect("WebSocket keyword query should build"); - assert!(keyword_query - .sql() - .contains("JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode'))")); - } -} diff --git a/crates/aether-data/adapters/mysql/src/usage/snapshots.rs b/crates/aether-data/adapters/mysql/src/usage/snapshots.rs deleted file mode 100644 index a929c089a..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/snapshots.rs +++ /dev/null @@ -1,1043 +0,0 @@ -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 deleted file mode 100644 index c65f027ae..000000000 --- a/crates/aether-data/adapters/mysql/src/usage/tests.rs +++ /dev/null @@ -1,1338 +0,0 @@ -use super::{MysqlUsageStorage, MysqlUsageWriteRepository}; -use crate::run_migrations; -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() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlUsageWriteRepository::new(pool); -} - -#[test] -fn mysql_usage_daily_heatmap_reads_imported_daily_aggregates() { - let source = include_str!("../usage.rs"); - assert!(source.contains("summarize_usage_daily_heatmap_from_daily_aggregates")); - assert!(source.contains("FROM stats_daily")); - assert!(source.contains("FROM stats_user_daily")); - assert!(source.contains("AS SIGNED) AS total_tokens")); - assert!(source.contains("CAST(COUNT(*) AS SIGNED) AS requests")); - assert!(source.contains("summaries.entry(item.date.clone()).or_insert(item)")); -} - -#[test] -fn mysql_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() { - let source = include_str!("../usage.rs"); - assert!(source.contains("async fn summarize_usage_totals_by_user_ids")); - assert!(source.contains("FROM stats_user_daily")); - assert!(source.contains("MAX(`date`) AS latest_date")); - assert!(source.contains("AS SIGNED) AS request_count")); - assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS request_count")); - assert!(source.contains("requested.cutoff_unix_secs")); -} - -#[test] -fn mysql_dashboard_reads_imported_daily_aggregates() { - let source = include_str!("../usage.rs"); - assert!(source.contains("summarize_dashboard_usage_from_daily_aggregates")); - assert!(source.contains("list_dashboard_daily_breakdown_from_daily_aggregates")); - assert!(source.contains("FROM stats_daily")); - assert!(source.contains("FROM stats_user_daily")); - assert!(source.contains("'aggregate' AS model")); - assert!(source.contains("AS SIGNED) AS total_requests")); - assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests")); - 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").replace("\r\n", "\n"); - 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( - "status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')" - )); - assert!(super::UPSERT_USAGE_SQL.contains("input_tokens = CASE")); - assert!(super::UPSERT_USAGE_SQL.contains("status_code = CASE")); - assert!(super::UPSERT_USAGE_SQL.contains("billing_status = CASE")); - assert!(super::UPSERT_USAGE_SQL.contains("finalized_at = CASE")); - assert!(super::UPSERT_USAGE_SQL.contains("updated_at_unix_secs = CASE")); - assert!(super::UPSERT_USAGE_SQL - .contains("WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status")); - assert!(super::UPSERT_USAGE_SQL.contains( - "WHEN status = 'streaming' AND VALUES(status) = 'streaming' AND VALUES(status_code) IS NULL THEN status_code" - )); -} - -#[test] -fn mysql_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() { - for field in [ - "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", - "request_metadata", - "candidate_id", - "candidate_index", - "key_name", - "planner_kind", - "route_family", - "route_kind", - "execution_path", - "local_execution_runtime_miss_reason", - ] { - let assignment = format!("{field} = CASE WHEN ("); - assert!( - super::UPSERT_USAGE_SQL.contains(&assignment), - "missing lifecycle guard for {field}" - ); - let preserve = format!("THEN {field} ELSE VALUES({field}) END"); - assert!( - super::UPSERT_USAGE_SQL.contains(&preserve), - "late lifecycle must preserve {field}" - ); - } - assert!(super::UPSERT_USAGE_SQL - .contains("OR (status = 'streaming' AND VALUES(status) = 'pending')")); -} - -#[tokio::test] -async fn mysql_stale_terminal_event_is_a_full_transaction_noop_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 stale terminal 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!("stale-user-{suffix}"); - let api_key_id = format!("stale-api-key-{suffix}"); - let provider_id = format!("stale-provider-{suffix}"); - let provider_key_id = format!("stale-provider-key-{suffix}"); - let request_id = format!("stale-request-{suffix}"); - seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await; - let repository = MysqlUsageWriteRepository::new(pool.clone()); - - let mut newer = sample_usage( - &request_id, - &user_id, - &api_key_id, - &provider_id, - &provider_key_id, - "completed", - "pending", - 2_000, - ); - newer.candidate_id = Some("candidate-new".to_string()); - newer.route_kind = Some("route-new".to_string()); - repository - .upsert(newer) - .await - .expect("newer terminal usage should upsert"); - - let counter_rows_before: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = ?") - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("counter rows should count"); - let routing_before: (Option, Option) = sqlx::query_as( - "SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = ?", - ) - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("routing snapshot should load"); - let settlement_before: (String, Option) = sqlx::query_as( - "SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = ?", - ) - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("settlement snapshot should load"); - - let mut stale = sample_usage( - &request_id, - &user_id, - &api_key_id, - &provider_id, - &provider_key_id, - "failed", - "void", - 1_999, - ); - stale.status_code = Some(503); - stale.total_cost_usd = Some(99.0); - stale.actual_total_cost_usd = Some(98.0); - stale.candidate_id = Some("candidate-stale".to_string()); - stale.route_kind = Some("route-stale".to_string()); - let stored = repository - .upsert(stale) - .await - .expect("stale terminal usage should be ignored"); - - assert_eq!(stored.status, "completed"); - assert_eq!(stored.billing_status, "pending"); - assert_eq!(stored.status_code, Some(200)); - assert_eq!(stored.total_cost_usd, 0.5); - assert_eq!(stored.routing_candidate_id(), Some("candidate-new")); - assert_eq!(stored.routing_route_kind(), Some("route-new")); - assert_eq!(stored.updated_at_unix_secs, 2_000); - - let counter_rows_after: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = ?") - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("counter rows should count"); - let routing_after: (Option, Option) = sqlx::query_as( - "SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = ?", - ) - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("routing snapshot should load"); - let settlement_after: (String, Option) = sqlx::query_as( - "SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = ?", - ) - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("settlement snapshot should load"); - assert_eq!(counter_rows_after, counter_rows_before); - assert_eq!(routing_after, routing_before); - assert_eq!(settlement_after, settlement_before); -} - -#[tokio::test] -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()) - else { - eprintln!("skipping mysql usage write smoke 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!("user-{suffix}"); - 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 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, 5); - assert_eq!( - record.request_metadata.as_ref().unwrap()["upstream_is_stream"], - true - ); - let upstream_is_stream: Option = - sqlx::query_scalar("SELECT upstream_is_stream FROM `usage` WHERE request_id = ?") - .bind(format!("request-{suffix}")) - .fetch_one(&pool) - .await - .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 = ?", - ) - .bind(&api_key_id) - .fetch_one(&pool) - .await - .expect("api key stats should load"); - 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 = ?", - ) - .bind(&provider_key_id) - .fetch_one(&pool) - .await - .expect("provider key stats should load"); - 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_is_not_persisted_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!(stored.request_headers.is_none()); - assert!(stored.request_body.is_none()); - assert!(stored.request_body_state.is_none()); - assert!(stored.request_body_ref.is_none()); - 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, 0); - 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!(sparse_stored.request_headers.is_none()); - assert!(sparse_stored.request_body.is_none()); - - 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!(cleared.request_body_state.is_none()); - assert!(cleared.provider_request_body.is_none()); -} - -#[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] -async fn mysql_usage_read_repository_reads_usage_contract_views_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 read smoke 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!("user-read-{suffix}"); - 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( - &first_request_id, - &user_id, - &api_key_id, - &provider_id, - &provider_key_id, - "completed", - "settled", - 1_000, - )) - .await - .expect("usage should upsert"); - writer - .upsert(sample_usage( - &second_request_id, - &user_id, - &api_key_id, - &provider_id, - &provider_key_id, - "failed", - "void", - 1_010, - )) - .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 - .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 == first_request_id) - .expect("usage should exist"); - 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() - .filter(|item| item.user_id.as_deref() == Some(&user_id)) - .count(), - 2 - ); -} - -#[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 stored_body: Option = - sqlx::query_scalar("SELECT CAST(request_body AS CHAR) FROM `usage` WHERE request_id = ?") - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("purged body should load"); - assert!(stored_body.is_none()); - let body_blobs: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = ?") - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("purged body blobs should count"); - assert_eq!(body_blobs, 0); - let body_audits: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM usage_http_audits WHERE request_id = ?") - .bind(&request_id) - .fetch_one(&pool) - .await - .expect("purged body refs should count"); - assert_eq!(body_audits, 0); - - 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, - api_key_id: &str, - provider_id: &str, - provider_key_id: &str, -) { - sqlx::query( - r#" -INSERT INTO users (id, auth_source, created_at, updated_at) -VALUES (?, 'local', 1, 1) -"#, - ) - .bind(user_id) - .execute(pool) - .await - .expect("user should seed"); - - sqlx::query( - r#" -INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) -VALUES (?, ?, ?, 1, 1) -"#, - ) - .bind(api_key_id) - .bind(user_id) - .bind(format!("hash-{api_key_id}")) - .execute(pool) - .await - .expect("api key should seed"); - - sqlx::query( - r#" -INSERT INTO providers (id, name, provider_type, created_at, updated_at) -VALUES (?, ?, 'openai', 1, 1) -"#, - ) - .bind(provider_id) - .bind(format!("Provider {provider_id}")) - .execute(pool) - .await - .expect("provider should seed"); - - sqlx::query( - r#" -INSERT INTO provider_api_keys (id, provider_id, name, created_at, updated_at) -VALUES (?, ?, ?, 1, 1) -"#, - ) - .bind(provider_key_id) - .bind(provider_id) - .bind(format!("Provider Key {provider_key_id}")) - .execute(pool) - .await - .expect("provider key should seed"); -} - -#[allow(clippy::too_many_arguments)] -fn sample_usage( - request_id: &str, - user_id: &str, - api_key_id: &str, - provider_id: &str, - provider_key_id: &str, - status: &str, - billing_status: &str, - updated_at: u64, -) -> UpsertUsageRecord { - UpsertUsageRecord { - request_id: request_id.to_string(), - user_id: Some(user_id.to_string()), - api_key_id: Some(api_key_id.to_string()), - username: Some("legacy-user".to_string()), - api_key_name: Some("legacy-key".to_string()), - provider_name: "Provider One".to_string(), - model: "model-1".to_string(), - target_model: Some("target-model".to_string()), - provider_id: Some(provider_id.to_string()), - provider_endpoint_id: Some("endpoint-1".to_string()), - provider_api_key_id: Some(provider_key_id.to_string()), - request_type: Some("chat".to_string()), - api_format: Some("openai".to_string()), - api_family: Some("chat".to_string()), - endpoint_kind: Some("chat".to_string()), - endpoint_api_format: Some("openai".to_string()), - provider_api_family: Some("chat".to_string()), - provider_endpoint_kind: Some("chat".to_string()), - has_format_conversion: Some(true), - is_stream: Some(false), - input_tokens: Some(2), - output_tokens: Some(3), - total_tokens: None, - cache_creation_input_tokens: None, - cache_creation_ephemeral_5m_input_tokens: Some(0), - cache_creation_ephemeral_1h_input_tokens: Some(0), - cache_read_input_tokens: Some(2), - cache_creation_cost_usd: Some(0.0), - cache_read_cost_usd: Some(0.1), - output_price_per_1m: Some(2.0), - total_cost_usd: Some(0.5), - actual_total_cost_usd: Some(0.4), - status_code: Some(200), - error_message: None, - error_category: None, - response_time_ms: Some(42), - first_byte_time_ms: Some(12), - status: status.to_string(), - billing_status: billing_status.to_string(), - request_headers: None, - request_body: None, - request_body_ref: None, - request_body_state: None, - provider_request_headers: None, - provider_request_body: None, - provider_request_body_ref: None, - provider_request_body_state: None, - response_headers: None, - response_body: None, - response_body_ref: None, - response_body_state: None, - client_response_headers: None, - client_response_body: None, - client_response_body_ref: None, - client_response_body_state: None, - candidate_id: Some("candidate-1".to_string()), - candidate_index: Some(1), - key_name: Some("key-one".to_string()), - planner_kind: Some("default".to_string()), - route_family: Some("chat".to_string()), - route_kind: Some("completion".to_string()), - execution_path: Some("remote".to_string()), - local_execution_runtime_miss_reason: None, - request_metadata: Some(serde_json::json!({ - "trace_id": "trace-1", - "upstream_is_stream": true, - })), - finalized_at_unix_secs: Some(updated_at), - created_at_unix_ms: Some(updated_at), - updated_at_unix_secs: updated_at, - } -} - -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/users.rs b/crates/aether-data/adapters/mysql/src/users.rs deleted file mode 100644 index 6d039f3e4..000000000 --- a/crates/aether-data/adapters/mysql/src/users.rs +++ /dev/null @@ -1,3354 +0,0 @@ -use async_trait::async_trait; -use chrono::{DateTime, TimeZone, Utc}; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::users::{ - is_valid_bcrypt_hash, last_oauth_unbind_denial, normalize_user_group_name, - BindUserOAuthLinkOutcome, BindUserOAuthLinkSessionExpectation, DeleteUserOAuthLinkOutcome, - LdapAuthUserProvisioningOutcome, ResolveOAuthLinkedUserOutcome, StoredUserAuthRecord, - StoredUserExportRow, StoredUserGroup, StoredUserGroupMember, StoredUserGroupMembership, - StoredUserOAuthLinkSummary, StoredUserPreferenceRecord, StoredUserSessionRecord, - StoredUserSummary, UpsertUserGroupRecord, UserExportListQuery, UserExportSortBy, - UserExportSummary, UserReadRepository, LAST_ACTIVE_ADMIN_DELETE_DENIED, - LAST_ACTIVE_ADMIN_UPDATE_DENIED, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const USER_SUMMARY_COLUMNS: &str = r#" -SELECT - id, - username, - email, - role, - is_active, - is_deleted -FROM users -"#; - -const MYSQL_LOCK_ACTIVE_ADMINS_SQL: &str = r#" -SELECT id -FROM users -WHERE LOWER(role) = 'admin' - AND is_active = 1 - AND is_deleted = 0 -ORDER BY id -FOR UPDATE -"#; - -const MYSQL_DELETE_USER_IF_WALLET_ABSENT_SQL: &str = r#" -DELETE FROM users -WHERE id = ? - AND NOT EXISTS ( - SELECT 1 - FROM wallets AS wallet - WHERE wallet.user_id = ? - OR EXISTS ( - SELECT 1 - FROM api_keys AS api_key - WHERE api_key.id = wallet.api_key_id - AND api_key.user_id = ? - ) - ) -"#; - -const MYSQL_DELETE_USER_API_KEYS_SQL: &str = "DELETE FROM api_keys WHERE user_id = ?"; - -const MYSQL_DELETE_USER_DEPENDENTS_SQL: &[&str] = &[ - "DELETE FROM usage_request_admissions WHERE subject_id = ?", - "DELETE FROM usage_cost_reservations WHERE subject_id = ?", - "DELETE FROM gemini_file_mappings WHERE user_id = ?", - "DELETE FROM api_key_provider_mappings WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - MYSQL_DELETE_USER_API_KEYS_SQL, - "DELETE FROM management_tokens WHERE user_id = ?", - "DELETE FROM user_sessions WHERE user_id = ?", - "DELETE FROM user_oauth_links WHERE user_id = ?", - "DELETE FROM user_group_members WHERE user_id = ?", - "DELETE FROM user_preferences WHERE user_id = ?", - "DELETE FROM user_invite_codes WHERE user_id = ?", - "DELETE FROM announcement_reads WHERE user_id = ?", -]; - -const MYSQL_PREPARE_USER_FACTS_FOR_DELETION_SQL: &[&str] = &[ - "UPDATE referral_rewards SET status = CASE WHEN status IN ('pending', 'failed', 'applying') THEN 'voided' ELSE status END, failure_reason = NULL, admin_note = NULL, updated_at = UNIX_TIMESTAMP() WHERE ? IN (inviter_user_id, invitee_user_id)", - "UPDATE referral_rewards SET failure_reason = NULL, admin_note = NULL, updated_at = UNIX_TIMESTAMP() WHERE admin_operator_id = ?", - "UPDATE user_referrals SET invite_code_snapshot = 'deleted-user', source_json = NULL, updated_at = UNIX_TIMESTAMP() WHERE ? IN (inviter_user_id, invitee_user_id)", - "UPDATE user_plan_entitlements SET status = CASE WHEN status = 'active' THEN 'revoked' ELSE status END, expires_at = LEAST(expires_at, UNIX_TIMESTAMP()), updated_at = UNIX_TIMESTAMP() WHERE user_id = ?", - "UPDATE wallets SET status = 'disabled', updated_at = UNIX_TIMESTAMP() WHERE user_id = ?", - "UPDATE wallets SET status = 'disabled', updated_at = UNIX_TIMESTAMP() WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - "UPDATE audit_logs SET description = 'deleted user event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE user_id = ?", - "UPDATE audit_logs SET description = 'deleted API key event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE user_id = ?)", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE wallet_transactions SET description = NULL WHERE operator_id = ?", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order WHERE history_order.user_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.user_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?) AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_orders SET gateway_response = NULL WHERE user_id = ?", - "UPDATE payment_orders SET gateway_response = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE user_id = ?", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE ? IN (requested_by, approved_by, processed_by)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE user_id = ?)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE redeem_code_batches SET description = NULL WHERE created_by = ?", -]; - -const MYSQL_ANONYMIZE_USER_HISTORY_SQL: &[&str] = &[ - "UPDATE request_candidates SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE video_tasks SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE `usage` SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE stats_user_daily SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_summary SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_model SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_api_format SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_model_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_model SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_model_provider SET username = NULL WHERE user_id = ?", -]; - -const MYSQL_ANONYMIZE_USER_API_KEY_HISTORY_SQL: &str = - "UPDATE stats_daily_api_key SET api_key_name = NULL WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)"; - -const USER_EXPORT_COLUMNS: &str = r#" -SELECT - id, - email, - email_verified, - username, - password_hash, - role, - auth_source, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - rate_limit, - rate_limit_mode, - model_capability_settings, - feature_settings, - is_active -FROM users -"#; - -const USER_AUTH_COLUMNS: &str = r#" -SELECT - id, - email, - email_verified, - username, - password_hash, - role, - auth_source, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - is_active, - is_deleted, - security_version, - created_at, - last_login_at -FROM users -"#; - -const USER_AUTH_COLUMNS_QUALIFIED: &str = r#" -SELECT - users.id AS id, - users.email AS email, - users.email_verified AS email_verified, - users.username AS username, - users.password_hash AS password_hash, - users.role AS role, - users.auth_source AS auth_source, - users.allowed_providers AS allowed_providers, - users.allowed_providers_mode AS allowed_providers_mode, - users.allowed_api_formats AS allowed_api_formats, - users.allowed_api_formats_mode AS allowed_api_formats_mode, - users.allowed_models AS allowed_models, - users.allowed_models_mode AS allowed_models_mode, - users.is_active AS is_active, - users.is_deleted AS is_deleted, - users.security_version AS security_version, - users.created_at AS created_at, - users.last_login_at AS last_login_at -FROM users -"#; - -const USER_OAUTH_LINK_SUMMARY_COLUMNS: &str = r#" -SELECT - user_oauth_links.provider_type, - oauth_providers.display_name, - user_oauth_links.provider_username, - user_oauth_links.provider_email, - user_oauth_links.linked_at, - user_oauth_links.last_login_at, - oauth_providers.is_enabled AS provider_enabled -FROM user_oauth_links -JOIN oauth_providers - ON oauth_providers.provider_type = user_oauth_links.provider_type -"#; - -const USER_PREFERENCES_COLUMNS: &str = r#" -SELECT - up.user_id, - up.avatar_url, - up.bio, - up.default_provider_id, - p.name AS default_provider_name, - up.theme, - up.language, - up.timezone, - up.email_notifications, - up.usage_alerts, - up.announcement_notifications -FROM user_preferences up -LEFT JOIN providers p - ON p.id = up.default_provider_id -"#; - -const USER_SESSION_COLUMNS: &str = r#" -SELECT - id, - user_id, - security_version, - client_device_id, - device_label, - refresh_token_hash, - prev_refresh_token_hash, - rotated_at, - last_seen_at, - expires_at, - revoked_at, - revoke_reason, - ip_address, - user_agent, - created_at, - updated_at -FROM user_sessions -"#; - -const USER_GROUP_COLUMNS: &str = r#" -SELECT - id, - name, - normalized_name, - description, - priority, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - rate_limit, - rate_limit_mode, - created_at, - updated_at -FROM user_groups -"#; - -const USER_GROUP_MEMBER_COLUMNS: &str = r#" -SELECT - user_group_members.group_id, - users.id AS user_id, - users.username, - users.email, - users.role, - users.is_active, - users.is_deleted, - user_group_members.created_at -FROM user_group_members -JOIN users ON users.id = user_group_members.user_id -"#; - -#[derive(Debug, Clone)] -pub struct MysqlUserReadRepository { - pool: MysqlPool, -} - -impl MysqlUserReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn fetch_summary_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_row).collect() - } - - async fn fetch_export_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_export_row).collect() - } - - async fn fetch_auth_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_auth_row).collect() - } - - async fn fetch_group_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_row).collect() - } - - async fn fetch_group_member_rows( - &self, - mut builder: QueryBuilder<'_, MySql>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_member_row).collect() - } - - async fn delete_local_auth_user_inner( - &self, - user_id: &str, - require_wallet_absent: bool, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let active_admin_ids = sqlx::query_scalar::<_, String>(MYSQL_LOCK_ACTIVE_ADMINS_SQL) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let target_security_state = - sqlx::query("SELECT role, is_active, is_deleted FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(target_security_state) = target_security_state else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let target_role = target_security_state - .try_get::("role") - .map_sql_err()?; - let target_is_active = target_security_state - .try_get::("is_active") - .map_sql_err()?; - let target_is_deleted = target_security_state - .try_get::("is_deleted") - .map_sql_err()?; - if target_role.eq_ignore_ascii_case("admin") - && target_is_active - && !target_is_deleted - && active_admin_ids.len() <= 1 - { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_DELETE_DENIED.to_string(), - )); - } - if require_wallet_absent { - let wallet_exists: Option = sqlx::query_scalar( - r#" -SELECT 1 -FROM wallets AS wallet -WHERE wallet.user_id = ? - OR EXISTS ( - SELECT 1 - FROM api_keys AS api_key - WHERE api_key.id = wallet.api_key_id - AND api_key.user_id = ? - ) -LIMIT 1 - "#, - ) - .bind(user_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if wallet_exists.is_some() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - } - for sql in MYSQL_PREPARE_USER_FACTS_FOR_DELETION_SQL { - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - for sql in MYSQL_ANONYMIZE_USER_HISTORY_SQL { - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query(MYSQL_ANONYMIZE_USER_API_KEY_HISTORY_SQL) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - for sql in MYSQL_DELETE_USER_DEPENDENTS_SQL { - if require_wallet_absent && *sql == MYSQL_DELETE_USER_API_KEYS_SQL { - continue; - } - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let result = if require_wallet_absent { - sqlx::query(MYSQL_DELETE_USER_IF_WALLET_ABSENT_SQL) - .bind(user_id) - .bind(user_id) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query("DELETE FROM users WHERE id = ?") - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - }; - if require_wallet_absent && result.rows_affected() == 0 { - // A wallet may have been inserted after the initial check. Do not - // commit the history/credential mutations when the guarded delete - // loses that race. - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - if require_wallet_absent { - sqlx::query(MYSQL_DELETE_USER_API_KEYS_SQL) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(result.rows_affected() > 0) - } -} - -#[async_trait] -impl UserReadRepository for MysqlUserReadRepository { - async fn list_users_by_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_SUMMARY_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push(") ORDER BY id ASC"); - self.fetch_summary_rows(builder).await - } - - async fn list_users_by_username_search( - &self, - username_search: &str, - ) -> Result, DataLayerError> { - let username_search = username_search.trim(); - if username_search.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_SUMMARY_COLUMNS); - builder - .push(" WHERE is_deleted = 0 AND LOWER(username) LIKE ") - .push_bind(format!("%{}%", username_search.to_ascii_lowercase())) - .push(" ORDER BY id ASC"); - self.fetch_summary_rows(builder).await - } - - async fn list_export_users(&self) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0 ORDER BY id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_users_page( - &self, - query: &UserExportListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0"); - if let Some(role) = query.role.as_deref() { - builder - .push(" AND LOWER(role) = ") - .push_bind(role.trim().to_ascii_lowercase()); - } - if let Some(is_active) = query.is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - if let Some(group_id) = query - .group_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = "); - builder.push_bind(group_id); - builder.push(")"); - } - 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(id) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(username) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(COALESCE(email, '')) LIKE ") - .push_bind(pattern) - .push(")"); - } - match query.sort_by { - UserExportSortBy::CreatedAt => { - builder - .push(" ORDER BY created_at ") - .push(if query.sort_order.is_desc() { - "DESC" - } else { - "ASC" - }) - .push(", id ASC"); - } - UserExportSortBy::Id => { - builder.push(" ORDER BY id ASC"); - } - } - - builder - .push(" LIMIT ") - .push_bind(i64::try_from(query.limit).map_err(|_| { - DataLayerError::InvalidInput(format!("invalid user export limit: {}", query.limit)) - })?) - .push(" OFFSET ") - .push_bind(i64::try_from(query.skip).map_err(|_| { - DataLayerError::InvalidInput(format!("invalid user export skip: {}", query.skip)) - })?); - self.fetch_export_rows(builder).await - } - - async fn count_export_users(&self, query: &UserExportListQuery) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(*) AS total FROM users"); - builder.push(" WHERE is_deleted = 0"); - if let Some(role) = query.role.as_deref() { - builder - .push(" AND LOWER(role) = ") - .push_bind(role.trim().to_ascii_lowercase()); - } - if let Some(is_active) = query.is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - if let Some(group_id) = query - .group_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = "); - builder.push_bind(group_id); - builder.push(")"); - } - 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(id) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(username) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(COALESCE(email, '')) LIKE ") - .push_bind(pattern) - .push(")"); - } - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - async fn summarize_export_users(&self) -> Result { - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END) AS active -FROM users -WHERE is_deleted = 0 -"#, - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - Ok(UserExportSummary { - total: row.try_get::("total").map_sql_err()?.max(0) as u64, - active: row - .try_get::, _>("active") - .map_sql_err()? - .unwrap_or(0) - .max(0) as u64, - }) - } - - async fn find_export_user_by_id( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder - .push(" WHERE is_deleted = 0 AND id = ") - .push_bind(user_id) - .push(" LIMIT 1"); - Ok(self.fetch_export_rows(builder).await?.into_iter().next()) - } - - async fn list_non_admin_export_users( - &self, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0 AND LOWER(role) != 'admin' ORDER BY id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_user_groups(&self) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder.push(" ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn find_user_group_by_id( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder - .push(" WHERE id = ") - .push_bind(group_id) - .push(" LIMIT 1"); - Ok(self.fetch_group_rows(builder).await?.into_iter().next()) - } - - async fn list_user_groups_by_ids( - &self, - group_ids: &[String], - ) -> Result, DataLayerError> { - if group_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for group_id in group_ids { - separated.push_bind(group_id); - } - } - builder.push(") ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn create_user_group( - &self, - record: UpsertUserGroupRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let id = uuid::Uuid::new_v4().to_string(); - let name = normalize_user_group_name(&record.name); - let normalized_name = name.to_ascii_lowercase(); - let result = sqlx::query( - r#" -INSERT INTO user_groups ( - id, name, normalized_name, description, priority, - allowed_providers, allowed_providers_mode, - allowed_api_formats, allowed_api_formats_mode, - allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&id) - .bind(name) - .bind(normalized_name) - .bind(record.description) - .bind(record.priority) - .bind(json_string_from_option_vec( - record.allowed_providers.as_ref(), - )) - .bind(record.allowed_providers_mode) - .bind(json_string_from_option_vec( - record.allowed_api_formats.as_ref(), - )) - .bind(record.allowed_api_formats_mode) - .bind(json_string_from_option_vec(record.allowed_models.as_ref())) - .bind(record.allowed_models_mode) - .bind(record.rate_limit) - .bind(record.rate_limit_mode) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - match result { - Ok(_) => self.find_user_group_by_id(&id).await, - Err(sqlx::Error::Database(err)) if err.is_unique_violation() => Err( - DataLayerError::InvalidInput("duplicate user group name".to_string()), - ), - Err(err) => Err(err).map_sql_err(), - } - } - - async fn update_user_group( - &self, - group_id: &str, - record: UpsertUserGroupRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let name = normalize_user_group_name(&record.name); - let normalized_name = name.to_ascii_lowercase(); - let result = sqlx::query( - r#" -UPDATE user_groups -SET name = ?, - normalized_name = ?, - description = ?, - priority = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(name) - .bind(normalized_name) - .bind(record.description) - .bind(record.priority) - .bind(json_string_from_option_vec( - record.allowed_providers.as_ref(), - )) - .bind(record.allowed_providers_mode) - .bind(json_string_from_option_vec( - record.allowed_api_formats.as_ref(), - )) - .bind(record.allowed_api_formats_mode) - .bind(json_string_from_option_vec(record.allowed_models.as_ref())) - .bind(record.allowed_models_mode) - .bind(record.rate_limit) - .bind(record.rate_limit_mode) - .bind(now) - .bind(group_id) - .execute(&self.pool) - .await; - match result { - Ok(result) if result.rows_affected() == 0 => Ok(None), - Ok(_) => self.find_user_group_by_id(group_id).await, - Err(sqlx::Error::Database(err)) if err.is_unique_violation() => Err( - DataLayerError::InvalidInput("duplicate user group name".to_string()), - ), - Err(err) => Err(err).map_sql_err(), - } - } - - /// Restore a group while holding its row lock. This keeps the snapshot - /// comparison and replacement atomic with respect to administrator edits. - async fn restore_user_group_if_matches( - &self, - expected: &StoredUserGroup, - restored: &StoredUserGroup, - ) -> Result { - if expected.id != restored.id || expected.id.trim().is_empty() { - return Ok(false); - } - - let mut tx = self.pool.begin().await.map_sql_err()?; - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder - .push(" WHERE id = ") - .push_bind(&expected.id) - .push(" FOR UPDATE"); - let row = builder - .build() - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_user_group_row(&row)?; - if ¤t != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE user_groups -SET name = ?, - normalized_name = ?, - description = ?, - priority = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - created_at = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&restored.name) - .bind(&restored.normalized_name) - .bind(&restored.description) - .bind(restored.priority) - .bind(json_string_from_option_vec( - restored.allowed_providers.as_ref(), - )) - .bind(&restored.allowed_providers_mode) - .bind(json_string_from_option_vec( - restored.allowed_api_formats.as_ref(), - )) - .bind(&restored.allowed_api_formats_mode) - .bind(json_string_from_option_vec( - restored.allowed_models.as_ref(), - )) - .bind(&restored.allowed_models_mode) - .bind(restored.rate_limit) - .bind(&restored.rate_limit_mode) - .bind(restored.created_at.map(|value| value.timestamp())) - .bind(restored.updated_at.map(|value| value.timestamp())) - .bind(&restored.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_user_group(&self, group_id: &str) -> Result { - let result = sqlx::query("DELETE FROM user_groups WHERE id = ?") - .bind(group_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn list_user_group_members( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_MEMBER_COLUMNS); - builder - .push(" WHERE user_group_members.group_id = ") - .push_bind(group_id) - .push(" ORDER BY users.username ASC, users.id ASC"); - self.fetch_group_member_rows(builder).await - } - - async fn replace_user_group_members( - &self, - group_id: &str, - user_ids: &[String], - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - // Membership rollback serializes on the owning user row. Lock every user that can be - // removed or inserted in deterministic order so it cannot race the per-user CAS restore. - let mut locked_user_ids = normalized_ids(user_ids); - let existing_user_ids = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_group_members WHERE group_id = ? ORDER BY user_id", - ) - .bind(group_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - locked_user_ids.extend(existing_user_ids); - locked_user_ids.sort(); - locked_user_ids.dedup(); - if !locked_user_ids.is_empty() { - let mut builder = QueryBuilder::::new("SELECT id FROM users WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for user_id in &locked_user_ids { - separated.push_bind(user_id); - } - } - builder.push(") ORDER BY id FOR UPDATE"); - builder.build().fetch_all(&mut *tx).await.map_sql_err()?; - } - sqlx::query("DELETE FROM user_group_members WHERE group_id = ?") - .bind(group_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for user_id in normalized_ids(user_ids) { - sqlx::query( - "INSERT IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.list_user_group_members(group_id).await - } - - async fn list_user_groups_for_user( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder - .push(" WHERE id IN (SELECT group_id FROM user_group_members WHERE user_id = ") - .push_bind(user_id) - .push(") ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn list_user_group_memberships_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - r#" -SELECT - user_group_members.user_id, - user_groups.id AS group_id, - user_groups.name AS group_name, - user_groups.priority AS group_priority, - user_group_members.created_at -FROM user_group_members -JOIN user_groups ON user_groups.id = user_group_members.group_id -WHERE user_group_members.user_id IN ( -"#, - ); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push( - ") ORDER BY user_group_members.user_id ASC, user_groups.name ASC, user_groups.id ASC", - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_membership_row).collect() - } - - async fn replace_user_groups_for_user( - &self, - user_id: &str, - group_ids: &[String], - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(Vec::new()); - } - sqlx::query("DELETE FROM user_group_members WHERE user_id = ?") - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for group_id in normalized_ids(group_ids) { - sqlx::query( - "INSERT IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.list_user_groups_for_user(user_id).await - } - - async fn restore_user_groups_if_matches( - &self, - user_id: &str, - expected_group_ids: &[String], - restored_group_ids: &[String], - ) -> Result { - let expected = normalized_ids(expected_group_ids); - let restored = normalized_ids(restored_group_ids); - let mut tx = self.pool.begin().await.map_sql_err()?; - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let current = sqlx::query_scalar::<_, String>( - "SELECT group_id FROM user_group_members WHERE user_id = ? ORDER BY group_id ASC FOR UPDATE", - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if current != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - if !restored.is_empty() { - let mut builder = QueryBuilder::::new( - "SELECT COUNT(*) AS count FROM user_groups WHERE id IN (", - ); - { - let mut separated = builder.separated(", "); - for group_id in &restored { - separated.push_bind(group_id); - } - } - builder.push(")"); - let count = builder - .build() - .fetch_one(&mut *tx) - .await - .map_sql_err()? - .try_get::("count") - .map_sql_err()?; - if count != restored.len() as i64 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - } - sqlx::query("DELETE FROM user_group_members WHERE user_id = ?") - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for group_id in restored { - sqlx::query( - "INSERT IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn add_user_to_group( - &self, - group_id: &str, - user_id: &str, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "INSERT IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(current_unix_secs()) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn find_user_auth_by_id( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE id = ") - .push_bind(user_id) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn list_user_auth_by_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push(") ORDER BY id ASC"); - self.fetch_auth_rows(builder).await - } - - async fn find_user_auth_by_identifier( - &self, - identifier: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE email = ") - .push_bind(identifier) - .push(" OR username = ") - .push_bind(identifier) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_user_auth_by_email( - &self, - email: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE email = ") - .push_bind(email) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_active_user_auth_by_email_ci( - &self, - email: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE LOWER(email) = LOWER(") - .push_bind(email) - .push(") AND is_deleted = 0 LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_user_auth_by_username( - &self, - username: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE username = ") - .push_bind(username) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn list_user_oauth_links( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_OAUTH_LINK_SUMMARY_COLUMNS); - builder - .push(" WHERE user_oauth_links.user_id = ") - .push_bind(user_id) - .push(" ORDER BY user_oauth_links.linked_at ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_oauth_link_summary_row).collect() - } - - async fn find_oauth_linked_user( - &self, - provider_type: &str, - provider_user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS_QUALIFIED); - builder - .push(" JOIN user_oauth_links ON users.id = user_oauth_links.user_id") - .push(" WHERE user_oauth_links.provider_type = ") - .push_bind(provider_type) - .push(" AND user_oauth_links.provider_user_id = ") - .push_bind(provider_user_id) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn resolve_enabled_oauth_linked_user( - &self, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - verified_email: Option<&str>, - touched_at: DateTime, - _provider_enabled_snapshot: bool, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let provider_enabled: Option = sqlx::query_scalar( - "SELECT is_enabled FROM oauth_providers WHERE provider_type = ? FOR UPDATE", - ) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_enabled != Some(true) { - tx.rollback().await.map_sql_err()?; - return Ok(ResolveOAuthLinkedUserOutcome::ProviderUnavailable); - } - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS_QUALIFIED} JOIN user_oauth_links ON users.id = user_oauth_links.user_id WHERE user_oauth_links.provider_type = ? AND user_oauth_links.provider_user_id = ? LIMIT 1 FOR UPDATE" - )) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(ResolveOAuthLinkedUserOutcome::NotLinked); - }; - let mut user = map_user_auth_row(&row)?; - sqlx::query( - "UPDATE user_oauth_links SET provider_username = COALESCE(?, provider_username), provider_email = COALESCE(?, provider_email), extra_data = COALESCE(?, extra_data), last_login_at = ? WHERE provider_type = ? AND provider_user_id = ?", - ) - .bind(provider_username) - .bind(provider_email) - .bind(optional_json_string(extra_data, "user_oauth_links.extra_data")?) - .bind(touched_at.timestamp()) - .bind(provider_type) - .bind(provider_user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if let Some(verified_email) = verified_email { - let result = sqlx::query( - "UPDATE users SET email_verified = 1, updated_at = ? WHERE id = ? AND email_verified = 0 AND LOWER(TRIM(email)) = LOWER(TRIM(?))", - ) - .bind(touched_at.timestamp()) - .bind(&user.id) - .bind(verified_email) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() == 1 { - user.email_verified = true; - } - } - tx.commit().await.map_sql_err()?; - Ok(ResolveOAuthLinkedUserOutcome::Linked(user)) - } - - async fn touch_oauth_link( - &self, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - touched_at: DateTime, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_oauth_links -SET provider_username = COALESCE(?, provider_username), - provider_email = COALESCE(?, provider_email), - extra_data = COALESCE(?, extra_data), - last_login_at = ? -WHERE provider_type = ? - AND provider_user_id = ? -"#, - ) - .bind(provider_username) - .bind(provider_email) - .bind(optional_json_string( - extra_data, - "user_oauth_links.extra_data", - )?) - .bind(touched_at.timestamp()) - .bind(provider_type) - .bind(provider_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn create_oauth_auth_user( - &self, - email: Option, - email_verified: bool, - username: String, - created_at: DateTime, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, rate_limit_mode, - is_active, is_deleted, created_at, updated_at, last_login_at -) -VALUES (?, ?, ?, ?, NULL, 'user', 'oauth', 'inherit', 'inherit', 'inherit', 'inherit', 1, 0, ?, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(created_at.timestamp()) - .bind(created_at.timestamp()) - .bind(created_at.timestamp()) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn find_oauth_link_owner( - &self, - provider_type: &str, - provider_user_id: &str, - ) -> Result, DataLayerError> { - sqlx::query_scalar( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ? AND provider_user_id = ? LIMIT 1", - ) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err() - } - - async fn has_user_oauth_provider_link( - &self, - user_id: &str, - provider_type: &str, - ) -> Result { - let owner: Option = sqlx::query_scalar( - "SELECT user_id FROM user_oauth_links WHERE user_id = ? AND provider_type = ? LIMIT 1", - ) - .bind(user_id) - .bind(provider_type) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(owner.is_some()) - } - - async fn count_user_oauth_links(&self, user_id: &str) -> Result { - let total: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links WHERE user_id = ?") - .bind(user_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(total.max(0) as u64) - } - - async fn has_oauth_links_for_provider( - &self, - provider_type: &str, - ) -> Result { - let exists: Option = - sqlx::query_scalar("SELECT 1 FROM user_oauth_links WHERE provider_type = ? LIMIT 1") - .bind(provider_type) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(exists.is_some()) - } - - async fn bind_user_oauth_link_if_provider_enabled( - &self, - user_id: &str, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - linked_at: DateTime, - _provider_enabled_snapshot: bool, - session_expectation: Option<&BindUserOAuthLinkSessionExpectation>, - ) -> Result { - let extra_data = optional_json_string(extra_data, "user_oauth_links.extra_data")?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let provider_enabled: Option = sqlx::query_scalar( - "SELECT is_enabled FROM oauth_providers WHERE provider_type = ? FOR UPDATE", - ) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_enabled.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::ProviderNotFound); - } - if provider_enabled != Some(true) { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::ProviderDisabled); - } - if let Some(expectation) = session_expectation { - let session_is_current: Option = sqlx::query_scalar( - r#" -SELECT 1 -FROM users -JOIN user_sessions - ON user_sessions.user_id = users.id -WHERE users.id = ? - AND users.is_active = 1 - AND users.is_deleted = 0 - AND users.security_version = ? - AND user_sessions.id = ? - AND user_sessions.security_version = ? - AND user_sessions.client_device_id = ? - AND user_sessions.revoked_at IS NULL - AND user_sessions.expires_at > GREATEST(?, UNIX_TIMESTAMP()) -FOR UPDATE -"#, - ) - .bind(user_id) - .bind(expectation.security_version) - .bind(&expectation.session_id) - .bind(expectation.security_version) - .bind(&expectation.client_device_id) - .bind(expectation.checked_at.timestamp()) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if session_is_current.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::SessionUnavailable); - } - } else { - let user_exists: Option = - sqlx::query_scalar("SELECT 1 FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::UserNotFound); - } - } - if let Some(owner) = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ? AND provider_user_id = ? LIMIT 1 FOR UPDATE", - ) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - { - tx.rollback().await.map_sql_err()?; - return Ok(if owner == user_id { - BindUserOAuthLinkOutcome::IdentityAlreadyBoundToUser - } else { - BindUserOAuthLinkOutcome::IdentityBoundToAnotherUser - }); - } - if sqlx::query_scalar::<_, i32>( - "SELECT 1 FROM user_oauth_links WHERE user_id = ? AND provider_type = ? LIMIT 1 FOR UPDATE", - ) - .bind(user_id) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .is_some() - { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::UserAlreadyLinkedProvider); - } - let inserted = sqlx::query( - r#" -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, provider_username, provider_email, - extra_data, linked_at, last_login_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(user_id) - .bind(provider_type) - .bind(provider_user_id) - .bind(provider_username) - .bind(provider_email) - .bind(extra_data.as_deref()) - .bind(linked_at.timestamp()) - .bind(linked_at.timestamp()) - .execute(&mut *tx) - .await; - match inserted { - Ok(_) => { - tx.commit().await.map_sql_err()?; - Ok(BindUserOAuthLinkOutcome::Bound) - } - Err(sqlx::Error::Database(err)) if err.is_unique_violation() => { - let owner = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ? AND provider_user_id = ? LIMIT 1", - ) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - tx.rollback().await.map_sql_err()?; - Ok(match owner { - Some(owner) if owner == user_id => { - BindUserOAuthLinkOutcome::IdentityAlreadyBoundToUser - } - Some(_) => BindUserOAuthLinkOutcome::IdentityBoundToAnotherUser, - None => BindUserOAuthLinkOutcome::UserAlreadyLinkedProvider, - }) - } - Err(err) => Err(DataLayerError::sql(err)), - } - } - - async fn upgrade_oauth_email_verification_if_matches( - &self, - user_id: &str, - verified_email: &str, - verified_at: DateTime, - ) -> Result { - let result = sqlx::query( - "UPDATE users SET email_verified = 1, updated_at = ? WHERE id = ? AND email_verified = 0 AND LOWER(TRIM(email)) = LOWER(TRIM(?))", - ) - .bind(verified_at.timestamp()) - .bind(user_id) - .bind(verified_email) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn delete_user_oauth_link( - &self, - user_id: &str, - provider_type: &str, - local_password_login_allowed: bool, - _enabled_provider_types_snapshot: &[String], - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let provider_exists: Option = sqlx::query_scalar( - "SELECT provider_type FROM oauth_providers WHERE provider_type = ? FOR UPDATE", - ) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - } - let user = - sqlx::query("SELECT auth_source, password_hash FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(user) = user else { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - }; - let auth_source = user.try_get::("auth_source").map_sql_err()?; - let password_hash = user - .try_get::, _>("password_hash") - .map_sql_err()?; - let provider_types = sqlx::query_scalar::<_, String>( - "SELECT provider_type FROM user_oauth_links WHERE user_id = ? FOR UPDATE", - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if !provider_types.iter().any(|value| value == provider_type) { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - } - let enabled_provider_types = sqlx::query_scalar::<_, String>( - r#" -SELECT user_oauth_links.provider_type -FROM user_oauth_links -JOIN oauth_providers - ON oauth_providers.provider_type = user_oauth_links.provider_type -WHERE user_oauth_links.user_id = ? - AND oauth_providers.is_enabled = 1 -FOR UPDATE -"#, - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let has_remaining_enabled_oauth_link = enabled_provider_types - .iter() - .any(|value| value != provider_type); - if !has_remaining_enabled_oauth_link { - if let Some(outcome) = last_oauth_unbind_denial( - &auth_source, - password_hash.as_deref(), - local_password_login_allowed, - ) { - tx.rollback().await.map_sql_err()?; - return Ok(outcome); - } - } - let result = - sqlx::query("DELETE FROM user_oauth_links WHERE user_id = ? AND provider_type = ?") - .bind(user_id) - .bind(provider_type) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - } - tx.commit().await.map_sql_err()?; - Ok(DeleteUserOAuthLinkOutcome::Deleted) - } - - async fn get_or_create_ldap_auth_user( - &self, - email: String, - username: String, - ldap_dn: Option, - ldap_username: Option, - logged_in_at: DateTime, - ) -> Result, DataLayerError> { - get_or_create_mysql_ldap_auth_user( - &self.pool, - email, - username, - ldap_dn, - ldap_username, - logged_in_at, - ) - .await - } - - async fn touch_auth_user_last_login( - &self, - user_id: &str, - logged_in_at: DateTime, - ) -> Result { - let result = sqlx::query("UPDATE users SET last_login_at = ?, updated_at = ? WHERE id = ?") - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn update_local_auth_user_profile( - &self, - user_id: &str, - email_present: bool, - email: Option, - email_verified: Option, - username: Option, - ) -> Result, DataLayerError> { - let now = chrono::Utc::now().timestamp(); - let result = sqlx::query( - "UPDATE users SET email = CASE WHEN ? THEN ? ELSE email END, email_verified = COALESCE(?, email_verified), username = COALESCE(?, username), updated_at = ? WHERE id = ?", - ) - .bind(email_present) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(now) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn restore_local_auth_user_state_if_matches( - &self, - expected_auth: &StoredUserAuthRecord, - restored_auth: &StoredUserAuthRecord, - expected_export: &StoredUserExportRow, - restored_export: &StoredUserExportRow, - expected_model_capability_settings: Option<&serde_json::Value>, - restored_model_capability_settings: Option, - expected_feature_settings: Option<&serde_json::Value>, - restored_feature_settings: Option, - ) -> Result { - if expected_auth.id != restored_auth.id - || expected_export.id != expected_auth.id - || restored_export.id != restored_auth.id - { - return Ok(false); - } - let mut tx = self.pool.begin().await.map_sql_err()?; - let active_admin_ids = sqlx::query_scalar::<_, String>(MYSQL_LOCK_ACTIVE_ADMINS_SQL) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let auth_row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(&expected_auth.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let export_row = sqlx::query(&format!( - "{USER_EXPORT_COLUMNS} WHERE id = ? LIMIT 1 FOR UPDATE" - )) - .bind(&expected_auth.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (Some(auth_row), Some(export_row)) = (auth_row, export_row) else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current_auth = map_user_auth_row(&auth_row)?; - let current_export = map_user_export_row(&export_row)?; - if !current_auth.matches_restore_state(expected_auth) - || !current_export.matches_restore_state(expected_export) - || current_export.rate_limit != expected_export.rate_limit - || current_export.rate_limit_mode != expected_export.rate_limit_mode - || current_export.model_capability_settings.as_ref() - != expected_model_capability_settings - || current_export.feature_settings.as_ref() != expected_feature_settings - { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let removes_active_admin = current_auth.role.eq_ignore_ascii_case("admin") - && current_auth.is_active - && !current_auth.is_deleted - && (!restored_auth.role.eq_ignore_ascii_case("admin") || !restored_auth.is_active); - if removes_active_admin && active_admin_ids.len() <= 1 { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_UPDATE_DENIED.to_string(), - )); - } - let security_state_changed = expected_auth.role != restored_auth.role - || expected_auth.is_active != restored_auth.is_active; - let result = sqlx::query( - r#" -UPDATE users -SET email = ?, - email_verified = ?, - username = ?, - role = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - model_capability_settings = ?, - feature_settings = ?, - is_active = ?, - security_version = security_version + CASE WHEN ? THEN 1 ELSE 0 END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(restored_auth.email.as_deref()) - .bind(restored_auth.email_verified) - .bind(&restored_auth.username) - .bind(&restored_auth.role) - .bind(optional_string_list_json( - restored_auth.allowed_providers.clone(), - "users.allowed_providers", - )?) - .bind(&restored_auth.allowed_providers_mode) - .bind(optional_string_list_json( - restored_auth.allowed_api_formats.clone(), - "users.allowed_api_formats", - )?) - .bind(&restored_auth.allowed_api_formats_mode) - .bind(optional_string_list_json( - restored_auth.allowed_models.clone(), - "users.allowed_models", - )?) - .bind(&restored_auth.allowed_models_mode) - .bind(restored_export.rate_limit) - .bind(&restored_export.rate_limit_mode) - .bind(optional_json_string( - restored_model_capability_settings.clone(), - "users.model_capability_settings", - )?) - .bind(optional_json_string( - restored_feature_settings.clone(), - "users.feature_settings", - )?) - .bind(restored_auth.is_active) - .bind(security_state_changed) - .bind(current_unix_secs()) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - if security_state_changed { - let now = current_unix_secs(); - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'user_security_state_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(now) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE api_keys SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE management_tokens SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn update_local_auth_user_password_hash( - &self, - user_id: &str, - password_hash: String, - updated_at: DateTime, - ) -> Result, DataLayerError> { - let result = sqlx::query( - "UPDATE users SET password_hash = ?, security_version = security_version + 1, updated_at = ? WHERE id = ?", - ) - .bind(password_hash) - .bind(updated_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn restore_local_auth_user_password_hash_if_matches( - &self, - user_id: &str, - expected_password_hash: Option<&str>, - password_hash: Option, - updated_at: DateTime, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE users -SET password_hash = ?, - security_version = security_version + 1, - updated_at = ? -WHERE id = ? - AND ((? IS NULL AND password_hash IS NULL) OR BINARY password_hash = BINARY ?) -"#, - ) - .bind(password_hash) - .bind(updated_at.timestamp()) - .bind(user_id) - .bind(expected_password_hash) - .bind(expected_password_hash) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn reset_local_auth_user_password_and_revoke_sessions( - &self, - user_id: &str, - password_hash: String, - changed_at: DateTime, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let updated = sqlx::query( - "UPDATE users SET password_hash = ?, security_version = security_version + 1, updated_at = ? WHERE id = ? AND is_deleted = 0", - ) - .bind(password_hash) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'admin_password_reset', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(changed_at.timestamp()) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn change_local_auth_password_and_revoke_sessions( - &self, - user_id: &str, - current_session_id: &str, - expected_password_hash: Option<&str>, - next_password_hash: String, - changed_at: DateTime, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let updated = sqlx::query( - r#" -UPDATE users -SET password_hash = ?, security_version = security_version + 1, updated_at = ? -WHERE id = ? - AND is_active = 1 - AND is_deleted = 0 - AND ((? IS NULL AND password_hash IS NULL) OR BINARY password_hash = BINARY ?) - AND EXISTS ( - SELECT 1 FROM user_sessions - WHERE user_id = ? AND id = ? AND revoked_at IS NULL AND expires_at > ? - ) -"#, - ) - .bind(next_password_hash) - .bind(changed_at.timestamp()) - .bind(user_id) - .bind(expected_password_hash) - .bind(expected_password_hash) - .bind(user_id) - .bind(current_session_id) - .bind(changed_at.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let revoked = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'password_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(changed_at.timestamp()) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if revoked.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn update_local_auth_user_admin_fields( - &self, - user_id: &str, - role: Option, - allowed_providers_present: bool, - allowed_providers: Option>, - allowed_api_formats_present: bool, - allowed_api_formats: Option>, - allowed_models_present: bool, - allowed_models: Option>, - rate_limit_present: bool, - rate_limit: Option, - is_active: Option, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let active_admin_ids = sqlx::query_scalar::<_, String>(MYSQL_LOCK_ACTIVE_ADMINS_SQL) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let current_security_state = - sqlx::query("SELECT role, is_active, is_deleted FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(current_security_state) = current_security_state else { - tx.rollback().await.map_sql_err()?; - return Ok(None); - }; - let current_role = current_security_state - .try_get::("role") - .map_sql_err()?; - let current_active = current_security_state - .try_get::("is_active") - .map_sql_err()?; - let current_deleted = current_security_state - .try_get::("is_deleted") - .map_sql_err()?; - let next_role = role.as_deref().unwrap_or(current_role.as_str()); - let next_active = is_active.unwrap_or(current_active); - if current_role.eq_ignore_ascii_case("admin") - && current_active - && !current_deleted - && (!next_role.eq_ignore_ascii_case("admin") || !next_active) - && active_admin_ids.len() <= 1 - { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_UPDATE_DENIED.to_string(), - )); - } - let security_state_changed = - !current_role.eq_ignore_ascii_case(next_role) || current_active != next_active; - let allowed_providers_mode = if allowed_providers - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_api_formats_mode = if allowed_api_formats - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_models_mode = if allowed_models - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let rate_limit_mode = if rate_limit.is_some() { - "custom" - } else { - "system" - }; - let result = sqlx::query( - r#" -UPDATE users -SET role = CASE WHEN ? THEN COALESCE(?, role) ELSE role END, - allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, - allowed_providers_mode = CASE WHEN ? THEN ? ELSE allowed_providers_mode END, - allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, - allowed_api_formats_mode = CASE WHEN ? THEN ? ELSE allowed_api_formats_mode END, - allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END, - allowed_models_mode = CASE WHEN ? THEN ? ELSE allowed_models_mode END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - rate_limit_mode = CASE WHEN ? THEN ? ELSE rate_limit_mode END, - is_active = CASE WHEN ? THEN ? ELSE is_active END, - security_version = security_version + CASE WHEN ? THEN 1 ELSE 0 END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(role.is_some()) - .bind(role) - .bind(allowed_providers_present) - .bind(optional_string_list_json( - allowed_providers, - "users.allowed_providers", - )?) - .bind(allowed_providers_present) - .bind(allowed_providers_mode) - .bind(allowed_api_formats_present) - .bind(optional_string_list_json( - allowed_api_formats, - "users.allowed_api_formats", - )?) - .bind(allowed_api_formats_present) - .bind(allowed_api_formats_mode) - .bind(allowed_models_present) - .bind(optional_string_list_json( - allowed_models, - "users.allowed_models", - )?) - .bind(allowed_models_present) - .bind(allowed_models_mode) - .bind(rate_limit_present) - .bind(rate_limit) - .bind(rate_limit_present) - .bind(rate_limit_mode) - .bind(is_active.is_some()) - .bind(is_active) - .bind(security_state_changed) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - if security_state_changed { - let revoked_at = chrono::Utc::now().timestamp(); - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'user_security_state_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(revoked_at) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE api_keys SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE management_tokens SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.find_user_auth_by_id(user_id).await - } - - async fn update_local_auth_user_policy_modes( - &self, - user_id: &str, - allowed_providers_mode: Option, - allowed_api_formats_mode: Option, - allowed_models_mode: Option, - rate_limit_mode: Option, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE users -SET allowed_providers_mode = CASE WHEN ? THEN ? ELSE allowed_providers_mode END, - allowed_api_formats_mode = CASE WHEN ? THEN ? ELSE allowed_api_formats_mode END, - allowed_models_mode = CASE WHEN ? THEN ? ELSE allowed_models_mode END, - rate_limit_mode = CASE WHEN ? THEN ? ELSE rate_limit_mode END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(allowed_providers_mode.is_some()) - .bind(allowed_providers_mode) - .bind(allowed_api_formats_mode.is_some()) - .bind(allowed_api_formats_mode) - .bind(allowed_models_mode.is_some()) - .bind(allowed_models_mode) - .bind(rate_limit_mode.is_some()) - .bind(rate_limit_mode) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn update_user_model_capability_settings( - &self, - user_id: &str, - settings: Option, - ) -> Result, DataLayerError> { - let normalized = normalize_optional_json_value(settings); - let result = sqlx::query( - "UPDATE users SET model_capability_settings = ?, updated_at = ? WHERE id = ?", - ) - .bind(optional_json_string( - normalized.clone(), - "users.model_capability_settings", - )?) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - Ok(normalized) - } - - async fn update_user_feature_settings( - &self, - user_id: &str, - settings: Option, - ) -> Result, DataLayerError> { - let normalized = normalize_optional_json_value(settings); - let result = - sqlx::query("UPDATE users SET feature_settings = ?, updated_at = ? WHERE id = ?") - .bind(optional_json_string( - normalized.clone(), - "users.feature_settings", - )?) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - Ok(normalized) - } - - async fn create_local_auth_user( - &self, - email: Option, - email_verified: bool, - username: String, - password_hash: String, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - let now = chrono::Utc::now().timestamp(); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, rate_limit_mode, - is_active, is_deleted, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, 'user', 'local', 'inherit', 'inherit', 'inherit', 'inherit', 1, 0, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(password_hash) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn create_local_auth_user_with_settings( - &self, - email: Option, - email_verified: bool, - username: String, - password_hash: String, - role: String, - allowed_providers: Option>, - allowed_api_formats: Option>, - allowed_models: Option>, - rate_limit: Option, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - let now = chrono::Utc::now().timestamp(); - let allowed_providers_mode = if allowed_providers - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_api_formats_mode = if allowed_api_formats - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_models_mode = if allowed_models - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let rate_limit_mode = if rate_limit.is_some() { - "custom" - } else { - "system" - }; - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers, allowed_providers_mode, - allowed_api_formats, allowed_api_formats_mode, - allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, - is_active, is_deleted, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 'local', ?, ?, ?, ?, ?, ?, ?, ?, 1, 0, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(password_hash) - .bind(role) - .bind(optional_string_list_json( - allowed_providers, - "users.allowed_providers", - )?) - .bind(allowed_providers_mode) - .bind(optional_string_list_json( - allowed_api_formats, - "users.allowed_api_formats", - )?) - .bind(allowed_api_formats_mode) - .bind(optional_string_list_json( - allowed_models, - "users.allowed_models", - )?) - .bind(allowed_models_mode) - .bind(rate_limit) - .bind(rate_limit_mode) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn delete_local_auth_user(&self, user_id: &str) -> Result { - self.delete_local_auth_user_inner(user_id, false).await - } - - async fn delete_local_auth_user_if_wallet_absent( - &self, - user_id: &str, - ) -> Result { - self.delete_local_auth_user_inner(user_id, true).await - } - - async fn count_active_admin_users(&self) -> Result { - let total: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM users -WHERE LOWER(role) = 'admin' - AND is_deleted = 0 - AND is_active = 1 -"#, - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(total.max(0) as u64) - } - - async fn read_user_preferences( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_PREFERENCES_COLUMNS); - builder.push(" WHERE up.user_id = ").push_bind(user_id); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_preference_row).transpose() - } - - async fn write_user_preferences( - &self, - preferences: &StoredUserPreferenceRecord, - ) -> Result, DataLayerError> { - let now = Utc::now().timestamp(); - sqlx::query( - r#" -INSERT INTO user_preferences ( - id, user_id, avatar_url, bio, default_provider_id, theme, language, timezone, - email_notifications, usage_alerts, announcement_notifications, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - avatar_url = VALUES(avatar_url), - bio = VALUES(bio), - default_provider_id = VALUES(default_provider_id), - theme = VALUES(theme), - language = VALUES(language), - timezone = VALUES(timezone), - email_notifications = VALUES(email_notifications), - usage_alerts = VALUES(usage_alerts), - announcement_notifications = VALUES(announcement_notifications), - updated_at = VALUES(updated_at) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&preferences.user_id) - .bind(preferences.avatar_url.as_deref()) - .bind(preferences.bio.as_deref()) - .bind(preferences.default_provider_id.as_deref()) - .bind(&preferences.theme) - .bind(&preferences.language) - .bind(&preferences.timezone) - .bind(preferences.email_notifications) - .bind(preferences.usage_alerts) - .bind(preferences.announcement_notifications) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.read_user_preferences(&preferences.user_id).await - } - - async fn find_user_session( - &self, - user_id: &str, - session_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(user_id) - .push(" AND id = ") - .push_bind(session_id) - .push(" LIMIT 1"); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_session_row).transpose() - } - - async fn list_user_sessions( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(user_id) - .push(" AND revoked_at IS NULL AND expires_at > ") - .push_bind(Utc::now().timestamp()) - .push(" ORDER BY last_seen_at DESC, created_at DESC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_session_row).collect() - } - - async fn create_user_session( - &self, - session: &StoredUserSessionRecord, - ) -> Result, DataLayerError> { - let now = session - .created_at - .or(session.updated_at) - .or(session.last_seen_at) - .unwrap_or_else(Utc::now); - let mut tx = self.pool.begin().await.map_sql_err()?; - let user_is_eligible: Option = sqlx::query_scalar( - "SELECT id FROM users WHERE id = ? AND is_active = 1 AND is_deleted = 0 AND security_version = ? FOR UPDATE", - ) - .bind(&session.user_id) - .bind(session.security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_is_eligible.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query( - r#" -UPDATE user_sessions -SET revoked_at = ?, revoke_reason = 'replaced_by_new_login', updated_at = ? -WHERE user_id = ? AND client_device_id = ? AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(now.timestamp()) - .bind(now.timestamp()) - .bind(&session.user_id) - .bind(&session.client_device_id) - .bind(now.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -INSERT INTO user_sessions ( - id, user_id, security_version, client_device_id, device_label, device_type, ip_address, user_agent, - refresh_token_hash, last_seen_at, expires_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, 'unknown', ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&session.id) - .bind(&session.user_id) - .bind(session.security_version) - .bind(&session.client_device_id) - .bind(session.device_label.as_deref()) - .bind(session.ip_address.as_deref()) - .bind(session.user_agent.as_deref()) - .bind(&session.refresh_token_hash) - .bind(session.last_seen_at.unwrap_or(now).timestamp()) - .bind(session.expires_at.unwrap_or(now).timestamp()) - .bind(session.created_at.unwrap_or(now).timestamp()) - .bind(session.updated_at.unwrap_or(now).timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(&session.user_id) - .push(" AND id = ") - .push_bind(&session.id) - .push(" LIMIT 1"); - let row = builder - .build() - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let created = row.as_ref().map(map_user_session_row).transpose()?; - tx.commit().await.map_sql_err()?; - Ok(created) - } - - async fn create_user_session_if_password_matches( - &self, - session: &StoredUserSessionRecord, - expected_password_hash: &str, - ) -> Result, DataLayerError> { - let now = session - .created_at - .or(session.updated_at) - .or(session.last_seen_at) - .unwrap_or_else(Utc::now); - let mut tx = self.pool.begin().await.map_sql_err()?; - let matched = sqlx::query_scalar::<_, String>( - r#" -SELECT password_hash FROM users -WHERE id = ? AND BINARY password_hash = BINARY ? AND LOWER(auth_source) = 'local' - AND is_active = 1 AND is_deleted = 0 AND security_version = ? -FOR UPDATE -"#, - ) - .bind(&session.user_id) - .bind(expected_password_hash) - .bind(session.security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if matched.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query("UPDATE users SET last_login_at = ? WHERE id = ?") - .bind(now.timestamp()) - .bind(&session.user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE user_sessions -SET revoked_at = ?, revoke_reason = 'replaced_by_new_login', updated_at = ? -WHERE user_id = ? AND client_device_id = ? AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(now.timestamp()) - .bind(now.timestamp()) - .bind(&session.user_id) - .bind(&session.client_device_id) - .bind(now.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -INSERT INTO user_sessions ( - id, user_id, security_version, client_device_id, device_label, device_type, ip_address, user_agent, - refresh_token_hash, last_seen_at, expires_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, 'unknown', ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&session.id) - .bind(&session.user_id) - .bind(session.security_version) - .bind(&session.client_device_id) - .bind(session.device_label.as_deref()) - .bind(session.ip_address.as_deref()) - .bind(session.user_agent.as_deref()) - .bind(&session.refresh_token_hash) - .bind(session.last_seen_at.unwrap_or(now).timestamp()) - .bind(session.expires_at.unwrap_or(now).timestamp()) - .bind(session.created_at.unwrap_or(now).timestamp()) - .bind(session.updated_at.unwrap_or(now).timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - self.find_user_session(&session.user_id, &session.id).await - } - - async fn touch_user_session( - &self, - user_id: &str, - session_id: &str, - touched_at: DateTime, - ip_address: Option<&str>, - user_agent: Option<&str>, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_sessions -SET last_seen_at = ?, ip_address = COALESCE(?, ip_address), - user_agent = COALESCE(?, user_agent), updated_at = ? -WHERE user_id = ? AND id = ? -"#, - ) - .bind(touched_at.timestamp()) - .bind(ip_address) - .bind(user_agent.map(|value| value.chars().take(1000).collect::())) - .bind(touched_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn update_user_session_device_label( - &self, - user_id: &str, - session_id: &str, - device_label: &str, - updated_at: DateTime, - ) -> Result { - let result = sqlx::query( - "UPDATE user_sessions SET device_label = ?, updated_at = ? WHERE user_id = ? AND id = ?", - ) - .bind(device_label.chars().take(120).collect::()) - .bind(updated_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn rotate_user_session_refresh_token( - &self, - user_id: &str, - session_id: &str, - expected_refresh_token_hash: &str, - next_refresh_token_hash: &str, - rotated_at: DateTime, - expires_at: DateTime, - ip_address: Option<&str>, - user_agent: Option<&str>, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_sessions -SET prev_refresh_token_hash = ?, rotated_at = ?, refresh_token_hash = ?, - expires_at = ?, last_seen_at = ?, ip_address = COALESCE(?, ip_address), - user_agent = COALESCE(?, user_agent), updated_at = ? -WHERE user_id = ? AND id = ? AND BINARY refresh_token_hash = BINARY ? - AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(expected_refresh_token_hash) - .bind(rotated_at.timestamp()) - .bind(next_refresh_token_hash) - .bind(expires_at.timestamp()) - .bind(rotated_at.timestamp()) - .bind(ip_address) - .bind(user_agent.map(|value| value.chars().take(1000).collect::())) - .bind(rotated_at.timestamp()) - .bind(user_id) - .bind(session_id) - .bind(expected_refresh_token_hash) - .bind(rotated_at.timestamp()) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn revoke_user_session( - &self, - user_id: &str, - session_id: &str, - revoked_at: DateTime, - reason: &str, - ) -> Result { - let result = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = ?, updated_at = ? WHERE user_id = ? AND id = ?", - ) - .bind(revoked_at.timestamp()) - .bind(reason.chars().take(100).collect::()) - .bind(revoked_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn revoke_all_user_sessions( - &self, - user_id: &str, - revoked_at: DateTime, - reason: &str, - ) -> Result { - let result = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = ?, updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(revoked_at.timestamp()) - .bind(reason.chars().take(100).collect::()) - .bind(revoked_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected()) - } - - async fn count_active_local_admin_users_with_valid_password( - &self, - ) -> Result { - let hashes = sqlx::query_scalar::<_, String>( - r#" -SELECT password_hash -FROM users -WHERE LOWER(role) = 'admin' - AND LOWER(auth_source) = 'local' - AND is_deleted = 0 - AND is_active = 1 - AND password_hash IS NOT NULL -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(hashes - .iter() - .filter(|hash| is_valid_bcrypt_hash(hash)) - .count() as u64) - } -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_string_list_json( - value: Option>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} could not be serialized as JSON: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_option_vec(value: Option<&Vec>) -> Option { - value.and_then(|items| serde_json::to_string(items).ok()) -} - -fn normalized_ids(values: &[String]) -> Vec { - values - .iter() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - .collect::>() - .into_iter() - .collect() -} - -fn current_unix_secs() -> i64 { - chrono::Utc::now().timestamp() -} - -fn optional_json_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} could not be serialized as JSON: {err}" - )) - }) - }) - .transpose() -} - -fn normalize_optional_json_value(value: Option) -> Option { - match value { - Some(serde_json::Value::Null) | None => None, - Some(value) => Some(value), - } -} - -async fn get_or_create_mysql_ldap_auth_user( - pool: &MysqlPool, - email: String, - username: String, - ldap_dn: Option, - ldap_username: Option, - logged_in_at: DateTime, -) -> Result, DataLayerError> { - let existing = - find_mysql_ldap_auth_user(pool, ldap_dn.as_deref(), ldap_username.as_deref(), &email) - .await?; - if let Some(existing) = existing { - if existing.is_deleted - || !existing.is_active - || !existing.auth_source.eq_ignore_ascii_case("ldap") - { - return Ok(None); - } - if existing.email.as_deref() != Some(email.as_str()) { - let taken: Option = - sqlx::query_scalar("SELECT 1 FROM users WHERE email = ? AND id <> ? LIMIT 1") - .bind(&email) - .bind(&existing.id) - .fetch_optional(pool) - .await - .map_sql_err()?; - if taken.is_some() { - return Ok(None); - } - } - sqlx::query("UPDATE users SET email = ?, email_verified = 1, ldap_dn = COALESCE(?, ldap_dn), ldap_username = COALESCE(?, ldap_username), last_login_at = ?, updated_at = ? WHERE id = ?") - .bind(&email) - .bind(ldap_dn.as_deref()) - .bind(ldap_username.as_deref()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(&existing.id) - .execute(pool) - .await - .map_sql_err()?; - let user = find_mysql_auth_by_id(pool, &existing.id) - .await? - .ok_or_else(|| { - DataLayerError::UnexpectedValue("updated LDAP user disappeared".to_string()) - })?; - return Ok(Some(LdapAuthUserProvisioningOutcome { - user, - created: false, - })); - } - - let base_username = ldap_username - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or(username.as_str()) - .trim() - .to_string(); - let mut candidate_username = base_username.clone(); - for _attempt in 0..3 { - let taken: Option = - sqlx::query_scalar("SELECT 1 FROM users WHERE username = ? LIMIT 1") - .bind(&candidate_username) - .fetch_optional(pool) - .await - .map_sql_err()?; - if taken.is_some() { - let suffix = uuid::Uuid::new_v4().simple().to_string(); - candidate_username = format!( - "{}_ldap_{}{}", - base_username, - logged_in_at.timestamp(), - &suffix[..4] - ); - continue; - } - let user_id = uuid::Uuid::new_v4().to_string(); - sqlx::query("INSERT INTO users (id, email, email_verified, username, password_hash, role, auth_source, ldap_dn, ldap_username, is_active, is_deleted, created_at, updated_at, last_login_at) VALUES (?, ?, 1, ?, NULL, 'user', 'ldap', ?, ?, 1, 0, ?, ?, ?)") - .bind(&user_id) - .bind(&email) - .bind(&candidate_username) - .bind(ldap_dn.as_deref()) - .bind(ldap_username.as_deref()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .execute(pool) - .await - .map_sql_err()?; - let user = find_mysql_auth_by_id(pool, &user_id) - .await? - .ok_or_else(|| { - DataLayerError::UnexpectedValue("created LDAP user disappeared".to_string()) - })?; - return Ok(Some(LdapAuthUserProvisioningOutcome { - user, - created: true, - })); - } - Ok(None) -} - -async fn find_mysql_ldap_auth_user( - pool: &MysqlPool, - ldap_dn: Option<&str>, - ldap_username: Option<&str>, - email: &str, -) -> Result, DataLayerError> { - if let Some(ldap_dn) = ldap_dn.filter(|value| !value.trim().is_empty()) { - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS} WHERE auth_source = 'ldap' AND ldap_dn = ? LIMIT 1" - )) - .bind(ldap_dn) - .fetch_optional(pool) - .await - .map_sql_err()?; - if let Some(row) = row.as_ref() { - return map_user_auth_row(row).map(Some); - } - } - if let Some(ldap_username) = ldap_username.filter(|value| !value.trim().is_empty()) { - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS} WHERE auth_source = 'ldap' AND ldap_username = ? LIMIT 1" - )) - .bind(ldap_username) - .fetch_optional(pool) - .await - .map_sql_err()?; - if let Some(row) = row.as_ref() { - return map_user_auth_row(row).map(Some); - } - } - let row = sqlx::query(&format!("{USER_AUTH_COLUMNS} WHERE email = ? LIMIT 1")) - .bind(email) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_auth_row).transpose() -} - -async fn find_mysql_auth_by_id( - pool: &MysqlPool, - user_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USER_AUTH_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(user_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_auth_row).transpose() -} - -fn optional_datetime_from_unix_secs(value: Option) -> Option> { - value.and_then(|value| Utc.timestamp_opt(value, 0).single()) -} - -fn map_user_row(row: &MySqlRow) -> Result { - StoredUserSummary::new( - row.try_get("id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_deleted").map_sql_err()?, - ) -} - -fn map_user_export_row(row: &MySqlRow) -> Result { - let feature_settings = optional_json_from_string( - row.try_get("feature_settings").map_sql_err()?, - "users.feature_settings", - )?; - StoredUserExportRow::new( - row.try_get("id").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("email_verified").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("password_hash").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("auth_source").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("rate_limit").map_sql_err()?, - optional_json_from_string( - row.try_get("model_capability_settings").map_sql_err()?, - "users.model_capability_settings", - )?, - row.try_get("is_active").map_sql_err()?, - ) - .map(|record| record.with_feature_settings(feature_settings)) - .and_then(|record| { - record.with_policy_modes( - row.try_get("allowed_providers_mode").map_sql_err()?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - row.try_get("allowed_models_mode").map_sql_err()?, - row.try_get("rate_limit_mode").map_sql_err()?, - ) - }) -} - -fn map_user_auth_row(row: &MySqlRow) -> Result { - StoredUserAuthRecord::new( - row.try_get("id").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("email_verified").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("password_hash").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("auth_source").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_deleted").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_login_at").map_sql_err()?), - ) - .and_then(|record| record.with_security_version(row.try_get("security_version").map_sql_err()?)) - .and_then(|record| { - record.with_policy_modes( - row.try_get("allowed_providers_mode").map_sql_err()?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - row.try_get("allowed_models_mode").map_sql_err()?, - ) - }) -} - -fn map_user_group_row(row: &MySqlRow) -> Result { - StoredUserGroup::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("normalized_name").map_sql_err()?, - row.try_get("description").map_sql_err()?, - row.try_get("priority").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "user_groups.allowed_providers", - )?, - row.try_get("allowed_providers_mode").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "user_groups.allowed_api_formats", - )?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "user_groups.allowed_models", - )?, - row.try_get("allowed_models_mode").map_sql_err()?, - row.try_get("rate_limit").map_sql_err()?, - row.try_get("rate_limit_mode").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("updated_at").map_sql_err()?), - ) -} - -fn map_user_group_member_row(row: &MySqlRow) -> Result { - Ok(StoredUserGroupMember { - group_id: row.try_get("group_id").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - username: row.try_get("username").map_sql_err()?, - email: row.try_get("email").map_sql_err()?, - role: row.try_get("role").map_sql_err()?, - is_active: row.try_get("is_active").map_sql_err()?, - is_deleted: row.try_get("is_deleted").map_sql_err()?, - created_at: optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - }) -} - -fn map_user_group_membership_row( - row: &MySqlRow, -) -> Result { - Ok(StoredUserGroupMembership { - user_id: row.try_get("user_id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - group_name: row.try_get("group_name").map_sql_err()?, - group_priority: row.try_get("group_priority").map_sql_err()?, - created_at: optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - }) -} - -fn map_oauth_link_summary_row( - row: &MySqlRow, -) -> Result { - StoredUserOAuthLinkSummary::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("provider_username").map_sql_err()?, - row.try_get("provider_email").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("linked_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_login_at").map_sql_err()?), - row.try_get("provider_enabled").map_sql_err()?, - ) -} - -fn map_user_preference_row(row: &MySqlRow) -> Result { - let user_id: String = row.try_get("user_id").map_sql_err()?; - if user_id.trim().is_empty() { - return Err(DataLayerError::UnexpectedValue( - "user_preferences.user_id is empty".to_string(), - )); - } - - Ok(StoredUserPreferenceRecord { - user_id, - avatar_url: row.try_get("avatar_url").map_sql_err()?, - bio: row.try_get("bio").map_sql_err()?, - default_provider_id: row.try_get("default_provider_id").map_sql_err()?, - default_provider_name: row.try_get("default_provider_name").map_sql_err()?, - theme: row.try_get("theme").map_sql_err()?, - language: row.try_get("language").map_sql_err()?, - timezone: row.try_get("timezone").map_sql_err()?, - email_notifications: row.try_get("email_notifications").map_sql_err()?, - usage_alerts: row.try_get("usage_alerts").map_sql_err()?, - announcement_notifications: row.try_get("announcement_notifications").map_sql_err()?, - }) -} - -fn map_user_session_row(row: &MySqlRow) -> Result { - StoredUserSessionRecord::new( - row.try_get("id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("client_device_id").map_sql_err()?, - row.try_get("device_label").map_sql_err()?, - row.try_get("refresh_token_hash").map_sql_err()?, - row.try_get("prev_refresh_token_hash").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("rotated_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_seen_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("expires_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("revoked_at").map_sql_err()?), - row.try_get("revoke_reason").map_sql_err()?, - row.try_get("ip_address").map_sql_err()?, - row.try_get("user_agent").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("updated_at").map_sql_err()?), - ) - .and_then(|record| record.with_security_version(row.try_get("security_version").map_sql_err()?)) -} - -#[cfg(test)] -mod tests { - use super::{ - MysqlUserReadRepository, MYSQL_ANONYMIZE_USER_API_KEY_HISTORY_SQL, - MYSQL_ANONYMIZE_USER_HISTORY_SQL, MYSQL_DELETE_USER_DEPENDENTS_SQL, - MYSQL_LOCK_ACTIVE_ADMINS_SQL, - }; - - #[test] - fn active_admin_mutations_use_a_deterministic_mysql_row_lock() { - let normalized = MYSQL_LOCK_ACTIVE_ADMINS_SQL - .split_whitespace() - .collect::>() - .join(" "); - assert!(normalized.contains("LOWER(role) = 'admin'")); - assert!(normalized.contains("is_active = 1")); - assert!(normalized.contains("is_deleted = 0")); - assert!(normalized.contains("ORDER BY id FOR UPDATE")); - assert!(MYSQL_DELETE_USER_DEPENDENTS_SQL - .iter() - .any(|sql| sql.starts_with("DELETE FROM management_tokens"))); - assert!(MYSQL_DELETE_USER_DEPENDENTS_SQL - .iter() - .any(|sql| sql.starts_with("DELETE FROM api_keys"))); - assert!(MYSQL_DELETE_USER_DEPENDENTS_SQL - .iter() - .any(|sql| sql.starts_with("DELETE FROM user_sessions"))); - assert_history_anonymization_contract(MYSQL_ANONYMIZE_USER_HISTORY_SQL); - assert!(MYSQL_ANONYMIZE_USER_API_KEY_HISTORY_SQL - .starts_with("UPDATE stats_daily_api_key SET api_key_name = NULL")); - assert!(MYSQL_ANONYMIZE_USER_API_KEY_HISTORY_SQL - .contains("SELECT id FROM api_keys WHERE user_id = ?")); - } - - fn assert_history_anonymization_contract(statements: &[&str]) { - const TABLES: &[&str] = &[ - "request_candidates", - "video_tasks", - "usage", - "stats_user_daily", - "stats_user_summary", - "stats_user_daily_model", - "stats_user_daily_provider", - "stats_user_daily_api_format", - "stats_user_daily_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", - ]; - - assert_eq!(statements.len(), TABLES.len()); - for table in TABLES { - let statement = statements - .iter() - .find(|sql| { - sql.starts_with(&format!("UPDATE {table} ")) - || sql.starts_with(&format!("UPDATE `{table}` ")) - }) - .unwrap_or_else(|| panic!("missing history anonymization for {table}")); - assert!(statement.contains("username = NULL")); - assert!(statement.ends_with("WHERE user_id = ?")); - } - for table in ["request_candidates", "video_tasks", "usage"] { - let statement = statements - .iter() - .find(|sql| { - sql.starts_with(&format!("UPDATE {table} ")) - || sql.starts_with(&format!("UPDATE `{table}` ")) - }) - .expect("identity snapshot table should be covered"); - assert!(statement.contains("api_key_name = NULL")); - } - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - let _repository = MysqlUserReadRepository::new(pool); - } -} diff --git a/crates/aether-data/adapters/mysql/src/video_tasks.rs b/crates/aether-data/adapters/mysql/src/video_tasks.rs deleted file mode 100644 index 5c132689d..000000000 --- a/crates/aether-data/adapters/mysql/src/video_tasks.rs +++ /dev/null @@ -1,1050 +0,0 @@ -use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::video_tasks::{ - StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount, - VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskStatusCount, - VideoTaskWriteRepository, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -const VIDEO_TASK_COLUMNS: &str = r#" -SELECT - id, - short_id, - request_id, - user_id, - api_key_id, - username, - api_key_name, - external_task_id, - provider_id, - endpoint_id, - key_id, - client_api_format, - provider_api_format, - format_converted, - model, - prompt, - original_request_body, - duration_seconds, - resolution, - aspect_ratio, - size, - status, - progress_percent, - progress_message, - retry_count, - poll_interval_seconds, - next_poll_at AS next_poll_at_unix_secs, - poll_count, - max_poll_count, - created_at AS created_at_unix_ms, - submitted_at AS submitted_at_unix_secs, - completed_at AS completed_at_unix_secs, - updated_at AS updated_at_unix_secs, - error_code, - error_message, - video_url, - request_metadata -FROM video_tasks -"#; - -#[derive(Debug, Clone)] -pub struct MysqlVideoTaskRepository { - pool: MysqlPool, -} - -impl MysqlVideoTaskRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } - - async fn find_by_id(&self, id: &str) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{VIDEO_TASK_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_id_for_user( - &self, - id: &str, - user_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE BINARY id = BINARY ? AND BINARY user_id = BINARY ? LIMIT 1" - )) - .bind(id) - .bind(user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_short_id( - &self, - short_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{VIDEO_TASK_COLUMNS} WHERE short_id = ? LIMIT 1")) - .bind(short_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_short_id_for_user( - &self, - short_id: &str, - user_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE BINARY short_id = BINARY ? AND BINARY user_id = BINARY ? LIMIT 1" - )) - .bind(short_id) - .bind(user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_user_external( - &self, - user_id: &str, - external_task_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE BINARY user_id = BINARY ? AND BINARY external_task_id = BINARY ? LIMIT 1" - )) - .bind(user_id) - .bind(external_task_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } -} - -#[async_trait] -impl VideoTaskReadRepository for MysqlVideoTaskRepository { - async fn find( - &self, - key: VideoTaskLookupKey<'_>, - ) -> Result, DataLayerError> { - match key { - VideoTaskLookupKey::Id(id) => self.find_by_id(id).await, - VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await, - VideoTaskLookupKey::UserExternal { - user_id, - external_task_id, - } => self.find_by_user_external(user_id, external_task_id).await, - } - } - - async fn find_for_user( - &self, - key: VideoTaskLookupKey<'_>, - user_id: &str, - ) -> Result, DataLayerError> { - match key { - VideoTaskLookupKey::Id(id) => self.find_by_id_for_user(id, user_id).await, - VideoTaskLookupKey::ShortId(short_id) => { - self.find_by_short_id_for_user(short_id, user_id).await - } - VideoTaskLookupKey::UserExternal { - user_id: lookup_user_id, - external_task_id, - } if lookup_user_id == user_id => { - self.find_by_user_external(user_id, external_task_id).await - } - VideoTaskLookupKey::UserExternal { .. } => Ok(None), - } - } - - async fn list_active(&self, limit: usize) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE status IN ('pending', 'submitted', 'queued', 'processing') ORDER BY updated_at DESC LIMIT ?" - )) - .bind(limit_i64(limit, "active video task limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_due( - &self, - now_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} 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 ?" - )) - .bind(u64_to_i64(now_unix_secs, "video task now")?) - .bind(limit_i64(limit, "due video task limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_page( - &self, - filter: &VideoTaskQueryFilter, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(VIDEO_TASK_COLUMNS); - push_filter(&mut builder, filter, None); - builder - .push(" ORDER BY created_at DESC, updated_at DESC LIMIT ") - .push_bind(limit_i64(limit, "video task page limit")?) - .push(" OFFSET ") - .push_bind(limit_i64(offset, "video task page offset")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_page_summary( - &self, - filter: &VideoTaskQueryFilter, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - self.list_page(filter, offset, limit).await - } - - async fn count(&self, filter: &VideoTaskQueryFilter) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - count_query(builder, &self.pool).await - } - - async fn count_by_status( - &self, - filter: &VideoTaskQueryFilter, - ) -> Result, DataLayerError> { - let mut builder = - QueryBuilder::::new("SELECT status, COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - builder.push(" GROUP BY status ORDER BY status ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(VideoTaskStatusCount { - status: VideoTaskStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - count: count_value(row.try_get("total").map_sql_err()?)?, - }) - }) - .collect() - } - - async fn count_distinct_users( - &self, - filter: &VideoTaskQueryFilter, - ) -> Result { - let mut builder = - QueryBuilder::::new("SELECT COUNT(DISTINCT user_id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - push_clause(&mut builder, "user_id IS NOT NULL"); - push_clause(&mut builder, "user_id <> ''"); - count_query(builder, &self.pool).await - } - - async fn top_models( - &self, - filter: &VideoTaskQueryFilter, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let mut builder = - QueryBuilder::::new("SELECT model, COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - push_clause(&mut builder, "model IS NOT NULL"); - push_clause(&mut builder, "model <> ''"); - builder - .push(" GROUP BY model ORDER BY total DESC, model ASC LIMIT ") - .push_bind(limit_i64(limit, "video task top models limit")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(VideoTaskModelCount { - model: row.try_get("model").map_sql_err()?, - count: count_value(row.try_get("total").map_sql_err()?)?, - }) - }) - .collect() - } - - async fn count_created_since( - &self, - filter: &VideoTaskQueryFilter, - created_since_unix_secs: u64, - ) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, Some(created_since_unix_secs)); - count_query(builder, &self.pool).await - } -} - -#[async_trait] -impl VideoTaskWriteRepository for MysqlVideoTaskRepository { - async fn upsert(&self, mut task: UpsertVideoTask) -> Result { - task.sanitize_for_persistence(); - let id = task.id.clone(); - let expected_identity = task.clone(); - bind_task(sqlx::query(upsert_sql()), task, true, false)? - .execute(&self.pool) - .await - .map_sql_err()?; - let stored = self.find_by_id(&id).await?.ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "video task {id} conflicts with persisted immutable identity" - )) - })?; - stored.ensure_immutable_identity_matches(&expected_identity)?; - Ok(stored) - } - - async fn update_if_active( - &self, - task: UpsertVideoTask, - ) -> Result, DataLayerError> { - let id = task.id.clone(); - let rows_affected = bind_task(sqlx::query(UPDATE_IF_ACTIVE_SQL), task, false, true)? - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - return Ok(None); - } - self.find_by_id(&id).await - } - - async fn claim_due( - &self, - now_unix_secs: u64, - claim_until_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - 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()); - } - - 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) - } -} - -const IMMUTABLE_IDENTITY_MATCH_SQL: &str = r#"BINARY id <=> BINARY VALUES(id) - AND BINARY short_id <=> BINARY VALUES(short_id) - AND BINARY request_id <=> BINARY VALUES(request_id) - AND BINARY user_id <=> BINARY VALUES(user_id) - AND BINARY api_key_id <=> BINARY VALUES(api_key_id) - AND BINARY external_task_id <=> BINARY VALUES(external_task_id) - AND BINARY provider_id <=> BINARY VALUES(provider_id) - AND BINARY endpoint_id <=> BINARY VALUES(endpoint_id) - AND BINARY key_id <=> BINARY VALUES(key_id) - AND BINARY client_api_format <=> BINARY VALUES(client_api_format) - AND BINARY provider_api_format <=> BINARY VALUES(provider_api_format) - AND format_converted <=> VALUES(format_converted) - AND BINARY model <=> BINARY VALUES(model) - AND duration_seconds <=> VALUES(duration_seconds) - AND BINARY resolution <=> BINARY VALUES(resolution) - AND BINARY aspect_ratio <=> BINARY VALUES(aspect_ratio) - AND BINARY size <=> BINARY VALUES(size)"#; - -const UPSERT_UPDATE_COLUMNS: &[&str] = &[ - "short_id", - "request_id", - "user_id", - "api_key_id", - "username", - "api_key_name", - "external_task_id", - "provider_id", - "endpoint_id", - "key_id", - "client_api_format", - "provider_api_format", - "format_converted", - "model", - "prompt", - "original_request_body", - "duration_seconds", - "resolution", - "aspect_ratio", - "size", - "status", - "progress_percent", - "progress_message", - "retry_count", - "poll_interval_seconds", - "next_poll_at", - "poll_count", - "max_poll_count", - "video_url", - "error_code", - "error_message", - "request_metadata", - "submitted_at", - "completed_at", - "updated_at", -]; - -fn upsert_sql() -> &'static str { - static SQL: std::sync::OnceLock = std::sync::OnceLock::new(); - SQL.get_or_init(|| { - let guarded_updates = UPSERT_UPDATE_COLUMNS - .iter() - .map(|column| { - format!( - " {column} = IF(({IMMUTABLE_IDENTITY_MATCH_SQL}), VALUES({column}), {column})" - ) - }) - .collect::>() - .join(",\n"); - format!( - r#" -INSERT INTO video_tasks ( - id, short_id, request_id, user_id, api_key_id, username, api_key_name, - external_task_id, provider_id, endpoint_id, key_id, client_api_format, - provider_api_format, format_converted, model, prompt, original_request_body, - duration_seconds, resolution, aspect_ratio, size, status, progress_percent, - progress_message, retry_count, poll_interval_seconds, next_poll_at, poll_count, - max_poll_count, video_url, error_code, error_message, request_metadata, - created_at, submitted_at, completed_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE -{guarded_updates} -"# - ) - }) -} - -const UPDATE_IF_ACTIVE_SQL: &str = r#" -UPDATE video_tasks SET - short_id = ?, - request_id = ?, - user_id = ?, - api_key_id = ?, - username = ?, - api_key_name = ?, - external_task_id = ?, - provider_id = ?, - endpoint_id = ?, - key_id = ?, - client_api_format = ?, - provider_api_format = ?, - format_converted = ?, - model = ?, - prompt = ?, - original_request_body = ?, - duration_seconds = ?, - resolution = ?, - aspect_ratio = ?, - size = ?, - status = ?, - progress_percent = ?, - progress_message = ?, - retry_count = ?, - poll_interval_seconds = ?, - next_poll_at = ?, - poll_count = ?, - max_poll_count = ?, - video_url = ?, - error_code = ?, - error_message = ?, - request_metadata = ?, - created_at = COALESCE(created_at, ?), - submitted_at = ?, - completed_at = ?, - updated_at = ? -WHERE id = ? - AND status IN ('pending', 'submitted', 'queued', 'processing') - AND BINARY short_id <=> BINARY ? - AND BINARY request_id <=> BINARY ? - AND BINARY user_id <=> BINARY ? - AND BINARY api_key_id <=> BINARY ? - AND BINARY external_task_id <=> BINARY ? - AND BINARY provider_id <=> BINARY ? - AND BINARY endpoint_id <=> BINARY ? - AND BINARY key_id <=> BINARY ? - AND BINARY client_api_format <=> BINARY ? - AND BINARY provider_api_format <=> BINARY ? - AND format_converted <=> ? - AND BINARY model <=> BINARY ? - AND duration_seconds <=> ? - AND BINARY resolution <=> BINARY ? - AND BINARY aspect_ratio <=> BINARY ? - AND BINARY size <=> BINARY ? -"#; - -fn bind_task<'q>( - query: sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments>, - mut task: UpsertVideoTask, - include_insert_id: bool, - include_update_id: bool, -) -> Result, DataLayerError> { - task.sanitize_for_persistence(); - let identity = task.clone(); - let original_request_body = json_to_string(&task.original_request_body)?; - let request_metadata = json_to_string(&task.request_metadata)?; - let query = if include_insert_id { - query.bind(task.id.clone()) - } else { - query - }; - let bound = query - .bind(task.short_id) - .bind(task.request_id) - .bind(task.user_id) - .bind(task.api_key_id) - .bind(task.username) - .bind(task.api_key_name) - .bind(task.external_task_id) - .bind(task.provider_id) - .bind(task.endpoint_id) - .bind(task.key_id) - .bind(task.client_api_format) - .bind(task.provider_api_format) - .bind(task.format_converted) - .bind(task.model) - .bind(task.prompt) - .bind(original_request_body) - .bind(optional_u32_to_i32( - task.duration_seconds, - "video task duration_seconds", - )?) - .bind(task.resolution) - .bind(task.aspect_ratio) - .bind(task.size) - .bind(status_to_database(task.status)) - .bind(i32::from(task.progress_percent)) - .bind(task.progress_message) - .bind(u32_to_i32(task.retry_count, "video task retry_count")?) - .bind(u32_to_i32( - task.poll_interval_seconds, - "video task poll_interval_seconds", - )?) - .bind(optional_u64_to_i64( - task.next_poll_at_unix_secs, - "video task next_poll_at", - )?) - .bind(u32_to_i32(task.poll_count, "video task poll_count")?) - .bind(u32_to_i32( - task.max_poll_count, - "video task max_poll_count", - )?) - .bind(task.video_url) - .bind(task.error_code) - .bind(task.error_message) - .bind(request_metadata) - .bind(u64_to_i64( - task.created_at_unix_ms, - "video task created_at", - )?) - .bind(optional_u64_to_i64( - task.submitted_at_unix_secs, - "video task submitted_at", - )?) - .bind(optional_u64_to_i64( - task.completed_at_unix_secs, - "video task completed_at", - )?) - .bind(u64_to_i64( - task.updated_at_unix_secs, - "video task updated_at", - )?); - if include_update_id { - bind_identity_guard(bound.bind(task.id), identity) - } else { - Ok(bound) - } -} - -fn bind_identity_guard<'q>( - query: sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments>, - identity: UpsertVideoTask, -) -> Result, DataLayerError> { - Ok(query - .bind(identity.short_id) - .bind(identity.request_id) - .bind(identity.user_id) - .bind(identity.api_key_id) - .bind(identity.external_task_id) - .bind(identity.provider_id) - .bind(identity.endpoint_id) - .bind(identity.key_id) - .bind(identity.client_api_format) - .bind(identity.provider_api_format) - .bind(identity.format_converted) - .bind(identity.model) - .bind(optional_u32_to_i32( - identity.duration_seconds, - "video task duration_seconds", - )?) - .bind(identity.resolution) - .bind(identity.aspect_ratio) - .bind(identity.size)) -} - -fn push_filter<'args>( - builder: &mut QueryBuilder<'args, MySql>, - filter: &'args VideoTaskQueryFilter, - created_since_unix_secs: Option, -) { - if let Some(user_id) = filter.user_id.as_deref() { - push_clause(builder, "BINARY user_id = BINARY "); - builder.push_bind(user_id); - } - if let Some(status) = filter.status { - push_clause(builder, "status = "); - builder.push_bind(status_to_database(status)); - } - if let Some(model_substring) = filter.model_substring.as_deref() { - push_clause(builder, "LOWER(model) LIKE "); - builder.push_bind(format!( - "%{}%", - escape_like_pattern(&model_substring.trim().to_ascii_lowercase()) - )); - builder.push(" ESCAPE '\\'"); - } - if let Some(client_api_format) = filter.client_api_format.as_deref() { - push_clause(builder, "client_api_format = "); - builder.push_bind(client_api_format); - } - if let Some(created_since_unix_secs) = created_since_unix_secs { - push_clause(builder, "created_at >= "); - builder.push_bind(created_since_unix_secs as i64); - } -} - -fn push_clause<'args>(builder: &mut QueryBuilder<'args, MySql>, clause: &str) { - let sql = builder.sql(); - if sql.contains(" WHERE ") || sql.contains("\nWHERE ") { - builder.push(" AND "); - } else { - builder.push(" WHERE "); - } - builder.push(clause); -} - -async fn count_query( - mut builder: QueryBuilder<'_, MySql>, - pool: &MysqlPool, -) -> Result { - let row = builder.build().fetch_one(pool).await.map_sql_err()?; - count_value(row.try_get("total").map_sql_err()?) -} - -fn count_value(value: i64) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid video task count result: {value}")) - }) -} - -fn map_video_task_row(row: &MySqlRow) -> Result { - StoredVideoTask::new( - row.try_get("id").map_sql_err()?, - row.try_get("short_id").map_sql_err()?, - row.try_get("request_id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("external_task_id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("endpoint_id").map_sql_err()?, - row.try_get("key_id").map_sql_err()?, - row.try_get("client_api_format").map_sql_err()?, - row.try_get("provider_api_format").map_sql_err()?, - row.try_get("format_converted").map_sql_err()?, - row.try_get("model").map_sql_err()?, - row.try_get("prompt").map_sql_err()?, - parse_json(row.try_get("original_request_body").ok().flatten())?, - row.try_get("duration_seconds").map_sql_err()?, - row.try_get("resolution").map_sql_err()?, - row.try_get("aspect_ratio").map_sql_err()?, - row.try_get("size").map_sql_err()?, - VideoTaskStatus::from_database(row.try_get::("status").map_sql_err()?.as_str())?, - row.try_get("progress_percent").map_sql_err()?, - row.try_get("progress_message").map_sql_err()?, - row.try_get("retry_count").map_sql_err()?, - row.try_get("poll_interval_seconds").map_sql_err()?, - row.try_get("next_poll_at_unix_secs").map_sql_err()?, - row.try_get("poll_count").map_sql_err()?, - row.try_get("max_poll_count").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("submitted_at_unix_secs").map_sql_err()?, - row.try_get("completed_at_unix_secs").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - row.try_get("error_code").map_sql_err()?, - row.try_get("error_message").map_sql_err()?, - row.try_get("video_url").map_sql_err()?, - parse_json(row.try_get("request_metadata").ok().flatten())?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("video task JSON field is invalid: {err}")) - }) - }) - .transpose() -} - -fn json_to_string(value: &Option) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "video task JSON field is unserializable: {err}" - )) - }) - }) - .transpose() -} - -fn status_to_database(status: VideoTaskStatus) -> &'static str { - match status { - VideoTaskStatus::Pending => "pending", - VideoTaskStatus::Submitted => "submitted", - VideoTaskStatus::Queued => "queued", - VideoTaskStatus::Processing => "processing", - VideoTaskStatus::Completed => "completed", - VideoTaskStatus::Failed => "failed", - VideoTaskStatus::Cancelled => "cancelled", - VideoTaskStatus::Expired => "expired", - VideoTaskStatus::Deleted => "deleted", - } -} - -fn escape_like_pattern(value: &str) -> String { - value - .replace('\\', "\\\\") - .replace('%', "\\%") - .replace('_', "\\_") -} - -fn limit_i64(value: usize, name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {name}: {value}"))) -} - -fn u64_to_i64(value: u64, name: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u64_to_i64(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u64_to_i64(value, name)).transpose() -} - -fn u32_to_i32(value: u32, name: &str) -> Result { - i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u32_to_i32(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u32_to_i32(value, name)).transpose() -} - -#[cfg(test)] -mod tests { - use super::{ - upsert_sql, MysqlVideoTaskRepository, IMMUTABLE_IDENTITY_MATCH_SQL, UPDATE_IF_ACTIVE_SQL, - UPSERT_UPDATE_COLUMNS, - }; - use crate::run_migrations; - use aether_data_contracts::repository::video_tasks::{ - UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository, - }; - use std::sync::Arc; - - #[test] - fn mysql_write_sql_atomically_guards_immutable_identity() { - assert!(IMMUTABLE_IDENTITY_MATCH_SQL.contains("BINARY id <=> BINARY VALUES(id)")); - for column in [ - "short_id", - "request_id", - "user_id", - "api_key_id", - "external_task_id", - "provider_id", - "endpoint_id", - "key_id", - "client_api_format", - "provider_api_format", - "model", - "resolution", - "aspect_ratio", - "size", - ] { - assert!( - IMMUTABLE_IDENTITY_MATCH_SQL - .contains(&format!("BINARY {column} <=> BINARY VALUES({column})")), - "upsert identity predicate should guard {column}" - ); - } - for column in ["format_converted", "duration_seconds"] { - assert!( - IMMUTABLE_IDENTITY_MATCH_SQL.contains(&format!("{column} <=> VALUES({column})")), - "upsert identity predicate should guard {column}" - ); - } - - let upsert = upsert_sql(); - assert!(!UPSERT_UPDATE_COLUMNS.contains(&"created_at")); - for column in UPSERT_UPDATE_COLUMNS { - assert!( - upsert.contains(&format!("{column} = IF((BINARY id <=> BINARY VALUES(id)")), - "upsert assignment should be conditional for {column}" - ); - } - for column in [ - "short_id", - "request_id", - "user_id", - "api_key_id", - "external_task_id", - "provider_id", - "endpoint_id", - "key_id", - "client_api_format", - "provider_api_format", - "model", - "resolution", - "aspect_ratio", - "size", - ] { - assert!( - UPDATE_IF_ACTIVE_SQL.contains(&format!("BINARY {column} <=> BINARY ?")), - "active update should guard {column}" - ); - } - for column in ["format_converted", "duration_seconds"] { - assert!( - UPDATE_IF_ACTIVE_SQL.contains(&format!("{column} <=> ?")), - "active update should guard {column}" - ); - } - assert!(UPDATE_IF_ACTIVE_SQL.contains("created_at = COALESCE(created_at, ?)")); - } - - #[tokio::test] - async fn repository_builds_from_lazy_pool() { - let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with( - "mysql://user:pass@localhost:3306/aether" - .parse() - .expect("mysql options should parse"), - ); - - 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 deleted file mode 100644 index 6418936a4..000000000 --- a/crates/aether-data/adapters/mysql/src/wallet.rs +++ /dev/null @@ -1,6778 +0,0 @@ -use async_trait::async_trait; -use chrono::Utc; -use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; - -use aether_data_contracts::repository::billing::{ - checked_plan_duration_days_from_snapshot, entitlements_have_replacement_selector, - entitlements_should_replace_existing, -}; -use aether_data_contracts::repository::wallet::{ - canonicalize_payment_method, canonicalize_wallet_refund_fields, - payment_callback_amount_matches_order, payment_callback_method_matches_order, - payment_callback_provider_matches_order, payment_order_is_failed_wallet_checkout_placeholder, - payment_order_is_uncertain_wallet_checkout_placeholder, - payment_order_refund_amounts_are_consistent, - payment_order_stripe_client_secret_cas_replacement, project_wallet_gateway_response, - project_wallet_recharge_gateway_response, redeem_code_payment_method, - redeem_code_refundable_amount, validate_admin_redeem_code_batch_input, - validate_manual_wallet_recharge, validate_payment_order_credit_amounts, - validate_plan_purchase_order_input, validate_plan_wallet_credit_entitlements, - validate_redeem_wallet_credit, validate_wallet_recharge_order_input, - wallet_recharge_checkout_claim_response, wallet_recharge_checkout_claim_token, - wallet_recharge_checkout_failed_response, wallet_recharge_checkout_uncertain_response, - wallet_recharge_order_is_checkout_placeholder, - wallet_recharge_order_is_reclaimable_placeholder, wallet_recharge_replay_matches, - wallet_recharge_response_is_checkout_placeholder, wallet_refund_proof_is_success, - AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, - AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery, - AdminWalletRefundRequestListQuery, CompareAndSwapPaymentOrderStripeClientSecretInput, - CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput, - CreateAdminRedeemCodeBatchResult, CreateManualWalletRechargeInput, - CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput, - CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput, - CreateWalletRefundRequestOutcome, CreatedAdminRedeemCodePlaintext, - CreditAdminPaymentOrderInput, DeleteAdminRedeemCodeBatchInput, - DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput, - FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome, ProcessAdminWalletRefundInput, - ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, ReclaimWalletRechargeCheckoutInput, - RedeemWalletCodeInput, RedeemWalletCodeOutcome, StoredAdminPaymentCallback, - StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, - StoredAdminRedeemCode, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, - StoredAdminRedeemCodePage, StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage, - StoredAdminWalletListItem, StoredAdminWalletListPage, StoredAdminWalletRefund, - StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem, - StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction, - StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger, - StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput, - UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome, - WalletReadRepository, WalletWriteRepository, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::MysqlPool; - -#[derive(Debug, Clone)] -pub struct MysqlWalletReadRepository { - pool: MysqlPool, -} - -impl MysqlWalletReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { pool } - } -} - -const ADMIN_WALLET_LIST_SELECT_SQL: &str = r#" -SELECT - 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 -"#; - -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, - payment_provider, payment_channel, order_kind, product_id, product_snapshot, - gateway_order_id, gateway_response, 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 1 = 1 -"#; - -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 -WHERE 1 = 1 -"#; - -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(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 ( - 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 -"#; - -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] -impl WalletReadRepository for MysqlWalletReadRepository { - async fn find( - &self, - key: WalletLookupKey<'_>, - ) -> Result, DataLayerError> { - 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( - &self, - user_id: &str, - limit_mode: &str, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE wallets SET limit_mode = ?, updated_at = ? WHERE user_id = ?") - .bind(limit_mode) - .bind(current_unix_secs_i64()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find(WalletLookupKey::UserId(user_id)).await - } - - async fn update_auth_api_key_wallet_limit_mode( - &self, - api_key_id: &str, - limit_mode: &str, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE wallets SET limit_mode = ?, updated_at = ? WHERE api_key_id = ?") - .bind(limit_mode) - .bind(current_unix_secs_i64()) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find(WalletLookupKey::ApiKeyId(api_key_id)).await - } - - async fn initialize_auth_user_wallet( - &self, - user_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_mysql_auth_wallet(&self.pool, Some(user_id), None, initial_gift_usd, unlimited) - .await - .map(|result| result.map(|(wallet, _created)| wallet)) - } - - async fn initialize_auth_user_wallet_with_outcome( - &self, - user_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_mysql_auth_wallet(&self.pool, Some(user_id), None, initial_gift_usd, unlimited) - .await - .map(|result| { - result.map(|(wallet, created)| InitializeAuthWalletOutcome { wallet, created }) - }) - } - - async fn initialize_auth_api_key_wallet( - &self, - api_key_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_mysql_auth_wallet( - &self.pool, - None, - Some(api_key_id), - initial_gift_usd, - unlimited, - ) - .await - .map(|result| result.map(|(wallet, _created)| wallet)) - } - - async fn initialize_auth_api_key_wallet_with_outcome( - &self, - api_key_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_mysql_auth_wallet( - &self.pool, - None, - Some(api_key_id), - initial_gift_usd, - unlimited, - ) - .await - .map(|result| { - result.map(|(wallet, created)| InitializeAuthWalletOutcome { wallet, created }) - }) - } - - async fn update_auth_user_wallet_snapshot( - &self, - user_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, - ) -> Result, DataLayerError> { - update_mysql_wallet_snapshot( - &self.pool, - "user_id", - user_id, - balance, - gift_balance, - limit_mode, - currency, - status, - total_recharged, - total_consumed, - total_refunded, - total_adjusted, - updated_at_unix_secs, - ) - .await?; - self.find(WalletLookupKey::UserId(user_id)).await - } - - async fn update_auth_api_key_wallet_snapshot( - &self, - api_key_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, - ) -> Result, DataLayerError> { - update_mysql_wallet_snapshot( - &self.pool, - "api_key_id", - api_key_id, - balance, - gift_balance, - limit_mode, - currency, - status, - total_recharged, - total_consumed, - total_refunded, - total_adjusted, - updated_at_unix_secs, - ) - .await?; - self.find(WalletLookupKey::ApiKeyId(api_key_id)).await - } - - async fn list_wallets_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - 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> { - 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 { - 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 { - 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 { - 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( - &self, - wallet_id: &str, - limit: usize, - offset: usize, - ) -> Result { - 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( - &self, - wallet_id: &str, - billing_timezone: &str, - ) -> Result, DataLayerError> { - let billing_date = current_billing_date(billing_timezone)?; - let sql = daily_usage_select_sql("AND billing_date = ? LIMIT 1"); - let row = sqlx::query(&sql) - .bind(wallet_id) - .bind(billing_timezone) - .bind(billing_date) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_daily_usage_row).transpose() - } - - async fn list_wallet_daily_usage_history( - &self, - wallet_id: &str, - billing_timezone: &str, - limit: usize, - ) -> Result { - let billing_date = current_billing_date(billing_timezone)?; - let total: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_timezone = ? - AND billing_date < ? -"#, - ) - .bind(wallet_id) - .bind(billing_timezone) - .bind(&billing_date) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let sql = daily_usage_select_sql("AND billing_date < ? ORDER BY billing_date DESC LIMIT ?"); - let rows = sqlx::query(&sql) - .bind(wallet_id) - .bind(billing_timezone) - .bind(billing_date) - .bind(i64_from_usize(limit, "wallet daily usage history limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_daily_usage_row) - .collect::, _>>()?; - Ok(StoredWalletDailyUsageLedgerPage { - items, - total: total.max(0) as u64, - }) - } - - async fn list_admin_wallet_refunds( - &self, - wallet_id: &str, - limit: usize, - offset: usize, - ) -> Result { - 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 { - 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> { - 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( - &self, - user_id: &str, - limit: usize, - offset: usize, - ) -> Result { - let total = read_count_row( - sqlx::query("SELECT COUNT(*) AS total FROM payment_orders WHERE user_id = ? AND order_kind = 'wallet_recharge'") - .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 = ? - AND order_kind = 'wallet_recharge' -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 { - 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 { - 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( - &self, - user_id: &str, - order_id: &str, - ) -> Result, DataLayerError> { - 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 = ? - AND order_kind = 'wallet_recharge' -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_wallet_recharge_order_by_order_no( - &self, - user_id: &str, - order_no: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql( - "WHERE user_id = ? AND order_no = ? AND order_kind = 'wallet_recharge' LIMIT 1", - ); - let row = sqlx::query(&sql) - .bind(user_id) - .bind(order_no) - .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( - &self, - user_id: &str, - product_id: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql( - r#" -WHERE user_id = ? - AND product_id = ? - AND order_kind = 'plan_purchase' - AND status = 'pending' - AND expires_at > ? -ORDER BY created_at DESC -LIMIT 1 -"#, - ); - let row = sqlx::query(&sql) - .bind(user_id) - .bind(product_id) - .bind(current_unix_secs_i64()) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_payment_order_row).transpose() - } - - async fn find_payment_order_by_order_no( - &self, - order_no: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE order_no = ? LIMIT 1"); - let row = sqlx::query(&sql) - .bind(order_no) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_payment_order_row).transpose() - } - - async fn find_wallet_refund( - &self, - wallet_id: &str, - refund_id: &str, - ) -> Result, DataLayerError> { - 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( - &self, - payment_method: Option<&str>, - limit: usize, - offset: usize, - ) -> Result { - 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 { - 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> { - 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 { - 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 }) - } -} - -#[async_trait] -impl WalletWriteRepository for MysqlWalletReadRepository { - async fn delete_wallet_if_unreferenced( - &self, - wallet_id: &str, - owner: WalletLookupKey<'_>, - ) -> Result { - if wallet_id.trim().is_empty() { - return Ok(false); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet compensation requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - let mut tx = self.pool.begin().await.map_sql_err()?; - let select_sql = format!( - r#" -SELECT id -FROM wallets -WHERE id = ? - AND {owner_clause} - AND balance = 0 - AND gift_balance = 0 - AND total_recharged = 0 - AND total_consumed = 0 - AND total_refunded = 0 - AND total_adjusted = 0 - AND limit_mode IN ('finite', 'unlimited') - AND currency = 'USD' - AND status = 'active' - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = wallets.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM `usage` u WHERE u.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = wallets.id - ) -LIMIT 1 -FOR UPDATE - "# - ); - let found = sqlx::query_scalar::<_, String>(&select_sql) - .bind(wallet_id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(found_id) = found else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let removed = sqlx::query("DELETE FROM wallets WHERE id = ?") - .bind(&found_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn delete_wallet_if_snapshot_matches_and_unreferenced( - &self, - expected: &StoredWalletSnapshot, - owner: WalletLookupKey<'_>, - ) -> Result { - if expected.id.trim().is_empty() { - return Ok(false); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet compensation requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - let mut tx = self.pool.begin().await.map_sql_err()?; - let select_sql = wallet_select_sql(&format!( - r#"WHERE id = ? - AND {owner_clause} - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = wallets.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM `usage` u WHERE u.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = wallets.id - ) -LIMIT 1 -FOR UPDATE"# - )); - let row = sqlx::query(&select_sql) - .bind(&expected.id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_wallet_row(&row)?; - if ¤t != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let removed = sqlx::query("DELETE FROM wallets WHERE id = ?") - .bind(&expected.id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn restore_wallet_if_snapshot_matches( - &self, - before: &StoredWalletSnapshot, - after: &StoredWalletSnapshot, - owner: WalletLookupKey<'_>, - ) -> Result { - if before.id.trim().is_empty() || after.id.trim().is_empty() { - return Ok(false); - } - if before.id != after.id { - return Err(DataLayerError::InvalidInput( - "wallet restore snapshots must reference the same wallet".to_string(), - )); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet restore requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - let owner_matches = match owner { - WalletLookupKey::UserId(user_id) => { - before.user_id.as_deref() == Some(user_id) - && after.user_id.as_deref() == Some(user_id) - && before.api_key_id.is_none() - && after.api_key_id.is_none() - } - WalletLookupKey::ApiKeyId(api_key_id) => { - before.api_key_id.as_deref() == Some(api_key_id) - && after.api_key_id.as_deref() == Some(api_key_id) - && before.user_id.is_none() - && after.user_id.is_none() - } - WalletLookupKey::WalletId(_) => false, - }; - if !owner_matches { - return Ok(false); - } - let before_updated_at = i64::try_from(before.updated_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput( - "wallet restore timestamp is outside the supported range".to_string(), - ) - })?; - - // Keep the row lock across compare and update so a concurrent wallet mutation cannot be - // mistaken for the import's own post-state. - let mut tx = self.pool.begin().await.map_sql_err()?; - let select_sql = wallet_select_sql(&format!( - "WHERE id = ? AND {owner_clause} LIMIT 1 FOR UPDATE" - )); - let row = sqlx::query(&select_sql) - .bind(&after.id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_wallet_row(&row)?; - if current != *after { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let updated = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - limit_mode = ?, - currency = ?, - status = ?, - total_recharged = ?, - total_consumed = ?, - total_refunded = ?, - total_adjusted = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(before.balance) - .bind(before.gift_balance) - .bind(&before.limit_mode) - .bind(&before.currency) - .bind(&before.status) - .bind(before.total_recharged) - .bind(before.total_consumed) - .bind(before.total_refunded) - .bind(before.total_adjusted) - .bind(before_updated_at) - .bind(&before.id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - if updated == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_provisional_auth_user_wallet( - &self, - wallet_id: &str, - user_id: &str, - ) -> Result { - if wallet_id.trim().is_empty() || user_id.trim().is_empty() { - return Ok(false); - } - let mut tx = self.pool.begin().await.map_sql_err()?; - let found_wallet_id = sqlx::query_scalar::<_, String>( - r#" -SELECT w.id -FROM wallets AS w -WHERE w.id = ? - AND w.user_id = ? - AND w.api_key_id IS NULL - AND w.balance = 0 - AND w.gift_balance >= 0 - AND w.total_recharged = 0 - AND w.total_consumed = 0 - AND w.total_refunded = 0 - AND w.total_adjusted = w.gift_balance - AND w.limit_mode IN ('finite', 'unlimited') - AND w.currency = 'USD' - AND w.status = 'active' - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = w.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = w.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = w.id - ) - AND NOT EXISTS (SELECT 1 FROM `usage` u WHERE u.wallet_id = w.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = w.id - ) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = w.id - ) - AND ( - (w.gift_balance = 0 AND NOT EXISTS ( - SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = w.id - )) - OR - (w.gift_balance > 0 - AND (SELECT COUNT(*) FROM wallet_transactions t WHERE t.wallet_id = w.id) = 1 - AND EXISTS ( - SELECT 1 FROM wallet_transactions t - WHERE t.wallet_id = w.id - AND t.category = 'gift' - AND t.reason_code = 'gift_initial' - AND t.amount = w.gift_balance - AND t.balance_before = 0 - AND t.balance_after = w.gift_balance - AND t.recharge_balance_before = 0 - AND t.recharge_balance_after = 0 - AND t.gift_balance_before = 0 - AND t.gift_balance_after = w.gift_balance - AND t.link_type = 'system_task' - AND t.link_id = w.user_id - AND t.operator_id IS NULL - )) - ) -LIMIT 1 -FOR UPDATE - "#, - ) - .bind(wallet_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(found_wallet_id) = found_wallet_id else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - sqlx::query("DELETE FROM wallet_transactions WHERE wallet_id = ?") - .bind(&found_wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let removed = sqlx::query("DELETE FROM wallets WHERE id = ? AND user_id = ?") - .bind(&found_wallet_id) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn create_wallet_recharge_order( - &self, - mut input: CreateWalletRechargeOrderInput, - ) -> Result { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - if !input.amount_usd.is_finite() || input.amount_usd <= 0.0 { - return Err(DataLayerError::InvalidInput( - "manual recharge amount must be finite and positive".to_string(), - )); - } - validate_wallet_recharge_order_input(&input).map_err(DataLayerError::InvalidInput)?; - if !input.amount_usd.is_finite() - || input.amount_usd <= 0.0 - || input - .pay_amount - .is_some_and(|value| !value.is_finite() || value <= 0.0) - || input - .exchange_rate - .is_some_and(|value| !value.is_finite() || value <= 0.0) - { - return Err(DataLayerError::InvalidInput( - "invalid wallet recharge numeric fields".to_string(), - )); - } - let projected_gateway_response = - project_wallet_recharge_gateway_response(&input.gateway_response) - .map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("wallet recharge expires_at overflow".to_string()) - })?; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - // `wallets.user_id` remains nullable to preserve deleted-user history; - // validate the live owner explicitly before creating a wallet/order. - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput("user not found".to_string())); - } - - let wallet_row = sqlx::query( - r#" -SELECT id, status -FROM wallets -WHERE user_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (wallet_id, wallet_status, created_wallet) = if let Some(row) = wallet_row { - ( - get::(&row, "id")?, - get::(&row, "status")?, - false, - ) - } else { - let requested_wallet_id = input - .preferred_wallet_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?) -ON DUPLICATE KEY UPDATE id = id -"#, - ) - .bind(&requested_wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = mysql_wallet_by_user_id_for_update(&mut tx, &input.user_id).await? - else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - }; - let wallet_id = get::(&row, "id")?; - let wallet_status = get::(&row, "status")?; - let created_wallet = - insert_result.rows_affected() > 0 && wallet_id == requested_wallet_id; - (wallet_id, wallet_status, created_wallet) - }; - if wallet_status != "active" { - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::WalletInactive); - } - - if let Some(existing_row) = - mysql_payment_order_by_order_no_for_update(&mut tx, &input.order_no).await? - { - let existing_user_id: Option = existing_row.try_get("user_id").map_sql_err()?; - let existing_kind: String = existing_row.try_get("order_kind").map_sql_err()?; - if existing_user_id.as_deref() == Some(input.user_id.as_str()) - && existing_kind == "wallet_recharge" - { - if !mysql_wallet_recharge_replay_matches(&existing_row, &wallet_id, &input)? { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge replay changes immutable order fields".to_string(), - )); - } - let existing = map_payment_order_row(&existing_row)?; - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::Existing(existing)); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment order number already belongs to another order".to_string(), - )); - } - - let order_id = uuid::Uuid::new_v4().to_string(); - let insert_result = sqlx::query( - r#" -INSERT INTO payment_orders ( - 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, fulfillment_status, - gateway_order_id, gateway_response, status, created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?, 'wallet_recharge', 'pending', ?, ?, 'pending', ?, ?) -ON DUPLICATE KEY UPDATE id = id -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(input.amount_usd) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(&input.payment_method) - .bind(input.payment_provider.as_deref()) - .bind(input.payment_channel.as_deref()) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(now) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = if insert_result.rows_affected() > 0 { - // A newly inserted row is identified by the generated id. A - // duplicate-key no-op has a different id and is resolved below. - if let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, &order_id).await? { - row - } else { - let Some(existing_row) = - mysql_payment_order_by_order_no_for_update(&mut tx, &input.order_no).await? - else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge order could not be created".to_string(), - )); - }; - let existing_user_id: Option = - existing_row.try_get("user_id").map_sql_err()?; - let existing_kind: String = existing_row.try_get("order_kind").map_sql_err()?; - let existing = map_payment_order_row(&existing_row)?; - if existing_user_id.as_deref() == Some(input.user_id.as_str()) - && existing_kind == "wallet_recharge" - { - if !mysql_wallet_recharge_replay_matches(&existing_row, &wallet_id, &input)? { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge replay changes immutable order fields".to_string(), - )); - } - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::Existing(existing)); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment order number already belongs to another order".to_string(), - )); - } - } else { - let Some(existing_row) = - mysql_payment_order_by_order_no_for_update(&mut tx, &input.order_no).await? - else { - if mysql_payment_order_by_gateway_order_id_for_update( - &mut tx, - &input.payment_method, - &input.gateway_order_id, - ) - .await? - .is_some() - { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment gateway order already belongs to another order".to_string(), - )); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge order could not be created".to_string(), - )); - }; - let existing_user_id: Option = existing_row.try_get("user_id").map_sql_err()?; - let existing_kind: String = existing_row.try_get("order_kind").map_sql_err()?; - let existing = map_payment_order_row(&existing_row)?; - if existing_user_id.as_deref() == Some(input.user_id.as_str()) - && existing_kind == "wallet_recharge" - { - if !mysql_wallet_recharge_replay_matches(&existing_row, &wallet_id, &input)? { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge replay changes immutable order fields".to_string(), - )); - } - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::Existing(existing)); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment order number already belongs to another order".to_string(), - )); - }; - tx.commit().await.map_sql_err()?; - Ok(CreateWalletRechargeOrderOutcome::Created( - map_payment_order_row(&row)?, - )) - } - - async fn update_wallet_recharge_checkout( - &self, - input: UpdateWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - if input.order_id.trim().is_empty() || input.gateway_order_id.trim().is_empty() { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout identifiers are required".to_string(), - )); - } - let projected_gateway_response = - match project_wallet_recharge_gateway_response(&input.gateway_response) { - Ok(value) => value, - Err(error) => return Ok(WalletMutationOutcome::Invalid(error)), - }; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_row) = - mysql_payment_order_by_id_for_update(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order_kind: Option = get(¤t_row, "order_kind")?; - if order_kind.as_deref() != Some("wallet_recharge") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not a wallet recharge".to_string(), - )); - } - let current = map_payment_order_row(¤t_row)?; - let current_is_checkout_placeholder = - wallet_recharge_order_is_checkout_placeholder(¤t); - let current_token = current - .gateway_response - .as_ref() - .and_then(wallet_recharge_checkout_claim_token); - let requested_token = wallet_recharge_checkout_claim_token(&projected_gateway_response); - if current_token.is_some() && current_token != requested_token { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - if current.status != "pending" { - if current.gateway_order_id.as_deref() == Some(input.gateway_order_id.as_str()) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(current)); - } - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is no longer pending".to_string(), - )); - } - let now = current_unix_secs_i64(); - if current - .expires_at_unix_secs - .is_none_or(|expires_at| expires_at <= now.max(0) as u64) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is expired".to_string(), - )); - } - // A newly-created row uses order_no as a temporary gateway id. Once - // the provider checkout is stored, do not let a concurrent request - // replace that checkout evidence. - if current.gateway_order_id.as_deref().is_some_and(|existing| { - existing != input.gateway_order_id.as_str() - && existing != current.order_no.as_str() - && !current_is_checkout_placeholder - }) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout is already bound".to_string(), - )); - } - if let Some(row) = mysql_payment_order_by_gateway_order_id_for_update( - &mut tx, - ¤t.payment_method, - &input.gateway_order_id, - ) - .await? - { - let existing_id: String = row.try_get("id").map_sql_err()?; - if existing_id != input.order_id { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment gateway order already belongs to another order".to_string(), - )); - } - } - let updated = sqlx::query( - "UPDATE payment_orders SET gateway_order_id = ?, gateway_response = ? WHERE id = ? AND status = 'pending' AND expires_at > ?", - ) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(&input.order_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is expired or no longer pending".to_string(), - )); - } - let updated = - map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn compare_and_swap_payment_order_stripe_client_secret( - &self, - input: CompareAndSwapPaymentOrderStripeClientSecretInput, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(false); - }; - let current = map_payment_order_row(&row)?; - let Some(replacement) = - payment_order_stripe_client_secret_cas_replacement(¤t, &input) - .map_err(DataLayerError::InvalidInput)? - else { - tx.commit().await.map_sql_err()?; - return Ok(false); - }; - let replacement = json_string(&replacement, "payment_orders.gateway_response")?; - let updated = sqlx::query("UPDATE payment_orders SET gateway_response = ? WHERE id = ?") - .bind(replacement) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(updated.rows_affected() == 1) - } - - async fn fail_wallet_recharge_checkout( - &self, - input: FailWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - if input.order_id.trim().is_empty() - || input.claim_token.trim().is_empty() - || input.claim_token.len() > 128 - { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout failure identifiers are required".to_string(), - )); - } - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if !wallet_recharge_order_is_checkout_placeholder(&order) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not a checkout placeholder".to_string(), - )); - } - let current_token = order - .gateway_response - .as_ref() - .and_then(wallet_recharge_checkout_claim_token); - if current_token != Some(input.claim_token.trim()) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - if order.status != "pending" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(order)); - } - let failed = if input.provider_request_may_have_succeeded { - wallet_recharge_checkout_uncertain_response( - order.gateway_response.as_ref(), - &input.reason, - current_unix_secs_i64().max(0) as u64, - ) - } else { - wallet_recharge_checkout_failed_response( - order.gateway_response.as_ref(), - &input.reason, - current_unix_secs_i64().max(0) as u64, - ) - }; - let failed = serde_json::to_string(&failed).map_err(|err| { - DataLayerError::UnexpectedValue(format!("payment_orders.gateway_response: {err}")) - })?; - let updated = sqlx::query( - "UPDATE payment_orders SET status = 'failed', gateway_response = ? WHERE id = ? AND status = 'pending'", - ) - .bind(failed) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - let updated = - map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn reclaim_wallet_recharge_checkout( - &self, - input: ReclaimWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64().max(0) as u64; - if input.order_id.trim().is_empty() - || input.claim_token.trim().is_empty() - || input.claim_token.len() > 128 - || input.expires_at_unix_secs <= now - { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout reclaim identifiers are invalid".to_string(), - )); - } - if !wallet_recharge_response_is_checkout_placeholder(&input.gateway_response) { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge reclaim response must be a placeholder".to_string(), - )); - } - let response = wallet_recharge_checkout_claim_response( - &input.gateway_response, - &input.claim_token, - now, - ) - .map_err(DataLayerError::InvalidInput)?; - let response = serde_json::to_string(&response).map_err(|err| { - DataLayerError::UnexpectedValue(format!("payment_orders.gateway_response: {err}")) - })?; - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("wallet recharge expires_at overflow".to_string()) - })?; - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order_kind: Option = get(&row, "order_kind")?; - let order = map_payment_order_row(&row)?; - if order_kind.as_deref() != Some("wallet_recharge") - || !wallet_recharge_order_is_reclaimable_placeholder(&order, now) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout is still in progress or already completed".to_string(), - )); - } - let updated = sqlx::query( - "UPDATE payment_orders SET gateway_order_id = ?, gateway_response = ?, status = 'pending', expires_at = ? WHERE id = ?", - ) - .bind(&order.order_no) - .bind(response) - .bind(expires_at) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout reclaim lost the order race".to_string(), - )); - } - let updated = - map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn create_plan_purchase_order( - &self, - mut input: CreatePlanPurchaseOrderInput, - ) -> Result { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - validate_plan_purchase_order_input(&input).map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("plan purchase expires_at overflow".to_string()) - })?; - let projected_gateway_response = project_wallet_gateway_response(&input.gateway_response) - .map_err(DataLayerError::InvalidInput)?; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - let product_snapshot = - json_string(&input.product_snapshot, "payment_orders.product_snapshot")?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - // MySQL's baseline schema does not declare a wallet owner foreign key. - // Lock and validate the user before any automatic wallet insert. - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput("user not found".to_string())); - } - - let wallet_row = sqlx::query( - r#" -SELECT id, status -FROM wallets -WHERE user_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (wallet_id, wallet_status) = if let Some(row) = wallet_row { - (get::(&row, "id")?, get::(&row, "status")?) - } else { - let requested_wallet_id = input - .preferred_wallet_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - // Resolve both the owner uniqueness race and a preferred wallet - // identifier collision through the owner row. This keeps plan - // checkout behavior deterministic across SQL backends. - let _insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?) -ON DUPLICATE KEY UPDATE id = id -"#, - ) - .bind(&requested_wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = mysql_wallet_by_user_id_for_update(&mut tx, &input.user_id).await? - else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - }; - (get::(&row, "id")?, get::(&row, "status")?) - }; - if wallet_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(CreatePlanPurchaseOrderOutcome::WalletInactive); - } - - let purchase_limit_scope = plan_purchase_limit_scope(&input.product_snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&input.product_snapshot); - let mut active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' -"#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? -"#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - active_count += sqlx::query_scalar::<_, i64>( - r#" - SELECT COUNT(*) - FROM payment_orders - WHERE user_id = ? - AND product_id = ? - AND order_kind = 'plan_purchase' - AND status = 'pending' - AND expires_at > ? - "#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - if active_count >= max_active_per_user { - tx.commit().await.map_sql_err()?; - return Ok(CreatePlanPurchaseOrderOutcome::ActivePlanLimitReached); - } - } - - let order_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO payment_orders ( - 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, - fulfillment_status, gateway_order_id, gateway_response, status, created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?, 'plan_purchase', ?, ?, 'pending', ?, ?, 'pending', ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(input.amount_usd) - .bind(input.pay_amount) - .bind(&input.pay_currency) - .bind(input.exchange_rate) - .bind(&input.payment_method) - .bind(input.payment_provider.as_deref()) - .bind(input.payment_channel.as_deref()) - .bind(&input.product_id) - .bind(product_snapshot) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(now) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = mysql_payment_order_by_id(&mut tx, &order_id).await?; - tx.commit().await.map_sql_err()?; - Ok(CreatePlanPurchaseOrderOutcome::Created( - map_payment_order_row(&row)?, - )) - } - - async fn create_wallet_refund_request( - &self, - input: CreateWalletRefundRequestInput, - ) -> Result { - if !input.amount_usd.is_finite() || input.amount_usd <= 0.0 { - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "refund amount must be finite and greater than zero".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - - let Some(wallet_row) = sqlx::query( - r#" -SELECT id, balance -FROM wallets -WHERE id = ? - AND user_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.wallet_id) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::WalletMissing); - }; - - if let Some(idempotency_key) = input.idempotency_key.as_deref() { - let existing = - mysql_refund_by_idempotency(&mut tx, &input.user_id, idempotency_key).await?; - if let Some(row) = existing { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::Duplicate(map_refund_row( - &row, - )?)); - } - } - - let wallet_recharge_balance: f64 = get(&wallet_row, "balance")?; - if !wallet_recharge_balance.is_finite() { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "wallet recharge balance is invalid".to_string(), - )); - } - let wallet_reserved_amount = sqlx::query_scalar::<_, Option>( - r#" -SELECT amount_usd -FROM refund_requests -WHERE wallet_id = ? - AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(&input.wallet_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()? - .into_iter() - .try_fold(0.0_f64, |total, amount| { - let amount = amount?; - if !amount.is_finite() || amount <= 0.0 { - return None; - } - let next = total + amount; - next.is_finite().then_some(next) - }); - let Some(wallet_reserved_amount) = wallet_reserved_amount else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "wallet refund reservation is invalid".to_string(), - )); - }; - if input.amount_usd > (wallet_recharge_balance - wallet_reserved_amount) { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::RefundAmountExceedsAvailableBalance); - } - - let mut payment_order_id = None; - let mut resolved_payment_method = None; - if let Some(order_id) = input.payment_order_id.as_deref() { - let Some(order_row) = sqlx::query( - r#" -SELECT id, status, payment_method, amount_usd, refunded_amount_usd, refundable_amount_usd -FROM payment_orders -WHERE id = ? - AND wallet_id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(order_id) - .bind(&input.wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::PaymentOrderNotFound); - }; - let status: String = get(&order_row, "status")?; - if status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::PaymentOrderNotRefundable); - } - let order_reserved_amount = sqlx::query_scalar::<_, Option>( - r#" -SELECT amount_usd -FROM refund_requests -WHERE payment_order_id = ? - AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(order_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()? - .into_iter() - .try_fold(0.0_f64, |total, amount| { - let amount = amount?; - if !amount.is_finite() || amount <= 0.0 { - return None; - } - let next = total + amount; - next.is_finite().then_some(next) - }); - let Some(order_reserved_amount) = order_reserved_amount else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "payment order refund reservation is invalid".to_string(), - )); - }; - let order_amount: f64 = get(&order_row, "amount_usd")?; - let refunded_amount: f64 = get(&order_row, "refunded_amount_usd")?; - let refundable_amount: f64 = get(&order_row, "refundable_amount_usd")?; - if !payment_order_refund_amounts_are_consistent( - order_amount, - refunded_amount, - refundable_amount, - ) { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "payment order refund amounts are invalid".to_string(), - )); - } - if input.amount_usd > (refundable_amount - order_reserved_amount) { - tx.commit().await.map_sql_err()?; - return Ok( - CreateWalletRefundRequestOutcome::RefundAmountExceedsAvailableOrderAmount, - ); - } - payment_order_id = Some(order_id.to_string()); - resolved_payment_method = Some(get::(&order_row, "payment_method")?); - } - - let canonical = canonicalize_wallet_refund_fields( - payment_order_id.as_deref(), - input.source_type.as_deref(), - input.source_id.as_deref(), - input.refund_mode.as_deref(), - resolved_payment_method.as_deref(), - ) - .map_err(DataLayerError::InvalidInput)?; - let source_type = canonical.source_type; - let source_id = canonical.source_id; - let refund_mode = canonical.refund_mode; - - let refund_id = uuid::Uuid::new_v4().to_string(); - let insert = sqlx::query( - r#" -INSERT INTO refund_requests ( - id, refund_no, wallet_id, user_id, payment_order_id, source_type, source_id, - refund_mode, amount_usd, status, reason, requested_by, idempotency_key, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending_approval', ?, ?, ?, ?, ?) -"#, - ) - .bind(&refund_id) - .bind(&input.refund_no) - .bind(&input.wallet_id) - .bind(&input.user_id) - .bind(payment_order_id.as_deref()) - .bind(&source_type) - .bind(source_id.as_deref()) - .bind(&refund_mode) - .bind(input.amount_usd) - .bind(input.reason.as_deref()) - .bind(&input.user_id) - .bind(input.idempotency_key.as_deref()) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await; - - if let Err(err) = insert { - if input.idempotency_key.is_some() - && err - .as_database_error() - .is_some_and(|database_error| database_error.is_unique_violation()) - { - tx.rollback().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::DuplicateRejected); - } - return Err(DataLayerError::sql(err)); - } - - let row = mysql_refund_by_id(&mut tx, &refund_id).await?; - tx.commit().await.map_sql_err()?; - Ok(CreateWalletRefundRequestOutcome::Created(map_refund_row( - &row, - )?)) - } - - async fn process_payment_callback( - &self, - mut input: ProcessPaymentCallbackInput, - ) -> Result { - input - .canonicalize_and_validate() - .map_err(DataLayerError::InvalidInput)?; - if input.callback_key.trim().is_empty() - || input.callback_key.chars().count() > 128 - || input.payload_hash.trim().is_empty() - || !input.amount_usd.is_finite() - || input.amount_usd <= 0.0 - || input - .pay_amount - .is_some_and(|value| !value.is_finite() || value <= 0.0) - || input - .exchange_rate - .is_some_and(|value| !value.is_finite() || value <= 0.0) - { - return Err(DataLayerError::InvalidInput( - "invalid payment callback numeric or identity fields".to_string(), - )); - } - let now = current_unix_secs_i64(); - let payload = json_string(&input.payload, "payment_callbacks.payload")?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - let candidate_callback_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO payment_callbacks ( - id, payment_order_id, payment_method, callback_key, order_no, gateway_order_id, - payload_hash, signature_valid, status, payload, error_message, created_at, processed_at -) -VALUES (?, NULL, ?, ?, ?, ?, ?, ?, 'received', NULL, NULL, ?, NULL) -ON DUPLICATE KEY UPDATE id = id -"#, - ) - .bind(&candidate_callback_id) - .bind(&input.payment_method) - .bind(&input.callback_key) - .bind(input.order_no.as_deref()) - .bind(input.gateway_order_id.as_deref()) - .bind(&input.payload_hash) - .bind(input.signature_valid) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - let callback_row = sqlx::query( - r#" -SELECT id, payment_order_id, payment_method, payload_hash, status, order_no, gateway_order_id -FROM payment_callbacks -WHERE callback_key = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.callback_key) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let callback_id: String = get(&callback_row, "id")?; - let duplicate = callback_id != candidate_callback_id; - let callback_order_no: Option = get(&callback_row, "order_no")?; - let callback_gateway_order_id: Option = get(&callback_row, "gateway_order_id")?; - let stored_method: String = get(&callback_row, "payment_method")?; - let stored_hash: Option = get(&callback_row, "payload_hash")?; - if !stored_method.eq_ignore_ascii_case(&input.payment_method) - || stored_hash.as_deref() != Some(input.payload_hash.as_str()) - { - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "callback key reused with different payment payload".to_string(), - }); - } - let status: String = get(&callback_row, "status")?; - if status == "processed" { - let order_id: Option = get(&callback_row, "payment_order_id")?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::DuplicateProcessed { order_id }); - } - - if !input.signature_valid { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "invalid callback signature", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "invalid callback signature".to_string(), - }); - } - - let lookup_order_no = input.order_no.clone().or_else(|| callback_order_no.clone()); - let lookup_gateway_order_id = input - .gateway_order_id - .clone() - .or_else(|| callback_gateway_order_id.clone()); - let order_row = if let Some(order_no) = lookup_order_no.as_deref() { - mysql_payment_order_by_order_no_for_update(&mut tx, order_no).await? - } else if let Some(gateway_order_id) = lookup_gateway_order_id.as_deref() { - mysql_payment_order_by_gateway_order_id_for_update( - &mut tx, - &input.payment_method, - gateway_order_id, - ) - .await? - } else { - None - }; - let Some(order_row) = order_row else { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order not found".to_string(), - }); - }; - - let order_id: String = get(&order_row, "id")?; - let order_no: String = get(&order_row, "order_no")?; - let order_wallet_id: String = get(&order_row, "wallet_id")?; - let order_payment_method: String = get(&order_row, "payment_method")?; - let order_payment_provider: Option = get(&order_row, "payment_provider")?; - let order_payment_channel: Option = get(&order_row, "payment_channel")?; - let order_pay_currency: Option = get(&order_row, "pay_currency")?; - let order_gateway_order_id: Option = get(&order_row, "gateway_order_id")?; - let order_kind: String = get(&order_row, "order_kind")?; - let order_amount_usd: f64 = get(&order_row, "amount_usd")?; - let order_pay_amount: Option = get(&order_row, "pay_amount")?; - let order_exchange_rate: Option = get(&order_row, "exchange_rate")?; - let order_status: String = get(&order_row, "status")?; - let expires_at_unix_secs: Option = get(&order_row, "expires_at_unix_secs")?; - let order_gateway_response = if order_status.eq_ignore_ascii_case("failed") { - optional_json( - get(&order_row, "gateway_response")?, - "payment_orders.gateway_response", - )? - } else { - None - }; - let failed_checkout_recoverable = payment_order_is_failed_wallet_checkout_placeholder( - &order_status, - &order_kind, - order_gateway_response.as_ref(), - ); - let uncertain_checkout = payment_order_is_uncertain_wallet_checkout_placeholder( - &order_status, - &order_kind, - order_gateway_response.as_ref(), - ); - if !order_amount_usd.is_finite() - || order_amount_usd <= 0.0 - || order_pay_amount.is_some_and(|value| !value.is_finite() || value <= 0.0) - { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order amount is invalid", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order amount is invalid".to_string(), - }); - } - - // A payment order may only credit a wallet owned by the same user. - // Lock the wallet (and its API-key owner row through the join) before - // any gateway binding, entitlement, wallet, or order mutation. Legacy - // rows with an ambiguous owner shape are deliberately rejected. - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(order_user_id) = order_user_id - .as_deref() - .filter(|value| !value.trim().is_empty()) - else { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order user missing", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order user missing".to_string(), - }); - }; - let Some(wallet_owner_row) = sqlx::query( - r#" -SELECT - w.user_id AS wallet_user_id, - w.api_key_id AS wallet_api_key_id, - api_keys.user_id AS api_key_user_id -FROM wallets AS w -LEFT JOIN api_keys ON api_keys.id = w.api_key_id -WHERE w.id = ? -LIMIT 1 -FOR UPDATE - "#, - ) - .bind(&order_wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet not found".to_string(), - }); - }; - let wallet_user_id: Option = get(&wallet_owner_row, "wallet_user_id")?; - let wallet_api_key_id: Option = get(&wallet_owner_row, "wallet_api_key_id")?; - let api_key_user_id: Option = get(&wallet_owner_row, "api_key_user_id")?; - let wallet_owner_matches = match ( - wallet_user_id.as_deref(), - wallet_api_key_id.as_deref(), - api_key_user_id.as_deref(), - ) { - (Some(wallet_user_id), None, _) if !wallet_user_id.trim().is_empty() => { - wallet_user_id == order_user_id - } - (None, Some(wallet_api_key_id), Some(api_key_user_id)) - if !wallet_api_key_id.trim().is_empty() && !api_key_user_id.trim().is_empty() => - { - api_key_user_id == order_user_id - } - _ => false, - }; - if !wallet_owner_matches { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order wallet owner mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order wallet owner mismatch".to_string(), - }); - } - - // The lookup identifier is not proof that the callback belongs to - // this order: order_no takes precedence over gateway_order_id. Check - // every identifier supplied by this delivery (and any persisted - // fallback from the callback row) before changing the order or - // wallet. Orders created before the gateway returns a provider - // transaction id store order_no as a placeholder; that value may be - // replaced by a verified callback, but a real id must never be - // rebound to another order. - if input - .order_no - .as_deref() - .is_some_and(|value| value != order_no) - || callback_order_no - .as_deref() - .is_some_and(|value| value != order_no) - { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order number mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order number mismatch".to_string(), - }); - } - let input_gateway_order_id = input - .gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - let callback_gateway_order_id = callback_gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - let stored_real_gateway_order_id = order_gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty() && *value != order_no); - let input_real_gateway_order_id = input_gateway_order_id.filter(|value| *value != order_no); - let callback_real_gateway_order_id = - callback_gateway_order_id.filter(|value| *value != order_no); - let effective_gateway_order_id = input_real_gateway_order_id - .or(callback_real_gateway_order_id) - .or(stored_real_gateway_order_id); - if let Some(expected_gateway_order_id) = stored_real_gateway_order_id { - if input_gateway_order_id.is_some_and(|value| value != expected_gateway_order_id) - || callback_gateway_order_id.is_some_and(|value| value != expected_gateway_order_id) - { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order mismatch".to_string(), - }); - } - } else if let (Some(input_gateway), Some(callback_gateway)) = - (input_real_gateway_order_id, callback_real_gateway_order_id) - { - if input_gateway != callback_gateway { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order identifier mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order identifier mismatch".to_string(), - }); - } - } - if stored_real_gateway_order_id.is_none() { - if let Some(gateway_order_id) = effective_gateway_order_id { - let conflicting_order_id: Option = sqlx::query_scalar( - "SELECT id FROM payment_orders WHERE payment_method = ? AND gateway_order_id = ? AND id <> ? LIMIT 1", - ) - .bind(&order_payment_method) - .bind(gateway_order_id) - .bind(&order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if conflicting_order_id.is_some() { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order belongs to another payment order", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order belongs to another payment order".to_string(), - }); - } - } - } - - let amount_matches = payment_callback_amount_matches_order( - order_amount_usd, - order_pay_amount, - order_pay_currency.as_deref(), - order_exchange_rate, - input.amount_usd, - input.pay_amount, - ); - if !amount_matches { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "callback amount mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "callback amount mismatch".to_string(), - }); - } - if !payment_callback_method_matches_order( - &order_payment_method, - order_payment_provider.as_deref(), - &input.payment_method, - input.payment_provider.as_deref(), - ) { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment method mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment method mismatch".to_string(), - }); - } - let payment_provider_matches = payment_callback_provider_matches_order( - &order_payment_method, - order_payment_provider.as_deref(), - &input.payment_method, - input.payment_provider.as_deref(), - ); - if !payment_provider_matches { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment provider mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment provider mismatch".to_string(), - }); - } - let currency_matches = match (input.pay_currency.as_deref(), order_pay_currency.as_deref()) - { - (Some(callback), Some(order)) => order.eq_ignore_ascii_case(callback), - (None, None) => input.pay_amount.is_none() && order_pay_amount.is_none(), - _ => false, - }; - if !currency_matches { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment currency mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment currency mismatch".to_string(), - }); - } - if let Some(expected_channel) = input.payment_channel.as_deref() { - let stored_channel = order_payment_channel.as_deref().or_else(|| { - (order_payment_provider.is_none() - && ["alipay", "wxpay"] - .iter() - .any(|method| method.eq_ignore_ascii_case(&order_payment_method))) - .then_some(order_payment_method.as_str()) - }); - if stored_channel.is_none_or(|value| !value.eq_ignore_ascii_case(expected_channel)) { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment channel mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment channel mismatch".to_string(), - }); - } - } - if order_status == "credited" { - mark_mysql_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::AlreadyCredited { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - }); - } - if !matches!(order_status.as_str(), "pending" | "paid") && !failed_checkout_recoverable { - let error = format!("payment order is not creditable: {order_status}"); - update_mysql_payment_callback_failure(&mut tx, &callback_id, &input, &payload, &error) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { duplicate, error }); - } - if (order_status == "pending" || (failed_checkout_recoverable && !uncertain_checkout)) - && expires_at_unix_secs.is_some_and(|value| value <= now) - { - sqlx::query("UPDATE payment_orders SET status = 'expired' WHERE id = ?") - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order expired", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order expired".to_string(), - }); - } - - if stored_real_gateway_order_id.is_none() { - if let Some(gateway_order_id) = effective_gateway_order_id { - if !mysql_bind_payment_gateway_order_id(&mut tx, &order_id, gateway_order_id) - .await? - { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order belongs to another payment order", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order belongs to another payment order".to_string(), - }); - } - } - } - - if order_kind == "plan_purchase" { - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(user_id) = order_user_id else { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order user missing", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order user missing".to_string(), - }); - }; - let product_id: Option = get(&order_row, "product_id")?; - let snapshot = optional_json( - get::>(&order_row, "product_snapshot")?, - "payment_orders.product_snapshot", - )? - .unwrap_or_else(|| serde_json::json!({})); - let plan_id = product_id.unwrap_or_else(|| { - snapshot - .get("id") - .and_then(|value| value.as_str()) - .unwrap_or("unknown") - .to_string() - }); - let entitlements = plan_entitlements_snapshot(&snapshot); - let existing_entitlement_id = sqlx::query_scalar::<_, String>( - "SELECT id FROM user_plan_entitlements WHERE payment_order_id = ? LIMIT 1", - ) - .bind(&order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if existing_entitlement_id.is_none() { - sqlx::query("SELECT id FROM wallets WHERE id = ? LIMIT 1 FOR UPDATE") - .bind(&order_wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let purchase_limit_scope = plan_purchase_limit_scope(&snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&snapshot); - let active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - if active_count >= max_active_per_user { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "plan purchase limit reached", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "plan purchase limit reached".to_string(), - }); - } - } - replace_matching_plan_entitlements_mysql(&mut tx, &user_id, &snapshot, now).await?; - sqlx::query( - r#" -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) -VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&user_id) - .bind(&plan_id) - .bind(&order_id) - .bind(now) - .bind(plan_expires_at_unix(&snapshot, now)?) - .bind(json_string( - &entitlements, - "user_plan_entitlements.entitlements_snapshot", - )?) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - apply_plan_wallet_credit_mysql( - &mut tx, - &order_wallet_id, - &order_id, - &input.payment_method, - &entitlements, - now, - ) - .await?; - } - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = COALESCE(?, gateway_order_id), - gateway_response = ?, - pay_amount = COALESCE(pay_amount, ?), - pay_currency = COALESCE(pay_currency, ?), - exchange_rate = COALESCE(exchange_rate, ?), - status = 'credited', - fulfillment_status = 'fulfilled', - fulfillment_error = NULL, - paid_at = COALESCE(paid_at, ?), - credited_at = ?, - refundable_amount_usd = 0 -WHERE id = ? - "#, - ) - .bind(effective_gateway_order_id) - .bind(json_string( - &input.gateway_response_projection(&order_no, effective_gateway_order_id), - "payment_orders.gateway_response", - )?) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(now) - .bind(now) - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated_order_row = mysql_payment_order_by_id(&mut tx, &order_id).await?; - mark_mysql_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Applied { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - order: map_payment_order_row(&updated_order_row)?, - }); - } - - let Some(wallet_row) = sqlx::query( - r#" -SELECT id, status, balance, gift_balance, total_recharged -FROM wallets -WHERE id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&order_wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet not found".to_string(), - }); - }; - let wallet_status: String = get(&wallet_row, "status")?; - if wallet_status != "active" { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet is not active", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet is not active".to_string(), - }); - } - - let before_recharge: f64 = get(&wallet_row, "balance")?; - let before_gift: f64 = get(&wallet_row, "gift_balance")?; - let total_recharged: f64 = get(&wallet_row, "total_recharged")?; - // Finite recharge balances may be negative: usage settlement permits a - // finite wallet to overdraft, and a later recharge must be able to - // restore that balance. Reject only malformed values and arithmetic - // overflow here. - if !before_recharge.is_finite() - || !before_gift.is_finite() - || before_gift < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - || !(total_recharged + order_amount_usd).is_finite() - || !(before_recharge + before_gift + order_amount_usd).is_finite() - { - update_mysql_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet balance is invalid", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet balance is invalid".to_string(), - }); - } - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + order_amount_usd; - let after_total = after_recharge + before_gift; - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = total_recharged + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(order_amount_usd) - .bind(now) - .bind(&order_wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_gateway', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&order_wallet_id) - .bind(order_amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&order_id) - .bind(format!("充值到账({})", input.payment_method)) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = COALESCE(?, gateway_order_id), - gateway_response = ?, - pay_amount = COALESCE(pay_amount, ?), - pay_currency = COALESCE(pay_currency, ?), - exchange_rate = COALESCE(exchange_rate, ?), - status = 'credited', - paid_at = COALESCE(paid_at, ?), - credited_at = ?, - refundable_amount_usd = amount_usd -WHERE id = ? -"#, - ) - .bind(effective_gateway_order_id) - .bind(json_string( - &input.gateway_response_projection(&order_no, effective_gateway_order_id), - "payment_orders.gateway_response", - )?) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(now) - .bind(now) - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated_order_row = mysql_payment_order_by_id(&mut tx, &order_id).await?; - mark_mysql_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - Ok(ProcessPaymentCallbackOutcome::Applied { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - order: map_payment_order_row(&updated_order_row)?, - }) - } - - async fn adjust_wallet_balance( - &self, - input: AdjustWalletBalanceInput, - ) -> Result, DataLayerError> { - if !input.amount_usd.is_finite() || input.amount_usd == 0.0 { - return Err(DataLayerError::InvalidInput( - "adjustment amount must be finite and non-zero".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_wallet_by_id_for_update(&mut tx, &input.wallet_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let before_recharge: f64 = get(&row, "balance")?; - let before_gift: f64 = get(&row, "gift_balance")?; - let before_total = before_recharge + before_gift; - let before_total_adjusted: f64 = get(&row, "total_adjusted")?; - if !before_recharge.is_finite() - || !before_gift.is_finite() - || !before_total.is_finite() - || !before_total_adjusted.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance is invalid".to_string(), - )); - } - let mut after_recharge = before_recharge; - let mut after_gift = before_gift; - apply_admin_balance_adjustment( - input.amount_usd, - &input.balance_type, - &mut after_recharge, - &mut after_gift, - ); - let after_total = after_recharge + after_gift; - let after_total_adjusted = before_total_adjusted + input.amount_usd; - if !after_recharge.is_finite() - || !after_gift.is_finite() - || !after_total.is_finite() - || !after_total_adjusted.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance overflow during admin adjustment".to_string(), - )); - } - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_adjusted = total_adjusted + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(input.amount_usd) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - - let transaction_id = uuid::Uuid::new_v4().to_string(); - let description = input - .description - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or("管理员调账") - .to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'adjust', 'adjust_admin', ?, ?, ?, ?, ?, ?, ?, 'admin_action', ?, ?, ?, ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(input.amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(after_gift) - .bind(&input.wallet_id) - .bind(input.operator_id.as_deref()) - .bind(&description) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - - Ok(Some(( - wallet, - StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "adjust".to_string(), - reason_code: "adjust_admin".to_string(), - amount: input.amount_usd, - balance_before: before_total, - balance_after: after_total, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: after_gift, - link_type: Some("admin_action".to_string()), - link_id: Some(input.wallet_id), - operator_id: input.operator_id, - operator_name: None, - operator_email: None, - description: Some(description), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }, - ))) - } - - async fn create_manual_wallet_recharge( - &self, - mut input: CreateManualWalletRechargeInput, - ) -> Result, DataLayerError> { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - if !input.amount_usd.is_finite() || input.amount_usd <= 0.0 { - return Err(DataLayerError::InvalidInput( - "manual recharge amount must be finite and positive".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(wallet_row) = mysql_wallet_by_id_for_update(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let before_recharge: f64 = get(&wallet_row, "balance")?; - let before_gift: f64 = get(&wallet_row, "gift_balance")?; - let before_total_recharged: f64 = get(&wallet_row, "total_recharged")?; - let (after_recharge, after_total_recharged) = validate_manual_wallet_recharge( - input.amount_usd, - before_recharge, - before_gift, - before_total_recharged, - ) - .map_err(DataLayerError::InvalidInput)?; - let user_id: Option = get(&wallet_row, "user_id")?; - let order_id = uuid::Uuid::new_v4().to_string(); - let gateway_response = json_string( - &serde_json::json!({ - "source": "manual", - "operator_id": input.operator_id, - "description": input.description, - }), - "payment_orders.gateway_response", - )?; - - sqlx::query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, status, gateway_response, - created_at, paid_at, credited_at -) -VALUES (?, ?, ?, ?, ?, 0, ?, ?, 'credited', ?, ?, ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&input.wallet_id) - .bind(user_id.as_deref()) - .bind(input.amount_usd) - .bind(input.amount_usd) - .bind(&input.payment_method) - .bind(&gateway_response) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_recharged) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let reason_code = if matches!( - input.payment_method.as_str(), - "card_code" | "gift_code" | "card_recharge" - ) { - "topup_card_code" - } else { - "topup_admin_manual" - }; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', ?, ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&input.wallet_id) - .bind(reason_code) - .bind(input.amount_usd) - .bind(before_recharge + before_gift) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&order_id) - .bind(input.operator_id.as_deref()) - .bind( - input - .description - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or("管理员手动充值"), - ) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let order = map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(Some((wallet, order))) - } - - async fn process_admin_wallet_refund( - &self, - input: ProcessAdminWalletRefundInput, - ) -> Result< - WalletMutationOutcome<( - StoredWalletSnapshot, - StoredAdminWalletRefund, - StoredAdminWalletTransaction, - )>, - DataLayerError, - > { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(refund_row) = - mysql_refund_by_id_and_wallet_for_update(&mut tx, &input.refund_id, &input.wallet_id) - .await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(&refund_row)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if !matches!(refund.status.as_str(), "approved" | "pending_approval") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status is not approvable".to_string(), - )); - } - - let Some(wallet_row) = mysql_wallet_by_id_for_update(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let before_recharge: f64 = get(&wallet_row, "balance")?; - let before_gift: f64 = get(&wallet_row, "gift_balance")?; - let before_total_refunded: f64 = get(&wallet_row, "total_refunded")?; - let amount_usd = refund.amount_usd; - let after_recharge = before_recharge - amount_usd; - let before_total = before_recharge + before_gift; - let after_total = after_recharge + before_gift; - let after_total_refunded = before_total_refunded + amount_usd; - if !before_recharge.is_finite() - || before_recharge < 0.0 - || !before_gift.is_finite() - || before_gift < 0.0 - || !before_total_refunded.is_finite() - || before_total_refunded < 0.0 - || !before_total.is_finite() - || !after_recharge.is_finite() - || !after_total.is_finite() - || !after_total_refunded.is_finite() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid".to_string(), - )); - } - if after_recharge < 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount exceeds refundable recharge balance".to_string(), - )); - } - - if let Some(payment_order_id) = refund.payment_order_id.as_deref() { - let Some(order_row) = - mysql_payment_order_by_id_for_update(&mut tx, payment_order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order not found".to_string(), - )); - }; - let order_wallet_id: String = get(&order_row, "wallet_id")?; - let order_status: String = get(&order_row, "status")?; - if order_wallet_id != input.wallet_id || order_status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not refundable for this wallet".to_string(), - )); - } - let order_amount: f64 = get(&order_row, "amount_usd")?; - let refunded_before: f64 = get(&order_row, "refunded_amount_usd")?; - let refundable_before: f64 = get(&order_row, "refundable_amount_usd")?; - let refunded_after = refunded_before + amount_usd; - let refundable_after = refundable_before - amount_usd; - if !payment_order_refund_amounts_are_consistent( - order_amount, - refunded_before, - refundable_before, - ) || amount_usd > refundable_before - || !refunded_after.is_finite() - || refunded_after < 0.0 - || refunded_after > order_amount - || !refundable_after.is_finite() - || refundable_after < 0.0 - || refundable_after > order_amount - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order refund amounts are invalid".to_string(), - )); - } - let result = sqlx::query( - r#" -UPDATE payment_orders -SET refunded_amount_usd = ?, - refundable_amount_usd = ? -WHERE id = ? -"#, - ) - .bind(refunded_after) - .bind(refundable_after) - .bind(payment_order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "payment order disappeared during refund processing".to_string(), - )); - } - } - - let result = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_refunded = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_refunded) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "wallet disappeared during refund processing".to_string(), - )); - } - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - - let transaction_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'refund', 'refund_out', ?, ?, ?, ?, ?, ?, ?, 'refund_request', ?, ?, '退款占款', ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(-amount_usd) - .bind(before_total) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.refund_id) - .bind(input.operator_id.as_deref()) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE refund_requests -SET status = 'processing', - approved_by = ?, - processed_by = ?, - processed_at = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let refund = map_refund_row(&mysql_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(( - wallet, - refund, - StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "refund".to_string(), - reason_code: "refund_out".to_string(), - amount: -amount_usd, - balance_before: before_total, - balance_after: after_recharge + before_gift, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: before_gift, - link_type: Some("refund_request".to_string()), - link_id: Some(input.refund_id.clone()), - operator_id: input.operator_id.clone(), - operator_name: None, - operator_email: None, - description: Some("退款占款".to_string()), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }, - ))) - } - - async fn complete_admin_wallet_refund( - &self, - input: CompleteAdminWalletRefundInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_refund) = - mysql_refund_by_id_and_wallet_for_update(&mut tx, &input.refund_id, &input.wallet_id) - .await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(¤t_refund)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if let (Some(existing_id), Some(incoming_id)) = ( - refund.gateway_refund_id.as_deref(), - input.gateway_refund_id.as_deref(), - ) { - if existing_id != incoming_id { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier conflicts with existing evidence".to_string(), - )); - } - } - if refund.status == "succeeded" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(refund)); - } - if refund.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status must be processing before completion".to_string(), - )); - } - // Preserve a processing proof for ordinary replays, but allow an - // explicit successful gateway proof to upgrade it at completion. - let selected_payout_proof = input - .payout_proof - .as_ref() - .filter(|proof| refund.payout_proof.is_none() || wallet_refund_proof_is_success(proof)) - .cloned() - .or_else(|| refund.payout_proof.clone()); - let payout_proof = selected_payout_proof - .as_ref() - .map(|value| json_string(value, "refund_requests.payout_proof")) - .transpose()?; - - let refund_update = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'succeeded', - gateway_refund_id = COALESCE(gateway_refund_id, ?), - payout_reference = COALESCE(payout_reference, ?), - payout_proof = ?, - completed_at = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(input.gateway_refund_id.as_deref()) - .bind(input.payout_reference.as_deref()) - .bind(payout_proof.as_deref()) - .bind(now) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if refund_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during refund completion".to_string(), - )); - } - let refund = map_refund_row(&mysql_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(refund)) - } - - async fn update_admin_wallet_refund_gateway( - &self, - input: UpdateAdminWalletRefundGatewayInput, - ) -> Result, DataLayerError> { - if input.gateway_refund_id.trim().is_empty() || input.gateway_refund_id.len() > 128 { - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier is invalid".to_string(), - )); - } - if input - .payout_proof - .as_ref() - .is_some_and(|proof| !proof.is_object()) - { - return Ok(WalletMutationOutcome::Invalid( - "gateway refund proof must be an object".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_row) = - mysql_refund_by_id_and_wallet_for_update(&mut tx, &input.refund_id, &input.wallet_id) - .await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let current = map_refund_row(¤t_row)?; - if !current.amount_usd.is_finite() || current.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if let Some(existing_id) = current.gateway_refund_id.as_deref() { - if existing_id != input.gateway_refund_id { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier conflicts with existing evidence".to_string(), - )); - } - } - if current.status == "succeeded" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(current)); - } - if current.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status must be processing before gateway update".to_string(), - )); - } - // Do not overwrite durable processing evidence with an arbitrary - // replay; only a terminal success proof may replace it. - let selected_payout_proof = input - .payout_proof - .as_ref() - .filter(|proof| current.payout_proof.is_none() || wallet_refund_proof_is_success(proof)) - .cloned() - .or_else(|| current.payout_proof.clone()); - let proof = selected_payout_proof - .as_ref() - .map(|value| json_string(value, "refund_requests.payout_proof")) - .transpose()?; - let gateway_update = sqlx::query( - r#" -UPDATE refund_requests -SET gateway_refund_id = COALESCE(gateway_refund_id, ?), - payout_proof = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(&input.gateway_refund_id) - .bind(proof.as_deref()) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if gateway_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during gateway evidence update".to_string(), - )); - } - let updated = map_refund_row(&mysql_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn fail_admin_wallet_refund( - &self, - input: FailAdminWalletRefundInput, - ) -> Result< - WalletMutationOutcome<( - StoredWalletSnapshot, - StoredAdminWalletRefund, - Option, - )>, - DataLayerError, - > { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(refund_row) = - mysql_refund_by_id_and_wallet_for_update(&mut tx, &input.refund_id, &input.wallet_id) - .await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(&refund_row)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - - if matches!(refund.status.as_str(), "pending_approval" | "approved") { - let refund_update = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'failed', - failure_reason = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(&input.reason) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if refund_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during refund failure".to_string(), - )); - } - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let refund = map_refund_row(&mysql_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((wallet, refund, None))); - } - - if refund.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "cannot fail refund in status: {}", - refund.status - ))); - } - - // Only an explicitly offline payout can be released without external - // settlement evidence. An original-channel refund may still be in - // flight between the provider request and the evidence update. - if refund.gateway_refund_id.is_some() - || refund.payout_proof.is_some() - || !refund - .refund_mode - .trim() - .eq_ignore_ascii_case("offline_payout") - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "cannot fail refund while gateway settlement is processing".to_string(), - )); - } - - let Some(wallet_row) = mysql_wallet_by_id_for_update(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let amount_usd = refund.amount_usd; - let before_recharge: f64 = get(&wallet_row, "balance")?; - let before_gift: f64 = get(&wallet_row, "gift_balance")?; - let before_total_refunded: f64 = get(&wallet_row, "total_refunded")?; - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + amount_usd; - let after_total = after_recharge + before_gift; - let after_total_refunded = before_total_refunded - amount_usd; - if !before_recharge.is_finite() - || before_recharge < 0.0 - || !before_gift.is_finite() - || before_gift < 0.0 - || !before_total_refunded.is_finite() - || before_total_refunded < 0.0 - || before_total_refunded < amount_usd - || !before_total.is_finite() - || !after_recharge.is_finite() - || !after_total.is_finite() - || !after_total_refunded.is_finite() - || after_total_refunded < 0.0 - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid for refund recovery".to_string(), - )); - } - - let mut order_amounts = None; - if let Some(payment_order_id) = refund.payment_order_id.as_deref() { - let Some(order_row) = - mysql_payment_order_by_id_for_update(&mut tx, payment_order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order not found".to_string(), - )); - }; - let order = map_payment_order_row(&order_row)?; - if order.wallet_id != input.wallet_id || order.status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not refundable for this wallet".to_string(), - )); - } - let refunded_before = order.refunded_amount_usd; - let refundable_before = order.refundable_amount_usd; - let refunded_after = refunded_before - amount_usd; - let refundable_after = refundable_before + amount_usd; - if !payment_order_refund_amounts_are_consistent( - order.amount_usd, - refunded_before, - refundable_before, - ) || refunded_before < amount_usd - || !refunded_after.is_finite() - || refunded_after < 0.0 - || !refundable_after.is_finite() - || refundable_after < 0.0 - || refundable_after > order.amount_usd - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order refund amounts are invalid".to_string(), - )); - } - order_amounts = Some(( - payment_order_id.to_string(), - refunded_after, - refundable_after, - )); - } - - let wallet_result = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_refunded = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_refunded) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if wallet_result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "wallet disappeared during refund recovery".to_string(), - )); - } - - let transaction_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'refund', 'refund_revert', ?, ?, ?, ?, ?, ?, ?, 'refund_request', ?, ?, '退款失败回补', ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.refund_id) - .bind(input.operator_id.as_deref()) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - if let Some((payment_order_id, refunded_after, refundable_after)) = order_amounts { - let result = sqlx::query( - r#" -UPDATE payment_orders -SET refunded_amount_usd = ?, - refundable_amount_usd = ? -WHERE id = ? -"#, - ) - .bind(refunded_after) - .bind(refundable_after) - .bind(payment_order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "payment order disappeared during refund recovery".to_string(), - )); - } - } - - let result = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'failed', - failure_reason = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(&input.reason) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during recovery".to_string(), - )); - } - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let refund = map_refund_row(&mysql_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(( - wallet, - refund, - Some(StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "refund".to_string(), - reason_code: "refund_revert".to_string(), - amount: amount_usd, - balance_before: before_total, - balance_after: after_total, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: before_gift, - link_type: Some("refund_request".to_string()), - link_id: Some(input.refund_id.clone()), - operator_id: input.operator_id.clone(), - operator_name: None, - operator_email: None, - description: Some("退款失败回补".to_string()), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }), - ))) - } - - async fn expire_admin_payment_order( - &self, - order_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, order_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "credited order cannot be expired".to_string(), - )); - } - if order.status == "expired" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, false))); - } - if order.status != "pending" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "only pending order can be expired: {}", - order.status - ))); - } - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - gateway_response.insert( - "expire_reason".to_string(), - serde_json::Value::String("admin_mark_expired".to_string()), - ); - gateway_response.insert( - "expired_at".to_string(), - serde_json::Value::String(Utc::now().to_rfc3339()), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - sqlx::query( - "UPDATE payment_orders SET status = 'expired', gateway_response = ? WHERE id = ?", - ) - .bind(gateway_response) - .bind(order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = map_payment_order_row(&mysql_payment_order_by_id(&mut tx, order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied((updated, true))) - } - - async fn fail_admin_payment_order( - &self, - order_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = mysql_payment_order_by_id_for_update(&mut tx, order_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "credited order cannot be failed".to_string(), - )); - } - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - gateway_response.insert( - "failure_reason".to_string(), - serde_json::Value::String("admin_mark_failed".to_string()), - ); - gateway_response.insert( - "failed_at".to_string(), - serde_json::Value::String(Utc::now().to_rfc3339()), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - sqlx::query( - "UPDATE payment_orders SET status = 'failed', gateway_response = ? WHERE id = ?", - ) - .bind(gateway_response) - .bind(order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = map_payment_order_row(&mysql_payment_order_by_id(&mut tx, order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn credit_admin_payment_order( - &self, - input: CreditAdminPaymentOrderInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(order_row) = - mysql_payment_order_by_id_for_update(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&order_row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, false))); - } - if matches!(order.status.as_str(), "failed" | "expired" | "refunded") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "payment order is not creditable: {}", - order.status - ))); - } - if order - .expires_at_unix_secs - .is_some_and(|value| value <= now as u64) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order expired".to_string(), - )); - } - let order_kind: String = get(&order_row, "order_kind")?; - let order_payment_provider: Option = get(&order_row, "payment_provider")?; - let order_payment_channel: Option = get(&order_row, "payment_channel")?; - if validate_payment_order_credit_amounts( - &order_kind, - &order.payment_method, - order_payment_provider.as_deref(), - order_payment_channel.as_deref(), - order.amount_usd, - order.pay_amount, - ) - .is_err() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order amount is invalid".to_string(), - )); - } - if order_kind == "plan_purchase" { - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(user_id) = order_user_id else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order user missing".to_string(), - )); - }; - let product_id: Option = get(&order_row, "product_id")?; - let snapshot = optional_json( - get::>(&order_row, "product_snapshot")?, - "payment_orders.product_snapshot", - )? - .unwrap_or_else(|| serde_json::json!({})); - let plan_id = product_id.unwrap_or_else(|| { - snapshot - .get("id") - .and_then(|value| value.as_str()) - .unwrap_or("unknown") - .to_string() - }); - let entitlements = plan_entitlements_snapshot(&snapshot); - let existing_entitlement_id = sqlx::query_scalar::<_, String>( - "SELECT id FROM user_plan_entitlements WHERE payment_order_id = ? LIMIT 1", - ) - .bind(&input.order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if existing_entitlement_id.is_none() { - let purchase_limit_scope = plan_purchase_limit_scope(&snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&snapshot); - let active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - if active_count >= max_active_per_user { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "plan purchase limit reached".to_string(), - )); - } - } - replace_matching_plan_entitlements_mysql(&mut tx, &user_id, &snapshot, now).await?; - sqlx::query( - r#" -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) -VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&user_id) - .bind(&plan_id) - .bind(&input.order_id) - .bind(now) - .bind(plan_expires_at_unix(&snapshot, now)?) - .bind(json_string( - &entitlements, - "user_plan_entitlements.entitlements_snapshot", - )?) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - apply_plan_wallet_credit_mysql( - &mut tx, - &order.wallet_id, - &input.order_id, - &order.payment_method, - &entitlements, - now, - ) - .await?; - } - - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - if let Some(serde_json::Value::Object(map)) = input.gateway_response_patch.clone() { - gateway_response.extend(map); - } - gateway_response.insert("manual_credit".to_string(), serde_json::Value::Bool(true)); - gateway_response.insert( - "credited_by".to_string(), - input - .operator_id - .clone() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - let next_gateway_order_id = input.gateway_order_id.clone().or(order.gateway_order_id); - let next_pay_amount = input.pay_amount.or(order.pay_amount); - let next_pay_currency = input.pay_currency.clone().or(order.pay_currency); - let next_exchange_rate = input.exchange_rate.or(order.exchange_rate); - let next_paid_at = order.paid_at_unix_secs.unwrap_or(now as u64) as i64; - - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = ?, - gateway_response = ?, - pay_amount = ?, - pay_currency = ?, - exchange_rate = ?, - status = 'credited', - fulfillment_status = 'fulfilled', - fulfillment_error = NULL, - paid_at = ?, - credited_at = ?, - refundable_amount_usd = 0 -WHERE id = ? -"#, - ) - .bind(next_gateway_order_id.as_deref()) - .bind(&gateway_response) - .bind(next_pay_amount) - .bind(next_pay_currency.as_deref()) - .bind(next_exchange_rate) - .bind(next_paid_at) - .bind(now) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let order = - map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, true))); - } - - let Some(wallet_row) = mysql_wallet_by_id_for_update(&mut tx, &order.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let wallet_status: String = get(&wallet_row, "status")?; - if wallet_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet is not active".to_string(), - )); - } - - let before_recharge: f64 = get(&wallet_row, "balance")?; - let before_gift: f64 = get(&wallet_row, "gift_balance")?; - let total_recharged: f64 = get(&wallet_row, "total_recharged")?; - if !before_recharge.is_finite() - || !before_gift.is_finite() - || before_gift < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - || !(total_recharged + order.amount_usd).is_finite() - || !(before_recharge + before_gift + order.amount_usd).is_finite() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid".to_string(), - )); - } - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + order.amount_usd; - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = total_recharged + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(order.amount_usd) - .bind(now) - .bind(&order.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_gateway', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&order.wallet_id) - .bind(order.amount_usd) - .bind(before_total) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.order_id) - .bind(format!("充值到账({})", order.payment_method)) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - if let Some(serde_json::Value::Object(map)) = input.gateway_response_patch { - gateway_response.extend(map); - } - gateway_response.insert("manual_credit".to_string(), serde_json::Value::Bool(true)); - gateway_response.insert( - "credited_by".to_string(), - input - .operator_id - .clone() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - let next_gateway_order_id = input.gateway_order_id.or(order.gateway_order_id); - let next_pay_amount = input.pay_amount.or(order.pay_amount); - let next_pay_currency = input.pay_currency.or(order.pay_currency); - let next_exchange_rate = input.exchange_rate.or(order.exchange_rate); - let next_paid_at = order.paid_at_unix_secs.unwrap_or(now as u64) as i64; - - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = ?, - gateway_response = ?, - pay_amount = ?, - pay_currency = ?, - exchange_rate = ?, - status = 'credited', - paid_at = ?, - credited_at = ?, - refundable_amount_usd = amount_usd -WHERE id = ? -"#, - ) - .bind(next_gateway_order_id.as_deref()) - .bind(&gateway_response) - .bind(next_pay_amount) - .bind(next_pay_currency.as_deref()) - .bind(next_exchange_rate) - .bind(next_paid_at) - .bind(now) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let order = - map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied((order, true))) - } - - async fn create_admin_redeem_code_batch( - &self, - input: CreateAdminRedeemCodeBatchInput, - ) -> Result { - validate_admin_redeem_code_batch_input(&input).map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let batch_id = uuid::Uuid::new_v4().to_string(); - let expires_at = input - .expires_at_unix_secs - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput( - "redeem code batch expires_at overflow".to_string(), - ) - }) - }) - .transpose()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - sqlx::query( - r#" -INSERT INTO redeem_code_batches ( - id, name, amount_usd, currency, balance_bucket, total_count, status, - description, created_by, expires_at, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) -"#, - ) - .bind(&batch_id) - .bind(&input.name) - .bind(input.amount_usd) - .bind(&input.currency) - .bind(&input.balance_bucket) - .bind(i64::try_from(input.total_count).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid redeem code count: {}", - input.total_count - )) - })?) - .bind(input.description.as_deref()) - .bind(input.created_by.as_deref()) - .bind(expires_at) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let mut codes = Vec::with_capacity(input.total_count); - for _ in 0..input.total_count { - let (code_id, code, masked_code, code_hash, prefix, suffix) = - generate_redeem_code_candidate(); - sqlx::query( - r#" -INSERT INTO redeem_codes ( - id, batch_id, code_hash, code_prefix, code_suffix, status, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, 'active', ?, ?) -"#, - ) - .bind(&code_id) - .bind(&batch_id) - .bind(&code_hash) - .bind(&prefix) - .bind(&suffix) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - codes.push(CreatedAdminRedeemCodePlaintext { - code_id, - code, - masked_code, - }); - } - - let batch = StoredAdminRedeemCodeBatch { - id: batch_id, - name: input.name, - amount_usd: input.amount_usd, - currency: input.currency, - balance_bucket: input.balance_bucket, - total_count: input.total_count as u64, - redeemed_count: 0, - active_count: input.total_count as u64, - status: "active".to_string(), - description: input.description, - created_by: input.created_by, - expires_at_unix_secs: input.expires_at_unix_secs, - created_at_unix_ms: timestamp(now, "redeem_code_batches.created_at")?, - updated_at_unix_secs: timestamp(now, "redeem_code_batches.updated_at")?, - }; - tx.commit().await.map_sql_err()?; - Ok(CreateAdminRedeemCodeBatchResult { batch, codes }) - } - - async fn disable_admin_redeem_code_batch( - &self, - input: DisableAdminRedeemCodeBatchInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_batch) = sqlx::query( - r#" -SELECT status -FROM redeem_code_batches -WHERE id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.batch_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let status: String = get(¤t_batch, "status")?; - if status != "disabled" { - sqlx::query( - r#" -UPDATE redeem_code_batches -SET status = 'disabled', - updated_at = ? -WHERE id = ? -"#, - ) - .bind(now) - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'disabled', - disabled_by = COALESCE(?, disabled_by), - updated_at = ? -WHERE batch_id = ? - AND status = 'active' -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - let batch = - map_redeem_batch_row(&mysql_redeem_batch_by_id(&mut tx, &input.batch_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(batch)) - } - - async fn delete_admin_redeem_code_batch( - &self, - input: DeleteAdminRedeemCodeBatchInput, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_batch) = sqlx::query( - r#" -SELECT status -FROM redeem_code_batches -WHERE id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.batch_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let status: String = get(¤t_batch, "status")?; - if status != "disabled" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "only disabled redeem code batch can be deleted".to_string(), - )); - } - - let redeemed_count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM redeem_codes -WHERE batch_id = ? - AND status = 'redeemed' -"#, - ) - .bind(&input.batch_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - if redeemed_count > 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "redeemed batch cannot be deleted".to_string(), - )); - } - - let batch = - map_redeem_batch_row(&mysql_redeem_batch_by_id(&mut tx, &input.batch_id).await?)?; - let _ = input.operator_id; - sqlx::query("DELETE FROM redeem_codes WHERE batch_id = ?") - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM redeem_code_batches WHERE id = ?") - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(batch)) - } - - async fn disable_admin_redeem_code( - &self, - input: DisableAdminRedeemCodeInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_code) = sqlx::query( - r#" -SELECT batch_id, status -FROM redeem_codes -WHERE id = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&input.code_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let batch_id: String = get(¤t_code, "batch_id")?; - let status: String = get(¤t_code, "status")?; - if status == "redeemed" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "redeemed code cannot be disabled".to_string(), - )); - } - if status != "disabled" { - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'disabled', - disabled_by = COALESCE(?, disabled_by), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(&input.code_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query("UPDATE redeem_code_batches SET updated_at = ? WHERE id = ?") - .bind(now) - .bind(&batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let code = map_redeem_code_row(&mysql_redeem_code_by_id(&mut tx, &input.code_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(code)) - } - - async fn redeem_wallet_code( - &self, - input: RedeemWalletCodeInput, - ) -> Result { - let Some(normalized) = normalize_redeem_code(&input.code) else { - return Ok(RedeemWalletCodeOutcome::InvalidCode); - }; - let code_hash = hash_redeem_code(&normalized); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - - let Some(code_row) = sqlx::query( - r#" -SELECT - codes.id AS code_id, - codes.status AS code_status, - codes.batch_id, - batches.name AS batch_name, - batches.status AS batch_status, - batches.balance_bucket, - batches.amount_usd, - batches.expires_at AS batch_expires_at -FROM redeem_codes AS codes -JOIN redeem_code_batches AS batches ON batches.id = codes.batch_id -WHERE codes.code_hash = ? -LIMIT 1 -FOR UPDATE -"#, - ) - .bind(&code_hash) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeNotFound); - }; - - let code_status: String = get(&code_row, "code_status")?; - match code_status.as_str() { - "disabled" => { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeDisabled); - } - "redeemed" => { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeRedeemed); - } - _ => {} - } - let batch_status: String = get(&code_row, "batch_status")?; - if batch_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::BatchDisabled); - } - let batch_expires_at: Option = get(&code_row, "batch_expires_at")?; - if batch_expires_at.is_some_and(|value| value <= now) { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeExpired); - } - - let code_id: String = get(&code_row, "code_id")?; - let batch_id: String = get(&code_row, "batch_id")?; - let batch_name: String = get(&code_row, "batch_name")?; - let balance_bucket: String = get(&code_row, "balance_bucket")?; - let amount_usd: f64 = get(&code_row, "amount_usd")?; - - let wallet_row = mysql_wallet_by_user_id_for_update(&mut tx, &input.user_id).await?; - let wallet_id = if let Some(row) = wallet_row.as_ref() { - let status: String = get(row, "status")?; - if status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::WalletInactive); - } - get(row, "id")? - } else { - uuid::Uuid::new_v4().to_string() - }; - - let (before_recharge, before_gift, before_total_recharged) = - if let Some(row) = wallet_row.as_ref() { - ( - get(row, "balance")?, - get(row, "gift_balance")?, - get(row, "total_recharged")?, - ) - } else { - sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?) -"#, - ) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - (0.0, 0.0, 0.0) - }; - let (after_recharge, after_gift, after_total_recharged) = validate_redeem_wallet_credit( - &balance_bucket, - amount_usd, - before_recharge, - before_gift, - before_total_recharged, - ) - .map_err(DataLayerError::UnexpectedValue)?; - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(after_total_recharged) - .bind(now) - .bind(&wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let payment_method = redeem_code_payment_method(&balance_bucket); - let order_id = uuid::Uuid::new_v4().to_string(); - let gateway_order_id = format!("card_{}", uuid::Uuid::new_v4().simple()); - let gateway_response = json_string( - &serde_json::json!({ - "source": "redeem_code", - "batch_id": batch_id, - "batch_name": batch_name, - "balance_bucket": balance_bucket, - }), - "payment_orders.gateway_response", - )?; - sqlx::query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, - exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, - gateway_order_id, gateway_response, status, created_at, paid_at, credited_at -) -VALUES (?, ?, ?, ?, ?, NULL, NULL, NULL, 0, ?, ?, ?, ?, 'credited', ?, ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(amount_usd) - .bind(redeem_code_refundable_amount(&balance_bucket, amount_usd)) - .bind(payment_method) - .bind(&gateway_order_id) - .bind(&gateway_response) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_card_code', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&wallet_id) - .bind(amount_usd) - .bind(before_recharge + before_gift) - .bind(after_recharge + after_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(after_gift) - .bind(&order_id) - .bind("兑换码充值") - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'redeemed', - redeemed_by_user_id = ?, - redeemed_wallet_id = ?, - redeemed_payment_order_id = ?, - redeemed_at = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&input.user_id) - .bind(&wallet_id) - .bind(&order_id) - .bind(now) - .bind(now) - .bind(&code_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE redeem_code_batches SET updated_at = ? WHERE id = ?") - .bind(now) - .bind(&batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let wallet = map_wallet_row(&mysql_wallet_by_id(&mut tx, &wallet_id).await?)?; - let order = map_payment_order_row(&mysql_payment_order_by_id(&mut tx, &order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(RedeemWalletCodeOutcome::Redeemed { - wallet, - order, - amount_usd, - batch_name, - }) - } -} - -fn daily_usage_select_sql(suffix: &'static str) -> String { - format!( - r#" -SELECT - id, billing_date, billing_timezone, total_cost_usd, total_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - first_finalized_at AS first_finalized_at_unix_secs, - last_finalized_at AS last_finalized_at_unix_secs, - aggregated_at AS aggregated_at_unix_secs -FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_timezone = ? - {suffix} -"# - ) -} - -fn current_billing_date(billing_timezone: &str) -> Result { - let timezone = billing_timezone.parse::().map_err(|err| { - DataLayerError::InvalidInput(format!("invalid wallet billing timezone: {err}")) - })?; - Ok(Utc::now().with_timezone(&timezone).date_naive().to_string()) -} - -fn map_wallet_row(row: &MySqlRow) -> Result { - StoredWalletSnapshot::new( - get(row, "id")?, - get(row, "user_id")?, - get(row, "api_key_id")?, - get(row, "balance")?, - get(row, "gift_balance")?, - get(row, "limit_mode")?, - get(row, "currency")?, - get(row, "status")?, - get(row, "total_recharged")?, - get(row, "total_consumed")?, - get(row, "total_refunded")?, - get(row, "total_adjusted")?, - get(row, "updated_at_unix_secs")?, - ) -} - -fn map_admin_wallet_list_item_row( - row: &MySqlRow, -) -> 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}")) - }) -} - -fn plan_entitlements_snapshot(snapshot: &serde_json::Value) -> serde_json::Value { - snapshot - .get("entitlements") - .or_else(|| snapshot.get("entitlements_json")) - .cloned() - .unwrap_or_else(|| serde_json::json!([])) -} - -fn plan_max_active_per_user(snapshot: &serde_json::Value) -> i64 { - snapshot - .get("max_active_per_user") - .and_then(|value| value.as_i64()) - .unwrap_or(1) - .max(1) -} - -fn plan_purchase_limit_scope(snapshot: &serde_json::Value) -> &str { - match snapshot - .get("purchase_limit_scope") - .and_then(|value| value.as_str()) - { - Some("lifetime") => "lifetime", - Some("unlimited") => "unlimited", - _ => "active_period", - } -} - -async fn replace_matching_plan_entitlements_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - user_id: &str, - snapshot: &serde_json::Value, - now: i64, -) -> Result<(), DataLayerError> { - let incoming_entitlements = plan_entitlements_snapshot(snapshot); - if !entitlements_have_replacement_selector(&incoming_entitlements) { - return Ok(()); - } - - let rows = sqlx::query( - r#" -SELECT id, entitlements_snapshot -FROM user_plan_entitlements -WHERE user_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(user_id) - .bind(now) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - - for row in rows { - let entitlements = optional_json( - get::>(&row, "entitlements_snapshot")?, - "user_plan_entitlements.entitlements_snapshot", - )? - .unwrap_or_else(|| serde_json::json!([])); - let should_replace = - entitlements_should_replace_existing(&incoming_entitlements, &entitlements); - if !should_replace { - continue; - } - let entitlement_id: String = get(&row, "id")?; - sqlx::query( - r#" -UPDATE user_plan_entitlements -SET status = 'replaced', - expires_at = CASE WHEN expires_at > ? THEN ? ELSE expires_at END, - updated_at = ? -WHERE id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(now) - .bind(now) - .bind(now) - .bind(entitlement_id) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - Ok(()) -} - -fn plan_expires_at_unix( - snapshot: &serde_json::Value, - starts_at_unix_secs: i64, -) -> Result { - let days = - checked_plan_duration_days_from_snapshot(snapshot).map_err(DataLayerError::InvalidInput)?; - let seconds = days.checked_mul(86_400).ok_or_else(|| { - DataLayerError::InvalidInput("plan duration exceeds the supported range".to_string()) - })?; - starts_at_unix_secs.checked_add(seconds).ok_or_else(|| { - DataLayerError::InvalidInput("plan expiration exceeds the supported range".to_string()) - }) -} - -async fn apply_plan_wallet_credit_mysql( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - wallet_id: &str, - order_id: &str, - payment_method: &str, - entitlements: &serde_json::Value, - now: i64, -) -> Result<(), DataLayerError> { - validate_plan_wallet_credit_entitlements(entitlements).map_err(DataLayerError::InvalidInput)?; - let credits = entitlements - .as_array() - .into_iter() - .flatten() - .filter(|item| { - item.get("type") - .and_then(|value| value.as_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("wallet_credit")) - }) - .filter_map(|item| { - let amount = item.get("amount_usd").and_then(|value| value.as_f64())?; - if amount <= 0.0 || !amount.is_finite() { - return None; - } - let bucket = item - .get("balance_bucket") - .and_then(|value| value.as_str()) - .unwrap_or("gift") - .trim() - .to_ascii_lowercase(); - Some((amount, bucket)) - }) - .collect::>(); - if credits.is_empty() { - return Ok(()); - } - let Some(wallet_row) = sqlx::query( - "SELECT id, status, balance, gift_balance, total_recharged FROM wallets WHERE id = ? LIMIT 1 FOR UPDATE", - ) - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - else { - return Err(DataLayerError::UnexpectedValue( - "wallet not found for plan wallet_credit".to_string(), - )); - }; - let status: String = get(&wallet_row, "status")?; - if status != "active" { - return Err(DataLayerError::UnexpectedValue( - "wallet is not active for plan wallet_credit".to_string(), - )); - } - let mut recharge_balance: f64 = get(&wallet_row, "balance")?; - let mut gift_balance: f64 = get(&wallet_row, "gift_balance")?; - let mut total_recharged: f64 = get(&wallet_row, "total_recharged")?; - if !recharge_balance.is_finite() - || !gift_balance.is_finite() - || gift_balance < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance is invalid for plan wallet_credit".to_string(), - )); - } - for (amount, bucket) in credits { - let before_recharge = recharge_balance; - let before_gift = gift_balance; - let before_total = before_recharge + before_gift; - let credits_recharge = bucket == "recharge"; - if credits_recharge { - recharge_balance += amount; - total_recharged += amount; - } else { - gift_balance += amount; - } - let after_total = recharge_balance + gift_balance; - if !before_total.is_finite() - || !recharge_balance.is_finite() - || !gift_balance.is_finite() - || !total_recharged.is_finite() - || !after_total.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance overflow for plan wallet_credit".to_string(), - )); - } - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(recharge_balance) - .bind(gift_balance) - .bind(total_recharged) - .bind(now) - .bind(wallet_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'plan_wallet_credit', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(wallet_id) - .bind(amount) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(recharge_balance) - .bind(before_gift) - .bind(gift_balance) - .bind(order_id) - .bind(format!("套餐附赠余额({payment_method})")) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - Ok(()) -} - -fn payment_gateway_response_map( - value: Option, -) -> serde_json::Map { - match value { - Some(serde_json::Value::Object(map)) => map, - _ => serde_json::Map::new(), - } -} - -fn normalize_redeem_code(value: &str) -> Option { - let normalized = value - .chars() - .filter(|ch| ch.is_ascii_alphanumeric()) - .map(|ch| ch.to_ascii_uppercase()) - .collect::(); - if normalized.len() < 16 { - None - } else { - Some(normalized) - } -} - -fn hash_redeem_code(normalized: &str) -> String { - use sha2::Digest; - - format!("{:x}", sha2::Sha256::digest(normalized.as_bytes())) -} - -fn format_redeem_code(normalized: &str) -> String { - normalized - .as_bytes() - .chunks(8) - .map(|chunk| std::str::from_utf8(chunk).unwrap_or_default()) - .collect::>() - .join("-") -} - -fn generate_redeem_code_candidate() -> (String, String, String, String, String, String) { - let normalized = uuid::Uuid::new_v4() - .simple() - .to_string() - .to_ascii_uppercase(); - let code = format_redeem_code(&normalized); - let code_id = uuid::Uuid::new_v4().to_string(); - let prefix = normalized.chars().take(4).collect::(); - let suffix = normalized - .chars() - .rev() - .take(4) - .collect::>() - .into_iter() - .rev() - .collect::(); - let masked_code = mask_redeem_code(&prefix, &suffix); - let code_hash = hash_redeem_code(&normalized); - (code_id, code, masked_code, code_hash, prefix, suffix) -} - -fn wallet_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -fn apply_admin_balance_adjustment( - amount_usd: f64, - balance_type: &str, - recharge_balance: &mut f64, - gift_balance: &mut f64, -) { - if amount_usd > 0.0 { - if balance_type.eq_ignore_ascii_case("gift") { - *gift_balance += amount_usd; - } else { - *recharge_balance += amount_usd; - } - return; - } - - let mut remaining = -amount_usd; - let consume_positive_bucket = |balance: &mut f64, to_consume: &mut f64| { - if *to_consume <= 0.0 { - return; - } - let available = (*balance).max(0.0); - let consumed = available.min(*to_consume); - *balance -= consumed; - *to_consume -= consumed; - }; - if balance_type.eq_ignore_ascii_case("gift") { - consume_positive_bucket(gift_balance, &mut remaining); - consume_positive_bucket(recharge_balance, &mut remaining); - } else { - consume_positive_bucket(recharge_balance, &mut remaining); - consume_positive_bucket(gift_balance, &mut remaining); - } - if remaining > 0.0 { - *recharge_balance -= remaining; - } -} - -async fn mysql_wallet_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - wallet_id: &str, -) -> Result { - let sql = wallet_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(wallet_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_wallet_by_id_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - wallet_id: &str, -) -> Result, DataLayerError> { - let sql = wallet_select_sql("WHERE id = ? LIMIT 1 FOR UPDATE"); - sqlx::query(&sql) - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_wallet_by_user_id_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - user_id: &str, -) -> Result, DataLayerError> { - let sql = wallet_select_sql("WHERE user_id = ? LIMIT 1 FOR UPDATE"); - sqlx::query(&sql) - .bind(user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -fn payment_order_select_sql(where_clause: &str) -> String { - format!( - 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, 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_clause} -"# - ) -} - -async fn mysql_payment_order_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - order_id: &str, -) -> Result { - let sql = payment_order_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(order_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_payment_order_by_id_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - order_id: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE id = ? LIMIT 1 FOR UPDATE"); - sqlx::query(&sql) - .bind(order_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_payment_order_by_order_no_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - order_no: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE order_no = ? LIMIT 1 FOR UPDATE"); - sqlx::query(&sql) - .bind(order_no) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_payment_order_by_gateway_order_id_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - payment_method: &str, - gateway_order_id: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql( - "WHERE payment_method = ? AND gateway_order_id = ? LIMIT 1 FOR UPDATE", - ); - sqlx::query(&sql) - .bind(payment_method) - .bind(gateway_order_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_bind_payment_gateway_order_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - order_id: &str, - gateway_order_id: &str, -) -> Result { - sqlx::query("SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - let bind_result = sqlx::query("UPDATE payment_orders SET gateway_order_id = ? WHERE id = ?") - .bind(gateway_order_id) - .bind(order_id) - .execute(&mut **tx) - .await; - match bind_result { - Ok(_) => { - sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(true) - } - Err(error) - if error - .as_database_error() - .is_some_and(|database_error| database_error.is_unique_violation()) => - { - sqlx::query("ROLLBACK TO SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(false) - } - Err(error) => { - let _ = sqlx::query("ROLLBACK TO SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await; - let _ = sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await; - Err(DataLayerError::sql(error)) - } - } -} - -fn refund_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -async fn mysql_refund_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - refund_id: &str, -) -> Result { - let sql = refund_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(refund_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_refund_by_id_and_wallet_for_update( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - refund_id: &str, - wallet_id: &str, -) -> Result, DataLayerError> { - let sql = refund_select_sql("WHERE id = ? AND wallet_id = ? LIMIT 1 FOR UPDATE"); - sqlx::query(&sql) - .bind(refund_id) - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn mysql_refund_by_idempotency( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - user_id: &str, - idempotency_key: &str, -) -> Result, DataLayerError> { - let sql = refund_select_sql("WHERE user_id = ? AND idempotency_key = ? LIMIT 1"); - sqlx::query(&sql) - .bind(user_id) - .bind(idempotency_key) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -fn redeem_batch_select_sql(where_clause: &str) -> String { - format!( - 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, - 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 -{where_clause} -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 -"# - ) -} - -async fn mysql_redeem_batch_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - batch_id: &str, -) -> Result { - let sql = redeem_batch_select_sql("WHERE batches.id = ?"); - sqlx::query(&sql) - .bind(batch_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -fn redeem_code_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -async fn mysql_redeem_code_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - code_id: &str, -) -> Result { - let sql = redeem_code_select_sql("WHERE codes.id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(code_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn update_mysql_payment_callback_failure( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - callback_id: &str, - input: &ProcessPaymentCallbackInput, - _payload: &str, - error: &str, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -UPDATE payment_callbacks -SET signature_valid = ?, - status = 'failed', - error_message = ?, - payload_hash = ?, - payload = NULL, - processed_at = ? -WHERE id = ? -"#, - ) - .bind(input.signature_valid) - .bind(error) - .bind(&input.payload_hash) - .bind(current_unix_secs_i64()) - .bind(callback_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -async fn mark_mysql_payment_callback_processed( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - callback_id: &str, - input: &ProcessPaymentCallbackInput, - _payload: &str, - order_id: &str, - order_no: &str, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -UPDATE payment_callbacks -SET payment_order_id = ?, - signature_valid = TRUE, - status = 'processed', - error_message = NULL, - payload_hash = ?, - payload = NULL, - processed_at = ?, - order_no = ?, - gateway_order_id = COALESCE(?, gateway_order_id) -WHERE id = ? -"#, - ) - .bind(order_id) - .bind(&input.payload_hash) - .bind(current_unix_secs_i64()) - .bind(order_no) - .bind(input.gateway_order_id.as_deref()) - .bind(callback_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -async fn update_mysql_wallet_snapshot( - pool: &MysqlPool, - owner_column: &str, - owner_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, -) -> Result<(), DataLayerError> { - let owner_predicate = match owner_column { - "user_id" => "user_id = ?", - "api_key_id" => "api_key_id = ?", - _ => { - return Err(DataLayerError::UnexpectedValue(format!( - "unsupported wallet owner column: {owner_column}" - ))); - } - }; - let sql = format!( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - limit_mode = ?, - currency = ?, - status = ?, - total_recharged = ?, - total_consumed = ?, - total_refunded = ?, - total_adjusted = ?, - updated_at = ? -WHERE {owner_predicate} -"# - ); - sqlx::query(&sql) - .bind(balance) - .bind(gift_balance) - .bind(limit_mode) - .bind(currency) - .bind(status) - .bind(total_recharged) - .bind(total_consumed) - .bind(total_refunded) - .bind(total_adjusted) - .bind( - updated_at_unix_secs - .map(|value| value as i64) - .unwrap_or_else(current_unix_secs_i64), - ) - .bind(owner_id) - .execute(pool) - .await - .map_sql_err()?; - Ok(()) -} - -async fn initialize_mysql_auth_wallet( - pool: &MysqlPool, - user_id: Option<&str>, - api_key_id: Option<&str>, - initial_gift_usd: f64, - unlimited: bool, -) -> Result, DataLayerError> { - let owner = user_id - .or(api_key_id) - .filter(|value| !value.trim().is_empty()); - if owner.is_none() || (user_id.is_some() && api_key_id.is_some()) { - return Err(DataLayerError::InvalidInput( - "wallet owner must be exactly one non-empty user or API-key id".to_string(), - )); - } - if !initial_gift_usd.is_finite() { - return Err(DataLayerError::InvalidInput( - "initial gift amount must be finite".to_string(), - )); - } - let gift_amount = if unlimited { - 0.0 - } else { - initial_gift_usd.max(0.0) - }; - let now = current_unix_secs_i64(); - let wallet = StoredWalletSnapshot::new( - uuid::Uuid::new_v4().to_string(), - user_id.map(str::to_string), - api_key_id.map(str::to_string), - 0.0, - gift_amount, - if unlimited { "unlimited" } else { "finite" }.to_string(), - "USD".to_string(), - "active".to_string(), - 0.0, - 0.0, - 0.0, - gift_amount, - now, - )?; - // Lock the owner row (when present) and keep the insert in the same - // transaction. This makes retries return the existing wallet instead of - // issuing another initial gift transaction. - let mut tx = pool.begin().await.map_sql_err()?; - let owner_column = if user_id.is_some() { - "user_id" - } else { - "api_key_id" - }; - let owner_value = owner.expect("validated wallet owner"); - - // Keep the ownership lock order identical to the guarded user-deletion - // path: users first, then api_keys/wallets. This closes the window where - // a deletion can observe no wallet while a concurrent initializer inserts - // one after the user row has been removed. - if let Some(user_id) = user_id { - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - } else { - let api_key_user_id: Option = - sqlx::query_scalar("SELECT user_id FROM api_keys WHERE id = ?") - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(api_key_user_id) = api_key_user_id else { - tx.rollback().await.map_sql_err()?; - return Ok(None); - }; - let user_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? FOR UPDATE") - .bind(&api_key_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - let api_key_exists: Option = - sqlx::query_scalar("SELECT id FROM api_keys WHERE id = ? AND user_id = ? FOR UPDATE") - .bind(owner_value) - .bind(&api_key_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if api_key_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - } - let existing_sql = format!( - "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 WHERE {owner_column} = ? LIMIT 1 FOR UPDATE" - ); - if let Some(row) = sqlx::query(&existing_sql) - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - { - let existing = map_wallet_row(&row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some((existing, false))); - } - - let insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, api_key_id, balance, gift_balance, limit_mode, currency, - status, total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, ?, 0, ?, ?, 'USD', 'active', 0, 0, 0, ?, ?, ?) -ON DUPLICATE KEY UPDATE id = id -"#, - ) - .bind(&wallet.id) - .bind(user_id) - .bind(api_key_id) - .bind(gift_amount) - .bind(&wallet.limit_mode) - .bind(gift_amount) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - // Re-read under lock. A duplicate owner insert means another caller won; - // return that row and never write a second gift transaction. - let existing_sql = format!( - "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 WHERE {owner_column} = ? LIMIT 1 FOR UPDATE" - ); - let Some(owner_row) = sqlx::query(&existing_sql) - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - }; - let owner_wallet_id: String = get(&owner_row, "id")?; - if insert_result.rows_affected() == 0 || owner_wallet_id != wallet.id { - let existing = map_wallet_row(&owner_row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some((existing, false))); - } - if gift_amount > 0.0 { - let link_id = user_id.or(api_key_id).unwrap_or_default(); - let description = if api_key_id.is_some() { - "独立余额 Key 初始赠款" - } else { - "用户初始赠款" - }; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'gift', 'gift_initial', ?, 0, ?, 0, 0, 0, ?, 'system_task', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&wallet.id) - .bind(gift_amount) - .bind(gift_amount) - .bind(gift_amount) - .bind(link_id) - .bind(description) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let wallet = map_wallet_row(&owner_row)?; - tx.commit().await.map_sql_err()?; - Ok(Some((wallet, true))) -} - -fn mysql_wallet_recharge_replay_matches( - row: &MySqlRow, - wallet_id: &str, - input: &CreateWalletRechargeOrderInput, -) -> Result { - let existing_wallet_id: String = get(row, "wallet_id")?; - let pay_currency: Option = get(row, "pay_currency")?; - let payment_method: String = get(row, "payment_method")?; - let payment_provider: Option = get(row, "payment_provider")?; - let payment_channel: Option = get(row, "payment_channel")?; - Ok(wallet_recharge_replay_matches( - &existing_wallet_id, - get(row, "amount_usd")?, - get(row, "pay_amount")?, - pay_currency.as_deref(), - get(row, "exchange_rate")?, - &payment_method, - payment_provider.as_deref(), - payment_channel.as_deref(), - wallet_id, - input, - )) -} - -fn map_payment_order_row(row: &MySqlRow) -> Result { - Ok(StoredAdminPaymentOrder { - id: get(row, "id")?, - order_no: get(row, "order_no")?, - wallet_id: get(row, "wallet_id")?, - user_id: get(row, "user_id")?, - amount_usd: get(row, "amount_usd")?, - pay_amount: get(row, "pay_amount")?, - pay_currency: get(row, "pay_currency")?, - exchange_rate: get(row, "exchange_rate")?, - refunded_amount_usd: get(row, "refunded_amount_usd")?, - refundable_amount_usd: get(row, "refundable_amount_usd")?, - payment_method: get(row, "payment_method")?, - payment_provider: get(row, "payment_provider")?, - order_kind: get(row, "order_kind")?, - gateway_order_id: get(row, "gateway_order_id")?, - gateway_response: optional_json( - get(row, "gateway_response")?, - "payment_orders.gateway_response", - )?, - status: get(row, "status")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "payment_orders.created_at", - )?, - paid_at_unix_secs: optional_timestamp( - get(row, "paid_at_unix_secs")?, - "payment_orders.paid_at", - )?, - credited_at_unix_secs: optional_timestamp( - get(row, "credited_at_unix_secs")?, - "payment_orders.credited_at", - )?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "payment_orders.expires_at", - )?, - }) -} - -fn map_payment_callback_row(row: &MySqlRow) -> Result { - Ok(StoredAdminPaymentCallback { - id: get(row, "id")?, - payment_order_id: get(row, "payment_order_id")?, - payment_method: get(row, "payment_method")?, - callback_key: get(row, "callback_key")?, - order_no: get(row, "order_no")?, - gateway_order_id: get(row, "gateway_order_id")?, - payload_hash: get(row, "payload_hash")?, - signature_valid: get(row, "signature_valid")?, - status: get(row, "status")?, - payload: optional_json(get(row, "payload")?, "payment_callbacks.payload")?, - error_message: get(row, "error_message")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "payment_callbacks.created_at", - )?, - processed_at_unix_secs: optional_timestamp( - get(row, "processed_at_unix_secs")?, - "payment_callbacks.processed_at", - )?, - }) -} - -fn map_wallet_transaction_row( - row: &MySqlRow, -) -> Result { - Ok(StoredAdminWalletTransaction { - 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")?, - created_at_unix_ms: optional_timestamp( - get(row, "created_at_unix_ms")?, - "wallet_transactions.created_at", - )?, - }) -} - -fn map_refund_row(row: &MySqlRow) -> Result { - Ok(StoredAdminWalletRefund { - 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")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "refund_requests.created_at", - )?, - updated_at_unix_secs: 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 map_redeem_batch_row(row: &MySqlRow) -> Result { - Ok(StoredAdminRedeemCodeBatch { - id: get(row, "id")?, - name: get(row, "name")?, - amount_usd: get(row, "amount_usd")?, - currency: get(row, "currency")?, - balance_bucket: get(row, "balance_bucket")?, - total_count: nonnegative_u64(get(row, "total_count")?, "redeem_code_batches.total_count")?, - redeemed_count: nonnegative_u64( - get(row, "redeemed_count")?, - "redeem_codes.redeemed_count", - )?, - active_count: nonnegative_u64(get(row, "active_count")?, "redeem_codes.active_count")?, - status: get(row, "status")?, - description: get(row, "description")?, - created_by: get(row, "created_by")?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "redeem_code_batches.expires_at", - )?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "redeem_code_batches.created_at", - )?, - updated_at_unix_secs: timestamp( - get(row, "updated_at_unix_secs")?, - "redeem_code_batches.updated_at", - )?, - }) -} - -fn map_redeem_code_row(row: &MySqlRow) -> Result { - let code_prefix: String = get(row, "code_prefix")?; - let code_suffix: String = get(row, "code_suffix")?; - Ok(StoredAdminRedeemCode { - id: get(row, "id")?, - batch_id: get(row, "batch_id")?, - batch_name: get(row, "batch_name")?, - masked_code: mask_redeem_code(&code_prefix, &code_suffix), - code_prefix, - code_suffix, - status: get(row, "status")?, - redeemed_by_user_id: get(row, "redeemed_by_user_id")?, - redeemed_by_user_name: get(row, "redeemed_by_user_name")?, - redeemed_wallet_id: get(row, "redeemed_wallet_id")?, - redeemed_payment_order_id: get(row, "redeemed_payment_order_id")?, - redeemed_order_no: get(row, "redeemed_order_no")?, - redeemed_at_unix_secs: optional_timestamp( - get(row, "redeemed_at_unix_secs")?, - "redeem_codes.redeemed_at", - )?, - disabled_by: get(row, "disabled_by")?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "redeem_code_batches.expires_at", - )?, - created_at_unix_ms: timestamp(get(row, "created_at_unix_ms")?, "redeem_codes.created_at")?, - updated_at_unix_secs: timestamp( - get(row, "updated_at_unix_secs")?, - "redeem_codes.updated_at", - )?, - }) -} - -fn map_daily_usage_row(row: &MySqlRow) -> Result { - Ok(StoredWalletDailyUsageLedger { - id: get(row, "id")?, - billing_date: get(row, "billing_date")?, - billing_timezone: get(row, "billing_timezone")?, - total_cost_usd: get(row, "total_cost_usd")?, - total_requests: nonnegative_u64( - get(row, "total_requests")?, - "wallet_daily_usage_ledgers.total_requests", - )?, - input_tokens: nonnegative_u64( - get(row, "input_tokens")?, - "wallet_daily_usage_ledgers.input_tokens", - )?, - output_tokens: nonnegative_u64( - get(row, "output_tokens")?, - "wallet_daily_usage_ledgers.output_tokens", - )?, - cache_creation_tokens: nonnegative_u64( - get(row, "cache_creation_tokens")?, - "wallet_daily_usage_ledgers.cache_creation_tokens", - )?, - cache_read_tokens: nonnegative_u64( - get(row, "cache_read_tokens")?, - "wallet_daily_usage_ledgers.cache_read_tokens", - )?, - first_finalized_at_unix_secs: optional_timestamp( - get(row, "first_finalized_at_unix_secs")?, - "wallet_daily_usage_ledgers.first_finalized_at", - )?, - last_finalized_at_unix_secs: optional_timestamp( - get(row, "last_finalized_at_unix_secs")?, - "wallet_daily_usage_ledgers.last_finalized_at", - )?, - aggregated_at_unix_secs: optional_timestamp( - get(row, "aggregated_at_unix_secs")?, - "wallet_daily_usage_ledgers.aggregated_at", - )?, - }) -} - -fn get(row: &MySqlRow, field: &str) -> Result -where - for<'r> T: sqlx::Decode<'r, sqlx::MySql> + sqlx::Type, -{ - 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, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn timestamp(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("{field_name} contains a negative timestamp")) - }) -} - -fn optional_timestamp(value: Option, field_name: &str) -> Result, DataLayerError> { - value.map(|value| timestamp(value, field_name)).transpose() -} - -fn nonnegative_u64(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("{field_name} contains a negative value")) - }) -} - -fn mask_redeem_code(prefix: &str, suffix: &str) -> String { - format!("{prefix}****{suffix}") -} - -#[cfg(test)] -mod tests; diff --git a/crates/aether-data/adapters/mysql/src/wallet/tests.rs b/crates/aether-data/adapters/mysql/src/wallet/tests.rs deleted file mode 100644 index a9a38b9b3..000000000 --- a/crates/aether-data/adapters/mysql/src/wallet/tests.rs +++ /dev/null @@ -1,630 +0,0 @@ -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, 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(" ") -} - -#[test] -fn mysql_gateway_order_uniqueness_preflights_before_persistent_changes() { - const UNIQUENESS_MIGRATION: &str = include_str!( - "../../migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql" - ); - - let executable_migration = UNIQUENESS_MIGRATION - .lines() - .filter(|line| !line.trim_start().starts_with("--")) - .collect::>() - .join("\n"); - let migration = compact_sql(&executable_migration); - let initial_cleanup = migration - .find("DROP TEMPORARY TABLE IF EXISTS aether_payment_gateway_order_uniqueness_preflight") - .expect("migration should clean up a same-session failed preflight"); - let create_preflight = migration - .find("CREATE TEMPORARY TABLE aether_payment_gateway_order_uniqueness_preflight") - .expect("migration should create a non-persistent conflict guard"); - let seed_preflight = migration - .find("INSERT INTO aether_payment_gateway_order_uniqueness_preflight (conflict_marker) VALUES (1)") - .expect("migration should seed the duplicate-key conflict guard"); - let conflict_probe = migration - .find("INSERT INTO aether_payment_gateway_order_uniqueness_preflight (conflict_marker) SELECT 1 FROM payment_orders") - .expect("migration should reject normalized historical conflicts"); - let final_cleanup = migration - .find("DROP TEMPORARY TABLE aether_payment_gateway_order_uniqueness_preflight;") - .expect("migration should remove the successful preflight guard"); - let first_persistent_update = migration - .find("UPDATE payment_orders SET payment_method") - .expect("migration should normalize payment order methods"); - let alter = migration - .find( - "MODIFY COLUMN gateway_order_id VARCHAR(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin NULL", - ) - .expect("migration should enforce binary gateway identifiers"); - let unique = migration - .find("ADD UNIQUE INDEX uq_payment_orders_payment_method_gateway_order_id") - .expect("migration should create composite uniqueness"); - - assert!(migration.contains("GROUP BY LOWER(TRIM(payment_method)), CONVERT(gateway_order_id USING utf8mb4) COLLATE utf8mb4_0900_bin HAVING COUNT(*) > 1 LIMIT 1")); - assert!(migration.contains("WHERE BINARY payment_method <> BINARY LOWER(TRIM(payment_method))")); - let preflight = &migration[..final_cleanup]; - assert!(!preflight.contains("UPDATE payment_orders")); - assert!(!preflight.contains("UPDATE payment_callbacks")); - assert!(!preflight.contains("ALTER TABLE")); - assert!( - initial_cleanup < create_preflight - && create_preflight < seed_preflight - && seed_preflight < conflict_probe - && conflict_probe < final_cleanup - && final_cleanup < first_persistent_update - && first_persistent_update < alter, - "the conflict probe must finish before any persistent UPDATE or ALTER" - ); - assert!( - alter < unique && !migration.contains("CREATE UNIQUE INDEX"), - "the collation and unique index must be one atomic ALTER TABLE" - ); -} - -#[tokio::test] -async fn mysql_wallet_read_repository_reads_wallet_contract_views() { - let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") - .ok() - .filter(|value| !value.trim().is_empty()) - else { - eprintln!("skipping mysql wallet read smoke 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_rows(&pool).await; - seed_rows(&pool).await; - - let repository = MysqlWalletReadRepository::new(pool); - let wallet = repository - .find(WalletLookupKey::UserId("user-1")) - .await - .expect("wallet find should query") - .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("wallet-read-smoke".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("wallet-read-smoke".to_string()), - owner_type: Some("unknown".to_string()), - limit: 10, - offset: 0, - }) - .await - .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 { - status: Some("credited".to_string()), - payment_method: Some("redeem_code".to_string()), - limit: 10, - offset: 0, - }) - .await - .expect("payment orders should list"); - assert_eq!(orders.total, 1); - assert_eq!( - orders.items[0].gateway_response.as_ref().unwrap()["ok"], - 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", 1, 1) - .await - .expect("refunds should list"); - 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"), 1, 1) - .await - .expect("callbacks should list"); - 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(), - status: Some("redeemed".to_string()), - limit: 10, - offset: 0, - }) - .await - .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'") - .bind(today) - .execute(repository.pool()) - .await - .expect("daily row should update"); - let daily = repository - .find_wallet_today_usage("wallet-1", "UTC") - .await - .expect("daily usage should query") - .expect("daily usage should exist"); - assert_eq!(daily.total_requests, 2); - - cleanup_rows(repository.pool()).await; -} - -impl MysqlWalletReadRepository { - fn pool(&self) -> &sqlx::MySqlPool { - &self.pool - } -} - -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 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) - .await - .expect("cleanup should succeed"); - } -} - -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), - ('user-2', 'Bob', 'bob@example.com', 'local', 1, 1) -"#, - ) - .execute(pool) - .await - .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, api_key_id, balance, gift_balance, status, total_recharged, - total_consumed, total_refunded, total_adjusted, created_at, updated_at -) VALUES - ('wallet-1', 'user-1', NULL, 10.0, 2.0, 'wallet-read-smoke', 20.0, 4.0, 1.0, 3.0, 1, 2), - ('wallet-2', 'user-2', NULL, 4.0, 1.0, 'wallet-read-smoke', 5.0, 2.0, 0.0, 0.0, 1, 3), - ('wallet-api-1', NULL, 'api-key-1', 7.0, 0.0, 'wallet-read-smoke', 7.0, 0.0, 0.0, 0.0, 1, 4) -"#, - ) - .execute(pool) - .await - .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, - 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 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), - ('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("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), - ('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("refunds should seed"); - - sqlx::query( - r#" -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), - ('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("transactions should seed"); - - sqlx::query( - r#" -INSERT INTO redeem_code_batches ( - 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 batches should seed"); - - sqlx::query( - r#" -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), - ('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 codes should seed"); - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES ( - 'daily-1', 'wallet-1', '2000-01-01', 'UTC', 1.25, 2, 10, 20, 3, 4, 11, 11, 11 -) -"#, - ) - .execute(pool) - .await - .expect("daily usage should seed"); -} diff --git a/crates/aether-data/adapters/postgres/src/quota.rs b/crates/aether-data/adapters/postgres/src/quota.rs index 04498cf75..84f0ab997 100644 --- a/crates/aether-data/adapters/postgres/src/quota.rs +++ b/crates/aether-data/adapters/postgres/src/quota.rs @@ -15,25 +15,21 @@ fn quota_snapshot_select() -> SelectQuery<'static> { DialectSql::common("billing_type").with_postgres("CAST(billing_type AS TEXT)"), ) .alias("billing_type"), - SelectColumn::expr(DialectSql::dialect( + SelectColumn::expr(DialectSql::common( "CAST(monthly_quota_usd AS DOUBLE PRECISION)", - "CAST(monthly_quota_usd AS REAL)", )) .alias("monthly_quota_usd"), - SelectColumn::expr(DialectSql::dialect( + SelectColumn::expr(DialectSql::common( "CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE PRECISION)", - "CAST(COALESCE(monthly_used_usd, 0) AS REAL)", )) .alias("monthly_used_usd"), SelectColumn::expr("quota_reset_day"), - SelectColumn::expr(DialectSql::dialect( + SelectColumn::expr(DialectSql::common( "CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT)", - "quota_last_reset_at", )) .alias("quota_last_reset_at_unix_secs"), - SelectColumn::expr(DialectSql::dialect( + SelectColumn::expr(DialectSql::common( "CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT)", - "quota_expires_at", )) .alias("quota_expires_at_unix_secs"), SelectColumn::expr("is_active"), diff --git a/crates/aether-data/adapters/sqlite/Cargo.toml b/crates/aether-data/adapters/sqlite/Cargo.toml deleted file mode 100644 index 931584177..000000000 --- a/crates/aether-data/adapters/sqlite/Cargo.toml +++ /dev/null @@ -1,24 +0,0 @@ -[package] -name = "aether-data-sqlite" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true -description = "SQLite repositories, pools, and migrations for Aether" - -[dependencies] -aether-ai-formats.workspace = true -aether-data-contracts.workspace = true -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 = ["sqlite", "runtime-tokio-rustls", "chrono", "migrate", "macros"] } -tracing.workspace = true -uuid.workspace = true - -[dev-dependencies] -tokio.workspace = true diff --git a/crates/aether-data/adapters/sqlite/migrations/20260403000000_baseline.sql b/crates/aether-data/adapters/sqlite/migrations/20260403000000_baseline.sql deleted file mode 100644 index 741957f48..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260403000000_baseline.sql +++ /dev/null @@ -1,1074 +0,0 @@ -CREATE TABLE IF NOT EXISTS users ( - id TEXT PRIMARY KEY, - external_id TEXT, - email TEXT UNIQUE, - username TEXT UNIQUE, - password_hash TEXT, - role TEXT, - auth_source TEXT NOT NULL DEFAULT 'local', - email_verified INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - is_deleted INTEGER NOT NULL DEFAULT 0, - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - model_capability_settings TEXT, - rate_limit INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - last_login_at INTEGER, - ldap_dn TEXT, - ldap_username TEXT -); - -CREATE TABLE IF NOT EXISTS api_keys ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - key_hash TEXT NOT NULL UNIQUE, - key_encrypted TEXT, - name TEXT, - key_prefix TEXT, - status TEXT NOT NULL DEFAULT 'active', - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - rate_limit INTEGER DEFAULT 100, - concurrent_limit INTEGER, - force_capabilities TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_locked INTEGER NOT NULL DEFAULT 0, - is_standalone INTEGER NOT NULL DEFAULT 0, - auto_delete_on_expiry INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - metadata TEXT, - expires_at INTEGER, - last_used_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS api_keys_user_id_idx ON api_keys (user_id); - -CREATE TABLE IF NOT EXISTS audit_logs ( - id TEXT PRIMARY KEY, - event_type TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - description TEXT NOT NULL, - ip_address TEXT, - user_agent TEXT, - request_id TEXT, - event_metadata TEXT, - status_code INTEGER, - error_message TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS audit_logs_created_at_idx ON audit_logs (created_at); -CREATE INDEX IF NOT EXISTS audit_logs_event_type_idx ON audit_logs (event_type); -CREATE INDEX IF NOT EXISTS audit_logs_request_id_idx ON audit_logs (request_id); -CREATE INDEX IF NOT EXISTS audit_logs_user_id_idx ON audit_logs (user_id); - -CREATE TABLE IF NOT EXISTS announcements ( - id TEXT PRIMARY KEY, - title TEXT NOT NULL, - content TEXT NOT NULL, - type TEXT NOT NULL DEFAULT 'info', - priority INTEGER NOT NULL DEFAULT 0, - author_id TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_pinned INTEGER NOT NULL DEFAULT 0, - start_time INTEGER, - end_time INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS announcements_author_id_idx ON announcements (author_id); -CREATE INDEX IF NOT EXISTS announcements_created_at_idx ON announcements (created_at); -CREATE INDEX IF NOT EXISTS announcements_is_active_idx ON announcements (is_active); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - announcement_id TEXT NOT NULL, - read_at INTEGER NOT NULL, - UNIQUE (user_id, announcement_id) -); -CREATE INDEX IF NOT EXISTS announcement_reads_announcement_id_idx ON announcement_reads (announcement_id); -CREATE INDEX IF NOT EXISTS announcement_reads_user_id_idx ON announcement_reads (user_id); - -CREATE TABLE IF NOT EXISTS management_tokens ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - name TEXT NOT NULL, - description TEXT, - token_hash TEXT NOT NULL UNIQUE, - token_prefix TEXT, - allowed_ips TEXT, - expires_at INTEGER, - last_used_at INTEGER, - last_used_ip TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (user_id, name) -); -CREATE INDEX IF NOT EXISTS management_tokens_user_id_idx ON management_tokens (user_id); - -CREATE TABLE IF NOT EXISTS user_preferences ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL UNIQUE, - avatar_url TEXT, - bio TEXT, - default_provider_id TEXT, - theme TEXT NOT NULL DEFAULT 'light', - language TEXT NOT NULL DEFAULT 'zh-CN', - timezone TEXT NOT NULL DEFAULT 'Asia/Shanghai', - email_notifications INTEGER NOT NULL DEFAULT 1, - usage_alerts INTEGER NOT NULL DEFAULT 1, - announcement_notifications INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS user_preferences_default_provider_id_idx - ON user_preferences (default_provider_id); -CREATE INDEX IF NOT EXISTS user_preferences_user_id_idx - ON user_preferences (user_id); - -CREATE TABLE IF NOT EXISTS user_sessions ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - client_device_id TEXT NOT NULL, - device_label TEXT, - device_type TEXT NOT NULL DEFAULT 'unknown', - browser_name TEXT, - browser_version TEXT, - os_name TEXT, - os_version TEXT, - device_model TEXT, - ip_address TEXT, - user_agent TEXT, - client_hints TEXT, - refresh_token_hash TEXT NOT NULL, - prev_refresh_token_hash TEXT, - rotated_at INTEGER, - last_seen_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - revoked_at INTEGER, - revoke_reason TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS user_sessions_user_active_idx - ON user_sessions (user_id, revoked_at, expires_at); -CREATE INDEX IF NOT EXISTS user_sessions_user_device_idx - ON user_sessions (user_id, client_device_id); -CREATE TABLE IF NOT EXISTS billing_rules ( - id TEXT PRIMARY KEY, - global_model_id TEXT, - model_id TEXT, - name TEXT NOT NULL, - task_type TEXT NOT NULL DEFAULT 'chat', - expression TEXT NOT NULL, - variables TEXT NOT NULL DEFAULT '{}', - dimension_mappings TEXT NOT NULL DEFAULT '{}', - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CHECK ( - (global_model_id IS NOT NULL AND model_id IS NULL) - OR (global_model_id IS NULL AND model_id IS NOT NULL) - ) -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_billing_rules_global_model_task - ON billing_rules (global_model_id, task_type) - WHERE is_enabled = 1 AND global_model_id IS NOT NULL; -CREATE UNIQUE INDEX IF NOT EXISTS uq_billing_rules_model_task - ON billing_rules (model_id, task_type) - WHERE is_enabled = 1 AND model_id IS NOT NULL; - -CREATE TABLE IF NOT EXISTS dimension_collectors ( - id TEXT PRIMARY KEY, - api_format TEXT NOT NULL, - task_type TEXT NOT NULL, - dimension_name TEXT NOT NULL, - source_type TEXT NOT NULL, - source_path TEXT, - value_type TEXT NOT NULL DEFAULT 'float', - transform_expression TEXT, - default_value TEXT, - priority INTEGER NOT NULL DEFAULT 0, - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CHECK ( - (source_type = 'computed' AND source_path IS NULL AND transform_expression IS NOT NULL) - OR (source_type <> 'computed' AND source_path IS NOT NULL) - ) -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_dimension_collectors_enabled - ON dimension_collectors (api_format, task_type, dimension_name, priority) - WHERE is_enabled = 1; - -CREATE TABLE IF NOT EXISTS providers ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - description TEXT, - website TEXT, - provider_type TEXT NOT NULL, - billing_type TEXT, - monthly_quota_usd REAL, - monthly_used_usd REAL, - quota_reset_day INTEGER, - quota_last_reset_at INTEGER, - quota_expires_at INTEGER, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - priority INTEGER NOT NULL DEFAULT 0, - provider_priority INTEGER NOT NULL DEFAULT 100, - keep_priority_on_conversion INTEGER NOT NULL DEFAULT 0, - enable_format_conversion INTEGER NOT NULL DEFAULT 1, - concurrent_limit INTEGER, - max_retries INTEGER, - proxy TEXT, - request_timeout REAL, - stream_first_byte_timeout REAL, - config TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - api_key TEXT, - encrypted_key TEXT, - auth_type TEXT NOT NULL DEFAULT 'api_key', - auth_config TEXT, - note TEXT, - internal_priority INTEGER NOT NULL DEFAULT 50, - capabilities TEXT, - api_formats TEXT, - auth_type_by_format TEXT, - allow_auth_channel_mismatch_formats TEXT, - rate_multipliers TEXT, - global_priority_by_format TEXT, - allowed_models TEXT, - expires_at INTEGER, - cache_ttl_minutes INTEGER NOT NULL DEFAULT 5, - max_probe_interval_minutes INTEGER NOT NULL DEFAULT 32, - proxy TEXT, - fingerprint TEXT, - concurrent_limit INTEGER, - learned_rpm_limit INTEGER, - concurrent_429_count INTEGER NOT NULL DEFAULT 0, - rpm_429_count INTEGER NOT NULL DEFAULT 0, - last_429_at INTEGER, - last_429_type TEXT, - adjustment_history TEXT, - utilization_samples TEXT, - last_probe_increase_at INTEGER, - last_rpm_peak INTEGER, - request_count INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - success_count INTEGER NOT NULL DEFAULT 0, - error_count INTEGER NOT NULL DEFAULT 0, - total_response_time_ms INTEGER NOT NULL DEFAULT 0, - last_used_at INTEGER, - auto_fetch_models INTEGER NOT NULL DEFAULT 0, - last_models_fetch_at INTEGER, - last_models_fetch_error TEXT, - locked_models TEXT, - model_include_patterns TEXT, - model_exclude_patterns TEXT, - upstream_metadata TEXT, - oauth_invalid_at INTEGER, - oauth_invalid_reason TEXT, - status_snapshot TEXT, - health_by_format TEXT, - circuit_breaker_by_format TEXT, - status TEXT NOT NULL DEFAULT 'active', - is_active INTEGER NOT NULL DEFAULT 1, - weight INTEGER NOT NULL DEFAULT 1, - rpm_limit INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id); -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id TEXT PRIMARY KEY, - pool_kind TEXT NOT NULL, - pool_id TEXT NOT NULL, - member_kind TEXT NOT NULL, - member_id TEXT NOT NULL, - capability TEXT NOT NULL, - scope_kind TEXT NOT NULL, - scope_id TEXT, - score REAL NOT NULL DEFAULT 0, - hard_state TEXT NOT NULL DEFAULT 'unknown', - score_version INTEGER NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at INTEGER, - last_scheduled_at INTEGER, - last_success_at INTEGER, - last_failure_at INTEGER, - failure_count INTEGER NOT NULL DEFAULT 0, - last_probe_attempt_at INTEGER, - last_probe_success_at INTEGER, - last_probe_failure_at INTEGER, - probe_failure_count INTEGER NOT NULL DEFAULT 0, - probe_status TEXT NOT NULL DEFAULT 'never', - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC); -CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id); -CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at); -CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at); - -CREATE TABLE IF NOT EXISTS gemini_file_mappings ( - id TEXT PRIMARY KEY, - file_name TEXT NOT NULL UNIQUE, - key_id TEXT NOT NULL, - user_id TEXT, - display_name TEXT, - mime_type TEXT, - source_hash TEXT, - created_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_key_id_idx ON gemini_file_mappings (key_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_user_id_idx ON gemini_file_mappings (user_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_expires_at_idx ON gemini_file_mappings (expires_at); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_source_hash_idx ON gemini_file_mappings (source_hash); - -CREATE TABLE IF NOT EXISTS request_candidates ( - id TEXT PRIMARY KEY, - request_id TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - candidate_index INTEGER NOT NULL, - retry_index INTEGER NOT NULL DEFAULT 0, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - status TEXT NOT NULL, - skip_reason TEXT, - is_cached INTEGER NOT NULL DEFAULT 0, - status_code INTEGER, - error_type TEXT, - error_message TEXT, - latency_ms INTEGER, - concurrent_requests INTEGER, - extra_data TEXT, - required_capabilities TEXT, - created_at INTEGER NOT NULL, - started_at INTEGER, - finished_at INTEGER, - UNIQUE (request_id, candidate_index, retry_index) -); -CREATE INDEX IF NOT EXISTS request_candidates_request_id_idx ON request_candidates (request_id); -CREATE INDEX IF NOT EXISTS request_candidates_provider_id_idx ON request_candidates (provider_id); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_id_idx ON request_candidates (endpoint_id); -CREATE INDEX IF NOT EXISTS request_candidates_status_idx ON request_candidates (status); -CREATE INDEX IF NOT EXISTS request_candidates_created_at_idx ON request_candidates (created_at); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_status_created_idx ON request_candidates (endpoint_id, status, created_at); - -CREATE TABLE IF NOT EXISTS video_tasks ( - id TEXT PRIMARY KEY, - short_id TEXT UNIQUE, - request_id TEXT NOT NULL UNIQUE, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - external_task_id TEXT, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - client_api_format TEXT, - provider_api_format TEXT, - format_converted INTEGER NOT NULL DEFAULT 0, - model TEXT, - prompt TEXT, - original_request_body TEXT, - duration_seconds INTEGER, - resolution TEXT, - aspect_ratio TEXT, - size TEXT, - status TEXT NOT NULL DEFAULT 'pending', - progress_percent INTEGER NOT NULL DEFAULT 0, - progress_message TEXT, - retry_count INTEGER NOT NULL DEFAULT 0, - poll_interval_seconds INTEGER NOT NULL DEFAULT 10, - next_poll_at INTEGER, - poll_count INTEGER NOT NULL DEFAULT 0, - max_poll_count INTEGER NOT NULL DEFAULT 360, - created_at INTEGER NOT NULL, - submitted_at INTEGER, - completed_at INTEGER, - updated_at INTEGER NOT NULL, - error_code TEXT, - error_message TEXT, - video_url TEXT, - request_metadata TEXT -); -CREATE INDEX IF NOT EXISTS video_tasks_external_id_idx ON video_tasks (external_task_id); -CREATE INDEX IF NOT EXISTS video_tasks_next_poll_idx ON video_tasks (next_poll_at); -CREATE INDEX IF NOT EXISTS video_tasks_request_id_idx ON video_tasks (request_id); -CREATE INDEX IF NOT EXISTS video_tasks_user_status_idx ON video_tasks (user_id, status); -CREATE INDEX IF NOT EXISTS video_tasks_api_key_id_idx ON video_tasks (api_key_id); -CREATE INDEX IF NOT EXISTS video_tasks_provider_id_idx ON video_tasks (provider_id); -CREATE INDEX IF NOT EXISTS video_tasks_endpoint_id_idx ON video_tasks (endpoint_id); -CREATE INDEX IF NOT EXISTS video_tasks_key_id_idx ON video_tasks (key_id); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - base_url TEXT NOT NULL, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - health_score REAL NOT NULL DEFAULT 1.0, - weight INTEGER NOT NULL DEFAULT 1, - header_rules TEXT, - body_rules TEXT, - max_retries INTEGER, - custom_path TEXT, - metadata TEXT, - config TEXT, - format_acceptance_config TEXT, - proxy TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_endpoints_provider_id_idx ON provider_endpoints (provider_id); - -CREATE TABLE IF NOT EXISTS models ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - global_model_id TEXT, - provider_model_name TEXT NOT NULL, - global_model_name TEXT, - api_format TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - is_available INTEGER NOT NULL DEFAULT 1, - price_per_request REAL, - tiered_pricing TEXT, - supports_vision INTEGER, - supports_function_calling INTEGER, - supports_streaming INTEGER, - supports_extended_thinking INTEGER, - supports_image_generation INTEGER, - provider_model_mappings TEXT, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS models_provider_id_idx ON models (provider_id); - -CREATE TABLE IF NOT EXISTS global_models ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - display_name TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - default_price_per_request REAL, - default_tiered_pricing TEXT, - supported_capabilities TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE TABLE IF NOT EXISTS system_configs ( - id TEXT PRIMARY KEY, - key TEXT NOT NULL UNIQUE, - value TEXT NOT NULL, - description TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - id TEXT PRIMARY KEY, - module_type TEXT NOT NULL UNIQUE, - enabled INTEGER NOT NULL DEFAULT 1, - config TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - provider_type TEXT PRIMARY KEY, - display_name TEXT NOT NULL, - client_id TEXT NOT NULL, - client_secret_encrypted TEXT, - authorization_url_override TEXT, - token_url_override TEXT, - userinfo_url_override TEXT, - scopes TEXT, - redirect_uri TEXT NOT NULL, - frontend_callback_url TEXT NOT NULL, - attribute_mapping TEXT, - extra_config TEXT, - is_enabled INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - server_url TEXT NOT NULL, - bind_dn TEXT NOT NULL, - bind_password_encrypted TEXT, - base_dn TEXT NOT NULL, - user_search_filter TEXT DEFAULT '(uid={username})' NOT NULL, - username_attr TEXT DEFAULT 'uid' NOT NULL, - email_attr TEXT DEFAULT 'mail' NOT NULL, - display_name_attr TEXT DEFAULT 'cn' NOT NULL, - is_enabled INTEGER NOT NULL DEFAULT 0, - is_exclusive INTEGER NOT NULL DEFAULT 0, - use_starttls INTEGER NOT NULL DEFAULT 0, - connect_timeout INTEGER NOT NULL DEFAULT 10, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - provider_type TEXT NOT NULL, - provider_user_id TEXT NOT NULL, - provider_username TEXT, - provider_email TEXT, - extra_data TEXT, - linked_at INTEGER NOT NULL, - last_login_at INTEGER -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_user_oauth_links_provider_user - ON user_oauth_links (provider_type, provider_user_id); -CREATE UNIQUE INDEX IF NOT EXISTS uq_user_oauth_links_user_provider - ON user_oauth_links (user_id, provider_type); -CREATE INDEX IF NOT EXISTS user_oauth_links_provider_type_idx ON user_oauth_links (provider_type); -CREATE INDEX IF NOT EXISTS user_oauth_links_user_id_idx ON user_oauth_links (user_id); -CREATE TABLE IF NOT EXISTS proxy_nodes ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - ip TEXT NOT NULL, - port INTEGER NOT NULL, - region TEXT, - status TEXT NOT NULL DEFAULT 'online', - registered_by TEXT, - last_heartbeat_at INTEGER, - heartbeat_interval INTEGER NOT NULL DEFAULT 30, - active_connections INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - avg_latency_ms REAL, - is_manual INTEGER NOT NULL DEFAULT 0, - proxy_url TEXT, - proxy_username TEXT, - proxy_password TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - remote_config TEXT, - config_version INTEGER NOT NULL DEFAULT 0, - hardware_info TEXT, - estimated_max_concurrency INTEGER, - tunnel_mode INTEGER NOT NULL DEFAULT 0, - tunnel_connected INTEGER NOT NULL DEFAULT 0, - tunnel_connected_at INTEGER, - failed_requests INTEGER NOT NULL DEFAULT 0, - dns_failures INTEGER NOT NULL DEFAULT 0, - stream_errors INTEGER NOT NULL DEFAULT 0, - proxy_metadata TEXT -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - node_id TEXT NOT NULL, - event_type TEXT NOT NULL, - detail TEXT, - created_at INTEGER NOT NULL -); -CREATE TABLE IF NOT EXISTS wallets ( - id TEXT PRIMARY KEY, - user_id TEXT UNIQUE, - api_key_id TEXT UNIQUE, - balance REAL NOT NULL DEFAULT 0, - gift_balance REAL NOT NULL DEFAULT 0, - limit_mode TEXT NOT NULL DEFAULT 'finite', - currency TEXT NOT NULL DEFAULT 'USD', - status TEXT NOT NULL DEFAULT 'active', - total_recharged REAL NOT NULL DEFAULT 0, - total_consumed REAL NOT NULL DEFAULT 0, - total_refunded REAL NOT NULL DEFAULT 0, - total_adjusted REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS wallets_api_key_id_idx ON wallets (api_key_id); -CREATE INDEX IF NOT EXISTS wallets_user_id_idx ON wallets (user_id); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - id TEXT PRIMARY KEY, - wallet_id TEXT NOT NULL, - category TEXT NOT NULL, - reason_code TEXT NOT NULL, - amount REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - recharge_balance_before REAL NOT NULL, - recharge_balance_after REAL NOT NULL, - gift_balance_before REAL NOT NULL, - gift_balance_after REAL NOT NULL, - link_type TEXT, - link_id TEXT, - operator_id TEXT, - description TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_wallet_created - ON wallet_transactions (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_category_created - ON wallet_transactions (category, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_reason_created - ON wallet_transactions (reason_code, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_link - ON wallet_transactions (link_type, link_id); -CREATE INDEX IF NOT EXISTS ix_wallet_transactions_operator_id - ON wallet_transactions (operator_id); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - id TEXT PRIMARY KEY, - wallet_id TEXT NOT NULL, - billing_date TEXT NOT NULL, - billing_timezone TEXT NOT NULL, - total_cost_usd REAL NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - first_finalized_at INTEGER, - last_finalized_at INTEGER, - aggregated_at INTEGER NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_daily_usage_wallet_date - ON wallet_daily_usage_ledgers (wallet_id, billing_timezone, billing_date); - -CREATE TABLE IF NOT EXISTS payment_orders ( - id TEXT PRIMARY KEY, - order_no TEXT NOT NULL UNIQUE, - wallet_id TEXT NOT NULL, - user_id TEXT, - amount_usd REAL NOT NULL, - pay_amount REAL, - pay_currency TEXT, - exchange_rate REAL, - refunded_amount_usd REAL NOT NULL DEFAULT 0, - refundable_amount_usd REAL NOT NULL DEFAULT 0, - payment_method TEXT NOT NULL, - gateway_order_id TEXT, - gateway_response TEXT, - status TEXT NOT NULL DEFAULT 'pending', - created_at INTEGER NOT NULL, - paid_at INTEGER, - credited_at INTEGER, - expires_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_payment_orders_wallet_created - ON payment_orders (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_user_created - ON payment_orders (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_status - ON payment_orders (status); -CREATE INDEX IF NOT EXISTS idx_payment_orders_gateway_order_id - ON payment_orders (gateway_order_id); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - id TEXT PRIMARY KEY, - payment_order_id TEXT, - payment_method TEXT NOT NULL, - callback_key TEXT NOT NULL UNIQUE, - order_no TEXT, - gateway_order_id TEXT, - payload_hash TEXT, - signature_valid INTEGER NOT NULL DEFAULT 0, - status TEXT NOT NULL DEFAULT 'received', - payload TEXT, - error_message TEXT, - created_at INTEGER NOT NULL, - processed_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_order - ON payment_callbacks (order_no); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_gateway_order - ON payment_callbacks (gateway_order_id); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_created - ON payment_callbacks (created_at); -CREATE INDEX IF NOT EXISTS ix_payment_callbacks_payment_order_id - ON payment_callbacks (payment_order_id); - -CREATE TABLE IF NOT EXISTS refund_requests ( - id TEXT PRIMARY KEY, - refund_no TEXT NOT NULL UNIQUE, - wallet_id TEXT NOT NULL, - user_id TEXT, - payment_order_id TEXT, - source_type TEXT NOT NULL, - source_id TEXT, - refund_mode TEXT NOT NULL, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending_approval', - reason TEXT, - requested_by TEXT, - approved_by TEXT, - processed_by TEXT, - gateway_refund_id TEXT, - payout_method TEXT, - payout_reference TEXT, - payout_proof TEXT, - failure_reason TEXT, - idempotency_key TEXT UNIQUE, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - processed_at INTEGER, - completed_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_refund_wallet_created - ON refund_requests (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_user_created - ON refund_requests (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_status - ON refund_requests (status); -CREATE INDEX IF NOT EXISTS ix_refund_requests_payment_order_id - ON refund_requests (payment_order_id); -CREATE INDEX IF NOT EXISTS ix_refund_requests_requested_by - ON refund_requests (requested_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_approved_by - ON refund_requests (approved_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_processed_by - ON refund_requests (processed_by); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - amount_usd REAL NOT NULL, - currency TEXT NOT NULL DEFAULT 'USD', - balance_bucket TEXT NOT NULL DEFAULT 'gift', - total_count INTEGER NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - description TEXT, - created_by TEXT, - expires_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_redeem_code_batches_status - ON redeem_code_batches (status, created_at); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - id TEXT PRIMARY KEY, - batch_id TEXT NOT NULL, - code_hash TEXT NOT NULL UNIQUE, - code_prefix TEXT NOT NULL, - code_suffix TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - redeemed_by_user_id TEXT, - redeemed_wallet_id TEXT, - redeemed_payment_order_id TEXT, - redeemed_at INTEGER, - disabled_by TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_batch_created - ON redeem_codes (batch_id, created_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_status - ON redeem_codes (status, updated_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_user - ON redeem_codes (redeemed_by_user_id, redeemed_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_order - ON redeem_codes (redeemed_payment_order_id); -CREATE TABLE IF NOT EXISTS "usage" ( - request_id TEXT PRIMARY KEY, - id TEXT, - user_id TEXT, - api_key_id TEXT, - provider_name TEXT NOT NULL DEFAULT 'unknown', - model TEXT NOT NULL DEFAULT 'unknown', - target_model TEXT, - provider_id TEXT, - provider_endpoint_id TEXT, - provider_api_key_id TEXT, - request_type TEXT, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - endpoint_api_format TEXT, - provider_api_family TEXT, - provider_endpoint_kind TEXT, - has_format_conversion INTEGER NOT NULL DEFAULT 0, - is_stream INTEGER NOT NULL DEFAULT 0, - upstream_is_stream INTEGER, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_5m_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_1h_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_cost_usd REAL NOT NULL DEFAULT 0, - cache_read_cost_usd REAL NOT NULL DEFAULT 0, - output_price_per_1m REAL, - status_code INTEGER, - error_message TEXT, - error_category TEXT, - response_time_ms INTEGER, - first_byte_time_ms INTEGER, - wallet_id TEXT, - status TEXT NOT NULL DEFAULT 'completed', - billing_status TEXT NOT NULL DEFAULT 'pending', - total_cost_usd REAL NOT NULL DEFAULT 0, - actual_total_cost_usd REAL NOT NULL DEFAULT 0, - request_metadata TEXT, - 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, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - wallet_gift_balance_before REAL, - wallet_gift_balance_after REAL, - finalized_at INTEGER, - created_at_unix_ms INTEGER NOT NULL DEFAULT 0, - updated_at_unix_secs INTEGER NOT NULL DEFAULT 0 -); -CREATE INDEX IF NOT EXISTS usage_api_key_id_idx ON "usage" (api_key_id); -CREATE INDEX IF NOT EXISTS usage_billing_status_idx ON "usage" (billing_status); -CREATE INDEX IF NOT EXISTS usage_created_at_idx ON "usage" (created_at_unix_ms); -CREATE INDEX IF NOT EXISTS usage_provider_api_key_id_idx ON "usage" (provider_api_key_id); -CREATE INDEX IF NOT EXISTS usage_provider_id_idx ON "usage" (provider_id); -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_settlement_snapshots ( - request_id TEXT PRIMARY KEY, - billing_status TEXT NOT NULL, - wallet_id TEXT, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - wallet_gift_balance_before REAL, - wallet_gift_balance_after REAL, - provider_monthly_used_usd REAL, - 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 TABLE IF NOT EXISTS stats_hourly ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL UNIQUE, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - aggregated_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, user_id) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, user_id, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL UNIQUE, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - fallback_count INTEGER NOT NULL DEFAULT 0, - unique_models INTEGER NOT NULL DEFAULT 0, - unique_providers INTEGER NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - aggregated_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", model) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - id TEXT PRIMARY KEY, - api_key_id TEXT NOT NULL, - "date" INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - api_key_name TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", api_key_id) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - error_category TEXT NOT NULL, - provider_name TEXT, - model TEXT, - count INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", error_category, provider_name, model) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - "date" INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - username TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", user_id) -); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260507120000_add_management_token_permissions.sql b/crates/aether-data/adapters/sqlite/migrations/20260507120000_add_management_token_permissions.sql deleted file mode 100644 index 2cc33bb3c..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260507120000_add_management_token_permissions.sql +++ /dev/null @@ -1,3 +0,0 @@ -ALTER TABLE management_tokens - ADD COLUMN permissions TEXT; - diff --git a/crates/aether-data/adapters/sqlite/migrations/20260508000000_add_proxy_node_metrics.sql b/crates/aether-data/adapters/sqlite/migrations/20260508000000_add_proxy_node_metrics.sql deleted file mode 100644 index b2b7d8c79..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260508000000_add_proxy_node_metrics.sql +++ /dev/null @@ -1,46 +0,0 @@ -ALTER TABLE proxy_node_events - ADD COLUMN event_metadata TEXT; - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m ( - node_id TEXT NOT NULL, - bucket_start_unix_secs INTEGER NOT NULL, - samples INTEGER NOT NULL DEFAULT 0, - uptime_samples INTEGER NOT NULL DEFAULT 0, - active_connections_sum INTEGER NOT NULL DEFAULT 0, - active_connections_max INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0, - connect_errors_delta INTEGER NOT NULL DEFAULT 0, - disconnects_delta INTEGER NOT NULL DEFAULT 0, - error_events_delta INTEGER NOT NULL DEFAULT 0, - ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_in_frames_delta INTEGER NOT NULL DEFAULT 0, - ws_out_frames_delta INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs) -); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h ( - node_id TEXT NOT NULL, - bucket_start_unix_secs INTEGER NOT NULL, - samples INTEGER NOT NULL DEFAULT 0, - uptime_samples INTEGER NOT NULL DEFAULT 0, - active_connections_sum INTEGER NOT NULL DEFAULT 0, - active_connections_max INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0, - connect_errors_delta INTEGER NOT NULL DEFAULT 0, - disconnects_delta INTEGER NOT NULL DEFAULT 0, - error_events_delta INTEGER NOT NULL DEFAULT 0, - ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_in_frames_delta INTEGER NOT NULL DEFAULT 0, - ws_out_frames_delta INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs) -); - -CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start - ON proxy_node_metrics_1m (bucket_start_unix_secs); - -CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start - ON proxy_node_metrics_1h (bucket_start_unix_secs); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260509000000_add_background_tasks.sql b/crates/aether-data/adapters/sqlite/migrations/20260509000000_add_background_tasks.sql deleted file mode 100644 index 60881617b..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260509000000_add_background_tasks.sql +++ /dev/null @@ -1,43 +0,0 @@ -CREATE TABLE IF NOT EXISTS background_task_runs ( - id TEXT PRIMARY KEY, - task_key TEXT NOT NULL, - kind TEXT NOT NULL, - "trigger" TEXT NOT NULL, - status TEXT NOT NULL, - attempt INTEGER NOT NULL DEFAULT 0, - max_attempts INTEGER NOT NULL DEFAULT 0, - owner_instance TEXT, - progress_percent INTEGER NOT NULL DEFAULT 0, - progress_message TEXT, - payload_json TEXT, - result_json TEXT, - error_message TEXT, - cancel_requested INTEGER NOT NULL DEFAULT 0, - created_by TEXT, - created_at_unix_secs INTEGER NOT NULL, - started_at_unix_secs INTEGER, - finished_at_unix_secs INTEGER, - updated_at_unix_secs INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key - ON background_task_runs (task_key); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_status - ON background_task_runs (status); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind - ON background_task_runs (kind); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at - ON background_task_runs (created_at_unix_secs DESC); - -CREATE TABLE IF NOT EXISTS background_task_events ( - id TEXT PRIMARY KEY, - run_id TEXT NOT NULL, - event_type TEXT NOT NULL, - message TEXT NOT NULL, - payload_json TEXT, - created_at_unix_secs INTEGER NOT NULL, - FOREIGN KEY (run_id) REFERENCES background_task_runs(id) ON DELETE CASCADE -); - -CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id - ON background_task_events (run_id, created_at_unix_secs ASC); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260509120000_add_user_groups.sql b/crates/aether-data/adapters/sqlite/migrations/20260509120000_add_user_groups.sql deleted file mode 100644 index 69b4b47d4..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260509120000_add_user_groups.sql +++ /dev/null @@ -1,101 +0,0 @@ -ALTER TABLE users ADD COLUMN allowed_providers_mode TEXT NOT NULL DEFAULT 'unrestricted'; -ALTER TABLE users ADD COLUMN allowed_api_formats_mode TEXT NOT NULL DEFAULT 'unrestricted'; -ALTER TABLE users ADD COLUMN allowed_models_mode TEXT NOT NULL DEFAULT 'unrestricted'; -ALTER TABLE users ADD COLUMN rate_limit_mode TEXT NOT NULL DEFAULT 'system'; - -UPDATE users -SET allowed_providers_mode = CASE WHEN allowed_providers IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_providers_mode = 'unrestricted'; - -UPDATE users -SET allowed_api_formats_mode = CASE WHEN allowed_api_formats IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_api_formats_mode = 'unrestricted'; - -UPDATE users -SET allowed_models_mode = CASE WHEN allowed_models IS NULL THEN 'unrestricted' ELSE 'specific' END -WHERE allowed_models_mode = 'unrestricted'; - -UPDATE users -SET rate_limit_mode = CASE WHEN rate_limit IS NULL THEN 'system' ELSE 'custom' END -WHERE rate_limit_mode = 'system'; - -CREATE TABLE IF NOT EXISTS user_groups ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - normalized_name TEXT NOT NULL UNIQUE, - description TEXT, - priority INTEGER NOT NULL DEFAULT 0, - allowed_providers TEXT, - allowed_providers_mode TEXT NOT NULL DEFAULT 'inherit', - allowed_api_formats TEXT, - allowed_api_formats_mode TEXT NOT NULL DEFAULT 'inherit', - allowed_models TEXT, - allowed_models_mode TEXT NOT NULL DEFAULT 'inherit', - rate_limit INTEGER, - rate_limit_mode TEXT NOT NULL DEFAULT 'inherit', - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS user_group_members ( - group_id TEXT NOT NULL REFERENCES user_groups(id) ON DELETE CASCADE, - user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, - created_at INTEGER NOT NULL, - PRIMARY KEY (group_id, user_id) -); - -CREATE INDEX IF NOT EXISTS user_group_members_user_id_idx - ON user_group_members (user_id); - -CREATE INDEX IF NOT EXISTS user_groups_priority_name_idx - ON user_groups (priority DESC, name ASC, id ASC); - -INSERT OR IGNORE 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 ( - '00000000-0000-0000-0000-000000000001', - 'Default', - 'default', - 'Default group for all users', - 0, - 'unrestricted', - 'unrestricted', - 'unrestricted', - 'system', - CAST(strftime('%s', 'now') AS INTEGER), - CAST(strftime('%s', 'now') AS INTEGER) -); - -INSERT OR IGNORE INTO system_configs ( - id, - key, - value, - description, - created_at, - updated_at -) -VALUES ( - '00000000-0000-0000-0000-000000000002', - 'default_user_group_id', - '"00000000-0000-0000-0000-000000000001"', - 'Default user group', - CAST(strftime('%s', 'now') AS INTEGER), - CAST(strftime('%s', 'now') AS INTEGER) -); - -INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at) -SELECT '00000000-0000-0000-0000-000000000001', id, CAST(strftime('%s', 'now') AS INTEGER) -FROM users -WHERE is_deleted = 0 - AND LOWER(role) <> 'admin'; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260510120000_normalize_empty_user_policy_modes.sql b/crates/aether-data/adapters/sqlite/migrations/20260510120000_normalize_empty_user_policy_modes.sql deleted file mode 100644 index ed2c25f17..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260510120000_normalize_empty_user_policy_modes.sql +++ /dev/null @@ -1,29 +0,0 @@ -UPDATE users -SET allowed_providers_mode = 'unrestricted' -WHERE allowed_providers_mode = 'specific' - AND ( - allowed_providers IS NULL - OR trim(allowed_providers) = '' - OR lower(trim(allowed_providers)) = 'null' - OR trim(allowed_providers) = '[]' - ); - -UPDATE users -SET allowed_api_formats_mode = 'unrestricted' -WHERE allowed_api_formats_mode = 'specific' - AND ( - allowed_api_formats IS NULL - OR trim(allowed_api_formats) = '' - OR lower(trim(allowed_api_formats)) = 'null' - OR trim(allowed_api_formats) = '[]' - ); - -UPDATE users -SET allowed_models_mode = 'unrestricted' -WHERE allowed_models_mode = 'specific' - AND ( - allowed_models IS NULL - OR trim(allowed_models) = '' - OR lower(trim(allowed_models)) = 'null' - OR trim(allowed_models) = '[]' - ); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260511120000_exclude_admins_from_default_user_group.sql b/crates/aether-data/adapters/sqlite/migrations/20260511120000_exclude_admins_from_default_user_group.sql deleted file mode 100644 index 404593495..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260511120000_exclude_admins_from_default_user_group.sql +++ /dev/null @@ -1,14 +0,0 @@ -DELETE FROM user_group_members -WHERE user_id IN ( - SELECT id - FROM users - WHERE LOWER(role) = 'admin' -) - AND ( - group_id = '00000000-0000-0000-0000-000000000001' - OR group_id IN ( - SELECT TRIM(value, '"') - FROM system_configs - WHERE key = 'default_user_group_id' - ) - ); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260511130000_add_pool_member_scores.sql b/crates/aether-data/adapters/sqlite/migrations/20260511130000_add_pool_member_scores.sql deleted file mode 100644 index 32bbc6f2c..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260511130000_add_pool_member_scores.sql +++ /dev/null @@ -1,37 +0,0 @@ -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id TEXT PRIMARY KEY, - pool_kind TEXT NOT NULL, - pool_id TEXT NOT NULL, - member_kind TEXT NOT NULL, - member_id TEXT NOT NULL, - capability TEXT NOT NULL, - scope_kind TEXT NOT NULL, - scope_id TEXT, - score REAL NOT NULL DEFAULT 0, - hard_state TEXT NOT NULL DEFAULT 'unknown', - score_version INTEGER NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at INTEGER, - last_scheduled_at INTEGER, - last_success_at INTEGER, - last_failure_at INTEGER, - failure_count INTEGER NOT NULL DEFAULT 0, - last_probe_attempt_at INTEGER, - last_probe_success_at INTEGER, - last_probe_failure_at INTEGER, - probe_failure_count INTEGER NOT NULL DEFAULT 0, - probe_status TEXT NOT NULL DEFAULT 'never', - updated_at INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx - ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC); - -CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx - ON pool_member_scores (pool_kind, pool_id, member_kind, member_id); - -CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx - ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at); - -CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx - ON pool_member_scores (updated_at); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260512000000_add_feature_settings.sql b/crates/aether-data/adapters/sqlite/migrations/20260512000000_add_feature_settings.sql deleted file mode 100644 index 2c1080a51..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260512000000_add_feature_settings.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE users ADD COLUMN feature_settings TEXT; -ALTER TABLE api_keys ADD COLUMN feature_settings TEXT; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260512090000_add_payment_gateway_plans.sql b/crates/aether-data/adapters/sqlite/migrations/20260512090000_add_payment_gateway_plans.sql deleted file mode 100644 index f8aca78a3..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260512090000_add_payment_gateway_plans.sql +++ /dev/null @@ -1,85 +0,0 @@ -ALTER TABLE payment_orders ADD COLUMN payment_provider TEXT; -ALTER TABLE payment_orders ADD COLUMN payment_channel TEXT; -ALTER TABLE payment_orders ADD COLUMN order_kind TEXT NOT NULL DEFAULT 'wallet_recharge'; -ALTER TABLE payment_orders ADD COLUMN product_id TEXT; -ALTER TABLE payment_orders ADD COLUMN product_snapshot TEXT; -ALTER TABLE payment_orders ADD COLUMN fulfillment_status TEXT NOT NULL DEFAULT 'pending'; -ALTER TABLE payment_orders ADD COLUMN fulfillment_error TEXT; - -CREATE INDEX IF NOT EXISTS idx_payment_orders_kind_status - ON payment_orders (order_kind, status); -CREATE INDEX IF NOT EXISTS idx_payment_orders_product - ON payment_orders (product_id); - -CREATE TABLE IF NOT EXISTS payment_gateway_configs ( - provider TEXT PRIMARY KEY, - enabled INTEGER NOT NULL DEFAULT 0, - endpoint_url TEXT NOT NULL, - callback_base_url TEXT, - merchant_id TEXT NOT NULL, - merchant_key_encrypted TEXT, - pay_currency TEXT NOT NULL DEFAULT 'CNY', - usd_exchange_rate REAL NOT NULL DEFAULT 7.2, - min_recharge_usd REAL NOT NULL DEFAULT 1, - channels_json TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS billing_plans ( - id TEXT PRIMARY KEY, - title TEXT NOT NULL, - description TEXT, - price_amount REAL NOT NULL, - price_currency TEXT NOT NULL DEFAULT 'CNY', - duration_unit TEXT NOT NULL, - duration_value INTEGER NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - sort_order INTEGER NOT NULL DEFAULT 0, - max_active_per_user INTEGER NOT NULL DEFAULT 1, - entitlements_json TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_billing_plans_enabled_sort - ON billing_plans (enabled, sort_order); - -CREATE TABLE IF NOT EXISTS user_plan_entitlements ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - plan_id TEXT NOT NULL, - payment_order_id TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - starts_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - entitlements_snapshot TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE, - FOREIGN KEY(plan_id) REFERENCES billing_plans(id) ON DELETE RESTRICT, - FOREIGN KEY(payment_order_id) REFERENCES payment_orders(id) ON DELETE RESTRICT -); - -CREATE INDEX IF NOT EXISTS idx_user_plan_entitlements_user_active - ON user_plan_entitlements (user_id, status, expires_at); -CREATE INDEX IF NOT EXISTS idx_user_plan_entitlements_order - ON user_plan_entitlements (payment_order_id); - -CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers ( - id TEXT PRIMARY KEY, - user_entitlement_id TEXT NOT NULL, - user_id TEXT NOT NULL, - request_id TEXT NOT NULL, - amount_usd REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - usage_date TEXT NOT NULL, - created_at INTEGER NOT NULL, - UNIQUE (user_entitlement_id, request_id), - FOREIGN KEY(user_entitlement_id) REFERENCES user_plan_entitlements(id) ON DELETE CASCADE, - FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE -); - -CREATE INDEX IF NOT EXISTS idx_entitlement_usage_user_date - ON entitlement_usage_ledgers (user_id, usage_date); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql b/crates/aether-data/adapters/sqlite/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql deleted file mode 100644 index 99fde62cf..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260512110000_add_billing_plan_purchase_limit_scope.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE billing_plans - ADD COLUMN purchase_limit_scope TEXT NOT NULL DEFAULT 'active_period'; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260516000000_add_routing_profiles.sql b/crates/aether-data/adapters/sqlite/migrations/20260516000000_add_routing_profiles.sql deleted file mode 100644 index 255c562fa..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260516000000_add_routing_profiles.sql +++ /dev/null @@ -1,45 +0,0 @@ -CREATE TABLE IF NOT EXISTS routing_groups ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - description TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_system_default INTEGER NOT NULL DEFAULT 0, - config_json TEXT NOT NULL, - version INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - published_at INTEGER, - UNIQUE (name) -); - -CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx - ON routing_groups (is_system_default, enabled); - -CREATE TABLE IF NOT EXISTS routing_group_bindings ( - id TEXT PRIMARY KEY NOT NULL, - group_id TEXT NOT NULL, - subject_type TEXT NOT NULL, - subject_id TEXT NOT NULL, - is_default INTEGER NOT NULL DEFAULT 0, - allow_explicit_select INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx - ON routing_group_bindings (group_id); -CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx - ON routing_group_bindings (subject_type, subject_id); - -CREATE TABLE IF NOT EXISTS routing_group_versions ( - id TEXT PRIMARY KEY NOT NULL, - group_id TEXT NOT NULL, - version INTEGER NOT NULL, - config_json TEXT NOT NULL, - created_at INTEGER NOT NULL, - created_by TEXT, - UNIQUE (group_id, version) -); - -CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx - ON routing_group_versions (group_id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260518000000_add_usage_counter_deltas.sql b/crates/aether-data/adapters/sqlite/migrations/20260518000000_add_usage_counter_deltas.sql deleted file mode 100644 index 61f06fd03..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260518000000_add_usage_counter_deltas.sql +++ /dev/null @@ -1,34 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_counter_deltas ( - id TEXT PRIMARY KEY NOT NULL, - request_id TEXT NOT NULL, - kind TEXT NOT NULL, - target_id TEXT NOT NULL, - request_count_delta INTEGER NOT NULL DEFAULT 0, - total_requests_delta INTEGER NOT NULL DEFAULT 0, - success_count_delta INTEGER NOT NULL DEFAULT 0, - error_count_delta INTEGER NOT NULL DEFAULT 0, - dns_failures_delta INTEGER NOT NULL DEFAULT 0, - stream_errors_delta INTEGER NOT NULL DEFAULT 0, - total_tokens_delta INTEGER NOT NULL DEFAULT 0, - total_cost_usd_delta REAL NOT NULL DEFAULT 0, - total_response_time_ms_delta INTEGER NOT NULL DEFAULT 0, - last_used_at_unix_secs INTEGER, - last_used_ip TEXT, - candidate_last_used_at_unix_secs INTEGER, - removed_last_used_at_unix_secs INTEGER, - usage_created_at_unix_secs INTEGER, - created_at INTEGER NOT NULL, - processed_at INTEGER -); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed - ON usage_counter_deltas (created_at, id); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed - ON usage_counter_deltas (processed_at, created_at, id); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind - ON usage_counter_deltas (request_id, kind, target_id); - -CREATE INDEX IF NOT EXISTS video_tasks_due_poll_idx - ON video_tasks (status, next_poll_at, updated_at); - -CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date - ON entitlement_usage_ledgers (user_entitlement_id, usage_date); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260519000000_add_referrals_privacy_required_announcements.sql b/crates/aether-data/adapters/sqlite/migrations/20260519000000_add_referrals_privacy_required_announcements.sql deleted file mode 100644 index 779f4eb65..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260519000000_add_referrals_privacy_required_announcements.sql +++ /dev/null @@ -1,69 +0,0 @@ -ALTER TABLE users ADD COLUMN privacy_policy_accepted_version TEXT; -ALTER TABLE users ADD COLUMN privacy_policy_accepted_at INTEGER; - -ALTER TABLE announcements ADD COLUMN requires_ack INTEGER NOT NULL DEFAULT 0; - -CREATE TABLE IF NOT EXISTS user_invite_codes ( - user_id TEXT PRIMARY KEY, - invite_code TEXT NOT NULL UNIQUE, - active INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE -); - -CREATE TABLE IF NOT EXISTS user_referrals ( - id TEXT PRIMARY KEY, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL UNIQUE, - invite_code_snapshot TEXT NOT NULL, - source_json TEXT, - first_paid_order_id TEXT, - first_paid_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(inviter_user_id) REFERENCES users(id) ON DELETE CASCADE, - FOREIGN KEY(invitee_user_id) REFERENCES users(id) ON DELETE CASCADE, - FOREIGN KEY(first_paid_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); - -CREATE INDEX IF NOT EXISTS idx_user_referrals_inviter - ON user_referrals (inviter_user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_user_referrals_created - ON user_referrals (created_at); -CREATE INDEX IF NOT EXISTS idx_user_referrals_invite_code - ON user_referrals (invite_code_snapshot); - -CREATE TABLE IF NOT EXISTS referral_rewards ( - id TEXT PRIMARY KEY, - referral_id TEXT NOT NULL, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL, - reward_type TEXT NOT NULL, - trigger_point TEXT NOT NULL, - source_order_id TEXT, - idempotency_key TEXT NOT NULL UNIQUE, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - wallet_transaction_id TEXT, - reversed_amount_usd REAL NOT NULL DEFAULT 0, - pending_reversal_amount_usd REAL NOT NULL DEFAULT 0, - failure_reason TEXT, - admin_operator_id TEXT, - admin_note TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(referral_id) REFERENCES user_referrals(id) ON DELETE CASCADE, - FOREIGN KEY(inviter_user_id) REFERENCES users(id) ON DELETE CASCADE, - FOREIGN KEY(invitee_user_id) REFERENCES users(id) ON DELETE CASCADE, - FOREIGN KEY(source_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); - -CREATE INDEX IF NOT EXISTS idx_referral_rewards_inviter_status - ON referral_rewards (inviter_user_id, status, created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_inviter_created - ON referral_rewards (inviter_user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_created - ON referral_rewards (created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_source_order - ON referral_rewards (source_order_id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260519120000_add_oauth_provider_icon_url.sql b/crates/aether-data/adapters/sqlite/migrations/20260519120000_add_oauth_provider_icon_url.sql deleted file mode 100644 index 9e5e9da05..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260519120000_add_oauth_provider_icon_url.sql +++ /dev/null @@ -1 +0,0 @@ -ALTER TABLE oauth_providers ADD COLUMN icon_url TEXT; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql b/crates/aether-data/adapters/sqlite/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql deleted file mode 100644 index be41d85e9..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260519130000_add_provider_api_key_default_pool_sort_index.sql +++ /dev/null @@ -1,2 +0,0 @@ -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort -ON provider_api_keys (provider_id, internal_priority, name, id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260520000000_add_api_key_ip_rules.sql b/crates/aether-data/adapters/sqlite/migrations/20260520000000_add_api_key_ip_rules.sql deleted file mode 100644 index e58e240c9..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260520000000_add_api_key_ip_rules.sql +++ /dev/null @@ -1 +0,0 @@ -ALTER TABLE api_keys ADD COLUMN ip_rules TEXT; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260520010000_preserve_usage_provider_identity.sql b/crates/aether-data/adapters/sqlite/migrations/20260520010000_preserve_usage_provider_identity.sql deleted file mode 100644 index d2c3d2de3..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260520010000_preserve_usage_provider_identity.sql +++ /dev/null @@ -1,18 +0,0 @@ --- Usage is a historical fact table. Backfill nullable provider_id snapshots --- from the unique provider name where the catalog row still exists. - -UPDATE "usage" -SET provider_id = ( - SELECT providers.id - FROM providers - WHERE providers.name = TRIM("usage".provider_name) - LIMIT 1 -) -WHERE provider_id IS NULL - AND TRIM(COALESCE(provider_name, '')) <> '' - AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending') - AND EXISTS ( - SELECT 1 - FROM providers - WHERE providers.name = TRIM("usage".provider_name) - ); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260524000000_add_scheduler_hotspot_indexes.sql b/crates/aether-data/adapters/sqlite/migrations/20260524000000_add_scheduler_hotspot_indexes.sql deleted file mode 100644 index f88877abf..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260524000000_add_scheduler_hotspot_indexes.sql +++ /dev/null @@ -1,16 +0,0 @@ -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active_priority_id - ON provider_api_keys (provider_id, is_active, internal_priority, id); - -CREATE INDEX IF NOT EXISTS pool_member_scores_scheduler_account_rank_idx - ON pool_member_scores ( - pool_kind, - pool_id, - capability, - scope_kind, - scope_id, - hard_state, - score DESC, - last_ranked_at DESC, - member_id, - id - ); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260527000000_add_provider_key_name_hotspot_index.sql b/crates/aether-data/adapters/sqlite/migrations/20260527000000_add_provider_key_name_hotspot_index.sql deleted file mode 100644 index 49f4d606d..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260527000000_add_provider_key_name_hotspot_index.sql +++ /dev/null @@ -1,2 +0,0 @@ -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id - ON provider_api_keys (provider_id, name, id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260528000000_normalize_endpoint_api_roots.sql b/crates/aether-data/adapters/sqlite/migrations/20260528000000_normalize_endpoint_api_roots.sql deleted file mode 100644 index aaf55e3e6..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260528000000_normalize_endpoint_api_roots.sql +++ /dev/null @@ -1,177 +0,0 @@ -WITH endpoint_url_parts AS ( - SELECT - e.id, - CASE - WHEN instr(e.base_url, '?') > 0 THEN rtrim(substr(e.base_url, 1, instr(e.base_url, '?') - 1), '/') - ELSE rtrim(e.base_url, '/') - END AS base_without_query, - CASE - WHEN instr(e.base_url, '?') > 0 THEN substr(e.base_url, instr(e.base_url, '?')) - ELSE '' - END AS query_suffix, - lower(trim(e.api_format)) AS normalized_api_format, - lower(trim(coalesce(e.custom_path, ''))) AS normalized_custom_path, - lower(CASE - WHEN instr(e.base_url, '?') > 0 THEN rtrim(substr(e.base_url, 1, instr(e.base_url, '?') - 1), '/') - ELSE rtrim(e.base_url, '/') - END) AS normalized_base, - lower(trim(coalesce(p.provider_type, ''))) AS provider_type - FROM provider_endpoints e - LEFT JOIN providers p ON p.id = e.provider_id - WHERE lower(trim(e.api_format)) IN ( - 'openai:chat', - 'openai:responses', - 'openai:responses:compact', - 'openai:embedding', - 'openai:rerank', - 'openai:image', - 'openai:video', - 'jina:embedding', - 'jina:rerank', - 'claude:messages', - 'gemini:generate_content', - 'gemini:embedding', - 'gemini:video' - ) -), -endpoint_api_root_updates AS ( - SELECT - id, - base_without_query - || CASE - WHEN normalized_api_format IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - THEN '/v1beta' - ELSE '/v1' - END - || query_suffix AS next_base_url - FROM endpoint_url_parts - WHERE provider_type NOT IN ( - 'codex', - 'chatgpt_web', - 'claude_code', - 'kiro', - 'gemini_cli', - 'vertex_ai', - 'antigravity', - 'grok', - 'windsurf' - ) - AND normalized_base NOT GLOB '*/v[0-9]' - AND normalized_base NOT GLOB '*/v[0-9][0-9]' - AND normalized_base NOT GLOB '*/v[0-9]/*' - AND normalized_base NOT GLOB '*/v[0-9][0-9]/*' - AND normalized_base NOT GLOB '*/v[0-9]beta*' - AND normalized_base NOT GLOB '*/v[0-9][0-9]beta*' - AND ( - ( - normalized_api_format IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND normalized_custom_path LIKE '/v1beta/%' - ) - OR ( - normalized_api_format NOT IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND normalized_custom_path LIKE '/v1/%' - ) - OR normalized_custom_path = '' - ) -) -UPDATE provider_endpoints -SET base_url = ( - SELECT next_base_url - FROM endpoint_api_root_updates - WHERE endpoint_api_root_updates.id = provider_endpoints.id -) -WHERE id IN (SELECT id FROM endpoint_api_root_updates); - -UPDATE provider_endpoints -SET custom_path = CASE - WHEN lower(trim(api_format)) = 'openai:chat' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/chat/completions' - THEN NULL - WHEN lower(trim(api_format)) = 'openai:responses' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/responses' - THEN NULL - WHEN lower(trim(api_format)) = 'openai:responses:compact' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/responses/compact' - THEN NULL - WHEN lower(trim(api_format)) = 'claude:messages' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/messages' - THEN NULL - WHEN lower(trim(api_format)) IN ('openai:embedding', 'jina:embedding') - AND lower(trim(coalesce(custom_path, ''))) = '/v1/embeddings' - THEN NULL - WHEN lower(trim(api_format)) IN ('openai:rerank', 'jina:rerank') - AND lower(trim(coalesce(custom_path, ''))) = '/v1/rerank' - THEN NULL - WHEN lower(trim(api_format)) = 'openai:image' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/images/generations' - THEN NULL - WHEN lower(trim(api_format)) = 'openai:video' - AND lower(trim(coalesce(custom_path, ''))) = '/v1/videos' - THEN NULL - WHEN lower(trim(api_format)) = 'gemini:generate_content' - AND lower(trim(coalesce(custom_path, ''))) = '/v1beta/models/{model}:{action}' - THEN NULL - WHEN lower(trim(api_format)) = 'gemini:embedding' - AND lower(trim(coalesce(custom_path, ''))) IN ('/v1beta/models/{model}:embedcontent', '/v1beta/models/{model}:{action}') - THEN NULL - WHEN lower(trim(api_format)) = 'gemini:video' - AND lower(trim(coalesce(custom_path, ''))) = '/v1beta/models/{model}:predictlongrunning' - THEN NULL - WHEN lower(trim(api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - THEN '/' || substr(trim(custom_path), 9) - ELSE '/' || substr(trim(custom_path), 5) - END -WHERE lower(trim(api_format)) IN ( - 'openai:chat', - 'openai:responses', - 'openai:responses:compact', - 'openai:embedding', - 'openai:rerank', - 'openai:image', - 'openai:video', - 'jina:embedding', - 'jina:rerank', - 'claude:messages', - 'gemini:generate_content', - 'gemini:embedding', - 'gemini:video' - ) - AND NOT EXISTS ( - SELECT 1 - FROM providers p - WHERE p.id = provider_endpoints.provider_id - AND lower(trim(coalesce(p.provider_type, ''))) IN ( - 'codex', - 'chatgpt_web', - 'claude_code', - 'kiro', - 'gemini_cli', - 'vertex_ai', - 'antigravity', - 'grok', - 'windsurf' - ) - ) - AND ( - ( - lower(trim(api_format)) IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND lower(trim(coalesce(custom_path, ''))) LIKE '/v1beta/%' - ) - OR ( - lower(trim(api_format)) NOT IN ('gemini:generate_content', 'gemini:embedding', 'gemini:video') - AND lower(trim(coalesce(custom_path, ''))) LIKE '/v1/%' - ) - ) - AND ( - lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9]' - OR lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9][0-9]' - OR lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9]/*' - OR lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9][0-9]/*' - OR lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9]beta*' - OR lower(rtrim(CASE WHEN instr(base_url, '?') > 0 THEN substr(base_url, 1, instr(base_url, '?') - 1) ELSE base_url END, '/')) GLOB '*/v[0-9][0-9]beta*' - ); - -UPDATE provider_endpoints -SET custom_path = NULL -WHERE custom_path IS NOT NULL - AND trim(custom_path) = ''; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260528020000_add_gateway_pressure_indexes.sql b/crates/aether-data/adapters/sqlite/migrations/20260528020000_add_gateway_pressure_indexes.sql deleted file mode 100644 index e574c6fb7..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260528020000_add_gateway_pressure_indexes.sql +++ /dev/null @@ -1,26 +0,0 @@ --- High-concurrency gateway read/cleanup paths. --- SQLite remains single-node/lightweight but benefits from the same bounded scans. - -CREATE INDEX IF NOT EXISTS idx_usage_created_id_desc - ON "usage" (created_at_unix_ms DESC, request_id ASC); - -CREATE INDEX IF NOT EXISTS idx_usage_user_created_id_desc - ON "usage" (user_id, created_at_unix_ms DESC, request_id ASC); - -CREATE INDEX IF NOT EXISTS idx_usage_api_format_created_id_desc - ON "usage" (api_format, created_at_unix_ms DESC, request_id ASC); - -CREATE INDEX IF NOT EXISTS idx_usage_status_created_id_desc - ON "usage" (status, created_at_unix_ms DESC, request_id ASC); - -CREATE INDEX IF NOT EXISTS idx_request_candidates_provider_created - ON request_candidates (provider_id, created_at DESC, id ASC); - -CREATE INDEX IF NOT EXISTS idx_request_candidates_api_key_created - ON request_candidates (api_key_id, created_at ASC, id ASC); - -CREATE INDEX IF NOT EXISTS idx_background_task_runs_status_created - ON background_task_runs (status, created_at_unix_secs DESC, updated_at_unix_secs DESC); - -CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind_created - ON background_task_runs (kind, created_at_unix_secs DESC, updated_at_unix_secs DESC); 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 deleted file mode 100644 index 7e321201d..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260725000000_normalize_imported_timestamp_storage.sql +++ /dev/null @@ -1,784 +0,0 @@ --- 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 deleted file mode 100644 index ecfcf2621..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260725010000_add_cross_driver_schema_parity.sql +++ /dev/null @@ -1,270 +0,0 @@ --- 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 deleted file mode 100644 index 0ed6b4983..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260725020000_add_advanced_stats_parity.sql +++ /dev/null @@ -1,447 +0,0 @@ -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 deleted file mode 100644 index c0201bc17..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260725030000_add_usage_routing_snapshot_parity.sql +++ /dev/null @@ -1,70 +0,0 @@ -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 deleted file mode 100644 index a7ad4e434..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260725040000_normalize_remaining_imported_timestamp_storage.sql +++ /dev/null @@ -1,821 +0,0 @@ --- 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/migrations/20260727000000_repair_routing_default_uniqueness.sql b/crates/aether-data/adapters/sqlite/migrations/20260727000000_repair_routing_default_uniqueness.sql deleted file mode 100644 index 805705e63..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260727000000_repair_routing_default_uniqueness.sql +++ /dev/null @@ -1,38 +0,0 @@ -UPDATE routing_groups -SET is_system_default = 0 -WHERE id IN ( - SELECT id - FROM ( - SELECT - id, - ROW_NUMBER() OVER (ORDER BY enabled DESC, updated_at DESC, id ASC) AS default_rank - FROM routing_groups - WHERE is_system_default = 1 - ) AS ranked_defaults - WHERE default_rank > 1 -); - -UPDATE routing_group_bindings -SET is_default = 0 -WHERE id IN ( - SELECT id - FROM ( - SELECT - id, - ROW_NUMBER() OVER ( - PARTITION BY subject_type, subject_id - ORDER BY created_at ASC, id ASC - ) AS default_rank - FROM routing_group_bindings - WHERE is_default = 1 - ) AS ranked_defaults - WHERE default_rank > 1 -); - -CREATE UNIQUE INDEX IF NOT EXISTS routing_groups_one_system_default_key - ON routing_groups (is_system_default) - WHERE is_system_default = 1; - -CREATE UNIQUE INDEX IF NOT EXISTS routing_group_bindings_subject_default_key - ON routing_group_bindings (subject_type, subject_id) - WHERE is_default = 1; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql b/crates/aether-data/adapters/sqlite/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql deleted file mode 100644 index f7d07d142..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260731000000_cleanup_duplicate_worker_boot_runs.sql +++ /dev/null @@ -1,22 +0,0 @@ --- Worker supervisors are registered as one logical row per task. Older binaries --- included the ephemeral gateway instance in the row id, leaving a permanently --- running row after every restart. Remove only those system-generated boot rows; --- current workers recreate the stable logical rows after migrations complete. --- The metadata predicate also replaces task-only rows written by early builds of --- this fix that still claimed an instance owner. Delete children explicitly so --- cleanup remains complete after imports performed with FK checks disabled. -DELETE FROM background_task_events -WHERE run_id IN ( - SELECT id - FROM background_task_runs - WHERE id LIKE 'boot:%' - AND owner_instance IS NOT NULL - AND created_by = 'system' - AND progress_message = 'worker booted' -); - -DELETE FROM background_task_runs -WHERE id LIKE 'boot:%' - AND owner_instance IS NOT NULL - AND created_by = 'system' - AND progress_message = 'worker booted'; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260814000000_add_usage_cost_reservations.sql b/crates/aether-data/adapters/sqlite/migrations/20260814000000_add_usage_cost_reservations.sql deleted file mode 100644 index 9253bdeac..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260814000000_add_usage_cost_reservations.sql +++ /dev/null @@ -1,34 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_cost_reservations ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - reservation_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - reserved_cost_units INTEGER NOT NULL CHECK (reserved_cost_units >= 0), - actual_cost_units INTEGER CHECK (actual_cost_units IS NULL OR actual_cost_units >= 0), - state TEXT NOT NULL CHECK (state IN ('reserved', 'finalized', 'released')), - reservation_expires_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - finalized_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CHECK (reservation_expires_at > admitted_at), - CHECK (retain_until >= reservation_expires_at), - CHECK ( - (state = 'reserved' AND actual_cost_units IS NULL AND finalized_at IS NULL) - OR (state = 'finalized' AND actual_cost_units IS NOT NULL AND finalized_at IS NOT NULL) - OR (state = 'released' AND actual_cost_units IS NOT NULL - AND actual_cost_units = 0 AND finalized_at IS NOT NULL) - ) -); - -CREATE INDEX IF NOT EXISTS usage_cost_reservations_request_id_idx - ON usage_cost_reservations (request_id); - -CREATE INDEX IF NOT EXISTS usage_cost_reservations_subject_admitted_at_idx - ON usage_cost_reservations (subject_id, admitted_at); - -CREATE INDEX IF NOT EXISTS usage_cost_reservations_reservation_expires_at_idx - ON usage_cost_reservations (reservation_expires_at); - -CREATE INDEX IF NOT EXISTS usage_cost_reservations_retain_until_token_idx - ON usage_cost_reservations (retain_until, reservation_token); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260815000000_add_usage_request_admissions.sql b/crates/aether-data/adapters/sqlite/migrations/20260815000000_add_usage_request_admissions.sql deleted file mode 100644 index 32f156149..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260815000000_add_usage_request_admissions.sql +++ /dev/null @@ -1,21 +0,0 @@ -CREATE TABLE IF NOT EXISTS usage_request_admissions ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - event_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - state TEXT NOT NULL CHECK (state IN ('active', 'released')), - released_at INTEGER, - created_at INTEGER NOT NULL, - CHECK (retain_until > admitted_at), - CHECK ( - (state = 'active' AND released_at IS NULL) - OR (state = 'released' AND released_at IS NOT NULL AND released_at >= admitted_at) - ) -); - -CREATE INDEX IF NOT EXISTS usage_request_admissions_subject_admitted_at_idx - ON usage_request_admissions (subject_id, admitted_at); - -CREATE INDEX IF NOT EXISTS usage_request_admissions_retain_until_token_idx - ON usage_request_admissions (retain_until, event_token); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260816000000_add_usage_policy_user_foreign_keys.sql b/crates/aether-data/adapters/sqlite/migrations/20260816000000_add_usage_policy_user_foreign_keys.sql deleted file mode 100644 index 14066f283..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260816000000_add_usage_policy_user_foreign_keys.sql +++ /dev/null @@ -1,117 +0,0 @@ --- SQLite cannot add a foreign key with ALTER TABLE. Rebuild both ledgers while --- preserving valid rows and dropping only records whose owning user was --- already deleted before the relationship became enforceable. -CREATE TABLE usage_cost_reservations_with_user_fk ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - reservation_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - reserved_cost_units INTEGER NOT NULL CHECK (reserved_cost_units >= 0), - actual_cost_units INTEGER CHECK (actual_cost_units IS NULL OR actual_cost_units >= 0), - state TEXT NOT NULL CHECK (state IN ('reserved', 'finalized', 'released')), - reservation_expires_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - finalized_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CONSTRAINT usage_cost_reservations_subject_id_fkey - FOREIGN KEY (subject_id) REFERENCES users (id) ON DELETE CASCADE, - CHECK (reservation_expires_at > admitted_at), - CHECK (retain_until >= reservation_expires_at), - CHECK ( - (state = 'reserved' AND actual_cost_units IS NULL AND finalized_at IS NULL) - OR (state = 'finalized' AND actual_cost_units IS NOT NULL AND finalized_at IS NOT NULL) - OR (state = 'released' AND actual_cost_units IS NOT NULL - AND actual_cost_units = 0 AND finalized_at IS NOT NULL) - ) -); - -INSERT INTO usage_cost_reservations_with_user_fk ( - request_id, - subject_id, - reservation_token, - admitted_at, - reserved_cost_units, - actual_cost_units, - state, - reservation_expires_at, - retain_until, - finalized_at, - created_at, - updated_at -) -SELECT - reservation.request_id, - reservation.subject_id, - reservation.reservation_token, - reservation.admitted_at, - reservation.reserved_cost_units, - reservation.actual_cost_units, - reservation.state, - reservation.reservation_expires_at, - reservation.retain_until, - reservation.finalized_at, - reservation.created_at, - reservation.updated_at -FROM usage_cost_reservations AS reservation -INNER JOIN users AS app_user ON app_user.id = reservation.subject_id; - -DROP TABLE usage_cost_reservations; -ALTER TABLE usage_cost_reservations_with_user_fk RENAME TO usage_cost_reservations; - -CREATE INDEX usage_cost_reservations_request_id_idx - ON usage_cost_reservations (request_id); -CREATE INDEX usage_cost_reservations_subject_admitted_at_idx - ON usage_cost_reservations (subject_id, admitted_at); -CREATE INDEX usage_cost_reservations_reservation_expires_at_idx - ON usage_cost_reservations (reservation_expires_at); -CREATE INDEX usage_cost_reservations_retain_until_token_idx - ON usage_cost_reservations (retain_until, reservation_token); - -CREATE TABLE usage_request_admissions_with_user_fk ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - event_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - state TEXT NOT NULL CHECK (state IN ('active', 'released')), - released_at INTEGER, - created_at INTEGER NOT NULL, - CONSTRAINT usage_request_admissions_subject_id_fkey - FOREIGN KEY (subject_id) REFERENCES users (id) ON DELETE CASCADE, - CHECK (retain_until > admitted_at), - CHECK ( - (state = 'active' AND released_at IS NULL) - OR (state = 'released' AND released_at IS NOT NULL AND released_at >= admitted_at) - ) -); - -INSERT INTO usage_request_admissions_with_user_fk ( - request_id, - subject_id, - event_token, - admitted_at, - retain_until, - state, - released_at, - created_at -) -SELECT - admission.request_id, - admission.subject_id, - admission.event_token, - admission.admitted_at, - admission.retain_until, - admission.state, - admission.released_at, - admission.created_at -FROM usage_request_admissions AS admission -INNER JOIN users AS app_user ON app_user.id = admission.subject_id; - -DROP TABLE usage_request_admissions; -ALTER TABLE usage_request_admissions_with_user_fk RENAME TO usage_request_admissions; - -CREATE INDEX usage_request_admissions_subject_admitted_at_idx - ON usage_request_admissions (subject_id, admitted_at); -CREATE INDEX usage_request_admissions_retain_until_token_idx - ON usage_request_admissions (retain_until, event_token); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql b/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql deleted file mode 100644 index 1ba4907fd..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql +++ /dev/null @@ -1,190 +0,0 @@ --- Preserve the explicit access granted by the original #741 implementation, --- which represented Codex Live as openai:responses. Runtime permission --- matching remains strict after this one-time data upgrade. - -UPDATE users -SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') -WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' - AND json_valid(allowed_api_formats) - AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' - AND EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(users.allowed_api_formats) THEN users.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'openai:responses' - ) - AND NOT EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(users.allowed_api_formats) THEN users.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'codex:live' - ); - -UPDATE user_groups -SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') -WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' - AND json_valid(allowed_api_formats) - AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' - AND EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(user_groups.allowed_api_formats) THEN user_groups.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'openai:responses' - ) - AND NOT EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(user_groups.allowed_api_formats) THEN user_groups.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'codex:live' - ); - -UPDATE api_keys -SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') -WHERE json_valid(allowed_api_formats) - AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' - AND EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(api_keys.allowed_api_formats) THEN api_keys.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'openai:responses' - ) - AND NOT EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(api_keys.allowed_api_formats) THEN api_keys.allowed_api_formats ELSE '[]' END - ) - WHERE value = 'codex:live' - ); - -UPDATE provider_api_keys -SET - api_formats = json_insert(api_formats, '$[#]', 'codex:live'), - updated_at = CAST(strftime('%s', 'now') AS INTEGER) -WHERE provider_id IN ( - SELECT id FROM providers - WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' - ) - AND json_valid(api_formats) - AND json_type(CASE WHEN json_valid(api_formats) THEN api_formats END) = 'array' - AND EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(provider_api_keys.api_formats) THEN provider_api_keys.api_formats ELSE '[]' END - ) - WHERE value = 'openai:responses' - ) - AND NOT EXISTS ( - SELECT 1 FROM json_each( - CASE WHEN json_valid(provider_api_keys.api_formats) THEN provider_api_keys.api_formats ELSE '[]' END - ) - WHERE value = 'codex:live' - ); - -UPDATE provider_api_keys -SET - auth_type_by_format = json_set( - auth_type_by_format, - '$."codex:live"', - json_extract(auth_type_by_format, '$."openai:responses"') - ), - updated_at = CAST(strftime('%s', 'now') AS INTEGER) -WHERE provider_id IN ( - SELECT id FROM providers - WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' - ) - AND json_valid(auth_type_by_format) - AND json_type(CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END) = 'object' - AND json_type( - CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END, - '$."openai:responses"' - ) IS NOT NULL - AND json_type( - CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END, - '$."codex:live"' - ) IS NULL; - -UPDATE provider_api_keys -SET - allow_auth_channel_mismatch_formats = json_insert( - allow_auth_channel_mismatch_formats, - '$[#]', - 'codex:live' - ), - updated_at = CAST(strftime('%s', 'now') AS INTEGER) -WHERE provider_id IN ( - SELECT id FROM providers - WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' - ) - AND json_valid(allow_auth_channel_mismatch_formats) - AND json_type( - CASE - WHEN json_valid(allow_auth_channel_mismatch_formats) - THEN allow_auth_channel_mismatch_formats - END - ) = 'array' - AND EXISTS ( - SELECT 1 FROM json_each( - CASE - WHEN json_valid(provider_api_keys.allow_auth_channel_mismatch_formats) - THEN provider_api_keys.allow_auth_channel_mismatch_formats - ELSE '[]' - END - ) - WHERE value = 'openai:responses' - ) - AND NOT EXISTS ( - SELECT 1 FROM json_each( - CASE - WHEN json_valid(provider_api_keys.allow_auth_channel_mismatch_formats) - THEN provider_api_keys.allow_auth_channel_mismatch_formats - ELSE '[]' - END - ) - WHERE value = 'codex:live' - ); - -UPDATE provider_api_keys -SET - rate_multipliers = json_set( - rate_multipliers, - '$."codex:live"', - json_extract(rate_multipliers, '$."openai:responses"') - ), - updated_at = CAST(strftime('%s', 'now') AS INTEGER) -WHERE provider_id IN ( - SELECT id FROM providers - WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' - ) - AND json_valid(rate_multipliers) - AND json_type(CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END) = 'object' - AND json_type( - CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END, - '$."openai:responses"' - ) IS NOT NULL - AND json_type( - CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END, - '$."codex:live"' - ) IS NULL; - -UPDATE provider_api_keys -SET - global_priority_by_format = json_set( - global_priority_by_format, - '$."codex:live"', - json_extract(global_priority_by_format, '$."openai:responses"') - ), - updated_at = CAST(strftime('%s', 'now') AS INTEGER) -WHERE provider_id IN ( - SELECT id FROM providers - WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' - ) - AND json_valid(global_priority_by_format) - AND json_type( - CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END - ) = 'object' - AND json_type( - CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END, - '$."openai:responses"' - ) IS NOT NULL - AND json_type( - CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END, - '$."codex:live"' - ) IS NULL; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql b/crates/aether-data/adapters/sqlite/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql deleted file mode 100644 index 4c6b8cdba..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260821120000_enforce_payment_gateway_order_uniqueness.sql +++ /dev/null @@ -1,19 +0,0 @@ --- A gateway transaction identifier may repeat across payment methods, but --- must never identify two orders in the same method. If historical conflicts --- exist, index creation intentionally fails without modifying financial data. --- Diagnose with: --- SELECT payment_method, gateway_order_id, COUNT(*) --- FROM payment_orders --- WHERE gateway_order_id IS NOT NULL --- GROUP BY payment_method, gateway_order_id --- HAVING COUNT(*) > 1; -UPDATE payment_orders -SET payment_method = lower(trim(payment_method)) -WHERE payment_method <> lower(trim(payment_method)); - -UPDATE payment_callbacks -SET payment_method = lower(trim(payment_method)) -WHERE payment_method <> lower(trim(payment_method)); - -CREATE UNIQUE INDEX uq_payment_orders_payment_method_gateway_order_id - ON payment_orders (payment_method, gateway_order_id); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260821130000_add_user_security_version.sql b/crates/aether-data/adapters/sqlite/migrations/20260821130000_add_user_security_version.sql deleted file mode 100644 index 5a97e0460..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260821130000_add_user_security_version.sql +++ /dev/null @@ -1,5 +0,0 @@ -ALTER TABLE users - ADD COLUMN security_version INTEGER NOT NULL DEFAULT 0; - -ALTER TABLE user_sessions - ADD COLUMN security_version INTEGER NOT NULL DEFAULT 0; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260827050000_anonymize_deleted_user_history.sql b/crates/aether-data/adapters/sqlite/migrations/20260827050000_anonymize_deleted_user_history.sql deleted file mode 100644 index 05d81afcc..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260827050000_anonymize_deleted_user_history.sql +++ /dev/null @@ -1,157 +0,0 @@ --- Retain legacy row contents. This migration only removes user foreign keys --- required by the current account-deletion flow. -ALTER TABLE user_plan_entitlements RENAME TO _aether_user_plan_entitlements_with_user_fk; -DROP INDEX IF EXISTS idx_user_plan_entitlements_user_active; -DROP INDEX IF EXISTS idx_user_plan_entitlements_order; - -CREATE TABLE user_plan_entitlements ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - plan_id TEXT NOT NULL, - payment_order_id TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - starts_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - entitlements_snapshot TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(plan_id) REFERENCES billing_plans(id) ON DELETE RESTRICT, - FOREIGN KEY(payment_order_id) REFERENCES payment_orders(id) ON DELETE RESTRICT -); - -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) -SELECT - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -FROM _aether_user_plan_entitlements_with_user_fk; - -CREATE INDEX idx_user_plan_entitlements_user_active - ON user_plan_entitlements (user_id, status, expires_at); -CREATE INDEX idx_user_plan_entitlements_order - ON user_plan_entitlements (payment_order_id); - -ALTER TABLE entitlement_usage_ledgers RENAME TO _aether_entitlement_usage_ledgers_with_user_fk; -DROP INDEX IF EXISTS idx_entitlement_usage_user_date; -DROP INDEX IF EXISTS idx_entitlement_usage_entitlement_date; - -CREATE TABLE entitlement_usage_ledgers ( - id TEXT PRIMARY KEY, - user_entitlement_id TEXT NOT NULL, - user_id TEXT NOT NULL, - request_id TEXT NOT NULL, - amount_usd REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - usage_date TEXT NOT NULL, - created_at INTEGER NOT NULL, - UNIQUE (user_entitlement_id, request_id), - FOREIGN KEY(user_entitlement_id) REFERENCES user_plan_entitlements(id) ON DELETE CASCADE -); - -INSERT INTO entitlement_usage_ledgers ( - id, user_entitlement_id, user_id, request_id, amount_usd, - balance_before, balance_after, usage_date, created_at -) -SELECT - id, user_entitlement_id, user_id, request_id, amount_usd, - balance_before, balance_after, usage_date, created_at -FROM _aether_entitlement_usage_ledgers_with_user_fk; - -CREATE INDEX idx_entitlement_usage_user_date - ON entitlement_usage_ledgers (user_id, usage_date); -CREATE INDEX idx_entitlement_usage_entitlement_date - ON entitlement_usage_ledgers (user_entitlement_id, usage_date); - -DROP TABLE _aether_entitlement_usage_ledgers_with_user_fk; -DROP TABLE _aether_user_plan_entitlements_with_user_fk; - -ALTER TABLE user_referrals RENAME TO _aether_user_referrals_with_user_fk; -DROP INDEX IF EXISTS idx_user_referrals_inviter; -DROP INDEX IF EXISTS idx_user_referrals_created; -DROP INDEX IF EXISTS idx_user_referrals_invite_code; - -CREATE TABLE user_referrals ( - id TEXT PRIMARY KEY, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL UNIQUE, - invite_code_snapshot TEXT NOT NULL, - source_json TEXT, - first_paid_order_id TEXT, - first_paid_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(first_paid_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); - -INSERT INTO user_referrals ( - id, inviter_user_id, invitee_user_id, invite_code_snapshot, source_json, - first_paid_order_id, first_paid_at, created_at, updated_at -) -SELECT - id, inviter_user_id, invitee_user_id, invite_code_snapshot, source_json, - first_paid_order_id, first_paid_at, created_at, updated_at -FROM _aether_user_referrals_with_user_fk; - -CREATE INDEX idx_user_referrals_inviter - ON user_referrals (inviter_user_id, created_at); -CREATE INDEX idx_user_referrals_created - ON user_referrals (created_at); -CREATE INDEX idx_user_referrals_invite_code - ON user_referrals (invite_code_snapshot); - -ALTER TABLE referral_rewards RENAME TO _aether_referral_rewards_with_user_fk; -DROP INDEX IF EXISTS idx_referral_rewards_inviter_status; -DROP INDEX IF EXISTS idx_referral_rewards_inviter_created; -DROP INDEX IF EXISTS idx_referral_rewards_created; -DROP INDEX IF EXISTS idx_referral_rewards_source_order; - -CREATE TABLE referral_rewards ( - id TEXT PRIMARY KEY, - referral_id TEXT NOT NULL, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL, - reward_type TEXT NOT NULL, - trigger_point TEXT NOT NULL, - source_order_id TEXT, - idempotency_key TEXT NOT NULL UNIQUE, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - wallet_transaction_id TEXT, - reversed_amount_usd REAL NOT NULL DEFAULT 0, - pending_reversal_amount_usd REAL NOT NULL DEFAULT 0, - failure_reason TEXT, - admin_operator_id TEXT, - admin_note TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - FOREIGN KEY(referral_id) REFERENCES user_referrals(id) ON DELETE CASCADE, - FOREIGN KEY(source_order_id) REFERENCES payment_orders(id) ON DELETE SET NULL -); - -INSERT INTO referral_rewards ( - id, referral_id, inviter_user_id, invitee_user_id, reward_type, - trigger_point, source_order_id, idempotency_key, amount_usd, status, - wallet_transaction_id, reversed_amount_usd, pending_reversal_amount_usd, - failure_reason, admin_operator_id, admin_note, created_at, updated_at -) -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, - trigger_point, source_order_id, idempotency_key, amount_usd, status, - wallet_transaction_id, reversed_amount_usd, pending_reversal_amount_usd, - failure_reason, admin_operator_id, admin_note, created_at, updated_at -FROM _aether_referral_rewards_with_user_fk; - -CREATE INDEX idx_referral_rewards_inviter_status - ON referral_rewards (inviter_user_id, status, created_at); -CREATE INDEX idx_referral_rewards_inviter_created - ON referral_rewards (inviter_user_id, created_at); -CREATE INDEX idx_referral_rewards_created - ON referral_rewards (created_at); -CREATE INDEX idx_referral_rewards_source_order - ON referral_rewards (source_order_id); - -DROP TABLE _aether_referral_rewards_with_user_fk; -DROP TABLE _aether_user_referrals_with_user_fk; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260831000000_enforce_ldap_config_singleton.sql b/crates/aether-data/adapters/sqlite/migrations/20260831000000_enforce_ldap_config_singleton.sql deleted file mode 100644 index a2333a451..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260831000000_enforce_ldap_config_singleton.sql +++ /dev/null @@ -1,11 +0,0 @@ --- LDAP configuration is a database-wide singleton. Preserve the row selected by the legacy --- reader (the smallest id), remove historical duplicates, and let the database arbitrate --- concurrent first creation. -DELETE FROM ldap_configs -WHERE id <> (SELECT MIN(id) FROM ldap_configs); - -ALTER TABLE ldap_configs -ADD COLUMN singleton_key INTEGER NOT NULL DEFAULT 1 CHECK (singleton_key = 1); - -CREATE UNIQUE INDEX ldap_configs_singleton_key_key -ON ldap_configs (singleton_key); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260831010000_add_proxy_node_tunnel_generation.sql b/crates/aether-data/adapters/sqlite/migrations/20260831010000_add_proxy_node_tunnel_generation.sql deleted file mode 100644 index 78f5333a9..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260831010000_add_proxy_node_tunnel_generation.sql +++ /dev/null @@ -1,20 +0,0 @@ -ALTER TABLE proxy_nodes - ADD COLUMN tunnel_generation TEXT NOT NULL DEFAULT ''; - -UPDATE proxy_nodes -SET tunnel_generation = lower(hex(randomblob(16))) -WHERE tunnel_generation = ''; - --- SQLite only permits a constant default when adding a NOT NULL column. Keep --- the upgrade compatible with legacy rows, then replace the temporary empty --- default for future legacy writers with a per-row random generation. This --- also covers importers that omit the newly added column. -CREATE TRIGGER IF NOT EXISTS proxy_nodes_fill_tunnel_generation -AFTER INSERT ON proxy_nodes -WHEN NEW.tunnel_generation IS NULL OR trim(NEW.tunnel_generation) = '' -BEGIN - UPDATE proxy_nodes - SET tunnel_generation = lower(hex(randomblob(16))) - WHERE id = NEW.id - AND (tunnel_generation IS NULL OR trim(tunnel_generation) = ''); -END; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql b/crates/aether-data/adapters/sqlite/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql deleted file mode 100644 index c1475a9f9..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260831020000_enforce_proxy_node_endpoint_uniqueness.sql +++ /dev/null @@ -1,9 +0,0 @@ --- A proxy endpoint has one stable node identity across manual and tunnel --- registrations. Index creation intentionally fails if historical duplicates --- exist; operators must resolve the conflicting identities explicitly. --- Diagnose with: --- SELECT ip, port, COUNT(*) --- FROM proxy_nodes --- GROUP BY ip, port --- HAVING COUNT(*) > 1; -CREATE UNIQUE INDEX uq_proxy_node_ip_port ON proxy_nodes (ip, port); diff --git a/crates/aether-data/adapters/sqlite/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql b/crates/aether-data/adapters/sqlite/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql deleted file mode 100644 index 1db6a846d..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260831030000_add_usage_counter_delta_tunnel_generation.sql +++ /dev/null @@ -1,2 +0,0 @@ -ALTER TABLE usage_counter_deltas - ADD COLUMN target_tunnel_generation TEXT; diff --git a/crates/aether-data/adapters/sqlite/migrations/20260903000000_add_routing_group_sort_order.sql b/crates/aether-data/adapters/sqlite/migrations/20260903000000_add_routing_group_sort_order.sql deleted file mode 100644 index 7cfdc1c1e..000000000 --- a/crates/aether-data/adapters/sqlite/migrations/20260903000000_add_routing_group_sort_order.sql +++ /dev/null @@ -1,4 +0,0 @@ -ALTER TABLE routing_groups ADD COLUMN sort_order INTEGER NOT NULL DEFAULT 0; - -CREATE INDEX IF NOT EXISTS routing_groups_enabled_sort_idx - ON routing_groups (enabled, sort_order, name, id); diff --git a/crates/aether-data/adapters/sqlite/src/announcements.rs b/crates/aether-data/adapters/sqlite/src/announcements.rs deleted file mode 100644 index 40f795a81..000000000 --- a/crates/aether-data/adapters/sqlite/src/announcements.rs +++ /dev/null @@ -1,480 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::announcements::*; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, push_limit, push_limit_offset, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const ANNOUNCEMENT_SELECT: &str = r#" -SELECT - a.id, - a.title, - a.content, - a.type, - a.priority, - a.is_active, - a.is_pinned, - a.requires_ack, - a.author_id, - u.username AS author_username, - a.start_time AS start_time_unix_secs, - a.end_time AS end_time_unix_secs, - a.created_at AS created_at_unix_ms, - a.updated_at AS updated_at_unix_secs -FROM announcements a -LEFT JOIN users u ON u.id = a.author_id -"#; - -#[derive(Debug, Clone)] -pub struct SqliteAnnouncementRepository { - pool: SqlitePool, -} - -impl SqliteAnnouncementRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn reload_by_id( - &self, - announcement_id: &str, - ) -> Result, DataLayerError> { - self.find_by_id(announcement_id).await - } - - fn apply_active_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - where_clause: &mut WhereClause, - active_only: bool, - now_unix_secs: u64, - ) -> Result<(), DataLayerError> { - if !active_only { - return Ok(()); - } - - let now = i64_from_u64(now_unix_secs, "announcements.now")?; - where_clause.push_next(builder); - builder - .push("a.is_active = 1 AND (a.start_time IS NULL OR a.start_time <= ") - .push_bind(now) - .push(") AND (a.end_time IS NULL OR a.end_time >= ") - .push_bind(now) - .push(")"); - Ok(()) - } -} - -#[async_trait] -impl AnnouncementReadRepository for SqliteAnnouncementRepository { - async fn find_by_id( - &self, - announcement_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(ANNOUNCEMENT_SELECT); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "a.id", - announcement_id.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_announcement_row).transpose() - } - - async fn list_announcements( - &self, - query: &AnnouncementListQuery, - ) -> Result { - let now_unix_secs = query.now_unix_secs.unwrap_or_else(current_unix_secs); - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(a.id) AS total FROM announcements a"); - let mut count_where = WhereClause::new(); - Self::apply_active_filter( - &mut count_builder, - &mut count_where, - query.active_only, - now_unix_secs, - )?; - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()? - .max(0) as u64; - - let mut list_builder = QueryBuilder::::new(ANNOUNCEMENT_SELECT); - let mut list_where = WhereClause::new(); - Self::apply_active_filter( - &mut list_builder, - &mut list_where, - query.active_only, - now_unix_secs, - )?; - list_builder - .push(" ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC"); - push_limit_offset(&mut list_builder, query.limit as i64, query.offset as i64); - let rows = list_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_announcement_row) - .collect::, _>>()?; - - Ok(StoredAnnouncementPage { items, total }) - } - - async fn count_unread_active_announcements( - &self, - user_id: &str, - now_unix_secs: u64, - ) -> Result { - let mut builder = - QueryBuilder::::new("SELECT COUNT(a.id) AS total FROM announcements a"); - let mut where_clause = WhereClause::new(); - Self::apply_active_filter(&mut builder, &mut where_clause, true, now_unix_secs)?; - where_clause.push_next(&mut builder); - builder - .push("NOT EXISTS (SELECT 1 FROM announcement_reads r WHERE r.user_id = ") - .push_bind(user_id.to_string()) - .push(" AND r.announcement_id = a.id)"); - let total = builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()? - .max(0) as u64; - Ok(total) - } - - async fn list_required_unread_active_announcements( - &self, - user_id: &str, - now_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{ANNOUNCEMENT_SELECT} -WHERE a.is_active = 1 - AND a.requires_ack = 1 - AND (a.start_time IS NULL OR a.start_time <= ?) - AND (a.end_time IS NULL OR a.end_time >= ?) - AND NOT EXISTS ( - SELECT 1 - FROM announcement_reads r - WHERE r.user_id = ? - AND r.announcement_id = a.id - ) -ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC -LIMIT ? -"# - )) - .bind(now_unix_secs as i64) - .bind(now_unix_secs as i64) - .bind(user_id) - .bind(limit as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_announcement_row).collect() - } -} - -#[async_trait] -impl AnnouncementWriteRepository for SqliteAnnouncementRepository { - async fn create_announcement( - &self, - record: CreateAnnouncementRecord, - ) -> Result { - record.validate()?; - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs() as i64; - sqlx::query( - r#" -INSERT INTO announcements ( - id, title, content, type, priority, author_id, is_active, is_pinned, - requires_ack, start_time, end_time, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&id) - .bind(record.title) - .bind(record.content) - .bind(record.kind) - .bind(record.priority) - .bind(record.author_id) - .bind(record.is_pinned) - .bind(record.requires_ack) - .bind(optional_i64_from_u64( - record.start_time_unix_secs, - "announcements.start_time", - )?) - .bind(optional_i64_from_u64( - record.end_time_unix_secs, - "announcements.end_time", - )?) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_id(&id) - .await? - .ok_or_else(|| DataLayerError::UnexpectedValue("created announcement missing".into())) - } - - async fn update_announcement( - &self, - record: UpdateAnnouncementRecord, - ) -> Result, DataLayerError> { - record.validate()?; - let id = record.announcement_id; - sqlx::query( - r#" -UPDATE announcements -SET title = COALESCE(?, title), - content = COALESCE(?, content), - type = COALESCE(?, type), - priority = COALESCE(?, priority), - is_active = COALESCE(?, is_active), - is_pinned = COALESCE(?, is_pinned), - requires_ack = COALESCE(?, requires_ack), - start_time = COALESCE(?, start_time), - end_time = COALESCE(?, end_time), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(record.title) - .bind(record.content) - .bind(record.kind) - .bind(record.priority) - .bind(record.is_active) - .bind(record.is_pinned) - .bind(record.requires_ack) - .bind(optional_i64_from_u64( - record.start_time_unix_secs, - "announcements.start_time", - )?) - .bind(optional_i64_from_u64( - record.end_time_unix_secs, - "announcements.end_time", - )?) - .bind(current_unix_secs() as i64) - .bind(&id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_id(&id).await - } - - async fn delete_announcement(&self, announcement_id: &str) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("DELETE FROM announcement_reads WHERE announcement_id = ?") - .bind(announcement_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let rows_affected = sqlx::query("DELETE FROM announcements WHERE id = ?") - .bind(announcement_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - tx.commit().await.map_sql_err()?; - Ok(rows_affected > 0) - } - - async fn mark_announcement_as_read( - &self, - user_id: &str, - announcement_id: &str, - read_at_unix_secs: u64, - ) -> Result { - let rows_affected = sqlx::query( - r#" -INSERT OR IGNORE INTO announcement_reads (id, user_id, announcement_id, read_at) -VALUES (?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(user_id) - .bind(announcement_id) - .bind(i64_from_u64( - read_at_unix_secs, - "announcement_reads.read_at", - )?) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| i64_from_u64(value, field_name)) - .transpose() -} - -fn map_announcement_row(row: &SqliteRow) -> Result { - StoredAnnouncement::new( - row.try_get("id").map_sql_err()?, - row.try_get("title").map_sql_err()?, - row.try_get("content").map_sql_err()?, - row.try_get("type").map_sql_err()?, - row.try_get("priority").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_pinned").map_sql_err()?, - row.try_get("requires_ack").map_sql_err()?, - row.try_get("author_id").map_sql_err()?, - row.try_get("author_username").map_sql_err()?, - row.try_get("start_time_unix_secs").map_sql_err()?, - row.try_get("end_time_unix_secs").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - ) -} - -#[cfg(test)] -mod tests { - use super::SqliteAnnouncementRepository; - use crate::run_migrations; - use aether_data_contracts::repository::announcements::{ - AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository, - CreateAnnouncementRecord, UpdateAnnouncementRecord, - }; - - #[tokio::test] - async fn sqlite_repository_reads_and_writes_announcements() { - 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_announcement_user(&pool).await; - - let repository = SqliteAnnouncementRepository::new(pool); - let created = repository - .create_announcement(CreateAnnouncementRecord { - title: "Initial".to_string(), - content: "Body".to_string(), - kind: "info".to_string(), - priority: 10, - is_pinned: true, - requires_ack: false, - author_id: "user-1".to_string(), - start_time_unix_secs: Some(100), - end_time_unix_secs: Some(300), - }) - .await - .expect("announcement should create"); - assert_eq!(created.author_username, Some("admin".to_string())); - assert!(created.is_active); - - let page = repository - .list_announcements(&AnnouncementListQuery { - active_only: true, - offset: 0, - limit: 10, - now_unix_secs: Some(200), - }) - .await - .expect("announcements should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].id, created.id); - - let unread = repository - .count_unread_active_announcements("user-1", 200) - .await - .expect("unread count should load"); - assert_eq!(unread, 1); - assert!(repository - .mark_announcement_as_read("user-1", &created.id, 210) - .await - .expect("read marker should insert")); - assert!(!repository - .mark_announcement_as_read("user-1", &created.id, 211) - .await - .expect("duplicate read marker should be ignored")); - assert_eq!( - repository - .count_unread_active_announcements("user-1", 200) - .await - .expect("unread count should reload"), - 0 - ); - - let updated = repository - .update_announcement(UpdateAnnouncementRecord { - announcement_id: created.id.clone(), - title: Some("Updated".to_string()), - content: None, - kind: None, - priority: Some(20), - is_active: Some(false), - is_pinned: Some(false), - requires_ack: Some(true), - start_time_unix_secs: None, - end_time_unix_secs: None, - }) - .await - .expect("announcement should update") - .expect("announcement should exist"); - assert_eq!(updated.title, "Updated"); - assert!(!updated.is_active); - - assert!(repository - .delete_announcement(&created.id) - .await - .expect("announcement should delete")); - assert!(repository - .find_by_id(&created.id) - .await - .expect("find should run") - .is_none()); - } - - async fn seed_announcement_user(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO users ( - id, email, username, role, auth_source, email_verified, is_active, is_deleted, created_at, updated_at -) -VALUES ('user-1', 'admin@example.com', 'admin', 'admin', 'local', 1, 1, 0, 1, 1) -"#, - ) - .execute(pool) - .await - .expect("user should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/audit.rs b/crates/aether-data/adapters/sqlite/src/audit.rs deleted file mode 100644 index 3fb126c4c..000000000 --- a/crates/aether-data/adapters/sqlite/src/audit.rs +++ /dev/null @@ -1,405 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, Row}; - -use aether_data_contracts::repository::audit::*; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -#[derive(Debug, Clone)] -pub struct SqliteAuditLogReadRepository { - pool: SqlitePool, -} - -impl SqliteAuditLogReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl AuditLogReadRepository for SqliteAuditLogReadRepository { - async fn list_admin_audit_logs( - &self, - query: &AuditLogListQuery, - ) -> Result { - let total = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM audit_logs AS a -LEFT JOIN users AS u ON a.user_id = u.id -WHERE a.created_at >= ? - AND (? IS NULL OR LOWER(u.username) LIKE LOWER(?) ESCAPE '\') - AND (? IS NULL OR a.event_type = ?) -"#, - ) - .bind(query.cutoff_unix_secs as i64) - .bind(query.username_pattern.as_deref()) - .bind(query.username_pattern.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query( - r#" -SELECT - a.id, - a.event_type, - a.user_id, - u.email AS user_email, - u.username AS user_username, - a.description, - a.ip_address, - a.status_code, - a.error_message, - a.event_metadata AS metadata, - a.created_at -FROM audit_logs AS a -LEFT JOIN users AS u ON a.user_id = u.id -WHERE a.created_at >= ? - AND (? IS NULL OR LOWER(u.username) LIKE LOWER(?) ESCAPE '\') - AND (? IS NULL OR a.event_type = ?) -ORDER BY a.created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.cutoff_unix_secs as i64) - .bind(query.username_pattern.as_deref()) - .bind(query.username_pattern.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.limit as i64) - .bind(query.offset as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let items = rows - .iter() - .map(map_sqlite_admin_audit_log_row) - .collect::, _>>()?; - - Ok(StoredAdminAuditLogPage { - items, - total: total.max(0) as u64, - }) - } - - async fn list_admin_suspicious_activities( - &self, - cutoff_unix_secs: u64, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT id, event_type, user_id, description, ip_address, event_metadata AS metadata, created_at -FROM audit_logs -WHERE created_at >= ? - AND event_type IN (?, ?, ?, ?) -ORDER BY created_at DESC -LIMIT 100 -"#, - ) - .bind(cutoff_unix_secs as i64) - .bind(SUSPICIOUS_EVENT_TYPES[0]) - .bind(SUSPICIOUS_EVENT_TYPES[1]) - .bind(SUSPICIOUS_EVENT_TYPES[2]) - .bind(SUSPICIOUS_EVENT_TYPES[3]) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - rows.iter() - .map(map_sqlite_suspicious_activity_row) - .collect() - } - - async fn read_admin_user_behavior_event_counts( - &self, - user_id: &str, - cutoff_unix_secs: u64, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT event_type, COUNT(*) AS count -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? -GROUP BY event_type -"#, - ) - .bind(user_id) - .bind(cutoff_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - Ok(rows - .iter() - .filter_map(|row| event_count_from_sqlite_row(row).ok()) - .collect()) - } - - async fn list_user_audit_logs( - &self, - user_id: &str, - query: &AuditLogListQuery, - ) -> Result { - let total = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? - AND (? IS NULL OR event_type = ?) -"#, - ) - .bind(user_id) - .bind(query.cutoff_unix_secs as i64) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query( - r#" -SELECT id, event_type, description, ip_address, status_code, created_at -FROM audit_logs -WHERE user_id = ? - AND created_at >= ? - AND (? IS NULL OR event_type = ?) -ORDER BY created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(user_id) - .bind(query.cutoff_unix_secs as i64) - .bind(query.event_type.as_deref()) - .bind(query.event_type.as_deref()) - .bind(query.limit as i64) - .bind(query.offset as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let items = rows - .iter() - .map(map_sqlite_user_audit_log_row) - .collect::, _>>()?; - - Ok(StoredUserAuditLogPage { - items, - total: total.max(0) as u64, - }) - } - - async fn delete_audit_logs_before( - &self, - cutoff_unix_secs: u64, - limit: usize, - ) -> Result { - let deleted = sqlx::query( - r#" -DELETE FROM audit_logs -WHERE id IN ( - SELECT id - FROM audit_logs - WHERE created_at < ? - ORDER BY created_at ASC, id ASC - LIMIT ? -) -"#, - ) - .bind(cutoff_unix_secs.min(i64::MAX as u64) as i64) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(deleted).unwrap_or(usize::MAX)) - } -} - -fn sqlite_created_at_unix_secs(row: &SqliteRow) -> Result { - let value = row.try_get::("created_at").map_sql_err()?; - Ok(value.max(0) as u64) -} - -fn map_sqlite_admin_audit_log_row(row: &SqliteRow) -> Result { - Ok(StoredAdminAuditLog { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - user_email: row.try_get("user_email").map_sql_err()?, - user_username: row.try_get("user_username").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - status_code: row.try_get("status_code").map_sql_err()?, - error_message: row.try_get("error_message").map_sql_err()?, - metadata: optional_json_from_text(row.try_get("metadata").map_sql_err()?)?, - created_at_unix_secs: sqlite_created_at_unix_secs(row)?, - }) -} - -fn map_sqlite_suspicious_activity_row( - row: &SqliteRow, -) -> Result { - Ok(StoredSuspiciousActivity { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - metadata: optional_json_from_text(row.try_get("metadata").map_sql_err()?)?, - created_at_unix_secs: sqlite_created_at_unix_secs(row)?, - }) -} - -fn map_sqlite_user_audit_log_row(row: &SqliteRow) -> Result { - Ok(StoredUserAuditLog { - id: row.try_get("id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - ip_address: row.try_get("ip_address").map_sql_err()?, - status_code: row.try_get("status_code").map_sql_err()?, - created_at_unix_secs: sqlite_created_at_unix_secs(row)?, - }) -} - -fn event_count_from_sqlite_row(row: &SqliteRow) -> Result<(String, u64), DataLayerError> { - let event_type = row.try_get("event_type").map_sql_err()?; - let count = row.try_get::("count").map_sql_err()?.max(0) as u64; - Ok((event_type, count)) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::run_migrations; - - #[tokio::test] - async fn sqlite_audit_log_repository_reads_monitoring_views() { - 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_sqlite_audit_logs(&pool).await; - - let repository = SqliteAuditLogReadRepository::new(pool); - let admin_page = repository - .list_admin_audit_logs(&AuditLogListQuery { - cutoff_unix_secs: 150, - username_pattern: Some("%ali%".to_string()), - event_type: Some("login_failed".to_string()), - limit: 10, - offset: 0, - }) - .await - .expect("admin audit logs should read"); - assert_eq!(admin_page.total, 1); - assert_eq!(admin_page.items[0].id, "audit-2"); - assert_eq!(admin_page.items[0].user_username.as_deref(), Some("alice")); - assert_eq!( - admin_page.items[0] - .metadata - .as_ref() - .and_then(|value| value.get("risk")) - .and_then(|value| value.as_str()), - Some("high") - ); - - let suspicious = repository - .list_admin_suspicious_activities(150) - .await - .expect("suspicious activities should read"); - assert_eq!(suspicious.len(), 1); - assert_eq!(suspicious[0].event_type, "login_failed"); - - let counts = repository - .read_admin_user_behavior_event_counts("user-1", 0) - .await - .expect("user behavior counts should read"); - assert_eq!(counts.get("login_failed"), Some(&1)); - assert_eq!(counts.get("request_success"), Some(&1)); - - let user_page = repository - .list_user_audit_logs( - "user-1", - &AuditLogListQuery { - cutoff_unix_secs: 0, - username_pattern: None, - event_type: Some("request_success".to_string()), - limit: 10, - offset: 0, - }, - ) - .await - .expect("user audit logs should read"); - assert_eq!(user_page.total, 1); - assert_eq!(user_page.items[0].id, "audit-1"); - assert_eq!(user_page.items[0].status_code, Some(200)); - - let deleted = repository - .delete_audit_logs_before(250, 1) - .await - .expect("audit cleanup should delete one old row"); - assert_eq!(deleted, 1); - let user_page = repository - .list_user_audit_logs( - "user-1", - &AuditLogListQuery { - cutoff_unix_secs: 0, - username_pattern: None, - event_type: Some("request_success".to_string()), - limit: 10, - offset: 0, - }, - ) - .await - .expect("user audit logs should read after cleanup"); - assert_eq!(user_page.total, 0); - } - - async fn seed_sqlite_audit_logs(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO users (id, email, username, role, auth_source, created_at, updated_at) -VALUES - ('user-1', 'alice@example.com', 'alice', 'user', 'local', 1, 1), - ('user-2', 'bob@example.com', 'bob', 'user', 'local', 1, 1) -"#, - ) - .execute(pool) - .await - .expect("users should insert"); - - sqlx::query( - r#" -INSERT INTO audit_logs ( - id, - event_type, - user_id, - description, - ip_address, - event_metadata, - status_code, - created_at -) -VALUES - ('audit-1', 'request_success', 'user-1', 'completed request', '127.0.0.1', NULL, 200, 100), - ('audit-2', 'login_failed', 'user-1', 'failed login', '127.0.0.2', '{"risk":"high"}', 401, 200), - ('audit-3', 'password_changed', 'user-2', 'other user changed password', '127.0.0.3', NULL, 200, 300) -"#, - ) - .execute(pool) - .await - .expect("audit logs should insert"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/auth.rs b/crates/aether-data/adapters/sqlite/src/auth.rs deleted file mode 100644 index c848067b7..000000000 --- a/crates/aether-data/adapters/sqlite/src/auth.rs +++ /dev/null @@ -1,2199 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::auth::{ - AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository, - AuthApiKeyWriteRepository, CompareAndSwapAuthApiKeyCiphertext, CreateStandaloneApiKeyRecord, - CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, - StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::{sqlite_real, SqlitePool}; - -const SNAPSHOT_COLUMNS: &str = r#" -SELECT - users.id AS user_id, - users.username, - users.email, - users.role AS user_role, - users.auth_source AS user_auth_source, - users.is_active AS user_is_active, - users.is_deleted AS user_is_deleted, - users.rate_limit AS user_rate_limit, - users.allowed_providers AS user_allowed_providers, - users.allowed_api_formats AS user_allowed_api_formats, - users.allowed_models AS user_allowed_models, - api_keys.id AS api_key_id, - api_keys.name AS api_key_name, - api_keys.is_active AS api_key_is_active, - api_keys.is_locked AS api_key_is_locked, - api_keys.is_standalone AS api_key_is_standalone, - api_keys.rate_limit AS api_key_rate_limit, - api_keys.concurrent_limit AS api_key_concurrent_limit, - api_keys.expires_at AS api_key_expires_at_unix_secs, - api_keys.allowed_providers AS api_key_allowed_providers, - api_keys.allowed_api_formats AS api_key_allowed_api_formats, - api_keys.allowed_models AS api_key_allowed_models, - api_keys.ip_rules AS api_key_ip_rules -FROM api_keys -JOIN users ON users.id = api_keys.user_id -"#; - -const EXPORT_COLUMNS: &str = r#" -SELECT - api_keys.user_id, - api_keys.id AS api_key_id, - api_keys.key_hash, - api_keys.key_encrypted, - api_keys.name, - api_keys.allowed_providers, - api_keys.allowed_api_formats, - api_keys.allowed_models, - api_keys.ip_rules, - api_keys.rate_limit, - api_keys.concurrent_limit, - api_keys.force_capabilities, - api_keys.feature_settings, - api_keys.is_active, - api_keys.expires_at AS expires_at_unix_secs, - api_keys.auto_delete_on_expiry, - api_keys.total_requests, - COALESCE(api_keys.total_tokens, 0) AS total_tokens, - CAST(COALESCE(api_keys.total_cost_usd, 0) AS REAL) AS total_cost_usd, - api_keys.last_used_at AS last_used_at_unix_secs, - api_keys.created_at AS created_at_unix_secs, - api_keys.updated_at AS updated_at_unix_secs, - api_keys.is_standalone -FROM api_keys -"#; - -const SQLITE_ANONYMIZE_API_KEY_HISTORY_SQL: &[&str] = &[ - "UPDATE request_candidates SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE video_tasks SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE usage SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE stats_daily_api_key SET api_key_name = NULL WHERE api_key_id = ?", - "UPDATE audit_logs SET description = 'deleted API key event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE api_key_id = ?", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.api_key_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_orders SET gateway_response = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id = ?)", -]; - -const SQLITE_DELETE_API_KEY_DEPENDENTS_SQL: &[&str] = - &["DELETE FROM api_key_provider_mappings WHERE api_key_id = ?"]; - -#[derive(Debug, Clone)] -pub struct SqliteAuthApiKeyReadRepository { - pool: SqlitePool, -} - -impl SqliteAuthApiKeyReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn fetch_snapshot_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_auth_api_key_snapshot_row).collect() - } - - async fn fetch_export_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_auth_api_key_export_row).collect() - } - - async fn reload_export_by_id( - &self, - api_key_id: &str, - ) -> Result, DataLayerError> { - Ok(self - .list_export_api_keys_by_ids(&[api_key_id.to_string()]) - .await? - .into_iter() - .next()) - } - - async fn create_api_key( - &self, - record: CreateApiKeyInsertRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let owner_exists: Option = - sqlx::query_scalar("SELECT id FROM users WHERE id = ? AND is_deleted = 0") - .bind(&record.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if owner_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query( - r#" -INSERT INTO api_keys ( - id, user_id, key_hash, key_encrypted, name, allowed_providers, - allowed_api_formats, allowed_models, ip_rules, rate_limit, concurrent_limit, - force_capabilities, feature_settings, is_active, expires_at, auto_delete_on_expiry, - total_requests, total_tokens, total_cost_usd, is_standalone, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.api_key_id) - .bind(&record.user_id) - .bind(&record.key_hash) - .bind(&record.key_encrypted) - .bind(&record.name) - .bind(json_string_from_string_list( - record.allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(json_string_from_string_list( - record.allowed_api_formats.as_ref(), - "api_keys.allowed_api_formats", - )?) - .bind(json_string_from_string_list( - record.allowed_models.as_ref(), - "api_keys.allowed_models", - )?) - .bind(json_string_from_string_list( - record.ip_rules.as_ref(), - "api_keys.ip_rules", - )?) - .bind(record.rate_limit) - .bind(record.concurrent_limit) - .bind(optional_json_to_string( - &record.force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(optional_json_to_string( - &record.feature_settings, - "api_keys.feature_settings", - )?) - .bind(record.is_active) - .bind(optional_i64_from_u64( - record.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(record.auto_delete_on_expiry) - .bind(i64_from_u64( - record.total_requests, - "api_keys.total_requests", - )?) - .bind(i64_from_u64(record.total_tokens, "api_keys.total_tokens")?) - .bind(record.total_cost_usd) - .bind(record.is_standalone) - .bind(now as i64) - .bind(now as i64) - .execute(&mut *tx) - .await - .map_sql_err()?; - let reload_sql = format!("{EXPORT_COLUMNS}\nWHERE api_keys.id = ?\nLIMIT 1"); - let row = sqlx::query(&reload_sql) - .bind(&record.api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::UnexpectedValue(format!( - "created api_keys row is missing: {}", - record.api_key_id - ))); - }; - let created = map_auth_api_key_export_row(&row)?; - tx.commit().await.map_sql_err()?; - Ok(Some(created)) - } -} - -struct CreateApiKeyInsertRecord { - user_id: String, - api_key_id: String, - key_hash: String, - key_encrypted: Option, - name: Option, - allowed_providers: Option>, - allowed_api_formats: Option>, - allowed_models: Option>, - ip_rules: Option>, - rate_limit: Option, - concurrent_limit: Option, - force_capabilities: Option, - feature_settings: Option, - is_active: bool, - expires_at_unix_secs: Option, - auto_delete_on_expiry: bool, - total_requests: u64, - total_tokens: u64, - total_cost_usd: f64, - is_standalone: bool, -} - -#[async_trait] -impl AuthApiKeyReadRepository for SqliteAuthApiKeyReadRepository { - async fn find_api_key_snapshot( - &self, - key: AuthApiKeyLookupKey<'_>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SNAPSHOT_COLUMNS); - match key { - AuthApiKeyLookupKey::KeyHash(key_hash) => { - builder - .push(" WHERE api_keys.key_hash = ") - .push_bind(key_hash); - } - AuthApiKeyLookupKey::ApiKeyId(api_key_id) => { - builder.push(" WHERE api_keys.id = ").push_bind(api_key_id); - } - AuthApiKeyLookupKey::UserApiKeyIds { - user_id, - api_key_id, - } => { - builder - .push(" WHERE api_keys.id = ") - .push_bind(api_key_id) - .push(" AND users.id = ") - .push_bind(user_id); - } - } - builder.push(" LIMIT 1"); - Ok(self.fetch_snapshot_rows(builder).await?.into_iter().next()) - } - - async fn list_api_key_snapshots_by_ids( - &self, - api_key_ids: &[String], - ) -> Result, DataLayerError> { - if api_key_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(SNAPSHOT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.id IN (", api_key_ids); - builder.push(" ORDER BY api_keys.id ASC"); - self.fetch_snapshot_rows(builder).await - } - - async fn list_export_api_keys_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.user_id IN (", user_ids); - builder - .push(" AND api_keys.is_standalone = 0 ORDER BY api_keys.user_id ASC, api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_api_keys_by_ids( - &self, - api_key_ids: &[String], - ) -> Result, DataLayerError> { - if api_key_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - push_in_clause(&mut builder, " WHERE api_keys.id IN (", api_key_ids); - builder.push(" ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_api_keys_by_name_search( - &self, - name_search: &str, - ) -> Result, DataLayerError> { - let name_search = name_search.trim(); - if name_search.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder - .push(" WHERE LOWER(COALESCE(api_keys.name, '')) LIKE ") - .push_bind(format!("%{}%", name_search.to_ascii_lowercase())) - .push(" ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_standalone_api_keys_page( - &self, - query: &StandaloneApiKeyExportListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder.push(" WHERE api_keys.is_standalone = 1"); - if let Some(is_active) = query.is_active { - builder - .push(" AND api_keys.is_active = ") - .push_bind(is_active); - } - builder - .push(" ORDER BY api_keys.id ASC LIMIT ") - .push_bind(i64::try_from(query.limit).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid standalone api key export limit: {}", - query.limit - )) - })?) - .push(" OFFSET ") - .push_bind(i64::try_from(query.skip).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid standalone api key export skip: {}", - query.skip - )) - })?); - self.fetch_export_rows(builder).await - } - - async fn count_export_standalone_api_keys( - &self, - is_active: Option, - ) -> Result { - let mut builder = QueryBuilder::::new( - "SELECT COUNT(*) AS total FROM api_keys WHERE is_standalone = 1", - ); - if let Some(is_active) = is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - async fn summarize_export_api_keys_by_user_ids( - &self, - user_ids: &[String], - now_unix_secs: u64, - ) -> Result { - if user_ids.is_empty() { - return Ok(AuthApiKeyExportSummary::default()); - } - let now_unix_secs = i64_from_u64(now_unix_secs, "api_keys.summary_now")?; - - let mut builder = QueryBuilder::::new( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 AND (expires_at IS NULL OR expires_at >= -"#, - ); - builder.push_bind(now_unix_secs); - builder.push( - r#") THEN 1 ELSE 0 END) AS active -FROM api_keys -"#, - ); - push_in_clause(&mut builder, " WHERE user_id IN (", user_ids); - builder.push(" AND is_standalone = 0"); - summarize_row(builder.build().fetch_one(&self.pool).await.map_sql_err()?) - } - - async fn summarize_export_non_standalone_api_keys( - &self, - now_unix_secs: u64, - ) -> Result { - summarize_api_keys(&self.pool, false, now_unix_secs).await - } - - async fn summarize_export_standalone_api_keys( - &self, - now_unix_secs: u64, - ) -> Result { - summarize_api_keys(&self.pool, true, now_unix_secs).await - } - - async fn find_export_standalone_api_key_by_id( - &self, - api_key_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder - .push(" WHERE api_keys.is_standalone = 1 AND api_keys.id = ") - .push_bind(api_key_id) - .push(" LIMIT 1"); - Ok(self.fetch_export_rows(builder).await?.into_iter().next()) - } - - async fn list_export_standalone_api_keys( - &self, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(EXPORT_COLUMNS); - builder.push(" WHERE api_keys.is_standalone = 1 ORDER BY api_keys.id ASC"); - self.fetch_export_rows(builder).await - } -} - -#[async_trait] -impl AuthApiKeyWriteRepository for SqliteAuthApiKeyReadRepository { - async fn touch_last_used_at(&self, api_key_id: &str) -> Result { - let now = current_unix_secs() as i64; - let rows_affected = sqlx::query( - r#" -UPDATE api_keys -SET last_used_at = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(now) - .bind(now) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn create_user_api_key( - &self, - record: CreateUserApiKeyRecord, - ) -> Result, DataLayerError> { - self.create_api_key(CreateApiKeyInsertRecord { - user_id: record.user_id, - api_key_id: record.api_key_id, - key_hash: record.key_hash, - key_encrypted: record.key_encrypted, - name: record.name, - allowed_providers: record.allowed_providers, - allowed_api_formats: record.allowed_api_formats, - allowed_models: record.allowed_models, - ip_rules: record.ip_rules, - rate_limit: Some(record.rate_limit), - concurrent_limit: record.concurrent_limit, - force_capabilities: record.force_capabilities, - feature_settings: record.feature_settings, - is_active: record.is_active, - expires_at_unix_secs: record.expires_at_unix_secs, - auto_delete_on_expiry: record.auto_delete_on_expiry, - total_requests: record.total_requests, - total_tokens: record.total_tokens, - total_cost_usd: record.total_cost_usd, - is_standalone: false, - }) - .await - } - - async fn create_standalone_api_key( - &self, - record: CreateStandaloneApiKeyRecord, - ) -> Result, DataLayerError> { - self.create_api_key(CreateApiKeyInsertRecord { - user_id: record.user_id, - api_key_id: record.api_key_id, - key_hash: record.key_hash, - key_encrypted: record.key_encrypted, - name: record.name, - allowed_providers: record.allowed_providers, - allowed_api_formats: record.allowed_api_formats, - allowed_models: record.allowed_models, - ip_rules: record.ip_rules, - rate_limit: record.rate_limit, - concurrent_limit: record.concurrent_limit, - force_capabilities: record.force_capabilities, - feature_settings: None, - is_active: record.is_active, - expires_at_unix_secs: record.expires_at_unix_secs, - auto_delete_on_expiry: record.auto_delete_on_expiry, - total_requests: record.total_requests, - total_tokens: record.total_tokens, - total_cost_usd: record.total_cost_usd, - is_standalone: true, - }) - .await - } - - async fn update_user_api_key_basic( - &self, - record: UpdateUserApiKeyBasicRecord, - ) -> Result, DataLayerError> { - self.update_user_api_key_basic_scoped(record, false).await - } - - async fn compare_and_swap_api_key_ciphertext( - &self, - mutation: &CompareAndSwapAuthApiKeyCiphertext, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = ? -WHERE CAST(id AS BLOB) = CAST(? AS BLOB) - AND CAST(user_id AS BLOB) = CAST(? AS BLOB) - AND CAST(key_hash AS BLOB) = CAST(? AS BLOB) - AND is_standalone = ? - AND CAST(key_encrypted AS BLOB) = CAST(? AS BLOB) -"#, - ) - .bind(&mutation.key_encrypted) - .bind(&mutation.api_key_id) - .bind(&mutation.user_id) - .bind(&mutation.key_hash) - .bind(mutation.is_standalone) - .bind(&mutation.expected_key_encrypted) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn update_user_api_key_basic_if_unlocked( - &self, - record: UpdateUserApiKeyBasicRecord, - ) -> Result, DataLayerError> { - self.update_user_api_key_basic_scoped(record, true).await - } - - async fn update_standalone_api_key_basic( - &self, - record: UpdateStandaloneApiKeyBasicRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs() as i64; - sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = CASE WHEN ? THEN ? ELSE key_encrypted END, - name = CASE WHEN ? THEN ? ELSE name END, - force_capabilities = CASE WHEN ? THEN ? ELSE force_capabilities END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, - allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, - allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, - allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END, - ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, - expires_at = CASE WHEN ? THEN ? ELSE expires_at END, - auto_delete_on_expiry = CASE WHEN ? THEN ? ELSE auto_delete_on_expiry END, - updated_at = ? -WHERE id = ? - AND is_standalone = 1 -"#, - ) - .bind(record.key_encrypted_present) - .bind(record.key_encrypted.as_deref()) - .bind(record.name_present) - .bind(record.name.as_deref()) - .bind(record.force_capabilities.is_some()) - .bind(optional_json_to_string( - &record.force_capabilities.clone().flatten(), - "api_keys.force_capabilities", - )?) - .bind(record.rate_limit_present) - .bind(record.rate_limit) - .bind(record.concurrent_limit_present) - .bind(record.concurrent_limit) - .bind(record.allowed_providers.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_providers, - "api_keys.allowed_providers", - )?) - .bind(record.allowed_api_formats.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_api_formats, - "api_keys.allowed_api_formats", - )?) - .bind(record.allowed_models.is_some()) - .bind(json_string_from_nested_string_list( - &record.allowed_models, - "api_keys.allowed_models", - )?) - .bind(record.ip_rules.is_some()) - .bind(json_string_from_nested_string_list( - &record.ip_rules, - "api_keys.ip_rules", - )?) - .bind(record.expires_at_present) - .bind(optional_i64_from_u64( - record.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(record.auto_delete_on_expiry_present) - .bind(record.auto_delete_on_expiry) - .bind(now) - .bind(&record.api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(&record.api_key_id).await - } - - async fn restore_api_key_if_matches( - &self, - expected: &StoredAuthApiKeyExportRecord, - restored: &StoredAuthApiKeyExportRecord, - ) -> Result { - if restored.api_key_id != expected.api_key_id - || restored.user_id != expected.user_id - || restored.key_hash != expected.key_hash - || restored.is_standalone != expected.is_standalone - { - return Ok(false); - } - - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let select_sql = format!("{EXPORT_COLUMNS} WHERE api_keys.id = ? LIMIT 1"); - let row = sqlx::query(&select_sql) - .bind(&expected.api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_auth_api_key_export_row(&row)?; - if current != *expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = ?, - name = ?, - allowed_providers = ?, - allowed_api_formats = ?, - allowed_models = ?, - ip_rules = ?, - rate_limit = ?, - concurrent_limit = ?, - force_capabilities = ?, - feature_settings = ?, - is_active = ?, - expires_at = ?, - auto_delete_on_expiry = ?, - total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - last_used_at = ?, - updated_at = ? -WHERE id = ? - AND user_id = ? - AND key_hash = ? - AND is_standalone = ? -"#, - ) - .bind(restored.key_encrypted.as_deref()) - .bind(restored.name.as_deref()) - .bind(json_string_from_string_list( - restored.allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(json_string_from_string_list( - restored.allowed_api_formats.as_ref(), - "api_keys.allowed_api_formats", - )?) - .bind(json_string_from_string_list( - restored.allowed_models.as_ref(), - "api_keys.allowed_models", - )?) - .bind(json_string_from_string_list( - restored.ip_rules.as_ref(), - "api_keys.ip_rules", - )?) - .bind(restored.rate_limit) - .bind(restored.concurrent_limit) - .bind(optional_json_to_string( - &restored.force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(optional_json_to_string( - &restored.feature_settings, - "api_keys.feature_settings", - )?) - .bind(restored.is_active) - .bind(optional_i64_from_u64( - restored.expires_at_unix_secs, - "api_keys.expires_at", - )?) - .bind(restored.auto_delete_on_expiry) - .bind(i64_from_u64( - restored.total_requests, - "api_keys.total_requests", - )?) - .bind(i64_from_u64( - restored.total_tokens, - "api_keys.total_tokens", - )?) - .bind(restored.total_cost_usd) - .bind(optional_i64_from_u64( - restored.last_used_at_unix_secs, - "api_keys.last_used_at", - )?) - .bind(current_unix_secs() as i64) - .bind(&restored.api_key_id) - .bind(&restored.user_id) - .bind(&restored.key_hash) - .bind(restored.is_standalone) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn set_user_api_key_active( - &self, - user_id: &str, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, Some(user_id), is_active, false, false) - .await - } - - async fn set_user_api_key_active_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, Some(user_id), is_active, false, true) - .await - } - - async fn set_standalone_api_key_active( - &self, - api_key_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_active(api_key_id, None, is_active, true, false) - .await - } - - async fn set_user_api_key_locked( - &self, - user_id: &str, - api_key_id: &str, - is_locked: bool, - ) -> Result { - let rows_affected = sqlx::query( - r#" -UPDATE api_keys -SET is_locked = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 -"#, - ) - .bind(is_locked) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn set_user_api_key_allowed_providers( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - ) -> Result, DataLayerError> { - self.set_user_api_key_allowed_providers_scoped( - user_id, - api_key_id, - allowed_providers, - false, - ) - .await - } - - async fn set_user_api_key_allowed_providers_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - ) -> Result, DataLayerError> { - self.set_user_api_key_allowed_providers_scoped(user_id, api_key_id, allowed_providers, true) - .await - } - - async fn set_user_api_key_force_capabilities( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_force_capabilities_scoped( - user_id, - api_key_id, - force_capabilities, - false, - ) - .await - } - - async fn set_user_api_key_force_capabilities_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_force_capabilities_scoped( - user_id, - api_key_id, - force_capabilities, - true, - ) - .await - } - - async fn set_user_api_key_feature_settings( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_feature_settings_scoped(user_id, api_key_id, feature_settings, false) - .await - } - - async fn set_user_api_key_feature_settings_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - self.set_user_api_key_feature_settings_scoped(user_id, api_key_id, feature_settings, true) - .await - } - - async fn set_api_key_usage_totals( - &self, - api_key_id: &str, - total_requests: u64, - total_tokens: u64, - total_cost_usd: f64, - ) -> Result, DataLayerError> { - if !total_cost_usd.is_finite() { - return Err(DataLayerError::InvalidInput( - "api_keys.total_cost_usd is not finite".to_string(), - )); - } - sqlx::query( - r#" -UPDATE api_keys -SET total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(i64_from_u64(total_requests, "api_keys.total_requests")?) - .bind(i64_from_u64(total_tokens, "api_keys.total_tokens")?) - .bind(total_cost_usd) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(api_key_id).await - } - - async fn delete_user_api_key( - &self, - user_id: &str, - api_key_id: &str, - ) -> Result { - self.delete_api_key(api_key_id, Some(user_id), false, false) - .await - } - - async fn delete_user_api_key_if_unlocked( - &self, - user_id: &str, - api_key_id: &str, - ) -> Result { - self.delete_api_key(api_key_id, Some(user_id), false, true) - .await - } - - async fn delete_standalone_api_key(&self, api_key_id: &str) -> Result { - self.delete_api_key(api_key_id, None, true, false).await - } - - async fn set_standalone_api_key_feature_settings( - &self, - api_key_id: &str, - feature_settings: Option, - ) -> Result, DataLayerError> { - sqlx::query( - r#" -UPDATE api_keys -SET feature_settings = ?, updated_at = ? -WHERE id = ? - AND is_standalone = 1 -"#, - ) - .bind(optional_json_to_string( - &feature_settings, - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_export_by_id(api_key_id).await - } -} - -impl SqliteAuthApiKeyReadRepository { - async fn update_user_api_key_basic_scoped( - &self, - record: UpdateUserApiKeyBasicRecord, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET key_encrypted = CASE WHEN ? THEN ? ELSE key_encrypted END, - name = CASE WHEN ? THEN ? ELSE name END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - concurrent_limit = CASE WHEN ? THEN ? ELSE concurrent_limit END, - ip_rules = CASE WHEN ? THEN ? ELSE ip_rules END, - feature_settings = CASE WHEN ? THEN ? ELSE feature_settings END, - updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(record.key_encrypted_present) - .bind(record.key_encrypted.as_deref()) - .bind(record.name_present) - .bind(record.name.as_deref()) - .bind(record.rate_limit_present) - .bind(record.rate_limit) - .bind(record.concurrent_limit_present) - .bind(record.concurrent_limit) - .bind(record.ip_rules.is_some()) - .bind(json_string_from_nested_string_list( - &record.ip_rules, - "api_keys.ip_rules", - )?) - .bind(record.feature_settings.is_some()) - .bind(optional_json_to_string( - &record.feature_settings.clone().flatten(), - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(&record.api_key_id) - .bind(&record.user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(&record.api_key_id).await - } - - async fn set_active( - &self, - api_key_id: &str, - user_id: Option<&str>, - is_active: bool, - is_standalone: bool, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new("UPDATE api_keys SET is_active = "); - builder - .push_bind(is_active) - .push(", updated_at = ") - .push_bind(current_unix_secs() as i64) - .push(" WHERE id = ") - .push_bind(api_key_id) - .push(" AND is_standalone = ") - .push_bind(is_standalone); - if let Some(user_id) = user_id { - builder.push(" AND user_id = ").push_bind(user_id); - } - if require_unlocked { - builder.push(" AND is_locked = ").push_bind(false); - } - let result = builder.build().execute(&self.pool).await.map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_allowed_providers_scoped( - &self, - user_id: &str, - api_key_id: &str, - allowed_providers: Option>, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET allowed_providers = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(json_string_from_string_list( - allowed_providers.as_ref(), - "api_keys.allowed_providers", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_force_capabilities_scoped( - &self, - user_id: &str, - api_key_id: &str, - force_capabilities: Option, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET force_capabilities = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(optional_json_to_string( - &force_capabilities, - "api_keys.force_capabilities", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn set_user_api_key_feature_settings_scoped( - &self, - user_id: &str, - api_key_id: &str, - feature_settings: Option, - require_unlocked: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE api_keys -SET feature_settings = ?, updated_at = ? -WHERE id = ? - AND user_id = ? - AND is_standalone = 0 - AND (? = 0 OR is_locked = 0) -"#, - ) - .bind(optional_json_to_string( - &feature_settings, - "api_keys.feature_settings", - )?) - .bind(current_unix_secs() as i64) - .bind(api_key_id) - .bind(user_id) - .bind(require_unlocked) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.reload_export_by_id(api_key_id).await - } - - async fn delete_api_key( - &self, - api_key_id: &str, - user_id: Option<&str>, - is_standalone: bool, - require_unlocked: bool, - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let matching_api_key = if let Some(user_id) = user_id { - if require_unlocked { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND user_id = ? AND is_standalone = 0 AND is_locked = 0", - ) - .bind(api_key_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND user_id = ? AND is_standalone = 0", - ) - .bind(api_key_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } - } else { - sqlx::query_scalar::<_, String>( - "SELECT id FROM api_keys WHERE id = ? AND is_standalone = 1", - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - }; - if matching_api_key.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - sqlx::query( - "UPDATE wallets SET status = 'disabled', updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE api_key_id = ? AND status <> 'disabled'", - ) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - for sql in SQLITE_ANONYMIZE_API_KEY_HISTORY_SQL { - sqlx::query(sql) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - for sql in SQLITE_DELETE_API_KEY_DEPENDENTS_SQL { - sqlx::query(sql) - .bind(api_key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let result = sqlx::query("DELETE FROM api_keys WHERE id = ? AND is_standalone = ?") - .bind(api_key_id) - .bind(is_standalone) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } -} - -fn push_in_clause<'args>( - builder: &mut QueryBuilder<'args, Sqlite>, - prefix: &str, - values: &'args [String], -) { - builder.push(prefix); - { - let mut separated = builder.separated(", "); - for value in values { - separated.push_bind(value); - } - } - builder.push(")"); -} - -async fn summarize_api_keys( - pool: &SqlitePool, - is_standalone: bool, - now_unix_secs: u64, -) -> Result { - let now_unix_secs = i64_from_u64(now_unix_secs, "api_keys.summary_now")?; - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 AND (expires_at IS NULL OR expires_at >= ?) THEN 1 ELSE 0 END) AS active -FROM api_keys -WHERE is_standalone = ? -"#, - ) - .bind(now_unix_secs) - .bind(is_standalone) - .fetch_one(pool) - .await - .map_sql_err()?; - summarize_row(row) -} - -fn summarize_row(row: SqliteRow) -> Result { - Ok(AuthApiKeyExportSummary { - total: row.try_get::("total").map_sql_err()?.max(0) as u64, - active: row - .try_get::, _>("active") - .map_sql_err()? - .unwrap_or(0) - .max(0) as u64, - }) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| i64_from_u64(value, field_name)) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_string_list( - value: Option<&Vec>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable string list: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_nested_string_list( - value: &Option>>, - field_name: &str, -) -> Result, DataLayerError> { - match value { - Some(Some(values)) => json_string_from_string_list(Some(values), field_name), - Some(None) | None => Ok(None), - } -} - -fn map_auth_api_key_snapshot_row( - row: &SqliteRow, -) -> Result { - let snapshot = StoredAuthApiKeySnapshot::new( - row.try_get("user_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("user_role").map_sql_err()?, - row.try_get("user_auth_source").map_sql_err()?, - row.try_get("user_is_active").map_sql_err()?, - row.try_get("user_is_deleted").map_sql_err()?, - optional_json_from_string( - row.try_get("user_allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("user_allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("user_allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("api_key_is_active").map_sql_err()?, - row.try_get("api_key_is_locked").map_sql_err()?, - row.try_get("api_key_is_standalone").map_sql_err()?, - row.try_get("api_key_rate_limit").map_sql_err()?, - row.try_get("api_key_concurrent_limit").map_sql_err()?, - row.try_get("api_key_expires_at_unix_secs").map_sql_err()?, - optional_json_from_string( - row.try_get("api_key_allowed_providers").map_sql_err()?, - "api_keys.allowed_providers", - )?, - optional_json_from_string( - row.try_get("api_key_allowed_api_formats").map_sql_err()?, - "api_keys.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("api_key_allowed_models").map_sql_err()?, - "api_keys.allowed_models", - )?, - )? - .with_api_key_ip_rules(optional_json_from_string( - row.try_get("api_key_ip_rules").map_sql_err()?, - "api_keys.ip_rules", - )?)?; - Ok(snapshot.with_user_rate_limit(row.try_get("user_rate_limit").map_sql_err()?)) -} - -fn map_auth_api_key_export_row( - row: &SqliteRow, -) -> Result { - let feature_settings = optional_json_from_string( - row.try_get("feature_settings").map_sql_err()?, - "api_keys.feature_settings", - )?; - StoredAuthApiKeyExportRecord::new( - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("key_hash").map_sql_err()?, - row.try_get("key_encrypted").map_sql_err()?, - row.try_get("name").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "api_keys.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "api_keys.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "api_keys.allowed_models", - )?, - row.try_get("rate_limit").map_sql_err()?, - row.try_get("concurrent_limit").map_sql_err()?, - optional_json_from_string( - row.try_get("force_capabilities").map_sql_err()?, - "api_keys.force_capabilities", - )?, - row.try_get("is_active").map_sql_err()?, - row.try_get("expires_at_unix_secs").map_sql_err()?, - row.try_get("auto_delete_on_expiry").map_sql_err()?, - row.try_get("total_requests").map_sql_err()?, - row.try_get("total_tokens").map_sql_err()?, - sqlite_real(row, "total_cost_usd")?, - row.try_get("is_standalone").map_sql_err()?, - ) - .and_then(|record| { - record.with_ip_rules(optional_json_from_string( - row.try_get("ip_rules").map_sql_err()?, - "api_keys.ip_rules", - )?) - }) - .map(|record| record.with_feature_settings(feature_settings)) - .and_then(|record| { - record.with_activity_timestamps( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - row.try_get("created_at_unix_secs").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - ) - }) -} - -#[cfg(test)] -mod tests { - use super::SqliteAuthApiKeyReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::auth::{ - AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository, - CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, - UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord, - }; - use serde_json::json; - use sqlx::Row; - - #[tokio::test] - async fn sqlite_repository_reads_auth_api_key_contract_views() { - 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_auth_api_key_rows(&pool).await; - - let repository = SqliteAuthApiKeyReadRepository::new(pool); - let snapshot = repository - .find_api_key_snapshot(AuthApiKeyLookupKey::KeyHash("hash-user")) - .await - .expect("snapshot lookup should run") - .expect("snapshot should exist"); - assert_eq!(snapshot.user_id, "user-1"); - assert_eq!( - snapshot.api_key_allowed_models, - Some(vec!["gpt-4.1".to_string()]) - ); - - let by_ids = repository - .list_api_key_snapshots_by_ids(&["key-user".to_string(), "key-standalone".to_string()]) - .await - .expect("snapshot list should run"); - assert_eq!(by_ids.len(), 2); - - let user_exports = repository - .list_export_api_keys_by_user_ids(&["user-1".to_string()]) - .await - .expect("user exports should load"); - assert_eq!(user_exports.len(), 1); - assert_eq!(user_exports[0].total_tokens, 456); - - let page = repository - .list_export_standalone_api_keys_page(&StandaloneApiKeyExportListQuery { - skip: 0, - limit: 10, - is_active: Some(true), - }) - .await - .expect("standalone page should load"); - assert_eq!(page.len(), 1); - assert_eq!(page[0].api_key_id, "key-standalone"); - - let count = repository - .count_export_standalone_api_keys(Some(true)) - .await - .expect("standalone count should load"); - assert_eq!(count, 1); - - let summary = repository - .summarize_export_api_keys_by_user_ids(&["user-1".to_string()], 100) - .await - .expect("summary should load"); - assert_eq!(summary.total, 1); - assert_eq!(summary.active, 1); - - assert_eq!( - repository - .find_export_standalone_api_key_by_id("key-standalone") - .await - .expect("standalone find should run") - .expect("standalone should exist") - .name, - Some("Standalone".to_string()) - ); - } - - #[tokio::test] - async fn sqlite_repository_writes_auth_api_key_contract_views() { - 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_auth_user(&pool).await; - - let repository = SqliteAuthApiKeyReadRepository::new(pool); - let missing_owner_key = repository - .create_user_api_key(CreateUserApiKeyRecord { - user_id: "missing-user".to_string(), - api_key_id: "key-missing-owner".to_string(), - key_hash: "hash-missing-owner".to_string(), - key_encrypted: Some("enc-missing-owner".to_string()), - name: Some("Missing Owner".to_string()), - allowed_providers: None, - allowed_api_formats: None, - allowed_models: None, - ip_rules: None, - rate_limit: 100, - concurrent_limit: None, - force_capabilities: None, - feature_settings: None, - is_active: true, - expires_at_unix_secs: None, - auto_delete_on_expiry: false, - total_requests: 0, - total_tokens: 0, - total_cost_usd: 0.0, - }) - .await - .expect("missing-owner creation should resolve"); - assert!(missing_owner_key.is_none()); - let user_key = repository - .create_user_api_key(CreateUserApiKeyRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-user".to_string(), - key_hash: "hash-created-user".to_string(), - key_encrypted: Some("enc-user".to_string()), - name: Some("Created User".to_string()), - allowed_providers: Some(vec!["openai".to_string()]), - allowed_api_formats: Some(vec!["openai:chat".to_string()]), - allowed_models: Some(vec!["gpt-4.1".to_string()]), - ip_rules: Some(vec!["203.0.113.10".to_string()]), - rate_limit: 100, - concurrent_limit: Some(5), - force_capabilities: Some(json!({"cache": true})), - feature_settings: Some(json!({"chat_pii_redaction": {"enabled": true}})), - is_active: true, - expires_at_unix_secs: Some(2_000_000_000), - auto_delete_on_expiry: false, - total_requests: 3, - total_tokens: 42, - total_cost_usd: 0.25, - }) - .await - .expect("user key should create") - .expect("user key should reload"); - assert_eq!(user_key.allowed_models, Some(vec!["gpt-4.1".to_string()])); - assert_eq!(user_key.total_tokens, 42); - assert_eq!( - user_key.feature_settings, - Some(json!({"chat_pii_redaction": {"enabled": true}})) - ); - - let updated_user_key = repository - .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-user".to_string(), - key_encrypted: Some("enc-user-rotated".to_string()), - key_encrypted_present: true, - name: Some("Updated User".to_string()), - name_present: true, - rate_limit: Some(150), - rate_limit_present: true, - concurrent_limit: Some(6), - concurrent_limit_present: true, - ip_rules: Some(Some(vec!["10.0.0.0/24".to_string()])), - feature_settings: Some(Some(json!({"compact": true}))), - }) - .await - .expect("user key should update") - .expect("user key should reload"); - assert_eq!(updated_user_key.name, Some("Updated User".to_string())); - assert_eq!( - updated_user_key.key_encrypted.as_deref(), - Some("enc-user-rotated") - ); - assert_eq!(updated_user_key.concurrent_limit, Some(6)); - assert_eq!( - updated_user_key.feature_settings, - Some(json!({"compact": true})) - ); - - // Rollback can explicitly restore nullable values, while a present zero remains a - // meaningful rate limit rather than being treated as an omitted field. - let cleared_user_key = repository - .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-user".to_string(), - key_encrypted: None, - key_encrypted_present: true, - name: None, - name_present: true, - rate_limit: None, - rate_limit_present: true, - concurrent_limit: None, - concurrent_limit_present: true, - ip_rules: None, - feature_settings: None, - }) - .await - .expect("nullable values should clear") - .expect("user key should remain"); - assert!(cleared_user_key.key_encrypted.is_none()); - assert!(cleared_user_key.name.is_none()); - assert!(cleared_user_key.rate_limit.is_none()); - assert!(cleared_user_key.concurrent_limit.is_none()); - - let zero_rate_limit = repository - .update_user_api_key_basic(UpdateUserApiKeyBasicRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-user".to_string(), - key_encrypted: None, - key_encrypted_present: false, - name: None, - name_present: false, - rate_limit: Some(0), - rate_limit_present: true, - concurrent_limit: None, - concurrent_limit_present: false, - ip_rules: None, - feature_settings: None, - }) - .await - .expect("zero rate limit should persist") - .expect("user key should remain"); - assert_eq!(zero_rate_limit.rate_limit, Some(0)); - - assert!(repository - .set_user_api_key_locked("user-1", "key-created-user", true) - .await - .expect("lock should update")); - let snapshot = repository - .find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-created-user")) - .await - .expect("snapshot should load") - .expect("snapshot should exist"); - assert!(snapshot.api_key_is_locked); - - assert!(repository - .update_user_api_key_basic_if_unlocked(UpdateUserApiKeyBasicRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-user".to_string(), - key_encrypted: None, - key_encrypted_present: false, - name: Some("must-not-change".to_string()), - name_present: true, - rate_limit: None, - rate_limit_present: false, - concurrent_limit: None, - concurrent_limit_present: false, - ip_rules: None, - feature_settings: Some(Some(json!({"must_not_change": true}))), - }) - .await - .expect("locked basic update should resolve") - .is_none()); - assert!(repository - .set_user_api_key_active_if_unlocked("user-1", "key-created-user", false) - .await - .expect("locked status update should resolve") - .is_none()); - assert!(repository - .set_user_api_key_allowed_providers_if_unlocked( - "user-1", - "key-created-user", - Some(vec!["must-not-change".to_string()]), - ) - .await - .expect("locked provider update should resolve") - .is_none()); - assert!(repository - .set_user_api_key_force_capabilities_if_unlocked( - "user-1", - "key-created-user", - Some(json!({"must_not_change": true})), - ) - .await - .expect("locked capability update should resolve") - .is_none()); - assert!(repository - .set_user_api_key_feature_settings_if_unlocked( - "user-1", - "key-created-user", - Some(json!({"must_not_change": true})), - ) - .await - .expect("locked feature update should resolve") - .is_none()); - assert!(!repository - .delete_user_api_key_if_unlocked("user-1", "key-created-user") - .await - .expect("locked delete should resolve")); - - let unchanged = repository - .reload_export_by_id("key-created-user") - .await - .expect("locked key should reload") - .expect("locked key should still exist"); - assert_ne!(unchanged.name.as_deref(), Some("must-not-change")); - assert!(unchanged.is_active); - assert_ne!( - unchanged.allowed_providers, - Some(vec!["must-not-change".to_string()]) - ); - assert_ne!( - unchanged.force_capabilities, - Some(json!({"must_not_change": true})) - ); - assert_eq!(unchanged.feature_settings, Some(json!({"compact": true}))); - - let active_user_key = repository - .set_user_api_key_active("user-1", "key-created-user", false) - .await - .expect("active flag should update") - .expect("user key should reload"); - assert!(!active_user_key.is_active); - - assert!(repository - .set_user_api_key_active("wrong-owner", "key-created-user", true) - .await - .expect("wrong-owner status update should resolve") - .is_none()); - - let provider_updated = repository - .set_user_api_key_allowed_providers( - "user-1", - "key-created-user", - Some(vec!["anthropic".to_string()]), - ) - .await - .expect("allowed providers should update") - .expect("user key should reload"); - assert_eq!( - provider_updated.allowed_providers, - Some(vec!["anthropic".to_string()]) - ); - - let capabilities_updated = repository - .set_user_api_key_force_capabilities( - "user-1", - "key-created-user", - Some(json!({"vision": true})), - ) - .await - .expect("force capabilities should update") - .expect("user key should reload"); - assert_eq!( - capabilities_updated.force_capabilities, - Some(json!({"vision": true})) - ); - - assert!(repository - .touch_last_used_at("key-created-user") - .await - .expect("touch should update")); - assert!(repository - .reload_export_by_id("key-created-user") - .await - .expect("touched key should reload") - .expect("touched key should exist") - .last_used_at_unix_secs - .is_some()); - - let standalone = repository - .create_standalone_api_key(CreateStandaloneApiKeyRecord { - user_id: "user-1".to_string(), - api_key_id: "key-created-standalone".to_string(), - key_hash: "hash-created-standalone".to_string(), - key_encrypted: Some("enc-standalone".to_string()), - name: Some("Created Standalone".to_string()), - allowed_providers: Some(vec!["openai".to_string()]), - allowed_api_formats: None, - allowed_models: None, - ip_rules: None, - rate_limit: None, - concurrent_limit: Some(2), - force_capabilities: None, - is_active: true, - expires_at_unix_secs: None, - auto_delete_on_expiry: false, - total_requests: 0, - total_tokens: 0, - total_cost_usd: 0.0, - }) - .await - .expect("standalone key should create") - .expect("standalone key should reload"); - assert!(standalone.is_standalone); - - let standalone = repository - .update_standalone_api_key_basic(UpdateStandaloneApiKeyBasicRecord { - api_key_id: "key-created-standalone".to_string(), - key_encrypted: Some("enc-standalone-rotated".to_string()), - key_encrypted_present: true, - name: Some("Updated Standalone".to_string()), - name_present: true, - force_capabilities: None, - rate_limit_present: true, - rate_limit: Some(20), - concurrent_limit_present: true, - concurrent_limit: None, - allowed_providers: Some(None), - allowed_api_formats: Some(Some(vec!["openai:responses".to_string()])), - allowed_models: Some(Some(vec!["gpt-4.1-mini".to_string()])), - ip_rules: None, - expires_at_present: true, - expires_at_unix_secs: Some(2_100_000_000), - auto_delete_on_expiry_present: true, - auto_delete_on_expiry: true, - }) - .await - .expect("standalone key should update") - .expect("standalone key should reload"); - assert_eq!(standalone.name, Some("Updated Standalone".to_string())); - assert_eq!( - standalone.key_encrypted.as_deref(), - Some("enc-standalone-rotated") - ); - assert_eq!(standalone.allowed_providers, None); - assert_eq!( - standalone.allowed_api_formats, - Some(vec!["openai:responses".to_string()]) - ); - assert_eq!(standalone.concurrent_limit, None); - assert!(standalone.auto_delete_on_expiry); - - let standalone = repository - .set_standalone_api_key_active("key-created-standalone", false) - .await - .expect("standalone active flag should update") - .expect("standalone key should reload"); - assert!(!standalone.is_active); - - assert!(repository - .delete_standalone_api_key("key-created-standalone") - .await - .expect("standalone key should delete")); - assert!(repository - .delete_user_api_key("user-1", "key-created-user") - .await - .expect("user key should delete")); - } - - #[tokio::test] - async fn sqlite_api_key_delete_is_owner_scoped_and_preserves_anonymized_facts() { - 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"); - sqlx::raw_sql( - r#" -INSERT INTO users ( - id, username, role, auth_source, is_active, is_deleted, created_at, updated_at -) VALUES - ('key-owner', 'key-owner', 'user', 'local', 1, 0, 1, 1), - ('other-owner', 'other-owner', 'user', 'local', 1, 0, 1, 1); - -INSERT INTO api_keys ( - id, user_id, key_hash, name, is_standalone, created_at, updated_at -) VALUES ( - 'key-to-delete', 'key-owner', 'key-to-delete-hash', 'private key name', 0, 1, 1 -); - -INSERT INTO wallets ( - id, api_key_id, balance, gift_balance, status, created_at, updated_at -) VALUES ( - 'key-wallet', 'key-to-delete', 12, 3, 'active', 1, 1 -); - -INSERT INTO request_candidates ( - id, request_id, user_id, api_key_id, username, api_key_name, - candidate_index, status, created_at -) VALUES ( - 'key-history-row', 'key-candidate-request', 'key-owner', 'key-to-delete', - 'key-owner', 'private key name', 0, 'success', 1 -); - -INSERT INTO video_tasks ( - id, request_id, user_id, api_key_id, username, api_key_name, created_at, updated_at -) VALUES ( - 'key-history-row', 'key-video-request', 'key-owner', 'key-to-delete', - 'key-owner', 'private key name', 1, 1 -); - -INSERT INTO usage ( - request_id, id, user_id, api_key_id, username, api_key_name -) VALUES ( - 'key-usage-request', 'key-history-row', 'key-owner', 'key-to-delete', - 'key-owner', 'private key name' -); - -INSERT INTO stats_daily_api_key ( - id, api_key_id, date, api_key_name, created_at, updated_at -) VALUES ( - 'key-history-row', 'key-to-delete', 1, 'private key name', 1, 1 -); - -INSERT INTO audit_logs ( - id, event_type, api_key_id, description, ip_address, user_agent, - event_metadata, error_message, created_at -) VALUES ( - 'key-audit', 'key_event', 'key-to-delete', 'private description', - '192.0.2.20', 'private agent', '{"private":true}', 'private error', 1 -); - -INSERT INTO api_key_provider_mappings ( - id, api_key_id, provider_id, created_at, updated_at -) VALUES ( - 'key-mapping', 'key-to-delete', 'provider-1', 1, 1 -); - -INSERT INTO payment_orders ( - id, order_no, wallet_id, amount_usd, payment_method, - gateway_response, status, created_at -) VALUES ( - 'key-order', 'key-order-no', 'key-wallet', 12, 'test', - '{"customer_email":"key@example.com"}', 'credited', 1 -); - -INSERT INTO payment_callbacks ( - id, payment_order_id, payment_method, callback_key, order_no, - payload_hash, signature_valid, status, payload, error_message, created_at -) VALUES ( - 'key-callback', 'key-order', 'test', 'key-callback-key', 'key-order-no', - 'key-payload-hash', 1, 'processed', - '{"customer_email":"key@example.com"}', 'private callback error', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("API key deletion fixtures should insert"); - - let repository = SqliteAuthApiKeyReadRepository::new(pool.clone()); - assert!(!repository - .delete_user_api_key("other-owner", "key-to-delete") - .await - .expect("wrong-owner delete should resolve")); - assert_eq!( - sqlx::query_scalar::<_, String>("SELECT status FROM wallets WHERE id = 'key-wallet'",) - .fetch_one(&pool) - .await - .expect("wallet status should load"), - "active" - ); - assert_eq!( - sqlx::query_scalar::<_, String>( - "SELECT api_key_name FROM request_candidates WHERE id = 'key-history-row'", - ) - .fetch_one(&pool) - .await - .expect("candidate key name should load"), - "private key name" - ); - - assert!(repository - .delete_user_api_key("key-owner", "key-to-delete") - .await - .expect("owner-scoped delete should succeed")); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM api_keys WHERE id = 'key-to-delete'", - ) - .fetch_one(&pool) - .await - .expect("API key count should load"), - 0 - ); - - let wallet = sqlx::query("SELECT api_key_id, status FROM wallets WHERE id = 'key-wallet'") - .fetch_one(&pool) - .await - .expect("wallet fact should remain"); - assert_eq!( - wallet - .try_get::, _>("api_key_id") - .expect("wallet API key id should decode") - .as_deref(), - Some("key-to-delete") - ); - assert_eq!( - wallet - .try_get::("status") - .expect("wallet status should decode"), - "disabled" - ); - - for table in [ - "request_candidates", - "video_tasks", - "usage", - "stats_daily_api_key", - ] { - let row = sqlx::query(&format!( - "SELECT api_key_id, api_key_name FROM {table} WHERE id = 'key-history-row'", - )) - .fetch_one(&pool) - .await - .unwrap_or_else(|error| panic!("{table} fact should remain: {error}")); - assert_eq!( - row.try_get::("api_key_id") - .expect("history API key id should decode"), - "key-to-delete", - "{table} API key id must remain stable" - ); - assert_eq!( - row.try_get::, _>("api_key_name") - .expect("history API key name should decode"), - None, - "{table} API key name must be removed" - ); - } - - let audit = sqlx::query( - "SELECT api_key_id, description, ip_address, user_agent, event_metadata, error_message FROM audit_logs WHERE id = 'key-audit'", - ) - .fetch_one(&pool) - .await - .expect("audit fact should remain"); - assert_eq!( - audit - .try_get::, _>("api_key_id") - .expect("audit API key id should decode") - .as_deref(), - Some("key-to-delete") - ); - assert_eq!( - audit - .try_get::("description") - .expect("audit description should decode"), - "deleted API key event" - ); - for column in [ - "ip_address", - "user_agent", - "event_metadata", - "error_message", - ] { - assert_eq!( - audit - .try_get::, _>(column) - .unwrap_or_else(|error| panic!("audit {column} should decode: {error}")), - None, - "audit {column} must be removed" - ); - } - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM api_key_provider_mappings WHERE id = 'key-mapping'", - ) - .fetch_one(&pool) - .await - .expect("mapping count should load"), - 0 - ); - - let order_gateway_response: Option = sqlx::query_scalar( - "SELECT gateway_response FROM payment_orders WHERE id = 'key-order'", - ) - .fetch_one(&pool) - .await - .expect("payment order should remain"); - assert_eq!(order_gateway_response, None); - let callback = sqlx::query( - "SELECT payload, error_message FROM payment_callbacks WHERE id = 'key-callback'", - ) - .fetch_one(&pool) - .await - .expect("payment callback should remain"); - assert_eq!( - callback - .try_get::, _>("payload") - .expect("callback payload should decode"), - None - ); - assert_eq!( - callback - .try_get::, _>("error_message") - .expect("callback error should decode"), - None - ); - } - - async fn seed_auth_api_key_rows(pool: &sqlx::SqlitePool) { - seed_auth_user(pool).await; - sqlx::query( - r#" -INSERT INTO api_keys ( - id, user_id, key_hash, key_encrypted, name, allowed_providers, - allowed_api_formats, allowed_models, rate_limit, concurrent_limit, - force_capabilities, is_active, expires_at, auto_delete_on_expiry, - total_requests, total_tokens, total_cost_usd, last_used_at, created_at, - updated_at, is_standalone -) VALUES - ( - 'key-user', 'user-1', 'hash-user', 'enc-user', 'User Key', '["openai"]', - '["openai:chat"]', '["gpt-4.1"]', 30, 2, '{"cache":true}', 1, 200, 0, - 123, 456, 1.25, 10, 1, 2, 0 - ), - ( - 'key-standalone', 'user-1', 'hash-standalone', 'enc-standalone', 'Standalone', NULL, - NULL, NULL, NULL, NULL, NULL, 1, NULL, 0, 0, 0, 0, NULL, 3, 4, 1 - ) -"#, - ) - .execute(pool) - .await - .expect("api keys should seed"); - } - - async fn seed_auth_user(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers, allowed_api_formats, allowed_models, rate_limit, - is_active, is_deleted, created_at, updated_at -) VALUES ( - 'user-1', 'user@example.com', 1, 'alice', NULL, 'user', 'local', - '["openai"]', '["openai:chat"]', '["gpt-4.1"]', 60, 1, 0, 1, 1 -) -"#, - ) - .execute(pool) - .await - .expect("user should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/auth_modules.rs b/crates/aether-data/adapters/sqlite/src/auth_modules.rs deleted file mode 100644 index 2c1261888..000000000 --- a/crates/aether-data/adapters/sqlite/src/auth_modules.rs +++ /dev/null @@ -1,597 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::auth_modules::*; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const OAUTH_PROVIDER_COLUMNS: &str = r#" -SELECT - provider_type, - display_name, - client_id, - client_secret_encrypted, - redirect_uri -FROM oauth_providers -"#; - -const LDAP_CONFIG_COLUMNS: &str = r#" -SELECT - server_url, - bind_dn, - bind_password_encrypted, - base_dn, - user_search_filter, - username_attr, - email_attr, - display_name_attr, - is_enabled, - is_exclusive, - use_starttls, - connect_timeout -FROM ldap_configs -"#; - -const UPDATE_LDAP_CONFIG_PRESERVE_PASSWORD_SQL: &str = r#" -UPDATE ldap_configs -SET - server_url = ?, - bind_dn = ?, - base_dn = ?, - user_search_filter = ?, - username_attr = ?, - email_attr = ?, - display_name_attr = ?, - is_enabled = ?, - is_exclusive = ?, - use_starttls = ?, - connect_timeout = ?, - updated_at = MAX(updated_at + 1, ?) -WHERE singleton_key = 1 - AND server_url IS ? - AND bind_dn IS ? - AND bind_password_encrypted IS ? - AND base_dn IS ? - AND user_search_filter IS ? - AND username_attr IS ? - AND email_attr IS ? - AND display_name_attr IS ? - AND is_enabled IS ? - AND is_exclusive IS ? - AND use_starttls IS ? - AND connect_timeout IS ? -"#; - -const UPDATE_LDAP_CONFIG_REPLACE_PASSWORD_SQL: &str = r#" -UPDATE ldap_configs -SET - server_url = ?, - bind_dn = ?, - bind_password_encrypted = ?, - base_dn = ?, - user_search_filter = ?, - username_attr = ?, - email_attr = ?, - display_name_attr = ?, - is_enabled = ?, - is_exclusive = ?, - use_starttls = ?, - connect_timeout = ?, - updated_at = MAX(updated_at + 1, ?) -WHERE singleton_key = 1 - AND server_url IS ? - AND bind_dn IS ? - AND bind_password_encrypted IS ? - AND base_dn IS ? - AND user_search_filter IS ? - AND username_attr IS ? - AND email_attr IS ? - AND display_name_attr IS ? - AND is_enabled IS ? - AND is_exclusive IS ? - AND use_starttls IS ? - AND connect_timeout IS ? -"#; - -const INSERT_LDAP_CONFIG_SQL: &str = r#" -INSERT INTO ldap_configs ( - singleton_key, - server_url, - bind_dn, - bind_password_encrypted, - base_dn, - user_search_filter, - username_attr, - email_attr, - display_name_attr, - is_enabled, - is_exclusive, - use_starttls, - connect_timeout, - created_at, - updated_at -) VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#; - -#[derive(Debug, Clone)] -pub struct SqliteAuthModuleReadRepository { - pool: SqlitePool, -} - -impl SqliteAuthModuleReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[derive(Debug, Clone)] -pub struct SqliteAuthModuleRepository { - pool: SqlitePool, -} - -impl SqliteAuthModuleRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -async fn list_enabled_oauth_providers( - pool: &SqlitePool, -) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(OAUTH_PROVIDER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "is_enabled", true); - builder.push(" ORDER BY provider_type ASC"); - let rows = builder.build().fetch_all(pool).await.map_sql_err()?; - rows.iter().map(map_oauth_row).collect() -} - -async fn get_ldap_config( - pool: &SqlitePool, -) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(LDAP_CONFIG_COLUMNS); - builder.push(" WHERE singleton_key = 1"); - let row = builder.build().fetch_optional(pool).await.map_sql_err()?; - row.as_ref().map(map_ldap_row).transpose() -} - -#[async_trait] -impl AuthModuleReadRepository for SqliteAuthModuleReadRepository { - async fn list_enabled_oauth_providers( - &self, - ) -> Result, DataLayerError> { - list_enabled_oauth_providers(&self.pool).await - } - - async fn get_ldap_config(&self) -> Result, DataLayerError> { - get_ldap_config(&self.pool).await - } -} - -#[async_trait] -impl AuthModuleReadRepository for SqliteAuthModuleRepository { - async fn list_enabled_oauth_providers( - &self, - ) -> Result, DataLayerError> { - list_enabled_oauth_providers(&self.pool).await - } - - async fn get_ldap_config(&self) -> Result, DataLayerError> { - get_ldap_config(&self.pool).await - } -} - -#[async_trait] -impl AuthModuleWriteRepository for SqliteAuthModuleRepository { - async fn compare_and_swap_ldap_config( - &self, - expected: Option<&StoredLdapModuleConfig>, - replacement: &StoredLdapModuleConfig, - bind_password_update: &LdapBindPasswordUpdate, - ) -> Result { - let persisted = - ldap_config_after_password_update(expected, replacement, bind_password_update)?; - let now = now_unix_secs(); - let Some(expected) = expected else { - let insert = sqlx::query(INSERT_LDAP_CONFIG_SQL) - .bind(&persisted.server_url) - .bind(&persisted.bind_dn) - .bind(persisted.bind_password_encrypted.as_deref()) - .bind(&persisted.base_dn) - .bind(persisted.user_search_filter.as_deref()) - .bind(persisted.username_attr.as_deref()) - .bind(persisted.email_attr.as_deref()) - .bind(persisted.display_name_attr.as_deref()) - .bind(persisted.is_enabled) - .bind(persisted.is_exclusive) - .bind(persisted.use_starttls) - .bind(persisted.connect_timeout) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await; - return match insert { - Ok(result) if result.rows_affected() == 1 => { - Ok(CompareAndSwapLdapConfigResult::Applied(persisted)) - } - Ok(_) => Ok(CompareAndSwapLdapConfigResult::Conflict), - Err(error) - if error - .as_database_error() - .is_some_and(|error| error.is_unique_violation()) => - { - Ok(CompareAndSwapLdapConfigResult::Conflict) - } - Err(error) => Err(DataLayerError::sql(error)), - }; - }; - - let rows_affected = match bind_password_update { - LdapBindPasswordUpdate::Preserve => { - sqlx::query(UPDATE_LDAP_CONFIG_PRESERVE_PASSWORD_SQL) - .bind(&replacement.server_url) - .bind(&replacement.bind_dn) - .bind(&replacement.base_dn) - .bind(replacement.user_search_filter.as_deref()) - .bind(replacement.username_attr.as_deref()) - .bind(replacement.email_attr.as_deref()) - .bind(replacement.display_name_attr.as_deref()) - .bind(replacement.is_enabled) - .bind(replacement.is_exclusive) - .bind(replacement.use_starttls) - .bind(replacement.connect_timeout) - .bind(now as i64) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - } - LdapBindPasswordUpdate::Set(_) | LdapBindPasswordUpdate::Clear => { - sqlx::query(UPDATE_LDAP_CONFIG_REPLACE_PASSWORD_SQL) - .bind(&replacement.server_url) - .bind(&replacement.bind_dn) - .bind(persisted.bind_password_encrypted.as_deref()) - .bind(&replacement.base_dn) - .bind(replacement.user_search_filter.as_deref()) - .bind(replacement.username_attr.as_deref()) - .bind(replacement.email_attr.as_deref()) - .bind(replacement.display_name_attr.as_deref()) - .bind(replacement.is_enabled) - .bind(replacement.is_exclusive) - .bind(replacement.use_starttls) - .bind(replacement.connect_timeout) - .bind(now as i64) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - } - }; - if rows_affected == 1 { - Ok(CompareAndSwapLdapConfigResult::Applied(persisted)) - } else { - Ok(CompareAndSwapLdapConfigResult::Conflict) - } - } - - async fn delete_ldap_config_if_matches( - &self, - expected: &StoredLdapModuleConfig, - ) -> Result { - let rows_affected = sqlx::query( - r#" -DELETE FROM ldap_configs -WHERE singleton_key = 1 - AND server_url IS ? - AND bind_dn IS ? - AND bind_password_encrypted IS ? - AND base_dn IS ? - AND user_search_filter IS ? - AND username_attr IS ? - AND email_attr IS ? - AND display_name_attr IS ? - AND is_enabled IS ? - AND is_exclusive IS ? - AND use_starttls IS ? - AND connect_timeout IS ? -"#, - ) - .bind(&expected.server_url) - .bind(&expected.bind_dn) - .bind(expected.bind_password_encrypted.as_deref()) - .bind(&expected.base_dn) - .bind(expected.user_search_filter.as_deref()) - .bind(expected.username_attr.as_deref()) - .bind(expected.email_attr.as_deref()) - .bind(expected.display_name_attr.as_deref()) - .bind(expected.is_enabled) - .bind(expected.is_exclusive) - .bind(expected.use_starttls) - .bind(expected.connect_timeout) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - async fn compare_and_swap_ldap_bind_password( - &self, - expected: &str, - replacement: &str, - ) -> Result { - let rows_affected = sqlx::query( - r#" -UPDATE ldap_configs -SET bind_password_encrypted = ?, updated_at = MAX(updated_at + 1, ?) -WHERE singleton_key = 1 - AND bind_password_encrypted = ? -"#, - ) - .bind(replacement) - .bind(now_unix_secs() as i64) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } -} - -fn ldap_config_after_password_update( - expected: Option<&StoredLdapModuleConfig>, - replacement: &StoredLdapModuleConfig, - bind_password_update: &LdapBindPasswordUpdate, -) -> Result { - let bind_password_encrypted = match bind_password_update { - LdapBindPasswordUpdate::Preserve => expected - .ok_or_else(|| { - DataLayerError::InvalidConfiguration( - "LDAP bind password cannot be preserved while creating the singleton" - .to_string(), - ) - })? - .bind_password_encrypted - .clone(), - LdapBindPasswordUpdate::Set(ciphertext) => Some(ciphertext.clone()), - LdapBindPasswordUpdate::Clear => None, - }; - Ok(StoredLdapModuleConfig { - bind_password_encrypted, - ..replacement.clone() - }) -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn map_oauth_row(row: &SqliteRow) -> Result { - StoredOAuthProviderModuleConfig::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("client_id").map_sql_err()?, - row.try_get("client_secret_encrypted").map_sql_err()?, - row.try_get("redirect_uri").map_sql_err()?, - ) -} - -fn map_ldap_row(row: &SqliteRow) -> Result { - Ok(StoredLdapModuleConfig { - server_url: row.try_get("server_url").map_sql_err()?, - bind_dn: row.try_get("bind_dn").map_sql_err()?, - bind_password_encrypted: row.try_get("bind_password_encrypted").map_sql_err()?, - base_dn: row.try_get("base_dn").map_sql_err()?, - user_search_filter: row.try_get("user_search_filter").map_sql_err()?, - username_attr: row.try_get("username_attr").map_sql_err()?, - email_attr: row.try_get("email_attr").map_sql_err()?, - display_name_attr: row.try_get("display_name_attr").map_sql_err()?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - is_exclusive: row.try_get("is_exclusive").map_sql_err()?, - use_starttls: row.try_get("use_starttls").map_sql_err()?, - connect_timeout: row.try_get("connect_timeout").map_sql_err()?, - }) -} - -#[cfg(test)] -mod tests { - use super::SqliteAuthModuleRepository; - use aether_data_contracts::repository::auth_modules::{ - AuthModuleReadRepository, AuthModuleWriteRepository, CompareAndSwapLdapConfigResult, - LdapBindPasswordUpdate, StoredLdapModuleConfig, - }; - - use crate::run_migrations; - - #[tokio::test] - async fn sqlite_repository_reads_and_writes_auth_module_configs() { - 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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES - ('github', 'GitHub', 'github-client', 'https://github.example.com/callback', - 'https://frontend.example.com/callback', 1, 1, 1), - ('disabled', 'Disabled', 'disabled-client', 'https://disabled.example.com/callback', - 'https://frontend.example.com/callback', 0, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("oauth providers should seed"); - - let repository = SqliteAuthModuleRepository::new(pool); - let oauth = repository - .list_enabled_oauth_providers() - .await - .expect("oauth providers should load"); - assert_eq!(oauth.len(), 1); - assert_eq!(oauth[0].provider_type, "github"); - - let ldap = StoredLdapModuleConfig { - server_url: "ldaps://ldap.example.com".to_string(), - bind_dn: "cn=admin,dc=example,dc=com".to_string(), - bind_password_encrypted: None, - base_dn: "dc=example,dc=com".to_string(), - user_search_filter: Some("(uid={username})".to_string()), - username_attr: Some("uid".to_string()), - email_attr: Some("mail".to_string()), - display_name_attr: Some("displayName".to_string()), - is_enabled: true, - is_exclusive: false, - use_starttls: true, - connect_timeout: Some(10), - }; - let stored = repository - .compare_and_swap_ldap_config( - None, - &ldap, - &LdapBindPasswordUpdate::Set("encrypted-password".to_string()), - ) - .await - .expect("ldap create CAS should execute"); - let CompareAndSwapLdapConfigResult::Applied(stored) = stored else { - panic!("initial LDAP create should apply"); - }; - assert_eq!(stored.server_url, "ldaps://ldap.example.com"); - assert_eq!( - stored.bind_password_encrypted.as_deref(), - Some("encrypted-password") - ); - - let competing_create = repository - .compare_and_swap_ldap_config( - None, - &ldap, - &LdapBindPasswordUpdate::Set("competing-password".to_string()), - ) - .await - .expect("competing LDAP create should execute"); - assert_eq!(competing_create, CompareAndSwapLdapConfigResult::Conflict); - - let preserve_replacement = StoredLdapModuleConfig { - server_url: "ldap://ldap.example.com".to_string(), - bind_password_encrypted: Some("stale-password-must-not-be-written".to_string()), - ..stored.clone() - }; - let updated = repository - .compare_and_swap_ldap_config( - Some(&stored), - &preserve_replacement, - &LdapBindPasswordUpdate::Preserve, - ) - .await - .expect("LDAP preserve CAS should execute"); - let CompareAndSwapLdapConfigResult::Applied(updated) = updated else { - panic!("fresh LDAP snapshot should update"); - }; - assert_eq!(updated.server_url, "ldap://ldap.example.com"); - assert_eq!( - updated.bind_password_encrypted.as_deref(), - Some("encrypted-password") - ); - - assert!(repository - .compare_and_swap_ldap_bind_password("encrypted-password", "rotated-password") - .await - .expect("LDAP password rotation should execute")); - let stale = repository - .compare_and_swap_ldap_config( - Some(&updated), - &updated, - &LdapBindPasswordUpdate::Preserve, - ) - .await - .expect("stale LDAP CAS should execute"); - assert_eq!(stale, CompareAndSwapLdapConfigResult::Conflict); - let rotated = repository - .get_ldap_config() - .await - .expect("rotated LDAP config should load") - .expect("rotated LDAP config should exist"); - assert_eq!( - rotated.bind_password_encrypted.as_deref(), - Some("rotated-password") - ); - - let set = repository - .compare_and_swap_ldap_config( - Some(&rotated), - &rotated, - &LdapBindPasswordUpdate::Set("replacement-password".to_string()), - ) - .await - .expect("LDAP password set CAS should execute"); - let CompareAndSwapLdapConfigResult::Applied(set) = set else { - panic!("fresh LDAP password set should apply"); - }; - let cleared = repository - .compare_and_swap_ldap_config(Some(&set), &set, &LdapBindPasswordUpdate::Clear) - .await - .expect("LDAP password clear CAS should execute"); - let CompareAndSwapLdapConfigResult::Applied(cleared) = cleared else { - panic!("fresh LDAP password clear should apply"); - }; - assert!(cleared.bind_password_encrypted.is_none()); - - let mismatched = StoredLdapModuleConfig { - base_dn: "dc=changed,dc=example".to_string(), - ..cleared.clone() - }; - assert!(!repository - .delete_ldap_config_if_matches(&mismatched) - .await - .expect("mismatched LDAP delete should execute")); - assert!(repository - .get_ldap_config() - .await - .expect("LDAP config should remain readable") - .is_some()); - assert!(repository - .delete_ldap_config_if_matches(&cleared) - .await - .expect("matching LDAP delete should execute")); - assert!(repository - .get_ldap_config() - .await - .expect("LDAP config should remain readable") - .is_none()); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/background_tasks.rs b/crates/aether-data/adapters/sqlite/src/background_tasks.rs deleted file mode 100644 index 624937200..000000000 --- a/crates/aether-data/adapters/sqlite/src/background_tasks.rs +++ /dev/null @@ -1,537 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::background_tasks::*; -use aether_data_query::{ - push_ci_contains, push_eq, push_limit, push_limit_offset, SqlDialect, WhereClause, -}; - -use crate::error::SqlResultExt; -use crate::{DataLayerError, SqlitePool}; - -const RUN_COLUMNS: &str = r#" -SELECT - id, - task_key, - kind, - "trigger", - status, - attempt, - max_attempts, - owner_instance, - progress_percent, - progress_message, - payload_json, - result_json, - error_message, - cancel_requested, - created_by, - created_at_unix_secs, - started_at_unix_secs, - finished_at_unix_secs, - updated_at_unix_secs -FROM background_task_runs -"#; - -const EVENT_COLUMNS: &str = r#" -SELECT - id, - run_id, - event_type, - message, - payload_json, - created_at_unix_secs -FROM background_task_events -"#; - -#[derive(Debug, Clone)] -pub struct SqliteBackgroundTaskRepository { - pool: SqlitePool, -} - -impl SqliteBackgroundTaskRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - fn apply_run_filter(builder: &mut QueryBuilder<'_, Sqlite>, query: &BackgroundTaskListQuery) { - let mut where_clause = WhereClause::new(); - if let Some(kind) = query.kind { - push_eq(builder, &mut where_clause, "kind", kind.as_database()); - } - if let Some(status) = query.status { - push_eq(builder, &mut where_clause, "status", status.as_database()); - } - if let Some(trigger) = query.trigger.as_deref() { - push_eq( - builder, - &mut where_clause, - &SqlDialect::Sqlite.quote_ident("trigger"), - trigger.to_string(), - ); - } - if let Some(task_key_substring) = query.task_key_substring.as_deref() { - push_ci_contains( - builder, - &mut where_clause, - SqlDialect::Sqlite, - "task_key", - task_key_substring, - ); - } - } -} - -#[async_trait] -impl BackgroundTaskReadRepository for SqliteBackgroundTaskRepository { - async fn find_run( - &self, - run_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(RUN_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", run_id.to_string()); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_run_row).transpose() - } - - async fn list_runs( - &self, - query: &BackgroundTaskListQuery, - ) -> Result { - let limit = query.limit.max(1); - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(id) AS total FROM background_task_runs"); - Self::apply_run_filter(&mut count_builder, query); - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let mut builder = QueryBuilder::::new(RUN_COLUMNS); - Self::apply_run_filter(&mut builder, query); - builder.push(" ORDER BY created_at_unix_secs DESC, updated_at_unix_secs DESC"); - push_limit_offset( - &mut builder, - i64_from_usize(limit, "run limit")?, - i64_from_usize(query.offset, "run offset")?, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let items = rows - .iter() - .map(map_run_row) - .collect::, _>>()?; - Ok(StoredBackgroundTaskRunPage { - items, - total: usize::try_from(total).unwrap_or_default(), - }) - } - - async fn list_events( - &self, - run_id: &str, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - let limit = limit.max(1); - let mut builder = QueryBuilder::::new(EVENT_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "run_id", - run_id.to_string(), - ); - builder.push(" ORDER BY created_at_unix_secs ASC, id ASC"); - push_limit_offset( - &mut builder, - i64_from_usize(limit, "event limit")?, - i64_from_usize(offset, "event offset")?, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_event_row).collect() - } - - async fn summarize_runs(&self) -> Result { - let total = sqlx::query_scalar::<_, i64>("SELECT COUNT(id) FROM background_task_runs") - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let running_count = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(id) FROM background_task_runs WHERE status = 'running'", - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let status_rows = sqlx::query( - "SELECT status, COUNT(id) AS total FROM background_task_runs GROUP BY status", - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let kind_rows = - sqlx::query("SELECT kind, COUNT(id) AS total FROM background_task_runs GROUP BY kind") - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - let mut by_status = std::collections::BTreeMap::new(); - for row in status_rows { - let key: String = row.try_get("status").map_sql_err()?; - let count: i64 = row.try_get("total").map_sql_err()?; - by_status.insert(key, u64::try_from(count).unwrap_or_default()); - } - let mut by_kind = std::collections::BTreeMap::new(); - for row in kind_rows { - let key: String = row.try_get("kind").map_sql_err()?; - let count: i64 = row.try_get("total").map_sql_err()?; - by_kind.insert(key, u64::try_from(count).unwrap_or_default()); - } - - Ok(BackgroundTaskSummary { - total: u64::try_from(total).unwrap_or_default(), - running_count: u64::try_from(running_count).unwrap_or_default(), - by_status, - by_kind, - }) - } -} - -#[async_trait] -impl BackgroundTaskWriteRepository for SqliteBackgroundTaskRepository { - async fn upsert_run( - &self, - mut run: UpsertBackgroundTaskRun, - ) -> Result { - run.sanitize_for_persistence(); - run.validate()?; - sqlx::query( - r#" -INSERT INTO background_task_runs ( - id, - task_key, - kind, - "trigger", - status, - attempt, - max_attempts, - owner_instance, - progress_percent, - progress_message, - payload_json, - result_json, - error_message, - cancel_requested, - created_by, - created_at_unix_secs, - started_at_unix_secs, - finished_at_unix_secs, - updated_at_unix_secs -) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) -ON CONFLICT(id) DO UPDATE SET - task_key = excluded.task_key, - kind = excluded.kind, - "trigger" = excluded."trigger", - status = excluded.status, - attempt = excluded.attempt, - max_attempts = excluded.max_attempts, - owner_instance = excluded.owner_instance, - progress_percent = excluded.progress_percent, - progress_message = excluded.progress_message, - payload_json = excluded.payload_json, - result_json = excluded.result_json, - error_message = excluded.error_message, - cancel_requested = excluded.cancel_requested, - created_by = excluded.created_by, - created_at_unix_secs = excluded.created_at_unix_secs, - started_at_unix_secs = excluded.started_at_unix_secs, - finished_at_unix_secs = excluded.finished_at_unix_secs, - updated_at_unix_secs = excluded.updated_at_unix_secs -"#, - ) - .bind(&run.id) - .bind(&run.task_key) - .bind(run.kind.as_database()) - .bind(&run.trigger) - .bind(run.status.as_database()) - .bind(i64::from(run.attempt)) - .bind(i64::from(run.max_attempts)) - .bind(run.owner_instance.as_deref()) - .bind(i32::from(run.progress_percent)) - .bind(run.progress_message.as_deref()) - .bind(run.payload_json.as_ref().map(serde_json::Value::to_string)) - .bind(run.result_json.as_ref().map(serde_json::Value::to_string)) - .bind(run.error_message.as_deref()) - .bind(run.cancel_requested) - .bind(run.created_by.as_deref()) - .bind(u64_to_i64( - run.created_at_unix_secs, - "created_at_unix_secs", - )?) - .bind(run.started_at_unix_secs.map(|value| value as i64)) - .bind(run.finished_at_unix_secs.map(|value| value as i64)) - .bind(u64_to_i64( - run.updated_at_unix_secs, - "updated_at_unix_secs", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.find_run(&run.id).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("background task run missing after upsert".to_string()) - }) - } - - async fn request_cancel( - &self, - run_id: &str, - updated_at_unix_secs: u64, - ) -> Result { - let affected = sqlx::query( - "UPDATE background_task_runs SET cancel_requested = 1, updated_at_unix_secs = ? WHERE id = ?", - ) - .bind(u64_to_i64(updated_at_unix_secs, "updated_at_unix_secs")?) - .bind(run_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(affected > 0) - } - - async fn upsert_event( - &self, - mut event: UpsertBackgroundTaskEvent, - ) -> Result { - event.sanitize_for_persistence(); - event.validate()?; - sqlx::query( - r#" -INSERT INTO background_task_events ( - id, run_id, event_type, message, payload_json, created_at_unix_secs -) VALUES (?, ?, ?, ?, ?, ?) -ON CONFLICT(id) DO UPDATE SET - run_id = excluded.run_id, - event_type = excluded.event_type, - message = excluded.message, - payload_json = excluded.payload_json, - created_at_unix_secs = excluded.created_at_unix_secs -"#, - ) - .bind(&event.id) - .bind(&event.run_id) - .bind(&event.event_type) - .bind(&event.message) - .bind( - event - .payload_json - .as_ref() - .map(serde_json::Value::to_string), - ) - .bind(u64_to_i64( - event.created_at_unix_secs, - "created_at_unix_secs", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - - let row = sqlx::query(&format!("{EVENT_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(&event.id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - map_event_row(&row) - } -} - -fn map_run_row(row: &SqliteRow) -> Result { - let kind: String = row.try_get("kind").map_sql_err()?; - let status: String = row.try_get("status").map_sql_err()?; - let attempt: i64 = row.try_get("attempt").map_sql_err()?; - let max_attempts: i64 = row.try_get("max_attempts").map_sql_err()?; - let progress_percent: i32 = row.try_get("progress_percent").map_sql_err()?; - let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?; - let started_at_unix_secs: Option = row.try_get("started_at_unix_secs").map_sql_err()?; - let finished_at_unix_secs: Option = row.try_get("finished_at_unix_secs").map_sql_err()?; - let updated_at_unix_secs: i64 = row.try_get("updated_at_unix_secs").map_sql_err()?; - - let mut run = StoredBackgroundTaskRun { - id: row.try_get("id").map_sql_err()?, - task_key: row.try_get("task_key").map_sql_err()?, - kind: BackgroundTaskKind::from_database(&kind)?, - trigger: row.try_get("trigger").map_sql_err()?, - status: BackgroundTaskStatus::from_database(&status)?, - attempt: u32::try_from(attempt).unwrap_or_default(), - max_attempts: u32::try_from(max_attempts).unwrap_or_default(), - owner_instance: row.try_get("owner_instance").map_sql_err()?, - progress_percent: u16::try_from(progress_percent).unwrap_or_default(), - progress_message: row.try_get("progress_message").map_sql_err()?, - payload_json: parse_optional_json(row.try_get("payload_json").map_sql_err()?)?, - result_json: parse_optional_json(row.try_get("result_json").map_sql_err()?)?, - error_message: row.try_get("error_message").map_sql_err()?, - cancel_requested: row.try_get("cancel_requested").map_sql_err()?, - created_by: row.try_get("created_by").map_sql_err()?, - created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(), - started_at_unix_secs: started_at_unix_secs.and_then(|value| u64::try_from(value).ok()), - finished_at_unix_secs: finished_at_unix_secs.and_then(|value| u64::try_from(value).ok()), - updated_at_unix_secs: u64::try_from(updated_at_unix_secs).unwrap_or_default(), - }; - run.sanitize_persisted_data(); - Ok(run) -} - -fn map_event_row(row: &SqliteRow) -> Result { - let created_at_unix_secs: i64 = row.try_get("created_at_unix_secs").map_sql_err()?; - let mut event = StoredBackgroundTaskEvent { - id: row.try_get("id").map_sql_err()?, - run_id: row.try_get("run_id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - message: row.try_get("message").map_sql_err()?, - payload_json: parse_optional_json(row.try_get("payload_json").map_sql_err()?)?, - created_at_unix_secs: u64::try_from(created_at_unix_secs).unwrap_or_default(), - }; - event.sanitize_persisted_data(); - Ok(event) -} - -fn parse_optional_json(value: Option) -> Result, DataLayerError> { - value - .map(|raw| { - serde_json::from_str::(&raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "invalid background task json payload: {err}" - )) - }) - }) - .transpose() -} - -fn i64_from_usize(value: usize, label: &str) -> Result { - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}")) - }) -} - -fn u64_to_i64(value: u64, label: &str) -> Result { - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("background task {label} overflow: {value}")) - }) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - use crate::run_migrations; - - #[tokio::test] - async fn sqlite_background_task_repository_round_trips_runs_and_events() { - 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 = SqliteBackgroundTaskRepository::new(pool); - let run = repository - .upsert_run(UpsertBackgroundTaskRun { - id: "run-1".to_string(), - task_key: "usage.cleanup".to_string(), - kind: BackgroundTaskKind::Scheduled, - trigger: "timer".to_string(), - status: BackgroundTaskStatus::Queued, - attempt: 0, - max_attempts: 3, - owner_instance: None, - progress_percent: 0, - progress_message: None, - payload_json: Some(json!({ - "partition": 7, - "refresh_token": "sensitive-refresh-token", - "nested": {"authorization": "Bearer sensitive"} - })), - result_json: None, - error_message: None, - cancel_requested: false, - created_by: Some("scheduler".to_string()), - created_at_unix_secs: 10, - started_at_unix_secs: None, - finished_at_unix_secs: None, - updated_at_unix_secs: 10, - }) - .await - .expect("background task run should upsert"); - assert_eq!(run.payload_json, Some(json!({"partition": 7}))); - - repository - .upsert_event(UpsertBackgroundTaskEvent { - id: "event-1".to_string(), - run_id: run.id.clone(), - event_type: "queued".to_string(), - message: "task queued".to_string(), - payload_json: Some(json!({ - "error_code": "provider_delete_failed", - "error": "sensitive provider detail" - })), - created_at_unix_secs: 11, - }) - .await - .expect("background task event should upsert"); - - let page = repository - .list_runs(&BackgroundTaskListQuery { - task_key_substring: Some("cleanup".to_string()), - kind: Some(BackgroundTaskKind::Scheduled), - status: Some(BackgroundTaskStatus::Queued), - trigger: Some("timer".to_string()), - offset: 0, - limit: 10, - }) - .await - .expect("background task runs should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].id, "run-1"); - - let events = repository - .list_events("run-1", 0, 10) - .await - .expect("background task events should list"); - assert_eq!(events.len(), 1); - assert_eq!(events[0].message, "queued"); - assert_eq!( - events[0].payload_json, - Some(json!({"error_code": "provider_delete_failed"})) - ); - - assert!(repository - .request_cancel("run-1", 20) - .await - .expect("background task cancellation should update")); - let cancelled = repository - .find_run("run-1") - .await - .expect("background task run should load") - .expect("background task run should exist"); - assert!(cancelled.cancel_requested); - assert_eq!(cancelled.updated_at_unix_secs, 20); - - let summary = repository - .summarize_runs() - .await - .expect("background task summary should load"); - assert_eq!(summary.total, 1); - assert_eq!(summary.by_status.get("queued"), Some(&1)); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/billing.rs b/crates/aether-data/adapters/sqlite/src/billing.rs deleted file mode 100644 index e13c08abf..000000000 --- a/crates/aether-data/adapters/sqlite/src/billing.rs +++ /dev/null @@ -1,2144 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, Row}; - -use aether_data_contracts::repository::billing::{ - AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, - AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput, - BillingPlanRecord, BillingPlanWriteInput, BillingReadRepository, - PaymentGatewayConfigCasWriteInput, PaymentGatewayConfigRecord, PaymentGatewayConfigWriteInput, - PaymentGatewaySecretCasUpdate, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord, - UserPlanEntitlementRecord, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, SqlitePool}; - -const MODEL_CONTEXT_COLUMNS: &str = r#" -SELECT - p.id AS provider_id, - p.billing_type AS provider_billing_type, - pak.id AS provider_api_key_id, - pak.rate_multipliers AS provider_api_key_rate_multipliers, - pak.cache_ttl_minutes AS provider_api_key_cache_ttl_minutes, - gm.id AS global_model_id, - gm.name AS global_model_name, - gm.config AS global_model_config, - CAST(gm.default_price_per_request AS REAL) AS default_price_per_request, - gm.default_tiered_pricing AS default_tiered_pricing, - m.id AS model_id, - m.provider_model_name AS model_provider_model_name, - m.config AS model_config, - CAST(m.price_per_request AS REAL) AS model_price_per_request, - m.tiered_pricing AS model_tiered_pricing, - m.provider_model_mappings AS provider_model_mappings, - m.is_available AS model_is_available, - m.created_at AS model_created_at -FROM providers p -"#; - -#[derive(Debug, Clone)] -pub struct SqliteBillingReadRepository { - pool: SqlitePool, -} - -impl SqliteBillingReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl BillingReadRepository for SqliteBillingReadRepository { - async fn find_model_context( - &self, - provider_id: &str, - provider_api_key_id: Option<&str>, - global_model_name: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{MODEL_CONTEXT_COLUMNS} -INNER JOIN global_models gm - ON gm.is_active = 1 -LEFT JOIN models m - ON m.global_model_id = gm.id - AND m.provider_id = p.id - AND m.is_active = 1 -LEFT JOIN provider_api_keys pak - ON pak.id = ? - AND pak.provider_id = p.id -WHERE p.id = ? - AND ( - gm.name = ? - OR m.provider_model_name = ? - OR m.provider_model_mappings IS NOT NULL - ) -"# - )) - .bind(provider_api_key_id) - .bind(provider_id) - .bind(global_model_name) - .bind(global_model_name) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - rows.iter() - .filter_map(|row| match_rank(row, global_model_name).transpose()) - .collect::, _>>()? - .into_iter() - .min_by_key(|candidate| { - ( - candidate.rank, - !candidate.is_available, - candidate.pricing_rank, - candidate.created_at, - ) - }) - .map(|candidate| candidate.context) - .transpose() - } - - async fn find_model_context_by_model_id( - &self, - provider_id: &str, - provider_api_key_id: Option<&str>, - model_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - r#" -{MODEL_CONTEXT_COLUMNS} -INNER JOIN models m - ON m.id = ? - AND m.provider_id = p.id - AND m.is_active = 1 -INNER JOIN global_models gm - ON gm.id = m.global_model_id - AND gm.is_active = 1 -LEFT JOIN provider_api_keys pak - ON pak.id = ? - AND pak.provider_id = p.id -WHERE p.id = ? -LIMIT 1 -"# - )) - .bind(model_id) - .bind(provider_api_key_id) - .bind(provider_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_row).transpose() - } - - async fn admin_billing_enabled_default_value_exists( - &self, - api_format: &str, - task_type: &str, - dimension_name: &str, - existing_id: Option<&str>, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM dimension_collectors -WHERE api_format = ? - AND task_type = ? - AND dimension_name = ? - AND is_enabled = 1 - AND default_value IS NOT NULL - AND (? IS NULL OR id <> ?) - "#, - ) - .bind(api_format) - .bind(task_type) - .bind(dimension_name) - .bind(existing_id) - .bind(existing_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(Some(read_count_sqlite(&row)? > 0)) - } - - async fn create_admin_billing_rule( - &self, - input: &AdminBillingRuleWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -INSERT INTO billing_rules ( - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&id) - .bind(&input.name) - .bind(&input.task_type) - .bind(input.global_model_id.as_deref()) - .bind(input.model_id.as_deref()) - .bind(&input.expression) - .bind(json_to_string(&input.variables)?) - .bind(json_to_string(&input.dimension_mappings)?) - .bind(input.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - if let Err(err) = result { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))); - } - match find_admin_billing_rule_sqlite(&self.pool, &id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing rule missing".to_string(), - )), - } - } - - async fn list_admin_billing_rules( - &self, - task_type: Option<&str>, - is_enabled: Option, - page: u32, - page_size: u32, - ) -> Result, u64)>, DataLayerError> { - let total_row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM billing_rules -WHERE (? IS NULL OR task_type = ?) - AND (? IS NULL OR is_enabled = ?) - "#, - ) - .bind(task_type) - .bind(task_type) - .bind(is_enabled) - .bind(is_enabled) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total = read_count_sqlite(&total_row)?; - let offset = u64::from(page.saturating_sub(1) * page_size); - let rows = sqlx::query( - r#" -SELECT - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM billing_rules -WHERE (? IS NULL OR task_type = ?) - AND (? IS NULL OR is_enabled = ?) -ORDER BY updated_at DESC, id DESC -LIMIT ? OFFSET ? - "#, - ) - .bind(task_type) - .bind(task_type) - .bind(is_enabled) - .bind(is_enabled) - .bind(i64::from(page_size)) - .bind( - i64::try_from(offset) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_admin_billing_rule_sqlite) - .collect::, _>>()?; - Ok(Some((items, total))) - } - - async fn find_admin_billing_rule( - &self, - rule_id: &str, - ) -> Result, DataLayerError> { - find_admin_billing_rule_sqlite(&self.pool, rule_id).await - } - - async fn update_admin_billing_rule( - &self, - rule_id: &str, - input: &AdminBillingRuleWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE billing_rules -SET name = ?, - task_type = ?, - global_model_id = ?, - model_id = ?, - expression = ?, - variables = ?, - dimension_mappings = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&input.name) - .bind(&input.task_type) - .bind(input.global_model_id.as_deref()) - .bind(input.model_id.as_deref()) - .bind(&input.expression) - .bind(json_to_string(&input.variables)?) - .bind(json_to_string(&input.dimension_mappings)?) - .bind(input.is_enabled) - .bind(current_unix_secs_i64()) - .bind(rule_id) - .execute(&self.pool) - .await; - let affected = match result { - Ok(result) => result.rows_affected(), - Err(err) => { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))) - } - }; - if affected == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match find_admin_billing_rule_sqlite(&self.pool, rule_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn create_admin_billing_collector( - &self, - input: &AdminBillingCollectorWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -INSERT INTO dimension_collectors ( - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&id) - .bind(&input.api_format) - .bind(&input.task_type) - .bind(&input.dimension_name) - .bind(&input.source_type) - .bind(input.source_path.as_deref()) - .bind(&input.value_type) - .bind(input.transform_expression.as_deref()) - .bind(input.default_value.as_deref()) - .bind(input.priority) - .bind(input.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - if let Err(err) = result { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))); - } - match find_admin_billing_collector_sqlite(&self.pool, &id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing collector missing".to_string(), - )), - } - } - - async fn list_admin_billing_collectors( - &self, - api_format: Option<&str>, - task_type: Option<&str>, - dimension_name: Option<&str>, - is_enabled: Option, - page: u32, - page_size: u32, - ) -> Result, u64)>, DataLayerError> { - let total_row = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM dimension_collectors -WHERE (? IS NULL OR api_format = ?) - AND (? IS NULL OR task_type = ?) - AND (? IS NULL OR dimension_name = ?) - AND (? IS NULL OR is_enabled = ?) - "#, - ) - .bind(api_format) - .bind(api_format) - .bind(task_type) - .bind(task_type) - .bind(dimension_name) - .bind(dimension_name) - .bind(is_enabled) - .bind(is_enabled) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total = read_count_sqlite(&total_row)?; - let offset = u64::from(page.saturating_sub(1) * page_size); - let rows = sqlx::query( - r#" -SELECT - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, - created_at AS created_at_unix_ms, updated_at AS updated_at_unix_secs -FROM dimension_collectors -WHERE (? IS NULL OR api_format = ?) - AND (? IS NULL OR task_type = ?) - AND (? IS NULL OR dimension_name = ?) - AND (? IS NULL OR is_enabled = ?) -ORDER BY updated_at DESC, priority DESC, id ASC -LIMIT ? OFFSET ? - "#, - ) - .bind(api_format) - .bind(api_format) - .bind(task_type) - .bind(task_type) - .bind(dimension_name) - .bind(dimension_name) - .bind(is_enabled) - .bind(is_enabled) - .bind(i64::from(page_size)) - .bind( - i64::try_from(offset) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_admin_billing_collector_sqlite) - .collect::, _>>()?; - Ok(Some((items, total))) - } - - async fn find_admin_billing_collector( - &self, - collector_id: &str, - ) -> Result, DataLayerError> { - find_admin_billing_collector_sqlite(&self.pool, collector_id).await - } - - async fn update_admin_billing_collector( - &self, - collector_id: &str, - input: &AdminBillingCollectorWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE dimension_collectors -SET api_format = ?, - task_type = ?, - dimension_name = ?, - source_type = ?, - source_path = ?, - value_type = ?, - transform_expression = ?, - default_value = ?, - priority = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&input.api_format) - .bind(&input.task_type) - .bind(&input.dimension_name) - .bind(&input.source_type) - .bind(input.source_path.as_deref()) - .bind(&input.value_type) - .bind(input.transform_expression.as_deref()) - .bind(input.default_value.as_deref()) - .bind(input.priority) - .bind(input.is_enabled) - .bind(current_unix_secs_i64()) - .bind(collector_id) - .execute(&self.pool) - .await; - let affected = match result { - Ok(result) => result.rows_affected(), - Err(err) => { - return Ok(AdminBillingMutationOutcome::Invalid(format!( - "Integrity error: {err}" - ))) - } - }; - if affected == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match find_admin_billing_collector_sqlite(&self.pool, collector_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn apply_admin_billing_preset( - &self, - preset: &str, - mode: &str, - collectors: &[AdminBillingCollectorWriteInput], - ) -> Result, DataLayerError> { - let mut created = 0_u64; - let mut updated = 0_u64; - let mut skipped = 0_u64; - let mut errors = Vec::new(); - - for collector in collectors { - let existing_id = match sqlx::query_scalar::<_, String>( - r#" -SELECT id -FROM dimension_collectors -WHERE api_format = ? - AND task_type = ? - AND dimension_name = ? - AND priority = ? - AND is_enabled = 1 -LIMIT 1 - "#, - ) - .bind(&collector.api_format) - .bind(&collector.task_type) - .bind(&collector.dimension_name) - .bind(collector.priority) - .fetch_optional(&self.pool) - .await - { - Ok(value) => value, - Err(err) => { - errors.push(format!( - "Failed to query collector: api_format={} task_type={} dim={}: {}", - collector.api_format, collector.task_type, collector.dimension_name, err - )); - continue; - } - }; - - if let Some(existing_id) = existing_id { - if mode == "overwrite" { - match sqlx::query( - r#" -UPDATE dimension_collectors -SET source_type = ?, - source_path = ?, - value_type = ?, - transform_expression = ?, - default_value = ?, - is_enabled = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(&collector.source_type) - .bind(collector.source_path.as_deref()) - .bind(&collector.value_type) - .bind(collector.transform_expression.as_deref()) - .bind(collector.default_value.as_deref()) - .bind(collector.is_enabled) - .bind(current_unix_secs_i64()) - .bind(&existing_id) - .execute(&self.pool) - .await - { - Ok(_) => updated += 1, - Err(err) => errors.push(format!( - "Failed to update collector {}: {}", - existing_id, err - )), - } - } else { - skipped += 1; - } - continue; - } - - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - match sqlx::query( - r#" -INSERT INTO dimension_collectors ( - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(id) - .bind(&collector.api_format) - .bind(&collector.task_type) - .bind(&collector.dimension_name) - .bind(&collector.source_type) - .bind(collector.source_path.as_deref()) - .bind(&collector.value_type) - .bind(collector.transform_expression.as_deref()) - .bind(collector.default_value.as_deref()) - .bind(collector.priority) - .bind(collector.is_enabled) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - { - Ok(_) => created += 1, - Err(err) => errors.push(format!( - "Failed to create collector: api_format={} task_type={} dim={}: {}", - collector.api_format, collector.task_type, collector.dimension_name, err - )), - } - } - - Ok(AdminBillingMutationOutcome::Applied( - AdminBillingPresetApplyResult { - preset: preset.to_string(), - mode: mode.to_string(), - created, - updated, - skipped, - errors, - }, - )) - } - - async fn find_payment_gateway_config( - &self, - provider: &str, - ) -> Result, DataLayerError> { - 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.trim().to_ascii_lowercase()) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref() - .map(map_payment_gateway_config_sqlite) - .transpose() - } - - async fn compare_and_swap_payment_gateway_secret( - &self, - update: &PaymentGatewaySecretCasUpdate, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE payment_gateway_configs -SET merchant_key_encrypted = ? -WHERE provider = ? - AND merchant_key_encrypted IS ? - "#, - ) - .bind(&update.merchant_key_encrypted) - .bind(update.provider.trim().to_ascii_lowercase()) - .bind(&update.expected_merchant_key_encrypted) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn compare_and_swap_payment_gateway_config( - &self, - mutation: &PaymentGatewayConfigCasWriteInput, - ) -> Result, DataLayerError> { - let input = &mutation.input; - let provider = input.provider.trim().to_ascii_lowercase(); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let result = if mutation.expected_existing { - sqlx::query( - r#" -UPDATE payment_gateway_configs -SET - enabled = ?, - endpoint_url = ?, - callback_base_url = ?, - merchant_id = ?, - merchant_key_encrypted = CASE - WHEN ? THEN merchant_key_encrypted - ELSE ? - END, - pay_currency = ?, - usd_exchange_rate = ?, - min_recharge_usd = ?, - channels_json = ?, - updated_at = ? -WHERE provider = ? - AND merchant_key_encrypted IS ? - "#, - ) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .bind(input.preserve_existing_secret) - .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(&provider) - .bind(mutation.expected_merchant_key_encrypted.as_deref()) - .execute(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -INSERT INTO payment_gateway_configs ( - provider, enabled, endpoint_url, callback_base_url, merchant_id, - merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, - channels_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(provider) DO NOTHING - "#, - ) - .bind(&provider) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .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(&mut *tx) - .await - .map_sql_err()? - }; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(AdminBillingMutationOutcome::NotFound); - } - - 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 upsert_payment_gateway_config( - &self, - input: &PaymentGatewayConfigWriteInput, - ) -> Result, DataLayerError> { - let provider = input.provider.trim().to_ascii_lowercase(); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query( - r#" -INSERT INTO payment_gateway_configs ( - provider, enabled, endpoint_url, callback_base_url, merchant_id, - merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd, - channels_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(provider) DO UPDATE SET - enabled = excluded.enabled, - endpoint_url = excluded.endpoint_url, - callback_base_url = excluded.callback_base_url, - merchant_id = excluded.merchant_id, - 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, - channels_json = excluded.channels_json, - updated_at = excluded.updated_at - "#, - ) - .bind(&provider) - .bind(input.enabled) - .bind(&input.endpoint_url) - .bind(input.callback_base_url.as_deref()) - .bind(&input.merchant_id) - .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) - .bind(input.preserve_existing_secret) - .execute(&mut *tx) - .await - .map_sql_err()?; - 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( - &self, - include_disabled: bool, - ) -> Result>, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM billing_plans -WHERE (? = 1 OR enabled = 1) -ORDER BY sort_order ASC, price_amount ASC, id ASC - "#, - ) - .bind(include_disabled) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(Some( - rows.iter() - .map(map_billing_plan_sqlite) - .collect::, _>>()?, - )) - } - - async fn find_billing_plan( - &self, - plan_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM billing_plans -WHERE id = ? -LIMIT 1 - "#, - ) - .bind(plan_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_billing_plan_sqlite).transpose() - } - - async fn create_billing_plan( - &self, - input: &BillingPlanWriteInput, - ) -> Result, DataLayerError> { - let id = uuid::Uuid::new_v4().to_string(); - let now = current_unix_secs_i64(); - sqlx::query(BILLING_PLAN_INSERT_SQLITE) - .bind(&id) - .bind(&input.title) - .bind(input.description.as_deref()) - .bind(input.price_amount) - .bind(&input.price_currency) - .bind(&input.duration_unit) - .bind(input.duration_value) - .bind(input.enabled) - .bind(input.sort_order) - .bind(input.max_active_per_user) - .bind(&input.purchase_limit_scope) - .bind(json_to_string(&input.entitlements_json)?) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - match self.find_billing_plan(&id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Err(DataLayerError::UnexpectedValue( - "created billing plan missing".to_string(), - )), - } - } - - async fn update_billing_plan( - &self, - plan_id: &str, - input: &BillingPlanWriteInput, - ) -> Result, DataLayerError> { - let result = sqlx::query(BILLING_PLAN_UPDATE_SQLITE) - .bind(&input.title) - .bind(input.description.as_deref()) - .bind(input.price_amount) - .bind(&input.price_currency) - .bind(&input.duration_unit) - .bind(input.duration_value) - .bind(input.enabled) - .bind(input.sort_order) - .bind(input.max_active_per_user) - .bind(&input.purchase_limit_scope) - .bind(json_to_string(&input.entitlements_json)?) - .bind(current_unix_secs_i64()) - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match self.find_billing_plan(plan_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn set_billing_plan_enabled( - &self, - plan_id: &str, - enabled: bool, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE billing_plans SET enabled = ?, updated_at = ? WHERE id = ?") - .bind(enabled) - .bind(current_unix_secs_i64()) - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - match self.find_billing_plan(plan_id).await? { - Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)), - None => Ok(AdminBillingMutationOutcome::NotFound), - } - } - - async fn delete_billing_plan( - &self, - plan_id: &str, - ) -> Result, DataLayerError> { - let exists = - sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM billing_plans WHERE id = ?") - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - if exists == 0 { - return Ok(AdminBillingMutationOutcome::NotFound); - } - - let order_count = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM payment_orders -WHERE product_id = ? - AND order_kind = 'plan_purchase' - "#, - ) - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let entitlement_count = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE plan_id = ?", - ) - .bind(plan_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - if order_count > 0 || entitlement_count > 0 { - return Ok(AdminBillingMutationOutcome::Invalid( - "套餐已有订单或权益,不能删除,请停用该套餐".to_string(), - )); - } - - let result = sqlx::query("DELETE FROM billing_plans WHERE id = ?") - .bind(plan_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - Ok(AdminBillingMutationOutcome::NotFound) - } else { - Ok(AdminBillingMutationOutcome::Applied(())) - } - } - - async fn list_user_plan_entitlements( - &self, - user_id: &str, - ) -> Result>, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, user_id, plan_id, payment_order_id, status, - starts_at AS starts_at_unix_secs, expires_at AS expires_at_unix_secs, - entitlements_snapshot, created_at AS created_at_unix_secs, - updated_at AS updated_at_unix_secs -FROM user_plan_entitlements -WHERE user_id = ? - AND status = 'active' - AND expires_at > ? -ORDER BY expires_at ASC, created_at ASC - "#, - ) - .bind(user_id) - .bind(current_unix_secs_i64()) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(Some( - rows.iter() - .map(map_user_plan_entitlement_sqlite) - .collect::, _>>()?, - )) - } - - async fn revoke_user_plan_entitlement( - &self, - user_id: &str, - entitlement_id: &str, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let result = sqlx::query( - r#" -UPDATE user_plan_entitlements -SET status = 'revoked', - expires_at = CASE WHEN expires_at > ? THEN ? ELSE expires_at END, - updated_at = ? -WHERE id = ? - AND user_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(now) - .bind(now) - .bind(now) - .bind(entitlement_id) - .bind(user_id) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - Ok(AdminBillingMutationOutcome::NotFound) - } else { - Ok(AdminBillingMutationOutcome::Applied(())) - } - } - - async fn find_user_daily_quota_availability( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let now_unix_secs = current_unix_secs_i64(); - let rows = sqlx::query( - r#" -SELECT - user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json -FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id -WHERE user_plan_entitlements.user_id = ? - AND user_plan_entitlements.status = 'active' - AND user_plan_entitlements.starts_at <= ? - AND user_plan_entitlements.expires_at > ? -ORDER BY user_plan_entitlements.expires_at ASC, - user_plan_entitlements.created_at ASC, - user_plan_entitlements.id ASC - "#, - ) - .bind(user_id) - .bind(now_unix_secs) - .bind(now_unix_secs) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let now = chrono::Utc::now(); - let mut grants = Vec::new(); - for row in rows { - let entitlement_id: String = row.try_get("id").map_sql_err()?; - let entitlements = parse_json(row.try_get("entitlements_snapshot").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])); - let plan_entitlements = - parse_json(row.try_get("plan_entitlements_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])); - grants.extend(daily_quota_grants_from_entitlement( - &entitlement_id, - &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), - now, - )?); - } - - let mut total_quota_usd = 0.0; - let mut used_usd = 0.0; - let mut remaining_usd = 0.0; - let mut allow_wallet_overage = true; - for grant in &grants { - allow_wallet_overage &= grant.allow_wallet_overage; - let used = sqlx::query_scalar::<_, f64>( - r#" -SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) -FROM entitlement_usage_ledgers -WHERE user_entitlement_id = ? - AND usage_date = ? - "#, - ) - .bind(&grant.entitlement_id) - .bind(&grant.usage_date) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - total_quota_usd += grant.daily_quota_usd; - used_usd += used.min(grant.daily_quota_usd).max(0.0); - remaining_usd += (grant.daily_quota_usd - used).max(0.0); - } - let has_active_daily_quota = !grants.is_empty(); - Ok(Some(UserDailyQuotaAvailabilityRecord { - has_active_daily_quota, - total_quota_usd, - used_usd, - remaining_usd, - allow_wallet_overage: has_active_daily_quota && allow_wallet_overage, - })) - } -} - -const BILLING_PLAN_INSERT_SQLITE: &str = r#" -INSERT INTO billing_plans ( - id, title, description, price_amount, price_currency, duration_unit, - duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, - entitlements_json, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#; - -const BILLING_PLAN_UPDATE_SQLITE: &str = r#" -UPDATE billing_plans -SET title = ?, - description = ?, - price_amount = ?, - price_currency = ?, - duration_unit = ?, - duration_value = ?, - enabled = ?, - sort_order = ?, - max_active_per_user = ?, - purchase_limit_scope = ?, - entitlements_json = ?, - updated_at = ? -WHERE id = ? -"#; - -struct RankedContext { - rank: u8, - is_available: bool, - pricing_rank: u8, - created_at: i64, - context: Result, -} - -fn match_rank( - row: &SqliteRow, - requested_model: &str, -) -> Result, DataLayerError> { - let provider_model_name: Option = - row.try_get("model_provider_model_name").map_sql_err()?; - let global_model_name: String = row.try_get("global_model_name").map_sql_err()?; - let mappings: Option = row.try_get("provider_model_mappings").ok().flatten(); - - let rank = if provider_model_name.as_deref() == Some(requested_model) { - 0 - } else if mappings - .as_deref() - .is_some_and(|mappings| provider_model_mappings_match(mappings, requested_model)) - { - 1 - } else if global_model_name == requested_model { - 2 - } else { - return Ok(None); - }; - - let has_model_price = sqlite_optional_real(row, "model_price_per_request")?.is_some() - || row - .try_get::, _>("model_tiered_pricing") - .ok() - .flatten() - .is_some(); - let has_default_price = sqlite_optional_real(row, "default_price_per_request")?.is_some() - || row - .try_get::, _>("default_tiered_pricing") - .ok() - .flatten() - .is_some(); - let pricing_rank = if has_model_price { - 0 - } else if has_default_price { - 1 - } else { - 2 - }; - - Ok(Some(RankedContext { - rank, - is_available: row - .try_get::, _>("model_is_available") - .map_sql_err()? - .unwrap_or(false), - pricing_rank, - created_at: row - .try_get::, _>("model_created_at") - .map_sql_err()? - .unwrap_or(i64::MAX), - context: map_row(row), - })) -} - -fn provider_model_mappings_match(raw: &str, requested_model: &str) -> bool { - let Ok(value) = serde_json::from_str::(raw) else { - return raw == requested_model; - }; - json_mapping_matches(&value, requested_model) -} - -fn json_mapping_matches(value: &serde_json::Value, requested_model: &str) -> bool { - match value { - serde_json::Value::String(value) => value == requested_model, - serde_json::Value::Array(values) => values - .iter() - .any(|value| json_mapping_matches(value, requested_model)), - serde_json::Value::Object(map) => map - .get("name") - .is_some_and(|value| json_mapping_matches(value, requested_model)), - _ => false, - } -} - -fn map_row(row: &SqliteRow) -> Result { - StoredBillingModelContext::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_billing_type").map_sql_err()?, - row.try_get("provider_api_key_id").map_sql_err()?, - parse_json( - row.try_get("provider_api_key_rate_multipliers") - .ok() - .flatten(), - )?, - row.try_get::, _>("provider_api_key_cache_ttl_minutes") - .map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - row.try_get("global_model_name").map_sql_err()?, - parse_json(row.try_get("global_model_config").ok().flatten())?, - sqlite_optional_real(row, "default_price_per_request")?, - parse_json(row.try_get("default_tiered_pricing").ok().flatten())?, - row.try_get("model_id").map_sql_err()?, - row.try_get("model_provider_model_name").map_sql_err()?, - parse_json(row.try_get("model_config").ok().flatten())?, - sqlite_optional_real(row, "model_price_per_request")?, - parse_json(row.try_get("model_tiered_pricing").ok().flatten())?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON field is invalid: {err}")) - }) - }) - .transpose() -} - -fn current_unix_secs_i64() -> i64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs() as i64 -} - -fn json_to_string(value: &serde_json::Value) -> Result { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON encode failed: {err}")) - }) -} - -#[derive(Debug)] -struct DailyQuotaGrant { - entitlement_id: String, - daily_quota_usd: f64, - usage_date: String, - allow_wallet_overage: bool, -} - -fn daily_quota_usage_date( - reset_timezone: Option<&str>, - now: chrono::DateTime, -) -> Result { - let timezone = reset_timezone - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("Asia/Shanghai") - .parse::() - .map_err(|err| DataLayerError::InvalidInput(format!("invalid reset_timezone: {err}")))?; - Ok(now.with_timezone(&timezone).date_naive().to_string()) -} - -fn daily_quota_grants_from_entitlement( - entitlement_id: &str, - entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, - now: chrono::DateTime, -) -> Result, DataLayerError> { - let mut grants = Vec::new(); - let Some(items) = entitlements.as_array() else { - return Ok(grants); - }; - for item in items { - if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") { - continue; - } - let daily_quota_usd = item - .get("daily_quota_usd") - .and_then(serde_json::Value::as_f64) - .unwrap_or(0.0); - if !daily_quota_usd.is_finite() || daily_quota_usd <= 0.0 { - continue; - } - grants.push(DailyQuotaGrant { - entitlement_id: entitlement_id.to_string(), - daily_quota_usd, - usage_date: daily_quota_usage_date( - item.get("reset_timezone") - .and_then(serde_json::Value::as_str), - now, - )?, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), - }); - } - Ok(grants) -} - -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - -fn read_count_sqlite(row: &SqliteRow) -> Result { - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) -} - -fn map_payment_gateway_config_sqlite( - row: &SqliteRow, -) -> Result { - Ok(PaymentGatewayConfigRecord { - provider: row.try_get("provider").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - endpoint_url: row.try_get("endpoint_url").map_sql_err()?, - callback_base_url: row.try_get("callback_base_url").map_sql_err()?, - merchant_id: row.try_get("merchant_id").map_sql_err()?, - merchant_key_encrypted: row.try_get("merchant_key_encrypted").map_sql_err()?, - pay_currency: row.try_get("pay_currency").map_sql_err()?, - usd_exchange_rate: sqlite_optional_real(row, "usd_exchange_rate")?.unwrap_or(0.0), - min_recharge_usd: sqlite_optional_real(row, "min_recharge_usd")?.unwrap_or(0.0), - channels_json: parse_json(row.try_get("channels_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn map_billing_plan_sqlite(row: &SqliteRow) -> Result { - Ok(BillingPlanRecord { - id: row.try_get("id").map_sql_err()?, - title: row.try_get("title").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - price_amount: sqlite_optional_real(row, "price_amount")?.unwrap_or(0.0), - price_currency: row.try_get("price_currency").map_sql_err()?, - duration_unit: row.try_get("duration_unit").map_sql_err()?, - duration_value: row.try_get("duration_value").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - sort_order: row.try_get("sort_order").map_sql_err()?, - max_active_per_user: row.try_get("max_active_per_user").map_sql_err()?, - purchase_limit_scope: row - .try_get::, _>("purchase_limit_scope") - .map_sql_err()? - .unwrap_or_else(|| "active_period".to_string()), - entitlements_json: parse_json(row.try_get("entitlements_json").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn map_user_plan_entitlement_sqlite( - row: &SqliteRow, -) -> Result { - Ok(UserPlanEntitlementRecord { - id: row.try_get("id").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - plan_id: row.try_get("plan_id").map_sql_err()?, - payment_order_id: row.try_get("payment_order_id").map_sql_err()?, - status: row.try_get("status").map_sql_err()?, - starts_at_unix_secs: row - .try_get::("starts_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - expires_at_unix_secs: row - .try_get::("expires_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - entitlements_snapshot: parse_json(row.try_get("entitlements_snapshot").ok().flatten())? - .unwrap_or_else(|| serde_json::json!([])), - created_at_unix_secs: row - .try_get::("created_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -async fn find_admin_billing_rule_sqlite( - pool: &SqlitePool, - rule_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, name, task_type, global_model_id, model_id, expression, variables, - dimension_mappings, is_enabled, created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM billing_rules -WHERE id = ? - "#, - ) - .bind(rule_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_admin_billing_rule_sqlite).transpose() -} - -fn map_admin_billing_rule_sqlite( - row: &SqliteRow, -) -> Result { - Ok(AdminBillingRuleRecord { - id: row.try_get("id").map_sql_err()?, - name: row.try_get("name").map_sql_err()?, - task_type: row.try_get("task_type").map_sql_err()?, - global_model_id: row.try_get("global_model_id").map_sql_err()?, - model_id: row.try_get("model_id").map_sql_err()?, - expression: row.try_get("expression").map_sql_err()?, - variables: parse_required_json(row.try_get("variables").map_sql_err()?)?, - dimension_mappings: parse_required_json(row.try_get("dimension_mappings").map_sql_err()?)?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - created_at_unix_ms: row - .try_get::("created_at_unix_ms") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -async fn find_admin_billing_collector_sqlite( - pool: &SqlitePool, - collector_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, api_format, task_type, dimension_name, source_type, source_path, value_type, - transform_expression, default_value, priority, is_enabled, - created_at AS created_at_unix_ms, updated_at AS updated_at_unix_secs -FROM dimension_collectors -WHERE id = ? - "#, - ) - .bind(collector_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref() - .map(map_admin_billing_collector_sqlite) - .transpose() -} - -fn map_admin_billing_collector_sqlite( - row: &SqliteRow, -) -> Result { - Ok(AdminBillingCollectorRecord { - id: row.try_get("id").map_sql_err()?, - api_format: row.try_get("api_format").map_sql_err()?, - task_type: row.try_get("task_type").map_sql_err()?, - dimension_name: row.try_get("dimension_name").map_sql_err()?, - source_type: row.try_get("source_type").map_sql_err()?, - source_path: row.try_get("source_path").map_sql_err()?, - value_type: row.try_get("value_type").map_sql_err()?, - transform_expression: row.try_get("transform_expression").map_sql_err()?, - default_value: row.try_get("default_value").map_sql_err()?, - priority: row.try_get("priority").map_sql_err()?, - is_enabled: row.try_get("is_enabled").map_sql_err()?, - created_at_unix_ms: row - .try_get::("created_at_unix_ms") - .map_sql_err()? - .max(0) as u64, - updated_at_unix_secs: row - .try_get::("updated_at_unix_secs") - .map_sql_err()? - .max(0) as u64, - }) -} - -fn parse_required_json(raw: String) -> Result { - serde_json::from_str(&raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!("billing JSON field is invalid: {err}")) - }) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::SqliteBillingReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::billing::{ - AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, AdminBillingRuleWriteInput, - BillingPlanWriteInput, BillingReadRepository, PaymentGatewayConfigCasWriteInput, - PaymentGatewayConfigWriteInput, PaymentGatewaySecretCasUpdate, - }; - - #[tokio::test] - async fn sqlite_repository_reads_billing_model_context() { - 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_billing_context(&pool).await; - - let repository = SqliteBillingReadRepository::new(pool); - let by_alias = repository - .find_model_context("provider-1", Some("key-1"), "gpt-upstream-alias") - .await - .expect("context lookup should run") - .expect("context should exist"); - assert_eq!(by_alias.model_id.as_deref(), Some("model-1")); - assert_eq!(by_alias.provider_api_key_cache_ttl_minutes, Some(60)); - assert_eq!(by_alias.model_price_per_request, Some(0.01)); - - let by_model_id = repository - .find_model_context_by_model_id("provider-1", Some("key-1"), "model-1") - .await - .expect("model lookup should run") - .expect("context should exist"); - assert_eq!(by_model_id.global_model_name, "gpt-5"); - } - - #[tokio::test] - async fn sqlite_repository_manages_admin_billing_rules_and_collectors() { - 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 rule = match repository - .create_admin_billing_rule(&AdminBillingRuleWriteInput { - name: "Chat rule".to_string(), - task_type: "chat".to_string(), - global_model_id: Some("global-1".to_string()), - model_id: None, - expression: "total_tokens * 0.01".to_string(), - variables: json!({"rate": 0.01}), - dimension_mappings: json!({"tokens": "total_tokens"}), - is_enabled: true, - }) - .await - .expect("rule create should run") - { - AdminBillingMutationOutcome::Applied(rule) => rule, - other => panic!("unexpected rule create outcome: {other:?}"), - }; - assert_eq!(rule.variables["rate"], json!(0.01)); - let (rules, total) = repository - .list_admin_billing_rules(Some("chat"), Some(true), 1, 20) - .await - .expect("rule list should run") - .expect("rule list should be available"); - assert_eq!(total, 1); - assert_eq!(rules[0].id, rule.id); - - let updated_rule = match repository - .update_admin_billing_rule( - &rule.id, - &AdminBillingRuleWriteInput { - name: "Updated chat rule".to_string(), - task_type: "chat".to_string(), - global_model_id: Some("global-1".to_string()), - model_id: None, - expression: "total_tokens * 0.02".to_string(), - variables: json!({"rate": 0.02}), - dimension_mappings: json!({"tokens": "total_tokens"}), - is_enabled: false, - }, - ) - .await - .expect("rule update should run") - { - AdminBillingMutationOutcome::Applied(rule) => rule, - other => panic!("unexpected rule update outcome: {other:?}"), - }; - assert_eq!(updated_rule.name, "Updated chat rule"); - assert!(!updated_rule.is_enabled); - - let collector = match repository - .create_admin_billing_collector(&AdminBillingCollectorWriteInput { - api_format: "openai".to_string(), - task_type: "chat".to_string(), - dimension_name: "total_tokens".to_string(), - source_type: "usage".to_string(), - source_path: Some("$.usage.total_tokens".to_string()), - value_type: "float".to_string(), - transform_expression: None, - default_value: Some("1".to_string()), - priority: 10, - is_enabled: true, - }) - .await - .expect("collector create should run") - { - AdminBillingMutationOutcome::Applied(collector) => collector, - other => panic!("unexpected collector create outcome: {other:?}"), - }; - assert!(repository - .admin_billing_enabled_default_value_exists("openai", "chat", "total_tokens", None,) - .await - .expect("default value check should run") - .expect("default value check should be available")); - let (collectors, total) = repository - .list_admin_billing_collectors(Some("openai"), Some("chat"), None, Some(true), 1, 20) - .await - .expect("collector list should run") - .expect("collector list should be available"); - assert_eq!(total, 1); - assert_eq!(collectors[0].id, collector.id); - - let preset = match repository - .apply_admin_billing_preset( - "openai-chat", - "overwrite", - &[AdminBillingCollectorWriteInput { - api_format: "openai".to_string(), - task_type: "chat".to_string(), - dimension_name: "total_tokens".to_string(), - source_type: "usage".to_string(), - source_path: Some("$.usage.total_tokens".to_string()), - value_type: "float".to_string(), - transform_expression: Some("max(total_tokens, 1)".to_string()), - default_value: Some("1".to_string()), - priority: 10, - is_enabled: true, - }], - ) - .await - .expect("preset apply should run") - { - AdminBillingMutationOutcome::Applied(result) => result, - other => panic!("unexpected preset outcome: {other:?}"), - }; - assert_eq!(preset.updated, 1); - assert_eq!(preset.errors, Vec::::new()); - } - - #[tokio::test] - async fn sqlite_repository_revokes_active_user_plan_entitlement() { - 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 now = super::current_unix_secs_i64(); - sqlx::query( - r#" -INSERT INTO users ( - id, username, email, role, auth_source, password_hash, is_active, - is_deleted, created_at, updated_at -) VALUES ( - 'user-revoke', 'revoke-user', 'revoke@example.com', 'user', 'local', - 'hash', 1, 0, 1, 1 -); -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, created_at, updated_at -) VALUES ( - 'wallet-revoke', 'user-revoke', 5.0, 0.0, 'finite', 1, 1 -); -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, - duration_value, entitlements_json, created_at, updated_at -) VALUES ( - 'plan-revoke', 'Revocable Plan', 0.0, 'USD', 'month', 1, - '[{"type":"daily_quota","daily_quota_usd":10.0,"allow_wallet_overage":true}]', - 1, 1 -); -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-revoke', 'order-revoke', 'wallet-revoke', 'user-revoke', 0.0, 0.0, - 0.0, 'admin_manual', '{}', 'credited', 1 -); -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) VALUES ( - 'entitlement-revoke', 'user-revoke', 'plan-revoke', 'order-revoke', - 'active', ?, ?, - '[{"type":"daily_quota","daily_quota_usd":10.0,"allow_wallet_overage":false}]', - ?, ? -); -"#, - ) - .bind(now - 60) - .bind(now + 3600) - .bind(now - 60) - .bind(now - 60) - .execute(&pool) - .await - .expect("revocable entitlement should seed"); - let repository = SqliteBillingReadRepository::new(pool.clone()); - - let quota = repository - .find_user_daily_quota_availability("user-revoke") - .await - .expect("quota should load") - .expect("quota should be available"); - assert!(quota.has_active_daily_quota); - assert!(quota.allow_wallet_overage); - - let wrong_user = repository - .revoke_user_plan_entitlement("other-user", "entitlement-revoke") - .await - .expect("ownership check should run"); - assert_eq!(wrong_user, AdminBillingMutationOutcome::NotFound); - - let outcome = repository - .revoke_user_plan_entitlement("user-revoke", "entitlement-revoke") - .await - .expect("entitlement revoke should run"); - assert_eq!(outcome, AdminBillingMutationOutcome::Applied(())); - let active = repository - .list_user_plan_entitlements("user-revoke") - .await - .expect("entitlements should load") - .expect("entitlements should be available"); - assert!(active.is_empty()); - let quota = repository - .find_user_daily_quota_availability("user-revoke") - .await - .expect("quota should load") - .expect("quota should be available"); - assert!(!quota.has_active_daily_quota); - let status: String = sqlx::query_scalar( - "SELECT status FROM user_plan_entitlements WHERE id = 'entitlement-revoke'", - ) - .fetch_one(&pool) - .await - .expect("entitlement status should load"); - assert_eq!(status, "revoked"); - } - - #[tokio::test] - async fn sqlite_repository_deletes_unused_billing_plans_only() { - 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.clone()); - - let input = BillingPlanWriteInput { - title: "Daily Plan".to_string(), - description: None, - price_amount: 100.0, - price_currency: "CNY".to_string(), - duration_unit: "month".to_string(), - duration_value: 1, - enabled: true, - sort_order: 10, - max_active_per_user: 1, - purchase_limit_scope: "active_period".to_string(), - entitlements_json: json!([{ - "type": "daily_quota", - "daily_quota_usd": 50.0, - "reset_timezone": "Asia/Shanghai", - "allow_wallet_overage": false - }]), - }; - let removable = match repository - .create_billing_plan(&input) - .await - .expect("plan create should run") - { - AdminBillingMutationOutcome::Applied(plan) => plan, - other => panic!("unexpected plan create outcome: {other:?}"), - }; - assert_eq!( - repository - .delete_billing_plan(&removable.id) - .await - .expect("plan delete should run"), - AdminBillingMutationOutcome::Applied(()) - ); - assert!(repository - .find_billing_plan(&removable.id) - .await - .expect("plan lookup should run") - .is_none()); - - let referenced = match repository - .create_billing_plan(&input) - .await - .expect("plan create should run") - { - AdminBillingMutationOutcome::Applied(plan) => plan, - other => panic!("unexpected plan create outcome: {other:?}"), - }; - sqlx::query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, amount_usd, payment_method, order_kind, - product_id, fulfillment_status, status, created_at -) -VALUES ('order-1', 'order-no-1', 'wallet-1', 0, 'epay', 'plan_purchase', - ?, 'pending', 'pending', 1) - "#, - ) - .bind(&referenced.id) - .execute(&pool) - .await - .expect("payment order should seed"); - match repository - .delete_billing_plan(&referenced.id) - .await - .expect("plan delete should run") - { - AdminBillingMutationOutcome::Invalid(detail) => { - assert!(detail.contains("不能删除")); - } - other => panic!("unexpected referenced plan delete outcome: {other:?}"), - } - assert!(repository - .find_billing_plan(&referenced.id) - .await - .expect("plan lookup should run") - .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") - ); - } - - #[tokio::test] - async fn sqlite_gateway_cas_is_create_only_and_secret_exact() { - 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 input = PaymentGatewayConfigWriteInput { - provider: "stripe".to_string(), - enabled: true, - endpoint_url: "https://api.stripe.com".to_string(), - callback_base_url: None, - merchant_id: "merchant".to_string(), - merchant_key_encrypted: Some("legacy-ciphertext".to_string()), - preserve_existing_secret: false, - pay_currency: "USD".to_string(), - usd_exchange_rate: 1.0, - min_recharge_usd: 1.0, - channels_json: json!({"channels": []}), - }; - let create = PaymentGatewayConfigCasWriteInput { - input: input.clone(), - expected_existing: false, - expected_merchant_key_encrypted: None, - }; - assert!(matches!( - repository - .compare_and_swap_payment_gateway_config(&create) - .await - .expect("create should run"), - AdminBillingMutationOutcome::Applied(_) - )); - assert_eq!( - repository - .compare_and_swap_payment_gateway_config(&create) - .await - .expect("conflicting create should run"), - AdminBillingMutationOutcome::NotFound - ); - - let before = repository - .find_payment_gateway_config("stripe") - .await - .expect("lookup should run") - .expect("config should exist"); - assert!(!repository - .compare_and_swap_payment_gateway_secret(&PaymentGatewaySecretCasUpdate { - provider: "stripe".to_string(), - expected_merchant_key_encrypted: "LEGACY-ciphertext".to_string(), - merchant_key_encrypted: "v2-ciphertext".to_string(), - }) - .await - .expect("case-mismatched CAS should run")); - assert!(repository - .compare_and_swap_payment_gateway_secret(&PaymentGatewaySecretCasUpdate { - provider: "stripe".to_string(), - expected_merchant_key_encrypted: "legacy-ciphertext".to_string(), - merchant_key_encrypted: "v2-ciphertext".to_string(), - }) - .await - .expect("exact CAS should run")); - let after = repository - .find_payment_gateway_config("stripe") - .await - .expect("lookup should run") - .expect("config should exist"); - assert_eq!(after.updated_at_unix_secs, before.updated_at_unix_secs); - assert_eq!( - after.merchant_key_encrypted.as_deref(), - Some("v2-ciphertext") - ); - - let stale_update = PaymentGatewayConfigCasWriteInput { - input, - expected_existing: true, - expected_merchant_key_encrypted: Some("legacy-ciphertext".to_string()), - }; - assert_eq!( - repository - .compare_and_swap_payment_gateway_config(&stale_update) - .await - .expect("stale update should run"), - AdminBillingMutationOutcome::NotFound - ); - } - - async fn seed_billing_context(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers (id, name, provider_type, billing_type, is_active, created_at, updated_at) -VALUES ('provider-1', 'Provider One', 'openai', 'pay_as_you_go', 1, 1, 1) -"#, - ) - .execute(pool) - .await - .expect("provider should seed"); - sqlx::query( - r#" -INSERT INTO provider_api_keys ( - id, provider_id, name, rate_multipliers, cache_ttl_minutes, created_at, updated_at -) -VALUES ('key-1', 'provider-1', 'Primary', '{"openai:chat":0.8}', 60, 1, 1) -"#, - ) - .execute(pool) - .await - .expect("provider key should seed"); - sqlx::query( - r#" -INSERT INTO global_models ( - id, name, display_name, is_active, default_price_per_request, default_tiered_pricing, - config, created_at, updated_at -) -VALUES ( - 'global-1', 'gpt-5', 'GPT-5', 1, 0.02, - '{"tiers":[{"up_to":null,"input_price_per_1m":3.0}]}', - '{"streaming":true}', 1, 1 -) -"#, - ) - .execute(pool) - .await - .expect("global model should seed"); - sqlx::query( - r#" -INSERT INTO models ( - id, provider_id, global_model_id, provider_model_name, is_active, is_available, - price_per_request, provider_model_mappings, created_at, updated_at -) -VALUES ( - 'model-1', 'provider-1', 'global-1', 'gpt-upstream', 1, 1, 0.01, - '[{"name":"gpt-upstream-alias"}]', 1, 1 -) -"#, - ) - .execute(pool) - .await - .expect("model should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs b/crates/aether-data/adapters/sqlite/src/candidate_selection.rs deleted file mode 100644 index 86d1f6615..000000000 --- a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs +++ /dev/null @@ -1,1989 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; - -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::candidate_selection::{ - provider_model_mapping_api_format_covers, MinimalCandidateSelectionReadRepository, - StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow, - StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, - StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, - StoredRequestedModelCandidateRowsQuery, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const CANDIDATE_SELECTION_COLUMNS: &str = r#" -SELECT - p.id AS provider_id, - p.name AS provider_name, - p.provider_type AS provider_type, - p.provider_priority AS provider_priority, - p.is_active AS provider_is_active, - p.config AS provider_config, - pe.id AS endpoint_id, - COALESCE(pe.api_format, '') AS endpoint_api_format, - pe.api_family AS endpoint_api_family, - pe.endpoint_kind AS endpoint_kind, - pe.is_active AS endpoint_is_active, - pak.id AS key_id, - pak.name AS key_name, - pak.auth_type AS key_auth_type, - pak.auth_config AS key_auth_config, - pak.is_active AS key_is_active, - pak.api_formats AS key_api_formats, - pak.allowed_models AS key_allowed_models, - pak.capabilities AS key_capabilities, - pak.internal_priority AS key_internal_priority, - pak.global_priority_by_format AS key_global_priority_by_format, - pak.last_used_at AS key_last_used_at_unix_secs, - m.id AS model_id, - m.global_model_id AS global_model_id, - gm.name AS global_model_name, - gm.config AS global_model_config, - m.provider_model_name AS model_provider_model_name, - m.provider_model_mappings AS model_provider_model_mappings, - m.supports_streaming AS model_supports_streaming, - m.is_active AS model_is_active, - m.is_available AS model_is_available, - CASE - WHEN json_valid(p.config) THEN - CASE - WHEN json_type(p.config, '$.pool_advanced') IS NOT NULL THEN 1 - ELSE 0 - END - ELSE 0 - END AS provider_pool_enabled -FROM providers p -INNER JOIN provider_endpoints pe ON pe.provider_id = p.id -INNER JOIN provider_api_keys pak ON pak.provider_id = p.id -INNER JOIN models m ON m.provider_id = p.id -INNER JOIN global_models gm ON gm.id = m.global_model_id -WHERE p.is_active = 1 - AND pe.is_active = 1 - AND pak.is_active = 1 - AND m.is_active = 1 - AND m.is_available = 1 - AND gm.is_active = 1 -"#; - -const REQUESTED_MODEL_RAW_PAGE_SIZE: u32 = 256; -const REQUESTED_MODEL_RAW_SCAN_LIMIT: u32 = 2048; - -#[derive(Debug, Clone)] -pub struct SqliteMinimalCandidateSelectionReadRepository { - pool: SqlitePool, -} - -#[derive(Debug, Clone)] -struct CandidateSelectionRow { - row: StoredMinimalCandidateSelectionRow, - key_auth_config: Option, -} - -#[derive(Debug, Clone, Copy)] -enum SelectedRowsOrder { - WithGlobalModel, - WithoutGlobalModel, -} - -#[derive(Debug, Clone, Copy)] -enum SelectedRowsFilter<'a> { - None, - GlobalModel(&'a str), - RequestedModel(&'a str), -} - -#[derive(Debug, Clone, Copy)] -struct SqlPage { - limit: i64, - offset: i64, -} - -#[derive(Debug)] -struct ExactPageAccumulator { - rows: Vec, - offset: usize, - limit: usize, - target_len: usize, -} - -impl ExactPageAccumulator { - fn new(offset: u32, limit: u32) -> Self { - let offset = usize::try_from(offset).unwrap_or(usize::MAX); - let limit = usize::try_from(limit).unwrap_or(usize::MAX); - Self { - rows: Vec::new(), - offset, - limit, - target_len: offset.saturating_add(limit), - } - } - - fn is_full(&self) -> bool { - self.rows.len() >= self.target_len - } - - fn push_matching(&mut self, rows: I, mut predicate: F) - where - I: IntoIterator, - F: FnMut(&T) -> bool, - { - let remaining = self.target_len.saturating_sub(self.rows.len()); - self.rows.extend( - rows.into_iter() - .filter(|row| predicate(row)) - .take(remaining), - ); - } - - fn into_page(self) -> Vec { - self.rows - .into_iter() - .skip(self.offset) - .take(self.limit) - .collect() - } -} - -#[derive(Debug)] -struct RequestedModelRawPage { - rows: Vec, - raw_len: u32, -} - -impl SqliteMinimalCandidateSelectionReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn selected_rows_for_api_format( - &self, - api_format: &str, - ) -> Result, DataLayerError> { - self.load_selected_rows_for_api_format( - api_format, - SelectedRowsFilter::None, - SelectedRowsOrder::WithGlobalModel, - None, - ) - .await - } - - async fn load_selected_rows_for_api_format( - &self, - api_format: &str, - filter: SelectedRowsFilter<'_>, - order: SelectedRowsOrder, - page: Option, - ) -> Result, DataLayerError> { - let canonical_api_format = normalize_api_format(api_format); - let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = - sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); - let mut rows = Vec::new(); - - for storage_api_format in storage_aliases { - let mut builder = QueryBuilder::::new("WITH candidate_rows AS ("); - builder.push(CANDIDATE_SELECTION_COLUMNS); - push_candidate_sql_filters(&mut builder, &storage_api_format, &match_aliases); - match filter { - SelectedRowsFilter::None => {} - SelectedRowsFilter::GlobalModel(global_model_name) => { - builder.push(" AND gm.name = "); - builder.push_bind(global_model_name); - } - SelectedRowsFilter::RequestedModel(requested_model_name) => { - push_requested_model_sql_filter( - &mut builder, - requested_model_name, - &match_aliases, - ); - } - } - push_selected_rows_query_tail(&mut builder, order, page); - - let query_rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut items = query_rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - items.retain(|item| { - api_format_matches(&item.row.endpoint_api_format, &canonical_api_format) - && item.row.key_supports_api_format(&canonical_api_format) - && key_auth_channel_matches(item, &canonical_api_format) - }); - rows.extend(items.into_iter().map(|item| item.row)); - } - - let rows = match filter { - SelectedRowsFilter::RequestedModel(requested_model_name) => rows - .into_iter() - .filter(|row| { - row_matches_requested_model(row, requested_model_name, &canonical_api_format) - }) - .collect(), - _ => rows, - }; - Ok(dedupe_candidate_selection_rows(rows)) - } - - async fn load_requested_model_raw_page( - &self, - api_format: &str, - requested_model_name: &str, - page: SqlPage, - ) -> Result { - let canonical_api_format = normalize_api_format(api_format); - let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = - sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); - let mut builder = QueryBuilder::::new("WITH candidate_rows AS ("); - builder.push(CANDIDATE_SELECTION_COLUMNS); - push_candidate_sql_filters_for_aliases( - &mut builder, - &storage_aliases, - &match_aliases, - &canonical_api_format, - ); - push_requested_model_sql_filter(&mut builder, requested_model_name, &match_aliases); - push_selected_rows_query_tail(&mut builder, SelectedRowsOrder::WithGlobalModel, Some(page)); - - let query_rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let raw_len = u32::try_from(query_rows.len()).unwrap_or(u32::MAX); - let mut rows = query_rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - rows.retain(|item| { - api_format_matches(&item.row.endpoint_api_format, &canonical_api_format) - && item.row.key_supports_api_format(&canonical_api_format) - && key_auth_channel_matches(item, &canonical_api_format) - }); - - Ok(RequestedModelRawPage { rows, raw_len }) - } -} - -#[async_trait] -impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelectionReadRepository { - async fn list_for_exact_api_format( - &self, - api_format: &str, - ) -> Result, DataLayerError> { - self.selected_rows_for_api_format(api_format).await - } - - async fn list_for_exact_api_format_page( - &self, - query: &StoredApiFormatCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - let fetch_limit = query.offset.saturating_add(query.limit); - let mut rows = self - .load_selected_rows_for_api_format( - &query.api_format, - SelectedRowsFilter::None, - SelectedRowsOrder::WithGlobalModel, - Some(SqlPage { - limit: i64::from(fetch_limit), - offset: 0, - }), - ) - .await?; - sort_candidate_selection_rows(&mut rows, true); - Ok(rows - .into_iter() - .skip(query.offset as usize) - .take(query.limit as usize) - .collect()) - } - - async fn list_for_exact_api_format_and_global_model( - &self, - api_format: &str, - global_model_name: &str, - ) -> Result, DataLayerError> { - self.load_selected_rows_for_api_format( - api_format, - SelectedRowsFilter::GlobalModel(global_model_name), - SelectedRowsOrder::WithoutGlobalModel, - None, - ) - .await - } - - async fn list_for_exact_api_format_and_requested_model( - &self, - api_format: &str, - requested_model_name: &str, - ) -> Result, DataLayerError> { - self.load_selected_rows_for_api_format( - api_format, - SelectedRowsFilter::RequestedModel(requested_model_name), - SelectedRowsOrder::WithGlobalModel, - None, - ) - .await - } - - async fn list_for_exact_api_format_and_requested_model_page( - &self, - query: &StoredRequestedModelCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - - let mut exact_page = ExactPageAccumulator::new(query.offset, query.limit); - let mut raw_offset = 0_u32; - while !exact_page.is_full() && raw_offset < REQUESTED_MODEL_RAW_SCAN_LIMIT { - let raw_limit = REQUESTED_MODEL_RAW_PAGE_SIZE - .min(REQUESTED_MODEL_RAW_SCAN_LIMIT.saturating_sub(raw_offset)); - let raw_page = self - .load_requested_model_raw_page( - &query.api_format, - &query.requested_model_name, - SqlPage { - limit: i64::from(raw_limit), - offset: i64::from(raw_offset), - }, - ) - .await?; - exact_page.push_matching(raw_page.rows, |item| { - row_matches_requested_model( - &item.row, - &query.requested_model_name, - &query.api_format, - ) - }); - raw_offset = raw_offset.saturating_add(raw_page.raw_len); - if raw_page.raw_len < raw_limit || raw_page.raw_len == 0 { - break; - } - } - - let rows = exact_page - .into_page() - .into_iter() - .map(|item| item.row) - .collect(); - Ok(dedupe_candidate_selection_rows(rows)) - } - - async fn list_pool_key_rows_for_group( - &self, - query: &StoredPoolKeyCandidateRowsQuery, - ) -> Result, DataLayerError> { - if query.limit == 0 { - return Ok(Vec::new()); - } - let canonical_api_format = normalize_api_format(&query.api_format); - let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = - sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); - let mut rows = Vec::::new(); - - for storage_api_format in storage_aliases { - let mut builder = QueryBuilder::::new(CANDIDATE_SELECTION_COLUMNS); - push_candidate_sql_filters(&mut builder, &storage_api_format, &match_aliases); - builder.push(" AND p.id = "); - builder.push_bind(&query.provider_id); - builder.push(" AND pe.id = "); - builder.push_bind(&query.endpoint_id); - builder.push(" AND m.id = "); - builder.push_bind(&query.model_id); - push_pool_key_order(&mut builder, &query.order); - builder.push(" LIMIT "); - builder.push_bind(i64::from(query.limit)); - builder.push(" OFFSET "); - builder.push_bind(i64::from(query.offset)); - - let query_rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut items = query_rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - items.retain(|item| { - api_format_matches(&item.row.endpoint_api_format, &canonical_api_format) - && item.row.key_supports_api_format(&canonical_api_format) - && key_auth_channel_matches(item, &canonical_api_format) - }); - rows.extend(items); - } - - Ok(dedupe_candidate_selection_rows( - rows.into_iter().map(|item| item.row).collect(), - )) - } - - async fn list_pool_key_rows_for_group_key_ids( - &self, - query: &StoredPoolKeyCandidateRowsByKeyIdsQuery, - ) -> Result, DataLayerError> { - if query.key_ids.is_empty() { - return Ok(Vec::new()); - } - let key_order = query - .key_ids - .iter() - .enumerate() - .map(|(index, key_id)| (key_id.as_str(), index)) - .collect::>(); - let canonical_api_format = normalize_api_format(&query.api_format); - let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = - sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); - let mut rows = Vec::new(); - - for storage_api_format in storage_aliases { - let mut builder = QueryBuilder::::new(CANDIDATE_SELECTION_COLUMNS); - push_candidate_sql_filters(&mut builder, &storage_api_format, &match_aliases); - builder.push(" AND p.id = "); - builder.push_bind(&query.provider_id); - builder.push(" AND pe.id = "); - builder.push_bind(&query.endpoint_id); - builder.push(" AND m.id = "); - builder.push_bind(&query.model_id); - builder.push(" AND pak.id IN ("); - { - let mut separated = builder.separated(", "); - for key_id in &query.key_ids { - separated.push_bind(key_id); - } - } - builder.push(")"); - builder.push(" ORDER BY CASE pak.id"); - for (index, key_id) in query.key_ids.iter().enumerate() { - builder.push(" WHEN "); - builder.push_bind(key_id); - builder.push(" THEN "); - builder.push_bind(i64::try_from(index).map_err(|_| { - DataLayerError::UnexpectedValue("key id order index overflowed".to_string()) - })?); - } - builder.push(" ELSE "); - builder.push_bind(i64::try_from(query.key_ids.len()).map_err(|_| { - DataLayerError::UnexpectedValue("key id order length overflowed".to_string()) - })?); - builder.push(" END ASC, pak.id ASC"); - - let query_rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut items = query_rows - .iter() - .map(map_candidate_selection_row) - .collect::, _>>()?; - items.retain(|item| { - api_format_matches(&item.row.endpoint_api_format, &canonical_api_format) - && item.row.key_supports_api_format(&canonical_api_format) - && key_auth_channel_matches(item, &canonical_api_format) - }); - rows.extend(items.into_iter().map(|item| item.row)); - } - - let mut rows = dedupe_candidate_selection_rows(rows); - rows.sort_by(|left, right| { - key_order - .get(left.key_id.as_str()) - .cmp(&key_order.get(right.key_id.as_str())) - .then(left.key_id.cmp(&right.key_id)) - }); - Ok(rows) - } -} - -fn push_candidate_sql_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - storage_api_format: &str, - match_aliases: &[String], -) { - builder.push(" AND LOWER(COALESCE(pe.api_format, '')) = "); - builder.push_bind(storage_api_format.trim().to_ascii_lowercase()); - push_key_api_format_sql_filter(builder, match_aliases); - push_key_auth_channel_sql_filter(builder, storage_api_format); -} - -fn push_candidate_sql_filters_for_aliases( - builder: &mut QueryBuilder<'_, Sqlite>, - storage_api_formats: &[String], - match_aliases: &[String], - requested_api_format: &str, -) { - builder.push(" AND LOWER(COALESCE(pe.api_format, '')) IN ("); - push_bind_list(builder, storage_api_formats); - builder.push(")"); - push_key_api_format_sql_filter(builder, match_aliases); - push_key_auth_channel_sql_filter(builder, requested_api_format); -} - -fn push_key_api_format_sql_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - match_aliases: &[String], -) { - builder.push( - r#" - AND ( - pak.api_formats IS NULL - OR TRIM(pak.api_formats) = '' - OR CASE - WHEN json_valid(pak.api_formats) THEN - ( - ( - json_type(pak.api_formats) = 'array' - AND EXISTS ( - SELECT 1 - FROM json_each(pak.api_formats) AS fmt - WHERE LOWER(TRIM(CAST(fmt.value AS TEXT))) IN ( -"#, - ); - push_bind_list(builder, match_aliases); - builder.push( - r#" - ) - ) - ) - OR ( - json_type(pak.api_formats) = 'text' - AND LOWER(TRIM(CAST(json_extract(pak.api_formats, '$') AS TEXT))) IN ( -"#, - ); - push_bind_list(builder, match_aliases); - builder.push( - r#" - ) - ) - OR ( - json_type(pak.api_formats) = 'text' - AND EXISTS ( - SELECT 1 - FROM json_each( - CASE - WHEN json_valid(CAST(json_extract(pak.api_formats, '$') AS TEXT)) - THEN CAST(json_extract(pak.api_formats, '$') AS TEXT) - ELSE '[]' - END - ) AS fmt - WHERE LOWER(TRIM(CAST(fmt.value AS TEXT))) IN ( -"#, - ); - push_bind_list(builder, match_aliases); - builder.push( - r#" - ) - ) - ) - ) - ELSE 0 - END - OR LOWER(TRIM(pak.api_formats)) IN ( -"#, - ); - push_bind_list(builder, match_aliases); - builder.push( - r#" - ) - ) -"#, - ); -} - -fn push_key_auth_channel_sql_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - storage_api_format: &str, -) { - let api_format = normalize_api_format(storage_api_format); - builder.push( - r#" - AND ( - ( - LOWER(TRIM(p.provider_type)) = 'codex' - AND LOWER(TRIM(pak.auth_type)) = 'oauth' - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'chatgpt_web' - AND LOWER(TRIM(pak.auth_type)) IN ('oauth', 'bearer') - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" = 'openai:image' - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'claude_code' - AND LOWER(TRIM(pak.auth_type)) = 'oauth' - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" = 'claude:messages' - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'kiro' - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" = 'claude:messages' - AND ( - LOWER(TRIM(pak.auth_type)) = 'oauth' - OR ( - LOWER(TRIM(pak.auth_type)) = 'bearer' - AND pak.auth_config IS NOT NULL - AND TRIM(pak.auth_config) <> '' - ) - ) - ) - OR ( - LOWER(TRIM(p.provider_type)) IN ('gemini_cli', 'antigravity') - AND LOWER(TRIM(pak.auth_type)) = 'oauth' - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" = 'gemini:generate_content' - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'grok' - AND LOWER(TRIM(pak.auth_type)) = 'oauth' - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image') - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'windsurf' - AND LOWER(TRIM(pak.auth_type)) IN ('oauth', 'api_key', 'bearer') - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" = 'openai:chat' - ) - OR ( - LOWER(TRIM(p.provider_type)) = 'vertex_ai' - AND LOWER(TRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(TRIM(p.provider_type)) NOT IN ( - 'chatgpt_web', - 'claude_code', - 'codex', - 'gemini_cli', - 'grok', - 'vertex_ai', - 'antigravity', - 'kiro', - 'windsurf' - ) - AND LOWER(TRIM(pak.auth_type)) <> 'oauth' - ) - ) -"#, - ); -} - -fn push_requested_model_sql_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - requested_model_name: &str, - _match_aliases: &[String], -) { - builder.push( - r#" - AND ( - gm.name = "#, - ); - builder.push_bind(requested_model_name.to_string()); - builder.push( - r#" - OR m.provider_model_name = "#, - ); - builder.push_bind(requested_model_name.to_string()); - builder.push( - r#" - OR ( - m.provider_model_mappings IS NOT NULL - AND m.provider_model_mappings LIKE "#, - ); - builder.push_bind(format!( - "%{}%", - requested_model_name - .replace('\\', "\\\\") - .replace('%', "\\%") - .replace('_', "\\_") - )); - builder.push( - r#" - ESCAPE '\' - ) - ) -"#, - ); -} - -fn push_selected_rows_query_tail( - builder: &mut QueryBuilder<'_, Sqlite>, - order: SelectedRowsOrder, - page: Option, -) { - builder.push( - r#" -), -pool_rows AS ( - SELECT candidate.* - FROM candidate_rows candidate - WHERE candidate.provider_pool_enabled = 1 - AND NOT EXISTS ( - SELECT 1 - FROM candidate_rows other - WHERE other.provider_pool_enabled = 1 - AND other.provider_id = candidate.provider_id - AND other.endpoint_id = candidate.endpoint_id - AND other.model_id = candidate.model_id - AND ( - other.key_internal_priority < candidate.key_internal_priority - OR ( - other.key_internal_priority = candidate.key_internal_priority - AND other.key_id < candidate.key_id - ) - ) - ) -), -selected_rows AS ( - SELECT * FROM candidate_rows WHERE provider_pool_enabled = 0 - UNION ALL - SELECT * FROM pool_rows -) -SELECT * FROM selected_rows -"#, - ); - push_selected_rows_order(builder, order); - if let Some(page) = page { - builder.push(" LIMIT "); - builder.push_bind(page.limit); - builder.push(" OFFSET "); - builder.push_bind(page.offset); - } -} - -fn push_selected_rows_order(builder: &mut QueryBuilder<'_, Sqlite>, order: SelectedRowsOrder) { - builder.push(" ORDER BY "); - if matches!(order, SelectedRowsOrder::WithGlobalModel) { - builder.push("global_model_name ASC, "); - } - builder.push( - "provider_priority ASC, key_internal_priority ASC, provider_id ASC, endpoint_id ASC, key_id ASC, model_id ASC", - ); -} - -fn push_pool_key_order( - builder: &mut QueryBuilder<'_, Sqlite>, - order: &StoredPoolKeyCandidateOrder, -) { - match order { - StoredPoolKeyCandidateOrder::InternalPriority => { - builder.push(" ORDER BY pak.internal_priority ASC, pak.id ASC"); - } - StoredPoolKeyCandidateOrder::Lru => { - builder.push( - " ORDER BY pak.last_used_at IS NOT NULL ASC, pak.last_used_at ASC, pak.internal_priority ASC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::CacheAffinity => { - builder.push( - " ORDER BY pak.last_used_at IS NULL ASC, pak.last_used_at DESC, pak.internal_priority ASC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::SingleAccount => { - builder.push( - " ORDER BY pak.internal_priority ASC, pak.last_used_at IS NULL ASC, pak.last_used_at DESC, pak.id ASC", - ); - } - StoredPoolKeyCandidateOrder::LoadBalance { seed } => { - push_seeded_pool_key_order(builder, seed); - } - } -} - -fn push_seeded_pool_key_order(builder: &mut QueryBuilder<'_, Sqlite>, seed: &str) { - const HEX_DIGITS: &[u8; 16] = b"0123456789abcdef"; - const PLACEHOLDERS: &[u8; 16] = b"ghijklmnopqrstuv"; - - let mut digits_by_rank = HEX_DIGITS.map(char::from); - digits_by_rank.sort_by(|left, right| { - stable_pool_key_hash(seed, &left.to_string()) - .cmp(&stable_pool_key_hash(seed, &right.to_string())) - .then(left.cmp(right)) - }); - let mut rank_by_digit = ['0'; 16]; - for (rank, digit) in digits_by_rank.into_iter().enumerate() { - let digit_index = digit - .to_digit(16) - .expect("seeded pool-key rank input must be hexadecimal") - as usize; - rank_by_digit[digit_index] = char::from(HEX_DIGITS[rank]); - } - - // SQLite has no built-in hash; placeholders avoid cascading replacements - // while remapping every key-id nibble to a seed-derived rank. - builder.push(" ORDER BY "); - for _ in 0..(HEX_DIGITS.len() + PLACEHOLDERS.len()) { - builder.push("replace("); - } - builder.push("lower(hex(pak.id))"); - for (digit, placeholder) in HEX_DIGITS.iter().zip(PLACEHOLDERS) { - builder.push(format!( - ", '{}', '{}')", - char::from(*digit), - char::from(*placeholder) - )); - } - for (placeholder, rank) in PLACEHOLDERS.iter().zip(rank_by_digit) { - builder.push(format!(", '{}', ", char::from(*placeholder))); - builder.push_bind(rank.to_string()); - builder.push(")"); - } - builder.push(" ASC, pak.id ASC"); -} - -fn push_bind_list(builder: &mut QueryBuilder<'_, Sqlite>, values: &[String]) { - let mut separated = builder.separated(", "); - for value in values { - separated.push_bind(value.clone()); - } -} - -fn stable_pool_key_hash(seed: &str, key_id: &str) -> u64 { - let mut hash = 0xcbf29ce484222325u64; - for byte in seed - .as_bytes() - .iter() - .copied() - .chain(std::iter::once(b':')) - .chain(key_id.as_bytes().iter().copied()) - { - hash ^= u64::from(byte); - hash = hash.wrapping_mul(0x100000001b3); - } - hash -} - -fn row_matches_requested_model( - row: &StoredMinimalCandidateSelectionRow, - requested_model_name: &str, - api_format: &str, -) -> bool { - (row_has_available_provider_model(row, api_format) - && row.global_model_name == requested_model_name) - || (row_default_provider_model_name_available(row, api_format) - && row.model_provider_model_name == requested_model_name) - || row - .model_provider_model_mappings - .as_ref() - .is_some_and(|mappings| { - mappings.iter().any(|mapping| { - mapping_scope_matches(mapping, row, api_format) - && mapping.name == requested_model_name - }) - }) -} - -fn row_has_available_provider_model( - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - row_mapping_matches_scope(row, api_format) - || row_default_provider_model_name_available(row, api_format) -} - -fn row_default_provider_model_name_available( - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - let Some(mappings) = row.model_provider_model_mappings.as_ref() else { - return true; - }; - let mut has_explicit_default_mapping = false; - for mapping in mappings { - if mapping.name != row.model_provider_model_name { - continue; - } - has_explicit_default_mapping = true; - if mapping_scope_matches(mapping, row, api_format) { - return true; - } - } - !has_explicit_default_mapping -} - -fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool { - row.model_provider_model_mappings - .as_ref() - .is_some_and(|mappings| { - mappings - .iter() - .any(|mapping| mapping_scope_matches(mapping, row, api_format)) - }) -} - -fn mapping_scope_matches( - mapping: &StoredProviderModelMapping, - row: &StoredMinimalCandidateSelectionRow, - api_format: &str, -) -> bool { - mapping.api_formats.as_ref().is_none_or(|formats| { - formats.iter().any(|value| { - provider_model_mapping_api_format_covers(&row.provider_type, value, api_format) - }) - }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { - endpoint_ids - .iter() - .any(|endpoint_id| endpoint_id == &row.endpoint_id) - }) -} - -fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool { - let provider_type = row.row.provider_type.trim().to_ascii_lowercase(); - let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase(); - let api_format = normalize_api_format(api_format); - match provider_type.as_str() { - "codex" => { - auth_type == "oauth" - && matches!( - api_format.as_str(), - "openai:responses" - | "openai:responses:compact" - | "openai:search" - | "openai:image" - | "codex:live" - ) - } - "chatgpt_web" => { - matches!(auth_type.as_str(), "oauth" | "bearer") && api_format == "openai:image" - } - "claude_code" => auth_type == "oauth" && api_format == "claude:messages", - "kiro" => { - api_format == "claude:messages" - && (auth_type == "oauth" - || (auth_type == "bearer" - && row - .key_auth_config - .as_deref() - .is_some_and(|value| !value.trim().is_empty()))) - } - "gemini_cli" | "antigravity" => { - auth_type == "oauth" && api_format == "gemini:generate_content" - } - "grok" => { - auth_type == "oauth" - && matches!( - api_format.as_str(), - "openai:chat" | "openai:responses" | "claude:messages" | "openai:image" - ) - } - "windsurf" => { - matches!(auth_type.as_str(), "oauth" | "api_key" | "bearer") - && api_format == "openai:chat" - } - "vertex_ai" => vertex_key_auth_channel_matches(&auth_type, &api_format), - _ => auth_type != "oauth", - } -} - -fn vertex_key_auth_channel_matches(auth_type: &str, api_format: &str) -> bool { - matches!(auth_type, "api_key" | "service_account" | "vertex_ai") - && matches!(api_format, "gemini:generate_content" | "gemini:embedding") -} - -fn dedupe_candidate_selection_rows( - rows: Vec, -) -> Vec { - let mut seen = BTreeSet::new(); - rows.into_iter() - .filter(|row| { - seen.insert(( - row.endpoint_id.clone(), - row.key_id.clone(), - row.model_id.clone(), - )) - }) - .collect() -} - -fn sort_candidate_selection_rows( - rows: &mut [StoredMinimalCandidateSelectionRow], - include_global_model: bool, -) { - rows.sort_by(|left, right| { - let global_model_order = if include_global_model { - left.global_model_name.cmp(&right.global_model_name) - } else { - std::cmp::Ordering::Equal - }; - global_model_order - .then(left.provider_priority.cmp(&right.provider_priority)) - .then(left.key_internal_priority.cmp(&right.key_internal_priority)) - .then(left.provider_id.cmp(&right.provider_id)) - .then(left.endpoint_id.cmp(&right.endpoint_id)) - .then(left.key_id.cmp(&right.key_id)) - .then(left.model_id.cmp(&right.model_id)) - }); -} - -fn map_candidate_selection_row(row: &SqliteRow) -> Result { - let _provider_config = parse_json(row.try_get("provider_config").ok().flatten())?; - let global_model_config = parse_json(row.try_get("global_model_config").ok().flatten())?; - let global_model_mappings = global_model_config - .as_ref() - .and_then(|value| value.get("model_mappings").cloned()); - let global_model_supports_streaming = global_model_config - .as_ref() - .and_then(|value| value.get("streaming")) - .and_then(json_bool); - Ok(CandidateSelectionRow { - row: StoredMinimalCandidateSelectionRow { - provider_id: row.try_get("provider_id").map_sql_err()?, - provider_name: row.try_get("provider_name").map_sql_err()?, - provider_type: row.try_get("provider_type").map_sql_err()?, - provider_priority: row.try_get("provider_priority").map_sql_err()?, - provider_is_active: row.try_get("provider_is_active").map_sql_err()?, - endpoint_id: row.try_get("endpoint_id").map_sql_err()?, - endpoint_api_format: row.try_get("endpoint_api_format").map_sql_err()?, - endpoint_api_family: row.try_get("endpoint_api_family").map_sql_err()?, - endpoint_kind: row.try_get("endpoint_kind").map_sql_err()?, - endpoint_is_active: row.try_get("endpoint_is_active").map_sql_err()?, - key_id: row.try_get("key_id").map_sql_err()?, - key_name: row.try_get("key_name").map_sql_err()?, - key_auth_type: row.try_get("key_auth_type").map_sql_err()?, - key_is_active: row.try_get("key_is_active").map_sql_err()?, - key_api_formats: parse_stored_key_policy_string_list( - row.try_get("key_api_formats").map_sql_err()?, - "provider_api_keys.api_formats", - )?, - key_allowed_models: parse_stored_key_policy_string_list( - row.try_get("key_allowed_models").map_sql_err()?, - "provider_api_keys.allowed_models", - )?, - key_capabilities: parse_json(row.try_get("key_capabilities").ok().flatten())?, - key_internal_priority: row.try_get("key_internal_priority").map_sql_err()?, - key_global_priority_by_format: parse_json( - row.try_get("key_global_priority_by_format").ok().flatten(), - )?, - model_id: row.try_get("model_id").map_sql_err()?, - global_model_id: row.try_get("global_model_id").map_sql_err()?, - global_model_name: row.try_get("global_model_name").map_sql_err()?, - global_model_mappings: parse_string_list( - global_model_mappings, - "global_models.config.model_mappings", - )?, - global_model_supports_streaming, - model_provider_model_name: row.try_get("model_provider_model_name").map_sql_err()?, - model_provider_model_mappings: parse_provider_model_mappings(parse_json( - row.try_get("model_provider_model_mappings").ok().flatten(), - )?)?, - model_supports_streaming: row.try_get("model_supports_streaming").map_sql_err()?, - model_is_active: row.try_get("model_is_active").map_sql_err()?, - model_is_available: row.try_get("model_is_available").map_sql_err()?, - }, - key_auth_config: row.try_get("key_auth_config").map_sql_err()?, - }) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "candidate selection JSON field is invalid: {err}" - )) - }) - }) - .transpose() -} - -fn json_bool(value: &serde_json::Value) -> Option { - value.as_bool().or_else(|| { - value - .as_str() - .and_then(|value| value.trim().parse::().ok()) - }) -} - -fn parse_string_list( - value: Option, - field_name: &str, -) -> Result>, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - parse_string_list_value(&value, field_name) -} - -fn parse_stored_key_policy_string_list( - raw: Option, - field_name: &str, -) -> Result>, DataLayerError> { - let Some(raw) = raw else { - return Ok(None); - }; - let value = serde_json::from_str::(&raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}")) - })?; - parse_key_policy_string_list_value(&value, field_name) -} - -fn parse_key_policy_string_list_value( - value: &serde_json::Value, - field_name: &str, -) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains JSON null; use SQL NULL for an unset policy" - ))), - serde_json::Value::Array(array) => { - parse_key_policy_string_list_array(array, field_name).map(Some) - } - serde_json::Value::String(raw) => parse_embedded_key_policy_string_list(raw, field_name), - _ => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} is not a JSON array" - ))), - } -} - -fn parse_embedded_key_policy_string_list( - raw: &str, - field_name: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains an empty string" - ))); - } - if raw.eq_ignore_ascii_case("null") { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains stringified JSON null; use SQL NULL for an unset policy" - ))); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_key_policy_string_list_value(&decoded, field_name); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_key_policy_string_list_array( - array: &[serde_json::Value], - field_name: &str, -) -> Result, DataLayerError> { - let mut items = Vec::with_capacity(array.len()); - for item in array { - let Some(item) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains a non-string item" - ))); - }; - let item = item.trim(); - if item.is_empty() { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains an empty item" - ))); - } - items.push(item.to_string()); - } - Ok(items) -} - -fn parse_string_list_value( - value: &serde_json::Value, - field_name: &str, -) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(array) => parse_string_list_array(array, field_name).map(Some), - serde_json::Value::String(raw) => parse_embedded_string_list(raw, field_name), - _ => Err(DataLayerError::UnexpectedValue(format!( - "{field_name} is not a JSON array" - ))), - } -} - -fn parse_embedded_string_list( - raw: &str, - field_name: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_string_list_value(&decoded, field_name); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_string_list_array( - array: &[serde_json::Value], - field_name: &str, -) -> Result, DataLayerError> { - let mut items = Vec::with_capacity(array.len()); - for item in array { - let Some(item) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue(format!( - "{field_name} contains a non-string item" - ))); - }; - let item = item.trim(); - if !item.is_empty() { - items.push(item.to_string()); - } - } - Ok(items) -} - -fn parse_provider_model_mappings( - value: Option, -) -> Result>, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(array) => parse_provider_model_mappings_array(&array), - serde_json::Value::Object(object) => parse_provider_model_mapping_object_lenient(&object) - .map(|mapping| mapping.map(|value| vec![value])), - serde_json::Value::String(raw) => parse_embedded_provider_model_mappings(&raw), - _ => Err(DataLayerError::UnexpectedValue( - "models.provider_model_mappings is not a JSON array".to_string(), - )), - } -} - -fn parse_embedded_provider_model_mappings( - raw: &str, -) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_provider_model_mappings(Some(decoded)); - } - - Ok(Some(vec![StoredProviderModelMapping { - name: raw.to_string(), - priority: 1, - api_formats: None, - endpoint_ids: None, - operations: None, - }])) -} - -fn parse_provider_model_mappings_array( - array: &[serde_json::Value], -) -> Result>, DataLayerError> { - let mut mappings = Vec::with_capacity(array.len()); - for raw in array { - match raw { - serde_json::Value::Object(object) => { - if let Some(mapping) = parse_provider_model_mapping_object_lenient(object)? { - mappings.push(mapping); - } - } - serde_json::Value::String(raw) if !raw.trim().is_empty() => { - mappings.push(StoredProviderModelMapping { - name: raw.trim().to_string(), - priority: 1, - api_formats: None, - endpoint_ids: None, - operations: None, - }); - } - _ => {} - } - } - - if mappings.is_empty() { - Ok(None) - } else { - Ok(Some(mappings)) - } -} - -fn parse_provider_model_mapping_object_lenient( - object: &serde_json::Map, -) -> Result, DataLayerError> { - let Some(name) = object - .get("name") - .and_then(serde_json::Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - else { - return Ok(None); - }; - - let priority = object - .get("priority") - .and_then(serde_json::Value::as_i64) - .unwrap_or(1) - .max(1); - let api_formats = parse_string_list( - object.get("api_formats").cloned(), - "models.provider_model_mappings.api_formats", - )? - .map(|formats| { - formats - .into_iter() - .map(|value| normalize_api_format(&value)) - .collect() - }); - let endpoint_ids = parse_string_list( - object.get("endpoint_ids").cloned(), - "models.provider_model_mappings.endpoint_ids", - )?; - let operations = parse_string_list( - object.get("operations").cloned(), - "models.provider_model_mappings.operations", - )? - .and_then(normalize_request_operations); - - Ok(Some(StoredProviderModelMapping { - name: name.to_string(), - priority: i32::try_from(priority).map_err(|_| { - DataLayerError::UnexpectedValue(format!( - "invalid models.provider_model_mappings.priority: {priority}" - )) - })?, - api_formats, - endpoint_ids, - operations, - })) -} - -fn normalize_request_operations(values: Vec) -> Option> { - let operations = values - .into_iter() - .map(|value| value.trim().to_ascii_lowercase()) - .filter(|value| !value.is_empty()) - .collect::>(); - (!operations.is_empty()).then_some(operations) -} - -fn api_format_aliases(api_format: &str) -> Vec { - aether_ai_formats::api_format_storage_aliases(api_format) -} - -fn api_format_permission_aliases(api_format: &str) -> Vec { - aether_ai_formats::api_format_permission_storage_aliases(api_format) -} - -fn normalize_api_format(api_format: &str) -> String { - aether_ai_formats::normalize_api_format_alias(api_format) -} - -fn api_format_matches(left: &str, right: &str) -> bool { - aether_ai_formats::api_format_alias_matches(left, right) -} - -fn sql_match_aliases(api_formats: &[String]) -> Vec { - api_formats - .iter() - .map(|value| value.trim().to_ascii_lowercase()) - .collect() -} - -#[cfg(test)] -mod tests { - use super::{ - parse_stored_key_policy_string_list, provider_model_mapping_api_format_covers, - push_key_auth_channel_sql_filter, push_pool_key_order, vertex_key_auth_channel_matches, - ExactPageAccumulator, SqliteMinimalCandidateSelectionReadRepository, - REQUESTED_MODEL_RAW_SCAN_LIMIT, - }; - use crate::run_migrations; - use aether_data_contracts::repository::candidate_selection::{ - MinimalCandidateSelectionReadRepository, StoredPoolKeyCandidateOrder, - StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery, - }; - - #[test] - fn vertex_auth_matrix_rejects_retired_claude_format() { - for auth_type in ["api_key", "service_account", "vertex_ai"] { - assert!(!vertex_key_auth_channel_matches( - auth_type, - "claude:messages" - )); - assert!(vertex_key_auth_channel_matches( - auth_type, - "gemini:generate_content" - )); - assert!(vertex_key_auth_channel_matches( - auth_type, - "gemini:embedding" - )); - } - - let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); - push_key_auth_channel_sql_filter(&mut builder, "claude:messages"); - let sql = builder.sql(); - let vertex_clause = sql - .split_once("LOWER(TRIM(p.provider_type)) = 'vertex_ai'") - .and_then(|(_, suffix)| suffix.split_once("LOWER(TRIM(p.provider_type)) NOT IN")) - .map(|(clause, _)| clause) - .expect("Vertex auth clause should exist"); - assert!(!vertex_clause.contains("claude:messages")); - assert!(vertex_clause.contains("gemini:generate_content")); - assert!(vertex_clause.contains("gemini:embedding")); - } - - #[test] - fn malformed_key_policy_never_degrades_to_unrestricted() { - for raw in ["null", "\"null\"", "\"\"", "[\"openai:chat\",null]"] { - assert!(parse_stored_key_policy_string_list( - Some(raw.to_string()), - "provider_api_keys.api_formats", - ) - .is_err()); - } - assert_eq!( - parse_stored_key_policy_string_list( - Some("[\"openai:chat\"]".to_string()), - "provider_api_keys.api_formats", - ) - .expect("valid key policy should parse"), - Some(vec!["openai:chat".to_string()]) - ); - } - - #[test] - fn codex_auth_sql_allows_live_for_oauth_keys() { - let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); - push_key_auth_channel_sql_filter(&mut builder, "codex:live"); - let sql = builder.sql(); - let codex_clause = sql - .split_once("LOWER(TRIM(p.provider_type)) = 'codex'") - .and_then(|(_, suffix)| { - suffix.split_once("LOWER(TRIM(p.provider_type)) = 'chatgpt_web'") - }) - .map(|(clause, _)| clause) - .expect("Codex auth clause should exist"); - - assert!(codex_clause.contains("LOWER(TRIM(pak.auth_type)) = 'oauth'")); - assert!(codex_clause.contains("'codex:live'")); - } - - #[test] - fn sqlite_mapping_scope_keeps_legacy_responses_compatibility_codex_only() { - assert!(provider_model_mapping_api_format_covers( - "codex", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "openai", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "custom", - "openai:responses", - "codex:live" - )); - assert!(!provider_model_mapping_api_format_covers( - "codex", - "openai:chat", - "codex:live" - )); - } - - #[test] - fn exact_page_accumulator_continues_after_coarse_false_positives() { - let mut accumulator = ExactPageAccumulator::new(1, 2); - accumulator.push_matching(vec![("coarse-1", false), ("coarse-2", false)], |row| row.1); - assert!(!accumulator.is_full()); - - accumulator.push_matching( - vec![ - ("exact-1", true), - ("coarse-3", false), - ("exact-2", true), - ("exact-3", true), - ], - |row| row.1, - ); - - assert!(accumulator.is_full()); - assert_eq!( - accumulator.into_page(), - vec![("exact-2", true), ("exact-3", true)] - ); - assert_eq!(REQUESTED_MODEL_RAW_SCAN_LIMIT, 2048); - } - - #[test] - fn load_balance_pool_key_order_is_seeded_and_pageable_in_sql() { - let sql_for_seed = |seed: &str| { - let mut builder = - sqlx::QueryBuilder::::new("SELECT pak.id FROM provider_api_keys pak"); - push_pool_key_order( - &mut builder, - &StoredPoolKeyCandidateOrder::LoadBalance { - seed: seed.to_string(), - }, - ); - builder.push(" LIMIT "); - builder.push_bind(64_i64); - builder.push(" OFFSET "); - builder.push_bind(128_i64); - builder.sql().to_string() - }; - - let first_seed_sql = sql_for_seed("seed-a"); - let second_seed_sql = sql_for_seed("seed-b"); - - assert!(first_seed_sql.contains("lower(hex(pak.id))")); - assert!(first_seed_sql.contains("ASC, pak.id ASC LIMIT ? OFFSET ?")); - assert_eq!(first_seed_sql, second_seed_sql); - assert_eq!(first_seed_sql.matches('?').count(), 18); - } - - #[tokio::test] - async fn sqlite_repository_reads_candidate_selection_rows() { - 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_candidate_selection(&pool).await; - - let repository = SqliteMinimalCandidateSelectionReadRepository::new(pool); - let rows = repository - .list_for_exact_api_format("openai:chat") - .await - .expect("candidate rows should load"); - assert_eq!( - rows.iter() - .map(|row| row.key_id.as_str()) - .collect::>(), - vec!["key-windsurf-oauth", "key-1"] - ); - assert_eq!( - rows[1].global_model_mappings, - Some(vec!["alias-global".to_string()]) - ); - assert_eq!(rows[1].global_model_supports_streaming, Some(true)); - assert_eq!( - rows[1] - .model_provider_model_mappings - .as_ref() - .and_then(|mappings| mappings.first()) - .and_then(|mapping| mapping.operations.as_ref()), - Some(&vec!["compact".to_string()]) - ); - - let requested = repository - .list_for_exact_api_format_and_requested_model_page( - &StoredRequestedModelCandidateRowsQuery { - api_format: "openai:chat".to_string(), - requested_model_name: "alias-provider".to_string(), - offset: 0, - limit: 10, - }, - ) - .await - .expect("requested model rows should load"); - assert_eq!(requested.len(), 1); - - let windsurf_requested = repository - .list_for_exact_api_format_and_requested_model_page( - &StoredRequestedModelCandidateRowsQuery { - api_format: "openai:chat".to_string(), - requested_model_name: "claude-opus-4-7".to_string(), - offset: 0, - limit: 10, - }, - ) - .await - .expect("windsurf oauth rows should load"); - assert_eq!( - windsurf_requested - .iter() - .map(|row| row.key_id.as_str()) - .collect::>(), - vec!["key-windsurf-oauth"] - ); - - let pool_keys = repository - .list_pool_key_rows_for_group(&StoredPoolKeyCandidateRowsQuery { - api_format: "openai:chat".to_string(), - provider_id: "provider-1".to_string(), - endpoint_id: "endpoint-1".to_string(), - model_id: "model-1".to_string(), - selected_provider_model_name: "provider-model".to_string(), - order: StoredPoolKeyCandidateOrder::InternalPriority, - offset: 1, - limit: 1, - }) - .await - .expect("pool keys should load"); - assert_eq!(pool_keys.len(), 1); - assert_eq!(pool_keys[0].key_id, "key-2"); - - let image_rows = repository - .list_for_exact_api_format_and_requested_model_page( - &StoredRequestedModelCandidateRowsQuery { - api_format: "openai:image".to_string(), - requested_model_name: "gpt-image-2".to_string(), - offset: 0, - limit: 10, - }, - ) - .await - .expect("chatgpt web image rows should load"); - assert_eq!( - image_rows - .iter() - .map(|row| row.key_id.as_str()) - .collect::>(), - vec!["key-chatgpt-web-oauth", "key-chatgpt-web-bearer"] - ); - - let search_rows = repository - .list_for_exact_api_format_and_requested_model_page( - &StoredRequestedModelCandidateRowsQuery { - api_format: "openai:search".to_string(), - requested_model_name: "gpt-5.6-sol".to_string(), - offset: 0, - limit: 10, - }, - ) - .await - .expect("Codex Search rows should load through Responses permissions"); - assert_eq!(search_rows.len(), 1); - assert_eq!(search_rows[0].key_id, "key-codex-search"); - assert_eq!(search_rows[0].endpoint_api_format, "openai:search"); - } - - #[tokio::test] - async fn sqlite_requested_model_page_crosses_coarse_false_positive_windows() { - 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_requested_model_pagination(&pool).await; - - let repository = SqliteMinimalCandidateSelectionReadRepository::new(pool); - let rows = repository - .list_for_exact_api_format_and_requested_model_page( - &StoredRequestedModelCandidateRowsQuery { - api_format: "openai:chat".to_string(), - requested_model_name: "sqlite-page-target".to_string(), - offset: 1, - limit: 1, - }, - ) - .await - .expect("requested model page should cross the coarse-only window"); - - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].model_id, "model-pagination-exact-1"); - } - - #[tokio::test] - async fn sqlite_load_balance_pool_key_pages_use_stable_seeded_order() { - 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_candidate_selection(&pool).await; - - let repository = SqliteMinimalCandidateSelectionReadRepository::new(pool); - let load_page = |seed: &str, offset, limit| StoredPoolKeyCandidateRowsQuery { - api_format: "openai:chat".to_string(), - provider_id: "provider-1".to_string(), - endpoint_id: "endpoint-1".to_string(), - model_id: "model-1".to_string(), - selected_provider_model_name: "provider-model".to_string(), - order: StoredPoolKeyCandidateOrder::LoadBalance { - seed: seed.to_string(), - }, - offset, - limit, - }; - - let seed_a_first = repository - .list_pool_key_rows_for_group(&load_page("seed-a", 0, 1)) - .await - .expect("first load-balance page should load"); - let seed_a_second = repository - .list_pool_key_rows_for_group(&load_page("seed-a", 1, 1)) - .await - .expect("second load-balance page should load"); - let seed_a_replay = repository - .list_pool_key_rows_for_group(&load_page("seed-a", 0, 2)) - .await - .expect("replayed load-balance window should load"); - let seed_b = repository - .list_pool_key_rows_for_group(&load_page("seed-b", 0, 2)) - .await - .expect("alternate load-balance seed should load"); - - let seed_a_pages = seed_a_first - .iter() - .chain(&seed_a_second) - .map(|row| row.key_id.as_str()) - .collect::>(); - let seed_a_replay = seed_a_replay - .iter() - .map(|row| row.key_id.as_str()) - .collect::>(); - let seed_b = seed_b - .iter() - .map(|row| row.key_id.as_str()) - .collect::>(); - - assert_eq!(seed_a_pages, vec!["key-1", "key-2"]); - assert_eq!(seed_a_pages, seed_a_replay); - assert_eq!(seed_b, vec!["key-2", "key-1"]); - } - - async fn seed_candidate_selection(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers ( - id, name, provider_type, provider_priority, config, is_active, created_at, updated_at -) -VALUES ('provider-1', 'Provider One', 'custom', 10, '{"pool_advanced":{}}', 1, 1, 1); - -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, is_active, created_at, updated_at -) -VALUES ('endpoint-1', 'provider-1', 'Endpoint One', 'https://example.test', 'openai:chat', 1, 1, 1); - -INSERT INTO provider_api_keys ( - id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at -) -VALUES - ('key-1', 'provider-1', 'Key One', 'api_key', '["openai:chat"]', 10, 1, 1, 1), - ('key-2', 'provider-1', 'Key Two', 'api_key', '["openai:chat"]', 20, 1, 1, 1); - -INSERT INTO providers ( - id, name, provider_type, provider_priority, is_active, created_at, updated_at -) -VALUES ('provider-chatgpt-web', 'ChatGPT Web', 'chatgpt_web', 20, 1, 1, 1); - -INSERT INTO providers ( - id, name, provider_type, provider_priority, is_active, created_at, updated_at -) -VALUES ('provider-windsurf', 'Windsurf', 'windsurf', 15, 1, 1, 1); - -INSERT INTO providers ( - id, name, provider_type, provider_priority, is_active, created_at, updated_at -) -VALUES ('provider-codex-search', 'Codex Search', 'codex', 12, 1, 1, 1); - -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, is_active, created_at, updated_at -) -VALUES ( - 'endpoint-chatgpt-web', 'provider-chatgpt-web', 'ChatGPT Web Image', - 'https://chatgpt.com', 'openai:image', 1, 1, 1 -); - -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, is_active, created_at, updated_at -) -VALUES ( - 'endpoint-windsurf', 'provider-windsurf', 'Windsurf Chat', - 'https://server.codeium.com', 'openai:chat', 1, 1, 1 -); - -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, is_active, created_at, updated_at -) -VALUES ( - 'endpoint-codex-search', 'provider-codex-search', 'Codex Search', - 'https://chatgpt.com/backend-api/codex', 'openai:search', 1, 1, 1 -); - -INSERT INTO provider_api_keys ( - id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at -) -VALUES - ('key-chatgpt-web-oauth', 'provider-chatgpt-web', 'OAuth', 'oauth', '["openai:image"]', 10, 1, 1, 1), - ('key-chatgpt-web-bearer', 'provider-chatgpt-web', 'Bearer', 'bearer', '["openai:image"]', 20, 1, 1, 1), - ('key-chatgpt-web-api-key', 'provider-chatgpt-web', 'API Key', 'api_key', '["openai:image"]', 30, 1, 1, 1); - -INSERT INTO provider_api_keys ( - id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at -) -VALUES ( - 'key-windsurf-oauth', 'provider-windsurf', 'OAuth', 'oauth', '["openai:chat"]', 10, 1, 1, 1 -); - -INSERT INTO provider_api_keys ( - id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at -) -VALUES ( - 'key-codex-search', 'provider-codex-search', 'OAuth', 'oauth', - '["openai:responses"]', 10, 1, 1, 1 -); - -INSERT INTO global_models ( - id, name, config, is_active, created_at, updated_at -) -VALUES - ('global-1', 'gpt-5', '{"model_mappings":["alias-global"],"streaming":true}', 1, 1, 1), - ('global-image-1', 'gpt-image-2', NULL, 1, 1, 1), - ('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1), - ('global-codex-search-1', 'search-global', '{"streaming":false}', 1, 1, 1); - -INSERT INTO models ( - id, provider_id, global_model_id, provider_model_name, provider_model_mappings, - supports_streaming, is_active, is_available, created_at, updated_at -) -VALUES ( - 'model-1', 'provider-1', 'global-1', 'provider-model', - '[{"name":"alias-provider","api_formats":["openai:chat"],"operations":["COMPACT"],"priority":1}]', - 1, 1, 1, 1, 1 -), -( - 'model-chatgpt-web-image', 'provider-chatgpt-web', 'global-image-1', 'gpt-image-2', - NULL, 1, 1, 1, 1, 1 -), -( - 'model-windsurf-opus', 'provider-windsurf', 'global-windsurf-1', 'claude-opus-4-7', - NULL, NULL, 1, 1, 1, 1 -), -( - 'model-codex-search', 'provider-codex-search', 'global-codex-search-1', 'search-upstream', - '[{"name":"gpt-5.6-sol","api_formats":["openai:responses"],"priority":1}]', - 0, 1, 1, 1, 1 -); -"#, - ) - .execute(pool) - .await - .expect("candidate selection rows should seed"); - } - - async fn seed_requested_model_pagination(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers ( - id, name, provider_type, provider_priority, is_active, created_at, updated_at -) -VALUES ( - 'provider-pagination', 'Pagination Provider', 'custom', 10, 1, 1, 1 -); - -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, is_active, created_at, updated_at -) -VALUES ( - 'endpoint-pagination', 'provider-pagination', 'Pagination Endpoint', - 'https://example.test', 'openai:chat', 1, 1, 1 -); - -INSERT INTO provider_api_keys ( - id, provider_id, name, auth_type, api_formats, internal_priority, - is_active, created_at, updated_at -) -VALUES ( - 'key-pagination', 'provider-pagination', 'Pagination Key', 'api_key', - '["openai:chat"]', 10, 1, 1, 1 -); - -WITH RECURSIVE sequence(value) AS ( - SELECT 0 - UNION ALL - SELECT value + 1 FROM sequence WHERE value < 255 -) -INSERT INTO global_models ( - id, name, is_active, created_at, updated_at -) -SELECT - printf('global-pagination-false-%03d', value), - printf('a-pagination-false-%03d', value), - 1, 1, 1 -FROM sequence; - -INSERT INTO global_models ( - id, name, is_active, created_at, updated_at -) -VALUES - ('global-pagination-exact-0', 'z-pagination-exact-0', 1, 1, 1), - ('global-pagination-exact-1', 'z-pagination-exact-1', 1, 1, 1); - -WITH RECURSIVE sequence(value) AS ( - SELECT 0 - UNION ALL - SELECT value + 1 FROM sequence WHERE value < 255 -) -INSERT INTO models ( - id, provider_id, global_model_id, provider_model_name, provider_model_mappings, - is_active, is_available, created_at, updated_at -) -SELECT - printf('model-pagination-false-%03d', value), - 'provider-pagination', - printf('global-pagination-false-%03d', value), - 'upstream-false', - '[{"name":"sqlite-page-target-noise","api_formats":["openai:chat"],"priority":1}]', - 1, 1, 1, 1 -FROM sequence; - -INSERT INTO models ( - id, provider_id, global_model_id, provider_model_name, provider_model_mappings, - is_active, is_available, created_at, updated_at -) -VALUES - ( - 'model-pagination-exact-0', 'provider-pagination', 'global-pagination-exact-0', - 'upstream-exact-0', - '[{"name":"sqlite-page-target","api_formats":["openai:chat"],"priority":1}]', - 1, 1, 1, 1 - ), - ( - 'model-pagination-exact-1', 'provider-pagination', 'global-pagination-exact-1', - 'upstream-exact-1', - '[{"name":"sqlite-page-target","api_formats":["openai:chat"],"priority":1}]', - 1, 1, 1, 1 - ); -"#, - ) - .execute(pool) - .await - .expect("requested model pagination rows should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/candidates.rs b/crates/aether-data/adapters/sqlite/src/candidates.rs deleted file mode 100644 index c5b9bbc1f..000000000 --- a/crates/aether-data/adapters/sqlite/src/candidates.rs +++ /dev/null @@ -1,1294 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; - -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqliteConnection}; - -use aether_data_contracts::repository::candidates::{ - request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket, - RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, - StoredRequestCandidate, UpsertRequestCandidateRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_in, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const CANDIDATE_COLUMNS: &str = r#" -SELECT - id, - request_id, - user_id, - api_key_id, - username, - api_key_name, - candidate_index, - retry_index, - provider_id, - endpoint_id, - key_id, - status, - skip_reason, - is_cached, - status_code, - error_type, - error_message, - latency_ms, - concurrent_requests, - extra_data, - required_capabilities, - created_at AS created_at_unix_ms, - started_at AS started_at_unix_ms, - finished_at AS finished_at_unix_ms -FROM request_candidates -"#; - -#[derive(Debug, Clone)] -pub struct SqliteRequestCandidateRepository { - pool: SqlitePool, -} - -impl SqliteRequestCandidateRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl RequestCandidateReadRepository for SqliteRequestCandidateRepository { - async fn list_by_request_id( - &self, - request_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? ORDER BY candidate_index ASC, retry_index ASC, created_at ASC" - )) - .bind(request_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_attempted_by_request_id( - &self, - request_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE request_id = ? \ - AND (status IN ('streaming', 'success', 'failed', 'cancelled') \ - OR (status = 'pending' AND started_at IS NOT NULL)) \ - ORDER BY candidate_index ASC, retry_index ASC, created_at ASC" - )) - .bind(request_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_recent( - &self, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} ORDER BY created_at DESC LIMIT ?" - )) - .bind(limit_i64(limit, "recent request candidate limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_by_provider_id( - &self, - provider_id: &str, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{CANDIDATE_COLUMNS} WHERE provider_id = ? ORDER BY created_at DESC LIMIT ?" - )) - .bind(provider_id) - .bind(limit_i64(limit, "provider request candidate limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn list_finalized_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() || limit == 0 { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(CANDIDATE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(unix_secs_to_ms_i64(since_unix_secs)?) - .push(" AND status IN ('success', 'failed', 'skipped')") - .push(" ORDER BY created_at DESC LIMIT ") - .push_bind(limit_i64(limit, "finalized request candidate limit")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_candidate_row).collect() - } - - async fn count_finalized_statuses_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - "SELECT endpoint_id, status, COUNT(id) AS count FROM request_candidates", - ); - let mut where_clause = WhereClause::new(); - push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(unix_secs_to_ms_i64(since_unix_secs)?) - .push(" AND status IN ('success', 'failed', 'skipped')") - .push(" GROUP BY endpoint_id, status"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(PublicHealthStatusCount { - endpoint_id: row.try_get("endpoint_id").map_sql_err()?, - status: RequestCandidateStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - count: u64::try_from(row.try_get::("count").map_sql_err()?).map_err( - |_| { - DataLayerError::UnexpectedValue( - "public health status count out of range".to_string(), - ) - }, - )?, - }) - }) - .collect() - } - - async fn aggregate_finalized_timeline_by_endpoint_ids_since( - &self, - endpoint_ids: &[String], - since_unix_secs: u64, - until_unix_secs: u64, - segments: u32, - ) -> Result, DataLayerError> { - if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs { - return Ok(Vec::new()); - } - let since_ms = unix_secs_to_ms_i64(since_unix_secs)?; - let until_ms = unix_secs_to_ms_i64(until_unix_secs)?; - let mut builder = QueryBuilder::::new(CANDIDATE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids); - builder - .push(" AND created_at >= ") - .push_bind(since_ms) - .push(" AND created_at <= ") - .push_bind(until_ms) - .push(" AND status IN ('success', 'failed', 'skipped')"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - aggregate_timeline( - rows.iter() - .map(map_candidate_row) - .collect::, _>>()?, - since_unix_secs, - until_unix_secs, - segments, - ) - } -} - -#[async_trait] -impl RequestCandidateWriteRepository for SqliteRequestCandidateRepository { - async fn upsert( - &self, - mut candidate: UpsertRequestCandidateRecord, - ) -> Result { - candidate.sanitize_for_persistence(); - candidate.validate()?; - 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, - mut candidates: Vec, - ) -> Result { - if candidates.is_empty() { - return Ok(0); - } - for candidate in &mut candidates { - candidate.sanitize_for_persistence(); - 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( - &self, - created_before_unix_secs: u64, - limit: usize, - ) -> Result { - if limit == 0 { - return Ok(0); - } - let rows_affected = sqlx::query( - r#" -DELETE FROM request_candidates -WHERE id IN ( - SELECT id - FROM request_candidates - WHERE created_at < ? - ORDER BY created_at ASC, id ASC - LIMIT ? -) -"#, - ) - .bind(unix_secs_to_ms_i64(created_before_unix_secs)?) - .bind(limit_i64(limit, "request candidate delete limit")?) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -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( - connection: &mut SqliteConnection, - candidate: &StoredRequestCandidate, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -INSERT INTO request_candidates ( - id, request_id, user_id, api_key_id, username, api_key_name, - candidate_index, retry_index, provider_id, endpoint_id, key_id, status, - skip_reason, is_cached, status_code, error_type, error_message, latency_ms, - concurrent_requests, extra_data, required_capabilities, created_at, started_at, finished_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(request_id, candidate_index, retry_index) DO UPDATE SET - user_id = excluded.user_id, - api_key_id = excluded.api_key_id, - username = excluded.username, - api_key_name = excluded.api_key_name, - provider_id = excluded.provider_id, - endpoint_id = excluded.endpoint_id, - key_id = excluded.key_id, - status = CASE - WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped') - AND excluded.status IN ('available', 'unused', 'pending', 'streaming') - THEN request_candidates.status - WHEN request_candidates.status = 'pending' - AND excluded.status IN ('available', 'unused') - THEN request_candidates.status - WHEN request_candidates.status = 'streaming' - AND excluded.status IN ('available', 'unused', 'pending') - THEN request_candidates.status - ELSE excluded.status - END, - skip_reason = excluded.skip_reason, - is_cached = excluded.is_cached, - status_code = CASE - WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped') - AND excluded.status IN ('available', 'unused', 'pending', 'streaming') - THEN request_candidates.status_code - WHEN request_candidates.status = 'pending' - AND excluded.status IN ('available', 'unused') - THEN request_candidates.status_code - WHEN request_candidates.status = 'streaming' - AND excluded.status IN ('available', 'unused', 'pending') - THEN request_candidates.status_code - ELSE COALESCE(excluded.status_code, request_candidates.status_code) - END, - error_type = excluded.error_type, - error_message = NULL, - latency_ms = CASE - WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped') - AND excluded.status IN ('available', 'unused', 'pending', 'streaming') - THEN request_candidates.latency_ms - WHEN request_candidates.status = 'pending' - AND excluded.status IN ('available', 'unused') - THEN request_candidates.latency_ms - WHEN request_candidates.status = 'streaming' - AND excluded.status IN ('available', 'unused', 'pending') - THEN request_candidates.latency_ms - ELSE COALESCE(excluded.latency_ms, request_candidates.latency_ms) - END, - concurrent_requests = excluded.concurrent_requests, - extra_data = excluded.extra_data, - required_capabilities = excluded.required_capabilities, - created_at = excluded.created_at, - started_at = excluded.started_at, - finished_at = CASE - WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped') - AND excluded.status IN ('available', 'unused', 'pending', 'streaming') - THEN request_candidates.finished_at - WHEN request_candidates.status = 'pending' - AND excluded.status IN ('available', 'unused') - THEN request_candidates.finished_at - WHEN request_candidates.status = 'streaming' - AND excluded.status IN ('available', 'unused', 'pending') - THEN request_candidates.finished_at - ELSE COALESCE(excluded.finished_at, request_candidates.finished_at) - END -"#, - ) - .bind(&candidate.id) - .bind(&candidate.request_id) - .bind(&candidate.user_id) - .bind(&candidate.api_key_id) - .bind(&candidate.username) - .bind(&candidate.api_key_name) - .bind(to_i32(candidate.candidate_index)?) - .bind(to_i32(candidate.retry_index)?) - .bind(&candidate.provider_id) - .bind(&candidate.endpoint_id) - .bind(&candidate.key_id) - .bind(status_to_database(candidate.status)) - .bind(&candidate.skip_reason) - .bind(candidate.is_cached) - .bind(candidate.status_code.map(i32::from)) - .bind(&candidate.error_type) - .bind(&candidate.error_message) - .bind(candidate.latency_ms.map(to_i32_u64).transpose()?) - .bind(candidate.concurrent_requests.map(to_i32).transpose()?) - .bind(json_to_string(&candidate.extra_data)?) - .bind(json_to_string(&candidate.required_capabilities)?) - .bind(u64_to_i64( - candidate.created_at_unix_ms, - "request candidate created_at", - )?) - .bind(optional_u64_to_i64( - candidate.started_at_unix_ms, - "request candidate started_at", - )?) - .bind(optional_u64_to_i64( - candidate.finished_at_unix_ms, - "request candidate finished_at", - )?) - .execute(connection) - .await - .map_sql_err()?; - Ok(()) -} - -fn merge_candidate( - mut candidate: UpsertRequestCandidateRecord, - existing: Option, -) -> Result { - candidate.sanitize_for_persistence(); - let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| { - request_candidate_lifecycle_would_regress(value.status, candidate.status) - }); - let merged_status = if preserve_existing_lifecycle { - existing - .as_ref() - .map(|value| value.status) - .unwrap_or(candidate.status) - } else { - candidate.status - }; - let created_at_unix_ms = existing - .as_ref() - .map(|value| value.created_at_unix_ms) - .filter(|value| *value > 1000) - .or_else(|| candidate.created_at_unix_ms.filter(|value| *value > 1000)) - .or(candidate.started_at_unix_ms) - .or(candidate.finished_at_unix_ms) - .unwrap_or_else(current_unix_ms); - let id = existing - .as_ref() - .map(|value| value.id.clone()) - .unwrap_or(candidate.id); - let extra_data = merge_json_objects( - existing.as_ref().and_then(|value| value.extra_data.clone()), - candidate.extra_data, - ); - StoredRequestCandidate::new( - id, - candidate.request_id, - existing - .as_ref() - .and_then(|value| value.user_id.clone()) - .or(candidate.user_id), - existing - .as_ref() - .and_then(|value| value.api_key_id.clone()) - .or(candidate.api_key_id), - existing - .as_ref() - .and_then(|value| value.username.clone()) - .or(candidate.username), - existing - .as_ref() - .and_then(|value| value.api_key_name.clone()) - .or(candidate.api_key_name), - to_i32(candidate.candidate_index)?, - to_i32(candidate.retry_index)?, - existing - .as_ref() - .and_then(|value| value.provider_id.clone()) - .or(candidate.provider_id), - existing - .as_ref() - .and_then(|value| value.endpoint_id.clone()) - .or(candidate.endpoint_id), - existing - .as_ref() - .and_then(|value| value.key_id.clone()) - .or(candidate.key_id), - merged_status, - candidate.skip_reason.or_else(|| { - existing - .as_ref() - .and_then(|value| value.skip_reason.clone()) - }), - candidate - .is_cached - .unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)), - if preserve_existing_lifecycle { - existing - .as_ref() - .and_then(|value| value.status_code.map(i32::from)) - } else { - candidate.status_code.map(i32::from).or_else(|| { - existing - .as_ref() - .and_then(|value| value.status_code.map(i32::from)) - }) - }, - if preserve_existing_lifecycle { - existing.as_ref().and_then(|value| value.error_type.clone()) - } else { - candidate - .error_type - .or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())) - }, - None, - if preserve_existing_lifecycle { - match existing.as_ref().and_then(|value| value.latency_ms) { - Some(value) => Some(to_i32_u64(value)?), - None => None, - } - } else { - candidate.latency_ms.map(to_i32_u64).transpose()?.or( - match existing.as_ref().and_then(|value| value.latency_ms) { - Some(value) => Some(to_i32_u64(value)?), - None => None, - }, - ) - }, - candidate.concurrent_requests.map(to_i32).transpose()?.or( - match existing - .as_ref() - .and_then(|value| value.concurrent_requests) - { - Some(value) => Some(to_i32(value)?), - None => None, - }, - ), - extra_data, - candidate.required_capabilities.or_else(|| { - existing - .as_ref() - .and_then(|value| value.required_capabilities.clone()) - }), - u64_to_i64(created_at_unix_ms, "request candidate created_at")?, - existing - .as_ref() - .and_then(|value| value.started_at_unix_ms) - .or(candidate.started_at_unix_ms) - .map(|value| u64_to_i64(value, "request candidate started_at")) - .transpose()?, - if preserve_existing_lifecycle { - existing - .as_ref() - .and_then(|value| value.finished_at_unix_ms) - } else { - candidate.finished_at_unix_ms.or_else(|| { - existing - .as_ref() - .and_then(|value| value.finished_at_unix_ms) - }) - } - .map(|value| u64_to_i64(value, "request candidate finished_at")) - .transpose()?, - ) -} - -fn aggregate_timeline( - candidates: Vec, - since_unix_secs: u64, - until_unix_secs: u64, - segments: u32, -) -> Result, DataLayerError> { - let endpoint_ids = candidates - .iter() - .filter_map(|candidate| candidate.endpoint_id.clone()) - .collect::>(); - let span_ms = until_unix_secs - .saturating_sub(since_unix_secs) - .saturating_mul(1000) - .max(1); - let since_ms = since_unix_secs.saturating_mul(1000); - let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new(); - for candidate in candidates { - let Some(endpoint_id) = candidate.endpoint_id.clone() else { - continue; - }; - let offset = candidate.created_at_unix_ms.saturating_sub(since_ms); - let segment_idx = ((offset.saturating_mul(u64::from(segments))) / span_ms) - .min(u64::from(segments.saturating_sub(1))) as u32; - let bucket = buckets.entry((endpoint_id.clone(), segment_idx)).or_insert( - PublicHealthTimelineBucket { - endpoint_id, - segment_idx, - total_count: 0, - success_count: 0, - failed_count: 0, - min_created_at_unix_ms: Some(candidate.created_at_unix_ms), - max_created_at_unix_ms: Some(candidate.created_at_unix_ms), - }, - ); - bucket.total_count += 1; - if candidate.status == RequestCandidateStatus::Success { - bucket.success_count += 1; - } - if candidate.status == RequestCandidateStatus::Failed { - bucket.failed_count += 1; - } - bucket.min_created_at_unix_ms = bucket - .min_created_at_unix_ms - .map(|value| value.min(candidate.created_at_unix_ms)); - bucket.max_created_at_unix_ms = bucket - .max_created_at_unix_ms - .map(|value| value.max(candidate.created_at_unix_ms)); - } - for endpoint_id in endpoint_ids { - for segment_idx in 0..segments { - buckets.entry((endpoint_id.clone(), segment_idx)).or_insert( - PublicHealthTimelineBucket { - endpoint_id: endpoint_id.clone(), - segment_idx, - total_count: 0, - success_count: 0, - failed_count: 0, - min_created_at_unix_ms: None, - max_created_at_unix_ms: None, - }, - ); - } - } - Ok(buckets.into_values().collect()) -} - -fn map_candidate_row(row: &SqliteRow) -> Result { - StoredRequestCandidate::new( - row.try_get("id").map_sql_err()?, - row.try_get("request_id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("candidate_index").map_sql_err()?, - row.try_get("retry_index").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("endpoint_id").map_sql_err()?, - row.try_get("key_id").map_sql_err()?, - RequestCandidateStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - row.try_get("skip_reason").map_sql_err()?, - row.try_get("is_cached").map_sql_err()?, - row.try_get("status_code").map_sql_err()?, - row.try_get("error_type").map_sql_err()?, - row.try_get("error_message").map_sql_err()?, - row.try_get("latency_ms").map_sql_err()?, - row.try_get("concurrent_requests").map_sql_err()?, - parse_json(row.try_get("extra_data").ok().flatten())?, - parse_json(row.try_get("required_capabilities").ok().flatten())?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("started_at_unix_ms").map_sql_err()?, - row.try_get("finished_at_unix_ms").map_sql_err()?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "request_candidates JSON field is invalid: {err}" - )) - }) - }) - .transpose() -} - -fn json_to_string(value: &Option) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "request_candidates JSON field is unserializable: {err}" - )) - }) - }) - .transpose() -} - -fn merge_json_objects( - existing: Option, - overlay: Option, -) -> Option { - match (existing, overlay) { - ( - Some(serde_json::Value::Object(mut existing_object)), - Some(serde_json::Value::Object(overlay_object)), - ) => { - existing_object.extend(overlay_object); - Some(serde_json::Value::Object(existing_object)) - } - (_existing, Some(overlay)) => Some(overlay), - (existing, None) => existing, - } -} - -fn status_to_database(status: RequestCandidateStatus) -> &'static str { - match status { - RequestCandidateStatus::Available => "available", - RequestCandidateStatus::Unused => "unused", - RequestCandidateStatus::Pending => "pending", - RequestCandidateStatus::Streaming => "streaming", - RequestCandidateStatus::Success => "success", - RequestCandidateStatus::Failed => "failed", - RequestCandidateStatus::Cancelled => "cancelled", - RequestCandidateStatus::Skipped => "skipped", - } -} - -fn current_unix_ms() -> u64 { - chrono::Utc::now().timestamp_millis().max(0) as u64 -} - -fn unix_secs_to_ms_i64(value: u64) -> Result { - let value = value.checked_mul(1000).ok_or_else(|| { - DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string()) - })?; - i64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string()) - }) -} - -fn limit_i64(value: usize, name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {name}: {value}"))) -} - -fn to_i32(value: u32) -> Result { - i32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}")) - }) -} - -fn to_i32_u64(value: u64) -> Result { - i32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}")) - }) -} - -fn u64_to_i64(value: u64, name: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u64_to_i64(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u64_to_i64(value, name)).transpose() -} - -#[cfg(test)] -mod tests { - use super::SqliteRequestCandidateRepository; - use crate::run_migrations; - use aether_data_contracts::repository::candidates::{ - RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, - UpsertRequestCandidateRecord, - }; - use serde_json::json; - - #[tokio::test] - async fn sqlite_repository_writes_and_reads_request_candidates() { - 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 = SqliteRequestCandidateRepository::new(pool.clone()); - let created = repository - .upsert(sample_upsert( - "candidate-1", - RequestCandidateStatus::Pending, - Some(json!({"gateway_execution_runtime": true})), - 1_000_000, - )) - .await - .expect("candidate should insert"); - assert_eq!(created.request_id, "request-1"); - sqlx::query( - "UPDATE request_candidates SET skip_reason = ?, error_type = ?, error_message = ?, extra_data = ?, required_capabilities = ? WHERE request_id = ?", - ) - .bind("legacy skip reason with tenant-secret") - .bind("legacy_error_type_with_token") - .bind("Bearer legacy-secret") - .bind(r#"{"gateway_execution_runtime":true,"request_body":{"password":"secret"}}"#) - .bind(r#"{"streaming":true,"internal_capability":"secret"}"#) - .bind("request-1") - .execute(&pool) - .await - .expect("legacy diagnostics should be injected for the conflict test"); - - let updated = repository - .upsert(sample_upsert( - "candidate-replacement", - RequestCandidateStatus::Success, - Some(json!({"stream_completed": true})), - 1_000_500, - )) - .await - .expect("candidate should update"); - assert_eq!(updated.id, "candidate-1"); - assert_eq!( - updated.extra_data, - Some(json!({ - "gateway_execution_runtime": true, - "stream_completed": true - })) - ); - let raw = sqlx::query( - "SELECT skip_reason, error_type, error_message, extra_data, required_capabilities FROM request_candidates WHERE request_id = ?", - ) - .bind("request-1") - .fetch_one(&pool) - .await - .expect("raw candidate diagnostics should load"); - assert!( - sqlx::Row::try_get::, _>(&raw, "error_message") - .expect("error_message should decode") - .is_none() - ); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "skip_reason") - .expect("skip_reason should decode") - .as_deref(), - Some("unclassified_skip") - ); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "error_type") - .expect("error_type should decode") - .as_deref(), - Some("unclassified_error") - ); - let raw_extra = sqlx::Row::try_get::, _>(&raw, "extra_data") - .expect("extra_data should decode") - .and_then(|value| serde_json::from_str::(&value).ok()); - assert_eq!(raw_extra, updated.extra_data); - let raw_capabilities = - sqlx::Row::try_get::, _>(&raw, "required_capabilities") - .expect("required_capabilities should decode") - .and_then(|value| serde_json::from_str::(&value).ok()); - assert_eq!(raw_capabilities, Some(json!({"streaming": true}))); - - sqlx::query( - "UPDATE request_candidates SET skip_reason = ?, error_type = ? WHERE request_id = ?", - ) - .bind("pool_cooldown") - .bind(" FirstByteTimeout ") - .bind("request-1") - .execute(&pool) - .await - .expect("known legacy diagnostics should be injected for the regression test"); - let mut late = sample_upsert( - "candidate-late-terminal", - RequestCandidateStatus::Failed, - Some(json!({"cache_1h": true})), - 1_000_250, - ); - late.user_id = Some("attacker-user".to_string()); - late.api_key_id = Some("attacker-api-key".to_string()); - late.username = Some("mallory".to_string()); - late.api_key_name = Some("attacker-key".to_string()); - late.provider_id = Some("attacker-provider".to_string()); - late.endpoint_id = Some("attacker-endpoint".to_string()); - late.key_id = Some("attacker-provider-key".to_string()); - late.error_type = Some("upstream5xx".to_string()); - let late_terminal = repository - .upsert(late) - .await - .expect("late terminal candidate should not replace the first terminal fact"); - assert_eq!(late_terminal.id, "candidate-1"); - assert_eq!(late_terminal.status, RequestCandidateStatus::Success); - assert_eq!(late_terminal.user_id.as_deref(), Some("user-1")); - assert_eq!(late_terminal.api_key_id.as_deref(), Some("key-1")); - assert_eq!(late_terminal.provider_id.as_deref(), Some("provider-1")); - assert_eq!(late_terminal.endpoint_id.as_deref(), Some("endpoint-1")); - assert_eq!(late_terminal.key_id.as_deref(), Some("provider-key-1")); - assert_eq!(late_terminal.skip_reason.as_deref(), Some("pool_cooldown")); - assert_eq!( - late_terminal.error_type.as_deref(), - Some("first_byte_timeout") - ); - assert_eq!(late_terminal.finished_at_unix_ms, Some(1_000_502)); - assert_eq!( - late_terminal.extra_data, - Some(json!({ - "cache_1h": true, - "gateway_execution_runtime": true, - "stream_completed": true - })) - ); - let raw = sqlx::query( - "SELECT skip_reason, error_type FROM request_candidates WHERE request_id = ?", - ) - .bind("request-1") - .fetch_one(&pool) - .await - .expect("raw candidate classifications should load"); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "skip_reason") - .expect("skip_reason should decode") - .as_deref(), - Some("pool_cooldown") - ); - assert_eq!( - sqlx::Row::try_get::, _>(&raw, "error_type") - .expect("error_type should decode") - .as_deref(), - Some("first_byte_timeout") - ); - - assert_eq!( - repository - .list_by_request_id("request-1") - .await - .expect("request list should load") - .len(), - 1 - ); - assert_eq!( - repository - .count_finalized_statuses_by_endpoint_ids_since(&["endpoint-1".to_string()], 900) - .await - .expect("status counts should load")[0] - .count, - 1 - ); - assert_eq!( - repository - .aggregate_finalized_timeline_by_endpoint_ids_since( - &["endpoint-1".to_string()], - 900, - 1200, - 3, - ) - .await - .expect("timeline should load") - .len(), - 3 - ); - assert_eq!( - repository - .delete_created_before(2_000, 10) - .await - .expect("old candidates should delete"), - 1 - ); - } - - #[tokio::test] - 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!({"gateway_execution_runtime": 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 extra_data = match writer { - 0 => json!({"stream_completed": true}), - 1 => json!({"cache_1h": true}), - 2 => json!({"first_byte_time_ms": 2}), - 3 => json!({"pool_key_index": 3}), - 4 => json!({"priority_slot": 4}), - 5 => json!({"ranking_index": 5}), - 6 => json!({"phase": "provider_request"}), - 7 => json!({"provider_api_format": "openai:responses"}), - 8 => json!({"client_api_format": "openai:chat"}), - 9 => json!({"execution_strategy": "local_cross_format"}), - _ => unreachable!("writer index is bounded by WRITERS"), - }; - let mut candidate = sample_upsert( - format!("writer-{writer}").as_str(), - status, - Some(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)); - assert_eq!( - candidate.extra_data, - Some(json!({ - "cache_1h": true, - "client_api_format": "openai:chat", - "execution_strategy": "local_cross_format", - "first_byte_time_ms": 2, - "gateway_execution_runtime": true, - "phase": "provider_request", - "pool_key_index": 3, - "priority_slot": 4, - "provider_api_format": "openai:responses", - "ranking_index": 5, - "stream_completed": true - })) - ); - - 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:") - .await - .expect("sqlite pool should connect"); - run_migrations(&pool) - .await - .expect("sqlite migrations should run"); - let repository = SqliteRequestCandidateRepository::new(pool); - let request_id = "request-batch"; - - let mut pending = sample_upsert( - "batch-first", - RequestCandidateStatus::Pending, - Some(json!({"gateway_execution_runtime": true})), - 4_000_000, - ); - pending.request_id = request_id.to_string(); - let mut streaming = sample_upsert( - "batch-second", - RequestCandidateStatus::Streaming, - Some(json!({"stream_completed": true})), - 4_000_100, - ); - streaming.request_id = request_id.to_string(); - streaming.is_cached = None; - let mut success = sample_upsert( - "batch-third", - RequestCandidateStatus::Success, - Some(json!({"cache_1h": 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!({"first_byte_time_ms": 42})), - 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); - - 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!({ - "cache_1h": true, - "first_byte_time_ms": 42, - "gateway_execution_runtime": true, - "stream_completed": true - })) - ); - - let rollback_request_id = "request-batch-rollback"; - let mut valid = sample_upsert( - "rollback-valid", - RequestCandidateStatus::Pending, - None, - 5_000_000, - ); - valid.request_id = rollback_request_id.to_string(); - let mut invalid = sample_upsert( - "rollback-invalid", - RequestCandidateStatus::Success, - None, - 5_000_100, - ); - invalid.request_id = rollback_request_id.to_string(); - 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()); - } - - fn sample_upsert( - id: &str, - status: RequestCandidateStatus, - extra_data: Option, - created_at_unix_ms: u64, - ) -> UpsertRequestCandidateRecord { - UpsertRequestCandidateRecord { - id: id.to_string(), - request_id: "request-1".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, - 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/sqlite/src/error.rs b/crates/aether-data/adapters/sqlite/src/error.rs deleted file mode 100644 index 4f25f36cd..000000000 --- a/crates/aether-data/adapters/sqlite/src/error.rs +++ /dev/null @@ -1,11 +0,0 @@ -use aether_data_contracts::DataLayerError; - -pub(crate) trait SqlResultExt { - fn map_sql_err(self) -> Result; -} - -impl SqlResultExt for Result { - fn map_sql_err(self) -> Result { - self.map_err(DataLayerError::sql) - } -} diff --git a/crates/aether-data/adapters/sqlite/src/gemini_file_mappings.rs b/crates/aether-data/adapters/sqlite/src/gemini_file_mappings.rs deleted file mode 100644 index 71a9a2fee..000000000 --- a/crates/aether-data/adapters/sqlite/src/gemini_file_mappings.rs +++ /dev/null @@ -1,667 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::gemini_file_mappings::{ - GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository, - GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping, - StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_ci_contains_any, push_limit_offset, SqlDialect, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -#[derive(Debug, Clone)] -pub struct SqliteGeminiFileMappingRepository { - pool: SqlitePool, -} - -impl SqliteGeminiFileMappingRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn reload_by_file_name( - &self, - file_name: &str, - ) -> Result { - self.find_by_file_name(file_name).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("gemini file mapping missing after write".to_string()) - }) - } -} - -#[async_trait] -impl GeminiFileMappingReadRepository for SqliteGeminiFileMappingRepository { - async fn find_by_file_name( - &self, - file_name: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE file_name = ? -LIMIT 1 -"#, - ) - .bind(file_name) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn find_active_by_file_name_for_user( - &self, - file_name: &str, - user_id: &str, - now_unix_secs: u64, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE file_name = ? AND user_id = ? AND expires_at > ? -LIMIT 1 -"#, - ) - .bind(file_name) - .bind(user_id) - .bind(i64_from_u64( - now_unix_secs, - "gemini_file_mappings.owner_read_now", - )?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn find_active_by_file_name_for_owner( - &self, - file_name: &str, - key_id: &str, - user_id: &str, - now_unix_secs: u64, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE file_name = ? AND key_id = ? AND user_id = ? AND expires_at > ? -LIMIT 1 -"#, - ) - .bind(file_name) - .bind(key_id) - .bind(user_id) - .bind(i64_from_u64( - now_unix_secs, - "gemini_file_mappings.owner_read_now", - )?) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - - row.as_ref().map(map_row).transpose() - } - - async fn list_mappings( - &self, - query: &GeminiFileMappingListQuery, - ) -> Result { - let total = build_list_count_query(query) - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let rows = build_list_rows_query(query) - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows.iter().map(map_row).collect::, _>>()?; - Ok(StoredGeminiFileMappingListPage { - items, - total: usize::try_from(total).unwrap_or_default(), - }) - } - - async fn summarize_mappings( - &self, - now_unix_secs: u64, - ) -> Result { - let totals = sqlx::query( - r#" -SELECT - COUNT(*) AS total_mappings, - SUM(CASE WHEN expires_at > ? THEN 1 ELSE 0 END) AS active_mappings -FROM gemini_file_mappings -"#, - ) - .bind(now_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let total_mappings = - usize::try_from(totals.try_get::("total_mappings").map_sql_err()?) - .unwrap_or_default(); - let active_mappings = usize::try_from( - totals - .try_get::, _>("active_mappings") - .map_sql_err()? - .unwrap_or(0), - ) - .unwrap_or_default(); - let by_mime_type_rows = sqlx::query( - r#" -SELECT - COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type, - COUNT(*) AS count -FROM gemini_file_mappings -WHERE expires_at > ? -GROUP BY COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') -ORDER BY mime_type ASC -"#, - ) - .bind(now_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let by_mime_type = by_mime_type_rows - .iter() - .map(|row| { - Ok(GeminiFileMappingMimeTypeCount { - mime_type: row.try_get("mime_type").map_sql_err()?, - count: usize::try_from(row.try_get::("count").map_sql_err()?) - .unwrap_or_default(), - }) - }) - .collect::, DataLayerError>>()?; - Ok(GeminiFileMappingStats { - total_mappings, - active_mappings, - expired_mappings: total_mappings.saturating_sub(active_mappings), - by_mime_type, - }) - } -} - -#[async_trait] -impl GeminiFileMappingWriteRepository for SqliteGeminiFileMappingRepository { - async fn upsert( - &self, - record: UpsertGeminiFileMappingRecord, - ) -> Result { - record.validate()?; - sqlx::query( - r#" -INSERT INTO gemini_file_mappings ( - id, file_name, key_id, user_id, display_name, mime_type, source_hash, - created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(file_name) DO UPDATE SET - key_id = excluded.key_id, - user_id = excluded.user_id, - display_name = excluded.display_name, - mime_type = excluded.mime_type, - source_hash = excluded.source_hash, - expires_at = excluded.expires_at -"#, - ) - .bind(&record.id) - .bind(&record.file_name) - .bind(&record.key_id) - .bind(&record.user_id) - .bind(&record.display_name) - .bind(&record.mime_type) - .bind(&record.source_hash) - .bind(current_unix_secs() as i64) - .bind(i64_from_u64( - record.expires_at_unix_secs, - "gemini_file_mappings.expires_at", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - self.reload_by_file_name(&record.file_name).await - } - - async fn upsert_if_owner_matches( - &self, - record: UpsertGeminiFileMappingRecord, - ) -> Result, DataLayerError> { - record.validate()?; - let rows_affected = sqlx::query( - r#" -INSERT INTO gemini_file_mappings ( - id, file_name, key_id, user_id, display_name, mime_type, source_hash, - created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(file_name) DO UPDATE SET - display_name = excluded.display_name, - mime_type = excluded.mime_type, - source_hash = excluded.source_hash, - expires_at = excluded.expires_at -WHERE gemini_file_mappings.key_id = excluded.key_id - AND gemini_file_mappings.user_id IS excluded.user_id -"#, - ) - .bind(&record.id) - .bind(&record.file_name) - .bind(&record.key_id) - .bind(&record.user_id) - .bind(&record.display_name) - .bind(&record.mime_type) - .bind(&record.source_hash) - .bind(current_unix_secs() as i64) - .bind(i64_from_u64( - record.expires_at_unix_secs, - "gemini_file_mappings.expires_at", - )?) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - return Ok(None); - } - self.reload_by_file_name(&record.file_name).await.map(Some) - } - - async fn delete_by_file_name(&self, file_name: &str) -> Result { - let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE file_name = ?") - .bind(file_name) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_file_name_for_user( - &self, - file_name: &str, - user_id: &str, - ) -> Result { - let rows_affected = - sqlx::query("DELETE FROM gemini_file_mappings WHERE file_name = ? AND user_id = ?") - .bind(file_name) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_file_name_for_owner( - &self, - file_name: &str, - key_id: &str, - user_id: &str, - ) -> Result { - let rows_affected = sqlx::query( - "DELETE FROM gemini_file_mappings WHERE file_name = ? AND key_id = ? AND user_id = ?", - ) - .bind(file_name) - .bind(key_id) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn delete_by_id( - &self, - mapping_id: &str, - ) -> Result, DataLayerError> { - let existing = sqlx::query( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(mapping_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - let Some(existing) = existing else { - return Ok(None); - }; - sqlx::query("DELETE FROM gemini_file_mappings WHERE id = ?") - .bind(mapping_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(Some(map_row(&existing)?)) - } - - async fn delete_expired_before(&self, now_unix_secs: u64) -> Result { - let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE expires_at <= ?") - .bind(now_unix_secs as i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -fn build_list_count_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Sqlite> { - let mut builder = - QueryBuilder::::new("SELECT COUNT(*) AS total FROM gemini_file_mappings"); - let mut where_clause = WhereClause::new(); - apply_list_filters(&mut builder, &mut where_clause, query); - builder -} - -fn build_list_rows_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Sqlite> { - let mut builder = QueryBuilder::::new( - r#" -SELECT - id, - file_name, - key_id, - user_id, - display_name, - mime_type, - source_hash, - created_at AS created_at_unix_ms, - expires_at AS expires_at_unix_secs -FROM gemini_file_mappings -"#, - ); - let mut where_clause = WhereClause::new(); - apply_list_filters(&mut builder, &mut where_clause, query); - builder.push(" ORDER BY created_at DESC, file_name ASC"); - push_limit_offset( - &mut builder, - i64::try_from(query.limit).unwrap_or(i64::MAX), - i64::try_from(query.offset).unwrap_or(i64::MAX), - ); - builder -} - -fn apply_list_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - where_clause: &mut WhereClause, - query: &GeminiFileMappingListQuery, -) { - if let Some(user_id) = query.user_id.as_deref() { - where_clause.push_next(builder); - builder.push("user_id = "); - builder.push_bind(user_id.to_string()); - } - if !query.include_expired { - where_clause.push_next(builder); - builder.push("expires_at > "); - builder.push_bind(query.now_unix_secs as i64); - } - if let Some(search) = query - .search - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - push_ci_contains_any( - builder, - where_clause, - SqlDialect::Sqlite, - &["file_name", "COALESCE(display_name, '')"], - search, - ); - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn i64_from_u64(value: u64, field_name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) -} - -fn map_row(row: &SqliteRow) -> Result { - Ok(StoredGeminiFileMapping { - id: row.try_get("id").map_sql_err()?, - file_name: row.try_get("file_name").map_sql_err()?, - key_id: row.try_get("key_id").map_sql_err()?, - user_id: row.try_get("user_id").ok().flatten(), - display_name: row.try_get("display_name").ok().flatten(), - mime_type: row.try_get("mime_type").ok().flatten(), - source_hash: row.try_get("source_hash").ok().flatten(), - created_at_unix_ms: u64::try_from( - row.try_get::("created_at_unix_ms").map_sql_err()?, - ) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "gemini_file_mappings.created_at is invalid".to_string(), - ) - })?, - expires_at_unix_secs: u64::try_from( - row.try_get::("expires_at_unix_secs") - .map_sql_err()?, - ) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "gemini_file_mappings.expires_at is invalid".to_string(), - ) - })?, - }) -} - -#[cfg(test)] -mod tests { - use super::SqliteGeminiFileMappingRepository; - use crate::run_migrations; - use aether_data_contracts::repository::gemini_file_mappings::{ - GeminiFileMappingListQuery, GeminiFileMappingReadRepository, - GeminiFileMappingWriteRepository, UpsertGeminiFileMappingRecord, - }; - - #[tokio::test] - async fn sqlite_repository_round_trips_gemini_file_mappings() { - 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 = SqliteGeminiFileMappingRepository::new(pool); - let created = repository - .upsert(UpsertGeminiFileMappingRecord { - id: "mapping-1".to_string(), - file_name: "files/example.png".to_string(), - key_id: "key-1".to_string(), - user_id: Some("user-1".to_string()), - display_name: Some("Example".to_string()), - mime_type: Some("image/png".to_string()), - source_hash: Some("hash-1".to_string()), - expires_at_unix_secs: 300, - }) - .await - .expect("mapping should upsert"); - assert_eq!(created.id, "mapping-1"); - assert_eq!(created.mime_type, Some("image/png".to_string())); - - let updated = repository - .upsert(UpsertGeminiFileMappingRecord { - id: "mapping-replacement".to_string(), - file_name: "files/example.png".to_string(), - key_id: "key-2".to_string(), - user_id: Some("user-2".to_string()), - display_name: Some("Updated".to_string()), - mime_type: Some("image/jpeg".to_string()), - source_hash: Some("hash-2".to_string()), - expires_at_unix_secs: 500, - }) - .await - .expect("mapping should update"); - assert_eq!(updated.id, "mapping-1"); - assert_eq!(updated.key_id, "key-2"); - - let guarded_reassignment = repository - .upsert_if_owner_matches(UpsertGeminiFileMappingRecord { - id: "mapping-attacker".to_string(), - file_name: "files/example.png".to_string(), - key_id: "key-attacker".to_string(), - user_id: Some("user-attacker".to_string()), - display_name: Some("Attacker".to_string()), - mime_type: Some("application/octet-stream".to_string()), - source_hash: Some("hash-attacker".to_string()), - expires_at_unix_secs: 900, - }) - .await - .expect("guarded reassignment should run"); - assert!(guarded_reassignment.is_none()); - let after_reassignment = repository - .find_by_file_name("files/example.png") - .await - .expect("mapping should read") - .expect("mapping should remain"); - assert_eq!(after_reassignment.key_id, "key-2"); - assert_eq!(after_reassignment.user_id.as_deref(), Some("user-2")); - - let guarded_refresh = repository - .upsert_if_owner_matches(UpsertGeminiFileMappingRecord { - id: "mapping-refresh".to_string(), - file_name: "files/example.png".to_string(), - key_id: "key-2".to_string(), - user_id: Some("user-2".to_string()), - display_name: Some("Updated".to_string()), - mime_type: Some("image/jpeg".to_string()), - source_hash: Some("hash-refreshed".to_string()), - expires_at_unix_secs: 500, - }) - .await - .expect("same-owner refresh should run"); - assert!(guarded_refresh.is_some()); - - assert!(repository - .find_active_by_file_name_for_user("files/example.png", "user-2", 400) - .await - .expect("owner-scoped read should run") - .is_some()); - assert!(repository - .find_active_by_file_name_for_user("files/example.png", "user-1", 400) - .await - .expect("foreign owner read should run") - .is_none()); - assert!(repository - .find_active_by_file_name_for_owner("files/example.png", "key-2", "user-2", 400,) - .await - .expect("provider owner read should run") - .is_some()); - assert!(repository - .find_active_by_file_name_for_owner("files/example.png", "key-attacker", "user-2", 400,) - .await - .expect("foreign provider read should run") - .is_none()); - assert!(repository - .find_active_by_file_name_for_user("files/example.png", "user-2", 500) - .await - .expect("expired owner read should run") - .is_none()); - - assert!(!repository - .delete_by_file_name_for_owner("files/example.png", "key-attacker", "user-2") - .await - .expect("wrong-key owner delete should run")); - - let page = repository - .list_mappings(&GeminiFileMappingListQuery { - user_id: Some("user-2".to_string()), - include_expired: false, - search: Some("updated".to_string()), - offset: 0, - limit: 10, - now_unix_secs: 400, - }) - .await - .expect("mappings should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].file_name, "files/example.png"); - - assert!(!repository - .delete_by_file_name_for_user("files/example.png", "user-1") - .await - .expect("non-owner delete should run")); - assert!(repository - .find_by_file_name("files/example.png") - .await - .expect("mapping should remain") - .is_some()); - - let stats = repository - .summarize_mappings(400) - .await - .expect("stats should load"); - assert_eq!(stats.total_mappings, 1); - assert_eq!(stats.active_mappings, 1); - assert_eq!(stats.by_mime_type[0].mime_type, "image/jpeg"); - - assert_eq!( - repository - .delete_expired_before(600) - .await - .expect("expired mappings should delete"), - 1 - ); - assert!(repository - .find_by_file_name("files/example.png") - .await - .expect("find should run") - .is_none()); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/global_models.rs b/crates/aether-data/adapters/sqlite/src/global_models.rs deleted file mode 100644 index fb74a5507..000000000 --- a/crates/aether-data/adapters/sqlite/src/global_models.rs +++ /dev/null @@ -1,1486 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::global_models::{ - metadata_supports_embedding, AdminGlobalModelListQuery, AdminProviderModelListQuery, - 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::{sqlite_optional_real, SqlitePool}; - -const LIST_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#" -SELECT - id, - name, - display_name, - is_active, - CAST(default_price_per_request AS REAL) AS default_price_per_request, - default_tiered_pricing, - supported_capabilities, - config, - 0 AS usage_count -FROM global_models -"#; - -const COUNT_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#" -SELECT COUNT(id) 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, - 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, - COALESCE( - m.supports_vision, - CASE - WHEN json_extract(gm.config, '$.vision') IS NULL THEN NULL - WHEN LOWER(CAST(json_extract(gm.config, '$.vision') AS TEXT)) 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(CAST(json_extract(gm.config, '$.function_calling') AS TEXT)) 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(CAST(json_extract(gm.config, '$.streaming') AS TEXT)) 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, - COUNT(id) AS total_models, - COALESCE(SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END), 0) 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, - m.provider_model_mappings, - CAST(m.price_per_request AS REAL) AS 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, - COALESCE(m.is_available, 1) AS 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, - CAST(gm.default_price_per_request AS REAL) 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 -"#; - -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 REAL) AS default_price_per_request, - gm.default_tiered_pricing, - gm.supported_capabilities, - gm.config, - COALESCE(gm_stats.provider_count, 0) AS provider_count, - COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count, - COALESCE(gm.usage_count, 0) 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, - COUNT(DISTINCT m.provider_id) AS provider_count, - 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 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 COUNT(id) 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 SqliteGlobalModelReadRepository { - pool: SqlitePool, -} - -impl SqliteGlobalModelReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn create_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - sqlx::query( - r#" -INSERT INTO models ( - id, - provider_id, - global_model_id, - provider_model_name, - provider_model_mappings, - price_per_request, - tiered_pricing, - supports_vision, - supports_function_calling, - supports_streaming, - supports_extended_thinking, - supports_image_generation, - is_active, - is_available, - config, - created_at, - updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.provider_id) - .bind(&record.global_model_id) - .bind(&record.provider_model_name) - .bind(optional_json_to_string( - &record.provider_model_mappings, - "models.provider_model_mappings", - )?) - .bind(record.price_per_request) - .bind(optional_json_to_string( - &record.tiered_pricing, - "models.tiered_pricing", - )?) - .bind(record.supports_vision) - .bind(record.supports_function_calling) - .bind(record.supports_streaming) - .bind(record.supports_extended_thinking) - .bind(record.supports_image_generation) - .bind(record.is_active) - .bind(record.is_available) - .bind(optional_json_to_string(&record.config, "models.config")?) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.get_admin_provider_model(&record.provider_id, &record.id) - .await - } - - pub async fn update_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let updated = sqlx::query( - r#" -UPDATE models -SET - global_model_id = ?, - provider_model_name = ?, - provider_model_mappings = ?, - price_per_request = ?, - tiered_pricing = ?, - supports_vision = ?, - supports_function_calling = ?, - supports_streaming = ?, - supports_extended_thinking = ?, - supports_image_generation = ?, - is_active = ?, - is_available = ?, - config = ?, - updated_at = ? -WHERE id = ? - AND provider_id = ? -"#, - ) - .bind(&record.global_model_id) - .bind(&record.provider_model_name) - .bind(optional_json_to_string( - &record.provider_model_mappings, - "models.provider_model_mappings", - )?) - .bind(record.price_per_request) - .bind(optional_json_to_string( - &record.tiered_pricing, - "models.tiered_pricing", - )?) - .bind(record.supports_vision) - .bind(record.supports_function_calling) - .bind(record.supports_streaming) - .bind(record.supports_extended_thinking) - .bind(record.supports_image_generation) - .bind(record.is_active) - .bind(record.is_available) - .bind(optional_json_to_string(&record.config, "models.config")?) - .bind(now as i64) - .bind(&record.id) - .bind(&record.provider_id) - .execute(&self.pool) - .await - .map_sql_err()?; - - if updated.rows_affected() == 0 { - return Ok(None); - } - - self.get_admin_provider_model(&record.provider_id, &record.id) - .await - } - - pub async fn delete_admin_provider_model( - &self, - provider_id: &str, - model_id: &str, - ) -> Result { - let deleted = sqlx::query( - r#" -DELETE FROM models -WHERE provider_id = ? - AND id = ? -"#, - ) - .bind(provider_id) - .bind(model_id) - .execute(&self.pool) - .await - .map_sql_err()?; - - Ok(deleted.rows_affected() > 0) - } - - pub async fn create_admin_global_model( - &self, - record: &CreateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let usage_count = - optional_admin_global_model_usage_count_i64(record.usage_count)?.unwrap_or_default(); - sqlx::query( - r#" -INSERT INTO global_models ( - id, - name, - display_name, - is_active, - default_price_per_request, - default_tiered_pricing, - supported_capabilities, - usage_count, - config, - created_at, - updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.name) - .bind(&record.display_name) - .bind(record.is_active) - .bind(record.default_price_per_request) - .bind(optional_json_to_string( - &record.default_tiered_pricing, - "global_models.default_tiered_pricing", - )?) - .bind(optional_json_to_string( - &record.supported_capabilities, - "global_models.supported_capabilities", - )?) - .bind(usage_count) - .bind(optional_json_to_string( - &record.config, - "global_models.config", - )?) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.get_admin_global_model_by_id(&record.id).await - } - - pub async fn update_admin_global_model( - &self, - record: &UpdateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let usage_count = optional_admin_global_model_usage_count_i64(record.usage_count)?; - let updated = sqlx::query( - r#" -UPDATE global_models -SET - display_name = ?, - is_active = ?, - default_price_per_request = ?, - default_tiered_pricing = ?, - supported_capabilities = ?, - config = ?, - usage_count = COALESCE(?, usage_count), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&record.display_name) - .bind(record.is_active) - .bind(record.default_price_per_request) - .bind(optional_json_to_string( - &record.default_tiered_pricing, - "global_models.default_tiered_pricing", - )?) - .bind(optional_json_to_string( - &record.supported_capabilities, - "global_models.supported_capabilities", - )?) - .bind(optional_json_to_string( - &record.config, - "global_models.config", - )?) - .bind(usage_count) - .bind(now as i64) - .bind(&record.id) - .execute(&self.pool) - .await - .map_sql_err()?; - - if updated.rows_affected() == 0 { - return Ok(None); - } - - self.get_admin_global_model_by_id(&record.id).await - } - - pub async fn delete_admin_global_model( - &self, - global_model_id: &str, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - - sqlx::query( - r#" -DELETE FROM models -WHERE global_model_id = ? -"#, - ) - .bind(global_model_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let deleted = sqlx::query( - r#" -DELETE FROM global_models -WHERE id = ? -"#, - ) - .bind(global_model_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - - Ok(deleted.rows_affected() > 0) - } -} - -#[async_trait] -impl GlobalModelReadRepository for SqliteGlobalModelReadRepository { - async fn list_public_models( - &self, - query: &PublicGlobalModelQuery, - ) -> Result { - 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> { - let row = sqlx::query( - r#" -SELECT - id, - name, - display_name, - is_active, - CAST(default_price_per_request AS REAL) AS default_price_per_request, - default_tiered_pricing, - supported_capabilities, - config, - 0 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> { - 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> { - 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 { - 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 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> { - 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> { - 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( - &self, - provider_id: &str, - model_id: &str, - ) -> Result, DataLayerError> { - 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> { - 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> { - 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> { - 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> { - 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> { - 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() - } -} - -#[async_trait] -impl GlobalModelWriteRepository for SqliteGlobalModelReadRepository { - async fn create_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - Self::create_admin_provider_model(self, record).await - } - - async fn update_admin_provider_model( - &self, - record: &UpsertAdminProviderModelRecord, - ) -> Result, DataLayerError> { - Self::update_admin_provider_model(self, record).await - } - - async fn delete_admin_provider_model( - &self, - provider_id: &str, - model_id: &str, - ) -> Result { - Self::delete_admin_provider_model(self, provider_id, model_id).await - } - - async fn create_admin_global_model( - &self, - record: &CreateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - Self::create_admin_global_model(self, record).await - } - - async fn update_admin_global_model( - &self, - record: &UpdateAdminGlobalModelRecord, - ) -> Result, DataLayerError> { - Self::update_admin_global_model(self, record).await - } - - async fn delete_admin_global_model( - &self, - global_model_id: &str, - ) -> Result { - Self::delete_admin_global_model(self, global_model_id).await - } -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_u64(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn first_tier_price(value: Option<&serde_json::Value>, key: &str) -> Option { - value - .and_then(|value| value.get("tiers")) - .and_then(serde_json::Value::as_array) - .and_then(|tiers| tiers.first()) - .and_then(|tier| tier.get(key)) - .and_then(serde_json::Value::as_f64) -} - -fn apply_public_model_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - 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<'_, Sqlite>, - 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<'_, Sqlite>, - 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, Sqlite> { - 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: &SqliteRow) -> Result { - StoredPublicGlobalModel::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - sqlite_optional_real(row, "default_price_per_request")?, - optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("supported_capabilities").map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "global_models.config")?, - row.try_get::("usage_count").map_sql_err()?.max(0) as u64, - ) -} - -fn map_admin_global_model_row(row: &SqliteRow) -> 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()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - sqlite_optional_real(row, "default_price_per_request")?, - optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("supported_capabilities").map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "global_models.config")?, - provider_count, - active_provider_count, - usage_count, - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "global_models.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "global_models.updated_at", - )?, - ) -} - -fn map_admin_provider_model_row( - row: &SqliteRow, -) -> Result { - StoredAdminProviderModel::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - row.try_get("provider_model_name").map_sql_err()?, - optional_json_from_string( - row.try_get("provider_model_mappings").map_sql_err()?, - "models.provider_model_mappings", - )?, - sqlite_optional_real(row, "price_per_request")?, - optional_json_from_string( - row.try_get("tiered_pricing").map_sql_err()?, - "models.tiered_pricing", - )?, - row.try_get("supports_vision").map_sql_err()?, - row.try_get("supports_function_calling").map_sql_err()?, - row.try_get("supports_streaming").map_sql_err()?, - row.try_get("supports_extended_thinking").map_sql_err()?, - row.try_get("supports_image_generation").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_available").map_sql_err()?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "models.config")?, - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "models.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "models.updated_at", - )?, - row.try_get("global_model_name").map_sql_err()?, - row.try_get("global_model_display_name").map_sql_err()?, - sqlite_optional_real(row, "global_model_default_price_per_request")?, - optional_json_from_string( - row.try_get("global_model_default_tiered_pricing") - .map_sql_err()?, - "global_models.default_tiered_pricing", - )?, - optional_json_from_string( - row.try_get("global_model_supported_capabilities") - .map_sql_err()?, - "global_models.supported_capabilities", - )?, - optional_json_from_string( - row.try_get("global_model_config").map_sql_err()?, - "global_models.config", - )?, - ) -} - -fn map_public_catalog_model_row( - row: &SqliteRow, -) -> Result { - let global_model_config = optional_json_from_string( - row.try_get("global_model_config").map_sql_err()?, - "global_models.config", - )?; - let global_model_supported_capabilities = optional_json_from_string( - row.try_get("global_model_supported_capabilities") - .map_sql_err()?, - "global_models.supported_capabilities", - )?; - let model_config = - optional_json_from_string(row.try_get("model_config").map_sql_err()?, "models.config")?; - let tiered_pricing = optional_json_from_string( - row.try_get("tiered_pricing").map_sql_err()?, - "models.tiered_pricing", - )?; - let default_tiered_pricing = optional_json_from_string( - row.try_get("default_tiered_pricing").map_sql_err()?, - "global_models.default_tiered_pricing", - )?; - let pricing = tiered_pricing.as_ref().or(default_tiered_pricing.as_ref()); - let global_model_is_active = row - .try_get::, _>("global_model_is_active") - .map_sql_err()? - .unwrap_or(true); - let model_is_active: bool = row.try_get("is_active").map_sql_err()?; - let provider_is_active: bool = row.try_get("provider_is_active").map_sql_err()?; - - StoredPublicCatalogModel::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_name").map_sql_err()?, - row.try_get("provider_model_name").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - global_model_config - .as_ref() - .and_then(|value| value.get("description")) - .and_then(serde_json::Value::as_str) - .map(ToString::to_string), - global_model_config - .as_ref() - .and_then(|value| value.get("icon_url")) - .and_then(serde_json::Value::as_str) - .map(ToString::to_string), - 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()?, - row.try_get("supports_function_calling").map_sql_err()?, - row.try_get("supports_streaming").map_sql_err()?, - metadata_supports_embedding( - global_model_supported_capabilities.as_ref(), - global_model_config.as_ref(), - model_config.as_ref(), - ), - model_is_active && provider_is_active && global_model_is_active, - ) -} - -fn map_provider_model_stats_row( - row: &SqliteRow, -) -> Result { - StoredProviderModelStats::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("total_models").map_sql_err()?, - row.try_get::, _>("active_models") - .map_sql_err()? - .unwrap_or(0), - ) -} - -fn map_active_global_model_row( - row: &SqliteRow, -) -> Result { - StoredProviderActiveGlobalModel::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("global_model_id").map_sql_err()?, - ) -} - -fn optional_admin_global_model_usage_count_i64( - value: Option, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput( - "global_models.usage_count exceeds i64 range".to_string(), - ) - }) - }) - .transpose() -} - -#[cfg(test)] -mod tests { - use super::SqliteGlobalModelReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::global_models::{ - AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord, - GlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, - PublicGlobalModelQuery, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord, - }; - use serde_json::json; - - #[tokio::test] - async fn sqlite_repository_reads_global_model_contract_views() { - 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_rows(&pool).await; - - let repository = SqliteGlobalModelReadRepository::new(pool); - let public = repository - .list_public_models(&PublicGlobalModelQuery { - offset: 0, - limit: 10, - is_active: Some(true), - search: Some("gpt".to_string()), - }) - .await - .expect("public models should load"); - assert_eq!(public.total, 1); - assert_eq!(public.items[0].name, "gpt-4.1"); - - let catalog = repository - .search_public_catalog_models(&PublicCatalogModelSearchQuery { - search: "provider".to_string(), - provider_id: Some("provider-1".to_string()), - limit: 10, - }) - .await - .expect("catalog search should load"); - assert_eq!(catalog.len(), 1); - assert_eq!(catalog[0].input_price_per_1m, Some(2.0)); - - let catalog_list = repository - .list_public_catalog_models(&PublicCatalogModelListQuery { - provider_id: None, - offset: 0, - limit: 10, - }) - .await - .expect("catalog list should load"); - assert_eq!(catalog_list.len(), 2); - assert_eq!(catalog_list[0].provider_id, "provider-1"); - assert_eq!(catalog_list[1].provider_id, "provider-3"); - - let admin_globals = repository - .list_admin_global_models(&AdminGlobalModelListQuery { - offset: 0, - limit: 10, - is_active: None, - search: None, - }) - .await - .expect("admin globals should load"); - assert_eq!(admin_globals.total, 1); - assert_eq!(admin_globals.items[0].provider_count, 3); - assert_eq!(admin_globals.items[0].active_provider_count, 2); - - let admin_models = repository - .list_admin_provider_models(&AdminProviderModelListQuery { - provider_id: "provider-1".to_string(), - is_active: Some(true), - offset: 0, - limit: 10, - }) - .await - .expect("admin provider models should load"); - assert_eq!(admin_models.len(), 1); - assert_eq!( - admin_models[0].global_model_name, - Some("gpt-4.1".to_string()) - ); - - let stats = repository - .list_provider_model_stats(&["provider-1".to_string()]) - .await - .expect("stats should load"); - assert_eq!(stats[0].total_models, 1); - assert_eq!(stats[0].active_models, 1); - - let refs = repository - .list_active_global_model_ids_by_provider_ids(&["provider-1".to_string()]) - .await - .expect("active refs should load"); - assert_eq!(refs[0].global_model_id, "global-1"); - } - - #[tokio::test] - async fn sqlite_repository_writes_global_models_and_provider_models() { - 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; - - let repository = SqliteGlobalModelReadRepository::new(pool); - let created_global = repository - .create_admin_global_model( - &CreateAdminGlobalModelRecord::new( - "global-write-1".to_string(), - "claude-3.7".to_string(), - "Claude 3.7".to_string(), - true, - Some(0.25), - Some(json!({"tiers":[{"input_price_per_1m":3.0}]})), - Some(json!(["chat", "vision"])), - Some(json!({"description":"write path"})), - ) - .expect("create global input should validate"), - ) - .await - .expect("global model should create") - .expect("created global model should return"); - assert_eq!(created_global.name, "claude-3.7"); - assert_eq!( - created_global.supported_capabilities, - Some(json!(["chat", "vision"])) - ); - - let updated_global = repository - .update_admin_global_model( - &UpdateAdminGlobalModelRecord::new( - "global-write-1".to_string(), - "Claude 3.7 Sonnet".to_string(), - false, - Some(0.35), - None, - Some(json!(["chat"])), - Some(json!({"description":"updated"})), - ) - .expect("update global input should validate"), - ) - .await - .expect("global model should update") - .expect("updated global model should return"); - assert_eq!(updated_global.display_name, "Claude 3.7 Sonnet"); - assert!(!updated_global.is_active); - - let created_provider_model = repository - .create_admin_provider_model( - &UpsertAdminProviderModelRecord::new( - "model-write-1".to_string(), - "provider-1".to_string(), - "global-write-1".to_string(), - "provider-claude-3.7".to_string(), - Some(json!(["claude-3.7", "claude-sonnet"])), - Some(0.75), - Some(json!({"tiers":[{"output_price_per_1m":15.0}]})), - Some(true), - Some(true), - Some(true), - Some(false), - Some(false), - true, - true, - Some(json!({"routing":"primary"})), - ) - .expect("create provider model input should validate"), - ) - .await - .expect("provider model should create") - .expect("created provider model should return"); - assert_eq!( - created_provider_model.global_model_name, - Some("claude-3.7".to_string()) - ); - assert_eq!( - created_provider_model.provider_model_mappings, - Some(json!(["claude-3.7", "claude-sonnet"])) - ); - - let updated_provider_model = repository - .update_admin_provider_model( - &UpsertAdminProviderModelRecord::new( - "model-write-1".to_string(), - "provider-1".to_string(), - "global-write-1".to_string(), - "provider-claude-3.7-v2".to_string(), - None, - Some(0.95), - None, - Some(false), - Some(true), - Some(false), - Some(true), - Some(false), - false, - true, - Some(json!({"routing":"secondary"})), - ) - .expect("update provider model input should validate"), - ) - .await - .expect("provider model should update") - .expect("updated provider model should return"); - assert_eq!( - updated_provider_model.provider_model_name, - "provider-claude-3.7-v2" - ); - assert!(!updated_provider_model.is_active); - - assert!(repository - .delete_admin_global_model("global-write-1") - .await - .expect("global model should delete")); - assert!(repository - .get_admin_provider_model("provider-1", "model-write-1") - .await - .expect("deleted provider model lookup should succeed") - .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( - r#" -INSERT INTO providers ( - id, name, provider_type, is_active, provider_priority, created_at, updated_at -) VALUES - ('provider-2', 'Inactive Provider', 'custom', 0, 1, 1, 1), - ('provider-3', 'Alpha Provider', 'custom', 1, 20, 1, 1) -"#, - ) - .execute(pool) - .await - .expect("extra providers should seed"); - sqlx::query( - r#" -INSERT INTO global_models ( - id, name, display_name, is_active, default_tiered_pricing, - supported_capabilities, usage_count, config, created_at, updated_at -) VALUES ( - 'global-1', 'gpt-4.1', 'GPT 4.1', 1, - '{"tiers":[{"input_price_per_1m":2.0,"output_price_per_1m":8.0}]}', - '["chat"]', 7, '{"description":"Flagship","icon_url":"https://example.com/icon.png"}', 2, 3 -) -"#, - ) - .execute(pool) - .await - .expect("global model should seed"); - sqlx::query( - r#" -INSERT INTO models ( - id, provider_id, global_model_id, provider_model_name, provider_model_mappings, - supports_vision, supports_function_calling, supports_streaming, is_active, - is_available, created_at, updated_at -) VALUES -( - 'model-1', 'provider-1', 'global-1', 'provider-gpt-4.1', '["gpt-4.1"]', - 1, 1, 1, 1, 1, 4, 5 -) -, -( - 'model-2', 'provider-2', 'global-1', 'inactive-provider-gpt-4.1', '["gpt-4.1"]', - 1, 1, 1, 1, 1, 6, 7 -), -( - 'model-3', 'provider-3', 'global-1', 'alpha-provider-gpt-4.1', '["gpt-4.1"]', - 1, 1, 1, 1, 1, 8, 9 -) -"#, - ) - .execute(pool) - .await - .expect("provider model should seed"); - } - - async fn seed_provider(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers ( - id, name, provider_type, is_active, provider_priority, created_at, updated_at -) VALUES ( - 'provider-1', 'Zulu Provider', 'custom', 1, 10, 1, 1 -) -"#, - ) - .execute(pool) - .await - .expect("provider should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/lib.rs b/crates/aether-data/adapters/sqlite/src/lib.rs deleted file mode 100644 index 252b9705c..000000000 --- a/crates/aether-data/adapters/sqlite/src/lib.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! SQLite repositories, pool primitives, and migrations. - -mod announcements; -mod audit; -mod auth; -mod auth_modules; -mod background_tasks; -mod billing; -mod candidate_selection; -mod candidates; -mod error; -mod gemini_file_mappings; -mod global_models; -mod management_tokens; -mod migrations; -mod oauth_providers; -mod pool; -mod pool_scores; -mod provider_catalog; -mod proxy_nodes; -mod quota; -mod routing_profiles; -mod settlement; -mod usage; -mod users; -mod video_tasks; -mod wallet; - -pub use aether_data_contracts::{DataLayerError, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; -pub use announcements::SqliteAnnouncementRepository; -pub use audit::SqliteAuditLogReadRepository; -pub use auth::SqliteAuthApiKeyReadRepository; -pub use auth_modules::{SqliteAuthModuleReadRepository, SqliteAuthModuleRepository}; -pub use background_tasks::SqliteBackgroundTaskRepository; -pub use billing::SqliteBillingReadRepository; -pub use candidate_selection::SqliteMinimalCandidateSelectionReadRepository; -pub use candidates::SqliteRequestCandidateRepository; -pub use gemini_file_mappings::SqliteGeminiFileMappingRepository; -pub use global_models::SqliteGlobalModelReadRepository; -pub use management_tokens::SqliteManagementTokenRepository; -pub use migrations::{pending_migrations, prepare_database_for_startup, run_migrations, MIGRATOR}; -pub use oauth_providers::SqliteOAuthProviderRepository; -pub use pool::{SqlitePool, SqlitePoolConfig, SqlitePoolFactory}; -pub use pool_scores::SqlitePoolMemberScoreRepository; -pub use provider_catalog::SqliteProviderCatalogReadRepository; -pub use proxy_nodes::SqliteProxyNodeReadRepository; -pub use quota::SqliteProviderQuotaRepository; -pub use routing_profiles::SqliteRoutingGroupRepository; -pub use settlement::SqliteSettlementRepository; -pub use usage::{SqliteUsageReadRepository, SqliteUsageWriteRepository}; -pub use users::SqliteUserReadRepository; -pub use video_tasks::SqliteVideoTaskRepository; -pub use wallet::SqliteWalletReadRepository; - -use sqlx::{sqlite::SqliteRow, Row}; - -pub fn sqlite_real(row: &SqliteRow, field: &str) -> Result { - match row.try_get::(field) { - Ok(value) => Ok(value), - Err(real_err) => match row.try_get::(field) { - Ok(value) => Ok(value as f64), - Err(_) => Err(DataLayerError::sql(real_err)), - }, - } -} - -pub fn sqlite_optional_real(row: &SqliteRow, field: &str) -> Result, DataLayerError> { - match row.try_get::, _>(field) { - Ok(value) => Ok(value), - Err(real_err) => match row.try_get::, _>(field) { - Ok(value) => Ok(value.map(|value| value as f64)), - Err(_) => Err(DataLayerError::sql(real_err)), - }, - } -} diff --git a/crates/aether-data/adapters/sqlite/src/management_tokens.rs b/crates/aether-data/adapters/sqlite/src/management_tokens.rs deleted file mode 100644 index e5b212ea8..000000000 --- a/crates/aether-data/adapters/sqlite/src/management_tokens.rs +++ /dev/null @@ -1,1112 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::management_tokens::{ - ActivateManagementTokenIfMatches, CreateManagementTokenRecord, ManagementTokenListQuery, - ManagementTokenReadRepository, ManagementTokenWriteRepository, RegenerateManagementTokenSecret, - StoredManagementToken, StoredManagementTokenListPage, StoredManagementTokenUserSummary, - StoredManagementTokenWithUser, UpdateManagementTokenRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, push_limit, push_limit_offset, push_optional_eq, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -#[derive(Debug, Clone)] -pub struct SqliteManagementTokenRepository { - pool: SqlitePool, -} - -impl SqliteManagementTokenRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn get_token( - &self, - token_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", token_id.to_string()); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_row).transpose() - } - - async fn get_token_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", token_id.to_string()); - push_optional_eq( - &mut builder, - &mut where_clause, - "user_id", - expected_user_id.map(ToOwned::to_owned), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_row).transpose() - } - - async fn update_management_token_scoped( - &self, - record: &UpdateManagementTokenRecord, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - record.validate()?; - let allowed_ips = json_to_string(record.allowed_ips.as_ref())?; - let permissions = json_to_string(record.permissions.as_ref())?; - let now = now_unix_secs(); - - sqlx::query(UPDATE_MANAGEMENT_TOKEN_SQL) - .bind(record.name.as_deref()) - .bind(record.clear_description) - .bind(record.description.as_deref()) - .bind(record.clear_allowed_ips) - .bind(allowed_ips) - .bind(permissions) - .bind(record.clear_expires_at) - .bind( - record - .expires_at_unix_secs - .and_then(|value| i64::try_from(value).ok()), - ) - .bind(record.is_active) - .bind(now as i64) - .bind(&record.token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_err(|err| map_sqlite_write_error(err, record.name.as_deref()))?; - self.get_token_scoped(&record.token_id, expected_user_id) - .await - } - - async fn delete_management_token_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - ) -> Result { - let result = sqlx::query( - "DELETE FROM management_tokens WHERE id = ? AND (? IS NULL OR user_id = ?)", - ) - .bind(token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn set_management_token_active_scoped( - &self, - token_id: &str, - expected_user_id: Option<&str>, - is_active: bool, - ) -> Result, DataLayerError> { - let result = sqlx::query( - "UPDATE management_tokens SET is_active = ?, updated_at = ? WHERE id = ? AND (? IS NULL OR user_id = ?)", - ) - .bind(is_active) - .bind(now_unix_secs() as i64) - .bind(token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token_scoped(token_id, expected_user_id).await - } - - async fn regenerate_management_token_secret_scoped( - &self, - mutation: &RegenerateManagementTokenSecret, - expected_user_id: Option<&str>, - ) -> Result, DataLayerError> { - mutation.validate()?; - let result = sqlx::query( - r#" -UPDATE management_tokens -SET token_hash = ?, token_prefix = ?, updated_at = ? -WHERE id = ? AND (? IS NULL OR user_id = ?) -"#, - ) - .bind(&mutation.token_hash) - .bind(mutation.token_prefix.as_deref()) - .bind(now_unix_secs() as i64) - .bind(&mutation.token_id) - .bind(expected_user_id) - .bind(expected_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token_scoped(&mutation.token_id, expected_user_id) - .await - } -} - -const UPDATE_MANAGEMENT_TOKEN_SQL: &str = r#" -UPDATE management_tokens -SET name = COALESCE(?, name), - description = CASE WHEN ? THEN NULL ELSE COALESCE(?, description) END, - allowed_ips = CASE WHEN ? THEN NULL ELSE COALESCE(?, allowed_ips) END, - permissions = COALESCE(?, permissions), - expires_at = CASE WHEN ? THEN NULL ELSE COALESCE(?, expires_at) END, - is_active = COALESCE(?, is_active), - updated_at = ? -WHERE id = ? AND (? IS NULL OR user_id = ?) -"#; - -const TOKEN_COLUMNS: &str = r#" -SELECT - id, - user_id, - name, - description, - token_prefix, - allowed_ips, - permissions, - expires_at AS expires_at_unix_secs, - last_used_at AS last_used_at_unix_secs, - last_used_ip, - COALESCE(usage_count, 0) AS usage_count, - is_active, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM management_tokens -"#; - -const LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL: &str = r#" -SELECT id -FROM users -WHERE id = ? - AND is_active = 1 - AND is_deleted = 0 - AND LOWER(role) = 'admin' - AND security_version = ? -"#; - -const LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL: &str = r#" -SELECT - id, - user_id, - token_hash, - name, - description, - token_prefix, - allowed_ips, - permissions, - expires_at AS expires_at_unix_secs, - last_used_at AS last_used_at_unix_secs, - last_used_ip, - COALESCE(usage_count, 0) AS usage_count, - is_active, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM management_tokens -WHERE id = ? -"#; - -const TOKEN_WITH_USER_COLUMNS: &str = r#" -SELECT - mt.id, - mt.user_id, - mt.name, - mt.description, - mt.token_prefix, - mt.allowed_ips, - mt.permissions, - mt.expires_at AS expires_at_unix_secs, - mt.last_used_at AS last_used_at_unix_secs, - mt.last_used_ip, - COALESCE(mt.usage_count, 0) AS usage_count, - mt.is_active, - mt.created_at AS created_at_unix_ms, - mt.updated_at AS updated_at_unix_secs, - u.id AS user_row_id, - u.email AS user_email, - u.username AS user_username, - u.role AS user_role -FROM management_tokens mt -JOIN users u ON u.id = mt.user_id -"#; - -#[async_trait] -impl ManagementTokenReadRepository for SqliteManagementTokenRepository { - async fn list_management_tokens( - &self, - query: &ManagementTokenListQuery, - ) -> Result { - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(mt.id) AS total FROM management_tokens mt"); - let mut count_where = WhereClause::new(); - apply_management_token_filters(&mut count_builder, &mut count_where, query); - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let mut list_builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut list_where = WhereClause::new(); - apply_management_token_filters(&mut list_builder, &mut list_where, query); - list_builder.push(" ORDER BY mt.created_at DESC, mt.id DESC"); - push_limit_offset( - &mut list_builder, - i64::try_from(query.limit).unwrap_or(i64::MAX), - i64::try_from(query.offset).unwrap_or(i64::MAX), - ); - let rows = list_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - Ok(StoredManagementTokenListPage { - items: rows - .iter() - .map(map_token_with_user_row) - .collect::, _>>()?, - total: usize::try_from(total.max(0)).unwrap_or(usize::MAX), - }) - } - - async fn get_management_token_with_user( - &self, - token_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "mt.id", - token_id.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_with_user_row).transpose() - } - - async fn get_management_token_with_user_by_hash( - &self, - token_hash: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(TOKEN_WITH_USER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "mt.token_hash", - token_hash.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_token_with_user_row).transpose() - } -} - -fn apply_management_token_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - where_clause: &mut WhereClause, - query: &ManagementTokenListQuery, -) { - push_optional_eq(builder, where_clause, "mt.user_id", query.user_id.clone()); - push_optional_eq(builder, where_clause, "mt.is_active", query.is_active); -} - -#[async_trait] -impl ManagementTokenWriteRepository for SqliteManagementTokenRepository { - async fn create_management_token( - &self, - record: &CreateManagementTokenRecord, - ) -> Result { - record.validate()?; - let now = now_unix_secs(); - sqlx::query( - r#" -INSERT INTO management_tokens ( - id, user_id, token_hash, token_prefix, name, description, allowed_ips, - permissions, expires_at, is_active, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&record.id) - .bind(&record.user_id) - .bind(&record.token_hash) - .bind(record.token_prefix.as_deref()) - .bind(&record.name) - .bind(record.description.as_deref()) - .bind(json_to_string(record.allowed_ips.as_ref())?) - .bind(json_to_string(record.permissions.as_ref())?) - .bind( - record - .expires_at_unix_secs - .and_then(|value| i64::try_from(value).ok()), - ) - .bind(record.is_active) - .bind(now as i64) - .bind(now as i64) - .execute(&self.pool) - .await - .map_err(|err| map_sqlite_write_error(err, Some(record.name.as_str())))?; - - self.get_token(&record.id).await?.ok_or_else(|| { - DataLayerError::UnexpectedValue("created management token missing".to_string()) - }) - } - - async fn update_management_token( - &self, - record: &UpdateManagementTokenRecord, - ) -> Result, DataLayerError> { - self.update_management_token_scoped(record, None).await - } - - async fn update_management_token_for_user( - &self, - record: &UpdateManagementTokenRecord, - user_id: &str, - ) -> Result, DataLayerError> { - self.update_management_token_scoped(record, Some(user_id)) - .await - } - - async fn delete_management_token(&self, token_id: &str) -> Result { - self.delete_management_token_scoped(token_id, None).await - } - - async fn delete_management_token_for_user( - &self, - token_id: &str, - user_id: &str, - ) -> Result { - self.delete_management_token_scoped(token_id, Some(user_id)) - .await - } - - async fn set_management_token_active( - &self, - token_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_management_token_active_scoped(token_id, None, is_active) - .await - } - - async fn set_management_token_active_for_user( - &self, - token_id: &str, - user_id: &str, - is_active: bool, - ) -> Result, DataLayerError> { - self.set_management_token_active_scoped(token_id, Some(user_id), is_active) - .await - } - - async fn activate_management_token_if_matches( - &self, - mutation: &ActivateManagementTokenIfMatches, - ) -> Result { - mutation.validate()?; - // BEGIN IMMEDIATE prevents a concurrent user downgrade or token mutation between the - // canonical pre-state reads and the activation write. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let eligible_user = - sqlx::query_scalar::<_, String>(LOCK_ELIGIBLE_MANAGEMENT_TOKEN_ADMIN_SQL) - .bind(&mutation.expected_token.user_id) - .bind(mutation.expected_user_security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if eligible_user.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let locked = sqlx::query(LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL) - .bind(&mutation.expected_token.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let snapshot_matches = match locked.as_ref() { - Some(row) => { - let token_hash: String = row.try_get("token_hash").map_sql_err()?; - let token = map_token_row(row)?; - mutation.matches_locked_token_snapshot(&token, &token_hash) - } - None => false, - }; - if !snapshot_matches { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE management_tokens -SET is_active = 1, updated_at = ? -WHERE id = ? - AND token_hash = ? - AND is_active = 0 - AND (expires_at IS NULL OR expires_at > ?) -"#, - ) - .bind(now_unix_secs() as i64) - .bind(&mutation.expected_token.id) - .bind(&mutation.token_hash) - .bind(i64::try_from(mutation.now_unix_secs).unwrap_or(i64::MAX)) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_inactive_management_token_if_matches( - &self, - mutation: &ActivateManagementTokenIfMatches, - ) -> Result { - mutation.validate()?; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let locked = sqlx::query(LOCK_MANAGEMENT_TOKEN_ACTIVATION_SNAPSHOT_SQL) - .bind(&mutation.expected_token.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let snapshot_matches = match locked.as_ref() { - Some(row) => { - let token_hash: String = row.try_get("token_hash").map_sql_err()?; - let token = map_token_row(row)?; - mutation.matches_locked_token_snapshot(&token, &token_hash) - } - None => false, - }; - if !snapshot_matches { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "DELETE FROM management_tokens WHERE id = ? AND token_hash = ? AND is_active = 0", - ) - .bind(&mutation.expected_token.id) - .bind(&mutation.token_hash) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn regenerate_management_token_secret( - &self, - mutation: &RegenerateManagementTokenSecret, - ) -> Result, DataLayerError> { - self.regenerate_management_token_secret_scoped(mutation, None) - .await - } - - async fn regenerate_management_token_secret_for_user( - &self, - mutation: &RegenerateManagementTokenSecret, - user_id: &str, - ) -> Result, DataLayerError> { - self.regenerate_management_token_secret_scoped(mutation, Some(user_id)) - .await - } - - async fn record_management_token_usage( - &self, - token_id: &str, - last_used_ip: Option<&str>, - ) -> Result, DataLayerError> { - let now = now_unix_secs(); - let result = sqlx::query( - r#" -UPDATE management_tokens -SET last_used_at = ?, - last_used_ip = ?, - usage_count = COALESCE(usage_count, 0) + 1, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(now as i64) - .bind(last_used_ip) - .bind(now as i64) - .bind(token_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get_token(token_id).await - } -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn non_negative_u64(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!( - "management_tokens.{field_name} must not be negative" - )) - }) -} - -fn optional_unix_secs(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| non_negative_u64(value, field_name)) - .transpose() -} - -fn json_to_string(value: Option<&serde_json::Value>) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "invalid management token JSON field: {err}" - )) - }) - }) - .transpose() -} - -fn json_from_string(value: Option) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "invalid management token JSON field: {err}" - )) - }) - }) - .transpose() -} - -fn map_sqlite_write_error(err: sqlx::Error, requested_name: Option<&str>) -> DataLayerError { - let message = err.to_string(); - if message.contains("management_tokens.user_id, management_tokens.name") { - return DataLayerError::InvalidInput( - requested_name - .map(|name| format!("已存在名为 '{}' 的 Token", name)) - .unwrap_or_else(|| "Management Token 名称已存在".to_string()), - ); - } - DataLayerError::sql(err) -} - -fn map_token_row(row: &SqliteRow) -> Result { - Ok(StoredManagementToken::new( - row.try_get("id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - )? - .with_display_fields( - row.try_get("description").map_sql_err()?, - row.try_get("token_prefix").map_sql_err()?, - json_from_string(row.try_get("allowed_ips").map_sql_err()?)?, - ) - .with_permissions(json_from_string(row.try_get("permissions").map_sql_err()?)?) - .with_runtime_fields( - optional_unix_secs( - row.try_get("expires_at_unix_secs").map_sql_err()?, - "expires_at", - )?, - optional_unix_secs( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - "last_used_at", - )?, - row.try_get("last_used_ip").map_sql_err()?, - non_negative_u64( - row.try_get::("usage_count").map_sql_err()?, - "usage_count", - )?, - row.try_get("is_active").map_sql_err()?, - ) - .with_timestamps( - optional_unix_secs( - row.try_get("created_at_unix_ms").map_sql_err()?, - "created_at", - )?, - optional_unix_secs( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "updated_at", - )?, - )) -} - -fn map_user_summary_row( - row: &SqliteRow, -) -> Result { - StoredManagementTokenUserSummary::new( - row.try_get("user_row_id").map_sql_err()?, - row.try_get("user_email").map_sql_err()?, - row.try_get("user_username").map_sql_err()?, - row.try_get("user_role").map_sql_err()?, - ) -} - -fn map_token_with_user_row( - row: &SqliteRow, -) -> Result { - Ok(StoredManagementTokenWithUser::new( - map_token_row(row)?, - map_user_summary_row(row)?, - )) -} - -#[cfg(test)] -mod tests { - use super::{ - non_negative_u64, optional_unix_secs, SqliteManagementTokenRepository, - UPDATE_MANAGEMENT_TOKEN_SQL, - }; - use crate::run_migrations; - use aether_data_contracts::repository::management_tokens::{ - ActivateManagementTokenIfMatches, CreateManagementTokenRecord, ManagementTokenListQuery, - ManagementTokenReadRepository, ManagementTokenWriteRepository, - RegenerateManagementTokenSecret, StoredManagementTokenUserSummary, - UpdateManagementTokenRecord, - }; - - #[test] - fn sqlite_management_token_mapping_rejects_negative_integer_state() { - assert!(optional_unix_secs(Some(-1), "expires_at").is_err()); - assert_eq!( - optional_unix_secs(None, "expires_at").expect("SQL NULL should remain optional"), - None - ); - assert!(non_negative_u64(-1, "usage_count").is_err()); - } - - #[test] - fn sqlite_management_token_updates_patch_only_explicit_fields() { - for clause in [ - "name = COALESCE(?, name)", - "allowed_ips = CASE WHEN ? THEN NULL ELSE COALESCE(?, allowed_ips) END", - "permissions = COALESCE(?, permissions)", - "expires_at = CASE WHEN ? THEN NULL ELSE COALESCE(?, expires_at) END", - "is_active = COALESCE(?, is_active)", - ] { - assert!(UPDATE_MANAGEMENT_TOKEN_SQL.contains(clause)); - } - } - - #[tokio::test] - async fn sqlite_repository_round_trips_management_tokens() { - 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"); - sqlx::query( - r#" -INSERT INTO users (id, email, username, role, is_active, created_at, updated_at) -VALUES ('user-1', 'user-1@example.com', 'user-1', 'admin', 1, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("seed user should insert"); - - let repository = SqliteManagementTokenRepository::new(pool); - let user = StoredManagementTokenUserSummary::new( - "user-1".to_string(), - Some("user-1@example.com".to_string()), - "user-1".to_string(), - "admin".to_string(), - ) - .expect("user summary should build"); - let created = repository - .create_management_token(&CreateManagementTokenRecord { - id: "token-1".to_string(), - user_id: "user-1".to_string(), - user, - token_hash: "hash-1".to_string(), - token_prefix: Some("ae_1234".to_string()), - name: "primary".to_string(), - description: Some("primary token".to_string()), - allowed_ips: Some(serde_json::json!(["127.0.0.1"])), - permissions: Some(serde_json::json!(["admin:usage:read"])), - expires_at_unix_secs: Some(1_800_000_000), - is_active: true, - }) - .await - .expect("token should create"); - assert_eq!(created.name, "primary"); - assert_eq!( - created.permissions, - Some(serde_json::json!(["admin:usage:read"])) - ); - - let page = repository - .list_management_tokens(&ManagementTokenListQuery { - user_id: Some("user-1".to_string()), - is_active: Some(true), - offset: 0, - limit: 10, - }) - .await - .expect("tokens should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].token.id, "token-1"); - - let by_hash = repository - .get_management_token_with_user_by_hash("hash-1") - .await - .expect("hash lookup should succeed") - .expect("token should exist"); - assert_eq!(by_hash.user.username, "user-1"); - - let pending_install = repository - .create_management_token(&CreateManagementTokenRecord { - id: "token-install".to_string(), - user_id: "user-1".to_string(), - user: by_hash.user.clone(), - token_hash: "hash-install".to_string(), - token_prefix: Some("ae_install".to_string()), - name: "pending install".to_string(), - description: None, - allowed_ips: Some(serde_json::json!(["127.0.0.1"])), - permissions: Some(serde_json::json!(["admin:proxy_nodes:write"])), - expires_at_unix_secs: Some(1_800_000_000), - is_active: false, - }) - .await - .expect("pending install token should create"); - let activation = ActivateManagementTokenIfMatches { - expected_token: pending_install, - token_hash: "hash-install".to_string(), - expected_user_security_version: 0, - now_unix_secs: 1_700_000_000, - }; - let mut wrong_secret = activation.clone(); - wrong_secret.token_hash = "hash-other".to_string(); - assert!(!repository - .activate_management_token_if_matches(&wrong_secret) - .await - .expect("wrong-secret activation should execute")); - let mut stale_snapshot = activation.clone(); - stale_snapshot.expected_token.description = Some("changed snapshot".to_string()); - assert!(!repository - .activate_management_token_if_matches(&stale_snapshot) - .await - .expect("stale-snapshot activation should execute")); - assert!(repository - .activate_management_token_if_matches(&activation) - .await - .expect("matching activation should execute")); - assert!(!repository - .activate_management_token_if_matches(&activation) - .await - .expect("already-active token must not activate again")); - - let cross_owner_update = UpdateManagementTokenRecord { - token_id: "token-1".to_string(), - name: Some("hijacked".to_string()), - description: None, - clear_description: false, - allowed_ips: None, - clear_allowed_ips: false, - permissions: None, - expires_at_unix_secs: None, - clear_expires_at: false, - is_active: None, - }; - assert!(repository - .update_management_token_for_user(&cross_owner_update, "user-2") - .await - .expect("owner-scoped update should execute") - .is_none()); - assert!(repository - .set_management_token_active_for_user("token-1", "user-2", false) - .await - .expect("owner-scoped toggle should execute") - .is_none()); - assert!(repository - .regenerate_management_token_secret_for_user( - &RegenerateManagementTokenSecret { - token_id: "token-1".to_string(), - token_hash: "hash-hijacked".to_string(), - token_prefix: Some("ae_hijacked".to_string()), - }, - "user-2", - ) - .await - .expect("owner-scoped regeneration should execute") - .is_none()); - assert!(!repository - .delete_management_token_for_user("token-1", "user-2") - .await - .expect("owner-scoped delete should execute")); - assert_eq!( - repository - .get_management_token_with_user_by_hash("hash-1") - .await - .expect("original hash lookup should succeed") - .expect("token should remain") - .token - .name, - "primary" - ); - assert!(repository - .get_management_token_with_user_by_hash("hash-hijacked") - .await - .expect("replacement hash lookup should succeed") - .is_none()); - - let updated = repository - .update_management_token(&UpdateManagementTokenRecord { - token_id: "token-1".to_string(), - name: Some("renamed".to_string()), - description: None, - clear_description: true, - allowed_ips: Some(serde_json::json!(["10.0.0.1"])), - clear_allowed_ips: false, - permissions: Some(serde_json::json!(["admin:usage:read", "admin:usage:write"])), - expires_at_unix_secs: None, - clear_expires_at: true, - is_active: Some(false), - }) - .await - .expect("update should succeed") - .expect("token should exist"); - assert_eq!(updated.name, "renamed"); - assert!(!updated.is_active); - assert_eq!(updated.description, None); - assert_eq!( - updated.permissions, - Some(serde_json::json!(["admin:usage:read", "admin:usage:write"])) - ); - assert_eq!(updated.expires_at_unix_secs, None); - - let toggled = repository - .set_management_token_active("token-1", true) - .await - .expect("toggle should succeed") - .expect("token should exist"); - assert!(toggled.is_active); - - let regenerated = repository - .regenerate_management_token_secret(&RegenerateManagementTokenSecret { - token_id: "token-1".to_string(), - token_hash: "hash-2".to_string(), - token_prefix: Some("ae_5678".to_string()), - }) - .await - .expect("regenerate should succeed") - .expect("token should exist"); - assert_eq!(regenerated.token_prefix.as_deref(), Some("ae_5678")); - assert!(repository - .get_management_token_with_user_by_hash("hash-1") - .await - .expect("old hash lookup should succeed") - .is_none()); - - let used = repository - .record_management_token_usage("token-1", Some("127.0.0.1")) - .await - .expect("usage should record") - .expect("token should exist"); - assert_eq!(used.usage_count, 1); - assert_eq!(used.last_used_ip.as_deref(), Some("127.0.0.1")); - - assert!(repository - .delete_management_token("token-1") - .await - .expect("delete should succeed")); - assert!(repository - .delete_management_token("token-install") - .await - .expect("install token delete should succeed")); - } - - #[tokio::test] - async fn sqlite_install_activation_rejects_changed_admin_identity_snapshot() { - 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"); - sqlx::query( - r#" -INSERT INTO users ( - id, email, username, role, is_active, is_deleted, security_version, created_at, updated_at -) -VALUES ('admin-install', 'admin@example.com', 'admin-install', 'admin', 1, 0, 7, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("admin user should insert"); - - let repository = SqliteManagementTokenRepository::new(pool.clone()); - let user = StoredManagementTokenUserSummary::new( - "admin-install".to_string(), - Some("admin@example.com".to_string()), - "admin-install".to_string(), - "admin".to_string(), - ) - .expect("user summary should build"); - - async fn create_pending( - repository: &SqliteManagementTokenRepository, - user: &StoredManagementTokenUserSummary, - id: &str, - ) -> aether_data_contracts::repository::management_tokens::StoredManagementToken { - repository - .create_management_token(&CreateManagementTokenRecord { - id: id.to_string(), - user_id: user.id.clone(), - user: user.clone(), - token_hash: format!("hash-{id}"), - token_prefix: Some("ae_install".to_string()), - name: id.to_string(), - description: Some("one-time tunnel install".to_string()), - allowed_ips: Some(serde_json::json!(["127.0.0.1"])), - permissions: Some(serde_json::json!(["admin:proxy_nodes:write"])), - expires_at_unix_secs: Some(1_800_000_000), - is_active: false, - }) - .await - .expect("pending install token should create") - } - - fn activation( - token: aether_data_contracts::repository::management_tokens::StoredManagementToken, - security_version: i64, - ) -> ActivateManagementTokenIfMatches { - ActivateManagementTokenIfMatches { - token_hash: format!("hash-{}", token.id), - expected_token: token, - expected_user_security_version: security_version, - now_unix_secs: 1_700_000_000, - } - } - - let role_activation = - activation(create_pending(&repository, &user, "install-role").await, 7); - sqlx::query("UPDATE users SET role = 'user' WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("role downgrade should update"); - assert!(!repository - .activate_management_token_if_matches(&role_activation) - .await - .expect("role-mismatched activation should execute")); - sqlx::query("UPDATE users SET role = 'admin' WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("admin role should restore"); - - let version_activation = activation( - create_pending(&repository, &user, "install-version").await, - 7, - ); - sqlx::query("UPDATE users SET security_version = 8 WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("security version should update"); - assert!(!repository - .activate_management_token_if_matches(&version_activation) - .await - .expect("version-mismatched activation should execute")); - - let inactive_activation = activation( - create_pending(&repository, &user, "install-inactive").await, - 8, - ); - sqlx::query("UPDATE users SET is_active = 0 WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("administrator should deactivate"); - assert!(!repository - .activate_management_token_if_matches(&inactive_activation) - .await - .expect("inactive-admin activation should execute")); - sqlx::query("UPDATE users SET is_active = 1 WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("administrator should reactivate"); - - let deleted_activation = activation( - create_pending(&repository, &user, "install-deleted").await, - 8, - ); - sqlx::query("UPDATE users SET is_deleted = 1 WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("administrator should be soft deleted"); - assert!(!repository - .activate_management_token_if_matches(&deleted_activation) - .await - .expect("deleted-admin activation should execute")); - sqlx::query("UPDATE users SET is_deleted = 0 WHERE id = 'admin-install'") - .execute(&pool) - .await - .expect("administrator deletion flag should restore"); - - let valid_activation = - activation(create_pending(&repository, &user, "install-valid").await, 8); - assert!(repository - .activate_management_token_if_matches(&valid_activation) - .await - .expect("matching administrator activation should execute")); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/migrations.rs b/crates/aether-data/adapters/sqlite/src/migrations.rs deleted file mode 100644 index d19bb3950..000000000 --- a/crates/aether-data/adapters/sqlite/src/migrations.rs +++ /dev/null @@ -1,995 +0,0 @@ -use sqlx::{ - migrate::{AppliedMigration, Migrate, MigrateError, Migrator}, - SqlitePool, -}; - -use aether_data_contracts::PendingMigrationInfo; - -pub static MIGRATOR: Migrator = sqlx::migrate!("./migrations"); - -pub async fn run_migrations(pool: &SqlitePool) -> Result<(), MigrateError> { - MIGRATOR.run(pool).await -} - -pub async fn pending_migrations( - pool: &SqlitePool, -) -> Result, MigrateError> { - 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) => { - 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)) -} - -pub async fn prepare_database_for_startup( - pool: &SqlitePool, -) -> Result, MigrateError> { - pending_migrations(pool).await -} - -fn is_missing_sqlx_migrations_table_error(err: &MigrateError) -> bool { - let message = err.to_string().to_ascii_lowercase(); - message.contains("_sqlx_migrations") - && (message.contains("no such table") - || message.contains("doesn't exist") - || message.contains("does not exist") - || message.contains("unknown table")) -} - -fn pending_migrations_from_applied( - applied_migrations: &[sqlx::migrate::AppliedMigration], -) -> Vec { - let applied_versions = applied_migrations - .iter() - .map(|migration| migration.version) - .collect::>(); - MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .filter(|migration| !applied_versions.contains(&migration.version)) - .map(|migration| PendingMigrationInfo { - version: migration.version, - description: migration.description.to_string(), - }) - .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 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() { - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("in-memory sqlite pool"); - - let pending = pending_migrations(&pool).await.expect("pending migrations"); - assert_eq!(pending.len(), MIGRATOR.iter().count()); - assert!(!pending.is_empty()); - - run_migrations(&pool).await.expect("run sqlite migrations"); - assert!(pending_migrations(&pool) - .await - .expect("pending migrations after run") - .is_empty()); - } - - #[tokio::test] - async fn legacy_proxy_node_inserts_receive_non_empty_tunnel_generations() { - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("in-memory sqlite pool"); - - run_migrations(&pool) - .await - .expect("sqlite migrations should run"); - // Simulate an older writer that does not know about tunnel_generation. - sqlx::query( - "INSERT INTO proxy_nodes (id, name, ip, port, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("legacy-node-a") - .bind("legacy node a") - .bind("127.0.0.1") - .bind(18080_i32) - .bind(1_i64) - .bind(1_i64) - .execute(&pool) - .await - .expect("legacy proxy node insert should succeed"); - sqlx::query( - "INSERT INTO proxy_nodes (id, name, ip, port, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("legacy-node-b") - .bind("legacy node b") - .bind("127.0.0.2") - .bind(18081_i32) - .bind(1_i64) - .bind(1_i64) - .execute(&pool) - .await - .expect("second legacy proxy node insert should succeed"); - - let generations = sqlx::query_scalar::<_, String>( - "SELECT tunnel_generation FROM proxy_nodes ORDER BY id", - ) - .fetch_all(&pool) - .await - .expect("proxy node generations should load"); - assert_eq!(generations.len(), 2); - assert!(generations.iter().all(|generation| !generation.is_empty())); - assert_ne!(generations[0], generations[1]); - } - - #[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 worker_boot_cleanup_removes_only_legacy_instance_rows_and_events() { - const CLEANUP_VERSION: i64 = 20260731000000; - const HASHED_LEGACY_ID: &str = - "boot:maintenance.request.candidate.cleanup:~0123456789abcdef0123"; - const OVERLONG_LEGACY_ID: &str = - "boot:maintenance.proxy.node.metrics.cleanup:gateway-instance-with-an-overlong-id"; - - assert_eq!(HASHED_LEGACY_ID.len(), 64); - assert!(OVERLONG_LEGACY_ID.len() > 64); - - 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 < CLEANUP_VERSION) - { - 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 background_task_runs ( - id, task_key, kind, "trigger", status, owner_instance, - progress_message, created_by, created_at_unix_secs, updated_at_unix_secs -) VALUES - ('boot:usage.queue.worker:gateway-a', 'usage.queue.worker', 'daemon', 'daemon', - 'running', 'gateway-a', 'worker booted', 'system', 1, 1), - ('boot:usage.queue.worker:gateway-b', 'usage.queue.worker', 'daemon', 'daemon', - 'running', 'gateway-b', 'worker booted', 'system', 2, 2), - ('boot:maintenance.request.candidate.cleanup:~0123456789abcdef0123', - 'maintenance.request.candidate.cleanup', 'scheduled', 'interval', - 'running', 'gateway-hash', 'worker booted', 'system', 3, 3), - ('boot:maintenance.proxy.node.metrics.cleanup:gateway-instance-with-an-overlong-id', - 'maintenance.proxy.node.metrics.cleanup', 'scheduled', 'interval', - 'running', 'gateway-overlong', 'worker booted', 'system', 4, 4), - ('boot:model.fetch.worker', 'model.fetch.worker', 'scheduled', 'interval', - 'running', 'gateway-early-fix', 'worker booted', 'system', 5, 5), - ('boot:usage.queue.worker', 'usage.queue.worker', 'daemon', 'daemon', - 'running', NULL, 'worker registered', 'system', 6, 6), - ('boot:ownerless-worker', 'ownerless.worker', 'daemon', 'daemon', - 'running', NULL, 'worker booted', 'system', 7, 7), - ('boot:custom-progress', 'custom.progress', 'daemon', 'daemon', - 'running', 'gateway-custom', 'worker healthy', 'system', 8, 8), - ('boot:user-request', 'user.request', 'on_demand', 'manual', - 'running', 'gateway-user', 'worker booted', 'admin', 9, 9); - -INSERT INTO background_task_events ( - id, run_id, event_type, message, created_at_unix_secs -) VALUES - ('legacy-event-a', 'boot:usage.queue.worker:gateway-a', 'worker_boot', 'legacy', 1), - ('legacy-event-b', 'boot:usage.queue.worker:gateway-b', 'worker_boot', 'legacy', 2), - ('legacy-event-hash', 'boot:maintenance.request.candidate.cleanup:~0123456789abcdef0123', - 'worker_boot', 'legacy hash', 3), - ('legacy-event-overlong', - 'boot:maintenance.proxy.node.metrics.cleanup:gateway-instance-with-an-overlong-id', - 'worker_boot', 'legacy overlong', 4), - ('early-fix-event', 'boot:model.fetch.worker', 'worker_boot', 'early fix', 5), - ('logical-event', 'boot:usage.queue.worker', 'worker_boot', 'logical', 6), - ('ownerless-event', 'boot:ownerless-worker', 'worker_boot', 'ownerless', 7), - ('custom-progress-event', 'boot:custom-progress', 'worker_boot', 'custom progress', 8), - ('manual-event', 'boot:user-request', 'manual', 'manual', 9); -"#, - ) - .execute(&pool) - .await - .expect("worker boot cleanup fixtures should insert"); - - let migration = MIGRATOR - .iter() - .find(|migration| migration.version == CLEANUP_VERSION) - .expect("worker boot cleanup migration should be embedded"); - sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect("worker boot cleanup migration should run"); - sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect("worker boot cleanup migration should be idempotent"); - - let remaining_runs = sqlx::query_as::<_, (String, Option, Option)>( - r#" -SELECT id, owner_instance, created_by -FROM background_task_runs -ORDER BY id -"#, - ) - .fetch_all(&pool) - .await - .expect("remaining worker task rows should load"); - assert_eq!( - remaining_runs, - vec![ - ( - "boot:custom-progress".to_string(), - Some("gateway-custom".to_string()), - Some("system".to_string()), - ), - ( - "boot:ownerless-worker".to_string(), - None, - Some("system".to_string()), - ), - ( - "boot:usage.queue.worker".to_string(), - None, - Some("system".to_string()), - ), - ( - "boot:user-request".to_string(), - Some("gateway-user".to_string()), - Some("admin".to_string()), - ), - ] - ); - - let remaining_events = - sqlx::query_scalar::<_, String>("SELECT id FROM background_task_events ORDER BY id") - .fetch_all(&pool) - .await - .expect("remaining worker task events should load"); - assert_eq!( - remaining_events, - vec![ - "custom-progress-event".to_string(), - "logical-event".to_string(), - "manual-event".to_string(), - "ownerless-event".to_string(), - ] - ); - } - - #[tokio::test] - async fn ldap_singleton_migration_keeps_legacy_min_id_and_enforces_one_fixed_row() { - const LDAP_SINGLETON_VERSION: i64 = 20260831000000; - - 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 < LDAP_SINGLETON_VERSION) - { - 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 ldap_configs ( - server_url, bind_dn, bind_password_encrypted, base_dn, created_at, updated_at -) VALUES - ('ldaps://first.example.com', 'cn=first', 'first-ciphertext', 'dc=first', 1, 1), - ('ldaps://second.example.com', 'cn=second', 'second-ciphertext', 'dc=second', 2, 2) -"#, - ) - .execute(&pool) - .await - .expect("legacy duplicate LDAP rows should seed"); - - let migration = MIGRATOR - .iter() - .find(|migration| migration.version == LDAP_SINGLETON_VERSION) - .expect("LDAP singleton migration should be embedded"); - sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect("LDAP singleton migration should run"); - - let surviving = sqlx::query_as::<_, (String, Option, i64)>( - "SELECT server_url, bind_password_encrypted, singleton_key FROM ldap_configs", - ) - .fetch_all(&pool) - .await - .expect("singleton LDAP row should load"); - assert_eq!( - surviving, - vec![( - "ldaps://first.example.com".to_string(), - Some("first-ciphertext".to_string()), - 1, - )] - ); - - let duplicate = sqlx::query( - r#" -INSERT INTO ldap_configs ( - server_url, bind_dn, bind_password_encrypted, base_dn, created_at, updated_at -) VALUES ('ldaps://third.example.com', 'cn=third', 'third-ciphertext', 'dc=third', 3, 3) -"#, - ) - .execute(&pool) - .await; - assert!( - duplicate.is_err(), - "a second singleton row must be rejected" - ); - - let invalid_key = sqlx::query("UPDATE ldap_configs SET singleton_key = 2") - .execute(&pool) - .await; - assert!( - invalid_key.is_err(), - "the singleton discriminator must remain fixed at one" - ); - } - - #[tokio::test] - async fn codex_live_permission_migration_is_scoped_and_idempotent() { - const MIGRATION_VERSION: i64 = 20260821000000; - - 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"); - - sqlx::raw_sql( - r#" -INSERT INTO users ( - id, username, allowed_api_formats, allowed_api_formats_mode, created_at, updated_at -) VALUES - ('user-specific', 'user-specific', '["openai:responses"]', 'specific', 1, 1), - ('user-unrestricted', 'user-unrestricted', '["openai:responses"]', 'unrestricted', 1, 1); - -INSERT INTO api_keys ( - id, user_id, key_hash, allowed_api_formats, created_at, updated_at -) VALUES ( - 'api-key-specific', 'user-specific', 'api-key-specific-hash', - '["openai:responses"]', 1, 1 -); - -INSERT INTO user_groups ( - id, name, normalized_name, allowed_api_formats, allowed_api_formats_mode, - created_at, updated_at -) VALUES ( - 'group-specific', 'Specific', 'specific', '["openai:responses"]', 'specific', 1, 1 -); - -INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES - ('provider-codex', 'Codex migration fixture', 'codex', 1, 1), - ('provider-openai', 'OpenAI migration fixture', 'openai', 1, 1); - -INSERT INTO provider_api_keys ( - id, provider_id, name, api_formats, auth_type_by_format, - allow_auth_channel_mismatch_formats, rate_multipliers, - global_priority_by_format, created_at, updated_at -) VALUES - ( - 'provider-key-codex', 'provider-codex', 'Codex key', - '["openai:responses"]', '{"openai:responses":"oauth"}', - '["openai:responses"]', '{"openai:responses":1.25}', - '{"openai:responses":17}', 1, 1 - ), - ( - 'provider-key-openai', 'provider-openai', 'OpenAI key', - '["openai:responses"]', '{"openai:responses":"api_key"}', - '["openai:responses"]', '{"openai:responses":2.0}', - '{"openai:responses":23}', 1, 1 - ); -"#, - ) - .execute(&pool) - .await - .expect("legacy Codex Live fixtures should insert"); - - let migration = MIGRATOR - .iter() - .find(|migration| migration.version == MIGRATION_VERSION) - .expect("Codex Live permission migration should be embedded"); - for _ in 0..2 { - sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect("Codex Live permission migration should run idempotently"); - } - - let user_specific: String = - sqlx::query_scalar("SELECT allowed_api_formats FROM users WHERE id = 'user-specific'") - .fetch_one(&pool) - .await - .expect("specific user formats should load"); - assert_eq!( - serde_json::from_str::(&user_specific).unwrap(), - serde_json::json!(["openai:responses", "codex:live"]) - ); - - let user_unrestricted: String = sqlx::query_scalar( - "SELECT allowed_api_formats FROM users WHERE id = 'user-unrestricted'", - ) - .fetch_one(&pool) - .await - .expect("unrestricted user formats should load"); - assert_eq!( - serde_json::from_str::(&user_unrestricted).unwrap(), - serde_json::json!(["openai:responses"]) - ); - - for (table, id_column, id) in [ - ("api_keys", "id", "api-key-specific"), - ("user_groups", "id", "group-specific"), - ] { - let sql = format!("SELECT allowed_api_formats FROM {table} WHERE {id_column} = ?"); - let formats: String = sqlx::query_scalar(sql.as_str()) - .bind(id) - .fetch_one(&pool) - .await - .expect("migrated downstream formats should load"); - assert_eq!( - serde_json::from_str::(&formats).unwrap(), - serde_json::json!(["openai:responses", "codex:live"]) - ); - } - - let codex_key = sqlx::query_as::<_, (String, String, String, String, String)>( - r#" -SELECT api_formats, auth_type_by_format, allow_auth_channel_mismatch_formats, - rate_multipliers, global_priority_by_format -FROM provider_api_keys -WHERE id = 'provider-key-codex' -"#, - ) - .fetch_one(&pool) - .await - .expect("migrated Codex key should load"); - assert_eq!( - serde_json::from_str::(&codex_key.0).unwrap(), - serde_json::json!(["openai:responses", "codex:live"]) - ); - assert_eq!( - serde_json::from_str::(&codex_key.1).unwrap(), - serde_json::json!({"openai:responses": "oauth", "codex:live": "oauth"}) - ); - assert_eq!( - serde_json::from_str::(&codex_key.2).unwrap(), - serde_json::json!(["openai:responses", "codex:live"]) - ); - assert_eq!( - serde_json::from_str::(&codex_key.3).unwrap(), - serde_json::json!({"openai:responses": 1.25, "codex:live": 1.25}) - ); - assert_eq!( - serde_json::from_str::(&codex_key.4).unwrap(), - serde_json::json!({"openai:responses": 17, "codex:live": 17}) - ); - - let openai_formats: String = sqlx::query_scalar( - "SELECT api_formats FROM provider_api_keys WHERE id = 'provider-key-openai'", - ) - .fetch_one(&pool) - .await - .expect("non-Codex key formats should load"); - assert_eq!( - serde_json::from_str::(&openai_formats).unwrap(), - serde_json::json!(["openai:responses"]) - ); - } - - #[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/oauth_providers.rs b/crates/aether-data/adapters/sqlite/src/oauth_providers.rs deleted file mode 100644 index a49ccf36a..000000000 --- a/crates/aether-data/adapters/sqlite/src/oauth_providers.rs +++ /dev/null @@ -1,653 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::oauth_providers::{ - OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig, - UpsertOAuthProviderConfigOutcome, UpsertOAuthProviderConfigRecord, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, push_limit, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -#[derive(Debug, Clone)] -pub struct SqliteOAuthProviderRepository { - pool: SqlitePool, -} - -impl SqliteOAuthProviderRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn get_provider( - &self, - provider_type: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(OAUTH_PROVIDER_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "provider_type", - provider_type.to_string(), - ); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_oauth_provider_row).transpose() - } -} - -const OAUTH_PROVIDER_COLUMNS: &str = r#" -SELECT - provider_type, - display_name, - client_id, - client_secret_encrypted, - authorization_url_override, - token_url_override, - userinfo_url_override, - scopes, - redirect_uri, - frontend_callback_url, - attribute_mapping, - extra_config, - icon_url, - is_enabled, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM oauth_providers -"#; - -const COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL: &str = r#" -SELECT COUNT(DISTINCT users.id) AS locked_count -FROM users -JOIN user_oauth_links - ON users.id = user_oauth_links.user_id -WHERE users.is_active = 1 - AND users.is_deleted = 0 - AND user_oauth_links.provider_type = ? - AND ( - ( - users.auth_source = 'oauth' - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links other_links - JOIN oauth_providers other_provider - ON other_links.provider_type = other_provider.provider_type - WHERE other_links.user_id = users.id - AND other_links.provider_type <> ? - AND other_provider.is_enabled = 1 - ) - ) OR ( - ? = 1 - AND users.auth_source = 'local' - AND users.role <> 'admin' - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links other_links - JOIN oauth_providers other_provider - ON other_links.provider_type = other_provider.provider_type - WHERE other_links.user_id = users.id - AND other_links.provider_type <> ? - AND other_provider.is_enabled = 1 - ) - ) - ) -"#; - -const COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL: &str = r#" -UPDATE oauth_providers -SET client_secret_encrypted = ? -WHERE provider_type = ? - AND client_secret_encrypted = ? -"#; - -#[async_trait] -impl OAuthProviderReadRepository for SqliteOAuthProviderRepository { - async fn list_oauth_provider_configs( - &self, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(OAUTH_PROVIDER_COLUMNS); - builder.push(" ORDER BY provider_type ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_oauth_provider_row).collect() - } - - async fn get_oauth_provider_config( - &self, - provider_type: &str, - ) -> Result, DataLayerError> { - self.get_provider(provider_type).await - } - - async fn count_locked_users_if_provider_disabled( - &self, - provider_type: &str, - ldap_exclusive: bool, - ) -> Result { - let row = sqlx::query(COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL) - .bind(provider_type) - .bind(provider_type) - .bind(ldap_exclusive) - .bind(provider_type) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let locked_count = row.try_get::("locked_count").map_sql_err()?; - usize::try_from(locked_count.max(0)).map_err(|_| { - DataLayerError::UnexpectedValue( - "oauth_providers.locked_user_count overflowed".to_string(), - ) - }) - } -} - -#[async_trait] -impl OAuthProviderWriteRepository for SqliteOAuthProviderRepository { - async fn upsert_oauth_provider_config_guarded( - &self, - record: &UpsertOAuthProviderConfigRecord, - ldap_exclusive: bool, - force_disable: bool, - _locked_users_snapshot: usize, - ) -> Result { - record.validate()?; - let now = now_unix_secs(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let existing_enabled: Option = - sqlx::query_scalar("SELECT is_enabled FROM oauth_providers WHERE provider_type = ?") - .bind(&record.provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if !force_disable && !record.is_enabled && existing_enabled == Some(true) { - let row = sqlx::query(COUNT_LOCKED_USERS_IF_PROVIDER_DISABLED_SQL) - .bind(&record.provider_type) - .bind(&record.provider_type) - .bind(ldap_exclusive) - .bind(&record.provider_type) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let affected_count = - usize::try_from(row.try_get::("locked_count").map_sql_err()?.max(0)) - .map_err(|_| { - DataLayerError::UnexpectedValue( - "oauth_providers.locked_user_count overflowed".to_string(), - ) - })?; - if affected_count > 0 { - tx.rollback().await.map_sql_err()?; - return Ok( - UpsertOAuthProviderConfigOutcome::DisableRequiresConfirmation { - affected_count, - }, - ); - } - } - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, - display_name, - client_id, - client_secret_encrypted, - authorization_url_override, - token_url_override, - userinfo_url_override, - scopes, - redirect_uri, - frontend_callback_url, - attribute_mapping, - extra_config, - icon_url, - is_enabled, - created_at, - updated_at -) VALUES ( - ?, ?, ?, - CASE ? WHEN 'set' THEN ? WHEN 'clear' THEN NULL ELSE NULL END, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? -) -ON CONFLICT(provider_type) DO UPDATE SET - display_name = excluded.display_name, - client_id = excluded.client_id, - client_secret_encrypted = CASE ? - WHEN 'set' THEN ? - WHEN 'clear' THEN NULL - ELSE oauth_providers.client_secret_encrypted - END, - authorization_url_override = excluded.authorization_url_override, - token_url_override = excluded.token_url_override, - userinfo_url_override = excluded.userinfo_url_override, - scopes = excluded.scopes, - redirect_uri = excluded.redirect_uri, - frontend_callback_url = excluded.frontend_callback_url, - attribute_mapping = excluded.attribute_mapping, - extra_config = excluded.extra_config, - icon_url = excluded.icon_url, - is_enabled = excluded.is_enabled, - updated_at = excluded.updated_at -"#, - ) - .bind(&record.provider_type) - .bind(&record.display_name) - .bind(&record.client_id) - .bind(record.client_secret_encrypted.mode_name()) - .bind(record.client_secret_encrypted.value()) - .bind(record.authorization_url_override.as_deref()) - .bind(record.token_url_override.as_deref()) - .bind(record.userinfo_url_override.as_deref()) - .bind(scopes_to_json_string(record.scopes.as_ref())?) - .bind(&record.redirect_uri) - .bind(&record.frontend_callback_url) - .bind(json_to_string(record.attribute_mapping.as_ref())?) - .bind(json_to_string(record.extra_config.as_ref())?) - .bind(record.icon_url.as_deref()) - .bind(record.is_enabled) - .bind(now as i64) - .bind(now as i64) - .bind(record.client_secret_encrypted.mode_name()) - .bind(record.client_secret_encrypted.value()) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = sqlx::query(&format!( - "{OAUTH_PROVIDER_COLUMNS} WHERE provider_type = ? LIMIT 1" - )) - .bind(&record.provider_type) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let provider = map_oauth_provider_row(&row)?; - tx.commit().await.map_sql_err()?; - Ok(UpsertOAuthProviderConfigOutcome::Upserted(provider)) - } - - async fn compare_and_swap_oauth_provider_client_secret( - &self, - provider_type: &str, - expected: &str, - replacement: &str, - ) -> Result { - let result = sqlx::query(COMPARE_AND_SWAP_OAUTH_PROVIDER_CLIENT_SECRET_SQL) - .bind(replacement) - .bind(provider_type) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn delete_oauth_provider_config_if_unlinked( - &self, - provider_type: &str, - has_links_snapshot: bool, - ) -> Result { - if has_links_snapshot { - return Ok(false); - } - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let provider_exists: Option = - sqlx::query_scalar("SELECT provider_type FROM oauth_providers WHERE provider_type = ?") - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "DELETE FROM oauth_providers WHERE provider_type = ? AND NOT EXISTS (SELECT 1 FROM user_oauth_links WHERE user_oauth_links.provider_type = oauth_providers.provider_type)", - ) - .bind(provider_type) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(result.rows_affected() == 1) - } -} - -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_unix_secs(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn json_to_string(value: Option<&serde_json::Value>) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("invalid OAuth provider JSON field: {err}")) - }) - }) - .transpose() -} - -fn json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn scopes_to_json_string(scopes: Option<&Vec>) -> Result, DataLayerError> { - json_to_string( - scopes - .map(|items| { - serde_json::Value::Array( - items - .iter() - .cloned() - .map(serde_json::Value::String) - .collect(), - ) - }) - .as_ref(), - ) -} - -fn parse_scopes(value: Option) -> Result>, DataLayerError> { - let Some(value) = json_from_string(value, "oauth_providers.scopes")? else { - return Ok(None); - }; - parse_scopes_value(&value) -} - -fn parse_scopes_value(value: &serde_json::Value) -> Result>, DataLayerError> { - match value { - serde_json::Value::Null => Ok(None), - serde_json::Value::Array(items) => parse_scopes_array(items).map(Some), - serde_json::Value::String(raw) => parse_embedded_scopes(raw), - _ => Err(DataLayerError::UnexpectedValue( - "oauth_providers.scopes is not a JSON array".to_string(), - )), - } -} - -fn parse_embedded_scopes(raw: &str) -> Result>, DataLayerError> { - let raw = raw.trim(); - if raw.is_empty() || raw.eq_ignore_ascii_case("null") { - return Ok(None); - } - - if let Ok(decoded) = serde_json::from_str::(raw) { - return parse_scopes_value(&decoded); - } - - Ok(Some(vec![raw.to_string()])) -} - -fn parse_scopes_array(items: &[serde_json::Value]) -> Result, DataLayerError> { - let mut scopes = Vec::with_capacity(items.len()); - for item in items { - let Some(scope) = item.as_str() else { - return Err(DataLayerError::UnexpectedValue( - "oauth_providers.scopes contains non-string value".to_string(), - )); - }; - let scope = scope.trim(); - if !scope.is_empty() { - scopes.push(scope.to_string()); - } - } - Ok(scopes) -} - -fn map_oauth_provider_row(row: &SqliteRow) -> Result { - Ok(StoredOAuthProviderConfig::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("client_id").map_sql_err()?, - row.try_get("redirect_uri").map_sql_err()?, - row.try_get("frontend_callback_url").map_sql_err()?, - )? - .with_config_fields( - row.try_get("client_secret_encrypted").map_sql_err()?, - row.try_get("authorization_url_override").map_sql_err()?, - row.try_get("token_url_override").map_sql_err()?, - row.try_get("userinfo_url_override").map_sql_err()?, - parse_scopes(row.try_get("scopes").map_sql_err()?)?, - json_from_string( - row.try_get("attribute_mapping").map_sql_err()?, - "oauth_providers.attribute_mapping", - )?, - json_from_string( - row.try_get("extra_config").map_sql_err()?, - "oauth_providers.extra_config", - )?, - row.try_get("icon_url").map_sql_err()?, - row.try_get("is_enabled").map_sql_err()?, - ) - .with_timestamps( - optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - optional_unix_secs(row.try_get("updated_at_unix_secs").map_sql_err()?), - )) -} - -#[cfg(test)] -mod tests { - use super::SqliteOAuthProviderRepository; - use crate::run_migrations; - use aether_data_contracts::repository::oauth_providers::{ - EncryptedSecretUpdate, OAuthProviderReadRepository, OAuthProviderWriteRepository, - UpsertOAuthProviderConfigOutcome, UpsertOAuthProviderConfigRecord, - }; - - fn sample_upsert(provider_type: &str) -> UpsertOAuthProviderConfigRecord { - UpsertOAuthProviderConfigRecord { - provider_type: provider_type.to_string(), - display_name: format!("{provider_type} display"), - client_id: format!("{provider_type}-client"), - client_secret_encrypted: EncryptedSecretUpdate::Preserve, - authorization_url_override: Some( - "https://connect.linux.do/oauth2/authorize".to_string(), - ), - token_url_override: Some("https://connect.linux.do/oauth2/token".to_string()), - userinfo_url_override: None, - scopes: Some(vec!["openid".to_string(), "profile".to_string()]), - redirect_uri: format!("https://{provider_type}.example.com/redirect"), - frontend_callback_url: "https://frontend.example.com/auth/callback".to_string(), - attribute_mapping: Some(serde_json::json!({"email": "email"})), - extra_config: Some(serde_json::json!({"team": true})), - icon_url: None, - is_enabled: true, - } - } - - #[tokio::test] - async fn sqlite_repository_round_trips_oauth_provider_configs() { - 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 = SqliteOAuthProviderRepository::new(pool.clone()); - let created = repository - .upsert_oauth_provider_config(&UpsertOAuthProviderConfigRecord { - client_secret_encrypted: EncryptedSecretUpdate::Set("secret-1".to_string()), - ..sample_upsert("linuxdo") - }) - .await - .expect("provider should upsert"); - assert_eq!(created.client_secret_encrypted.as_deref(), Some("secret-1")); - assert_eq!( - created.scopes, - Some(vec!["openid".to_string(), "profile".to_string()]) - ); - - let updated = repository - .upsert_oauth_provider_config(&UpsertOAuthProviderConfigRecord { - client_secret_encrypted: EncryptedSecretUpdate::Preserve, - display_name: "Linux.do".to_string(), - ..sample_upsert("linuxdo") - }) - .await - .expect("provider should update"); - assert_eq!(updated.display_name, "Linux.do"); - assert_eq!(updated.client_secret_encrypted.as_deref(), Some("secret-1")); - assert!( - repository - .compare_and_swap_oauth_provider_client_secret( - "linuxdo", - "secret-1", - "record-bound-v2", - ) - .await - .expect("client secret CAS should execute") - ); - let migrated = repository - .get_oauth_provider_config("linuxdo") - .await - .expect("provider should fetch") - .expect("provider should exist"); - assert_eq!(migrated.display_name, "Linux.do"); - assert_eq!(migrated.updated_at_unix_secs, updated.updated_at_unix_secs); - assert_eq!( - migrated.client_secret_encrypted.as_deref(), - Some("record-bound-v2") - ); - assert!(!repository - .compare_and_swap_oauth_provider_client_secret("linuxdo", "secret-1", "must-not-win",) - .await - .expect("stale client secret CAS should execute")); - - let listed = repository - .list_oauth_provider_configs() - .await - .expect("providers should list"); - assert_eq!(listed.len(), 1); - - let fetched = repository - .get_oauth_provider_config("linuxdo") - .await - .expect("provider should fetch") - .expect("provider should exist"); - assert_eq!( - fetched.attribute_mapping, - Some(serde_json::json!({"email": "email"})) - ); - - sqlx::query( - r#" -INSERT INTO users ( - id, email, username, role, auth_source, is_active, is_deleted, created_at, updated_at -) VALUES - ('user-oauth', 'oauth@example.com', 'oauth-user', 'user', 'oauth', 1, 0, 1, 1), - ('user-local', 'local@example.com', 'local-user', 'user', 'local', 1, 0, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("users should seed"); - sqlx::query( - r#" -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, linked_at -) VALUES - ('link-1', 'user-oauth', 'linuxdo', 'linuxdo-1', 1), - ('link-2', 'user-local', 'linuxdo', 'linuxdo-2', 1) -"#, - ) - .execute(&pool) - .await - .expect("oauth links should seed"); - assert_eq!( - repository - .count_locked_users_if_provider_disabled("linuxdo", false) - .await - .expect("locked users should count"), - 1 - ); - assert_eq!( - repository - .count_locked_users_if_provider_disabled("linuxdo", true) - .await - .expect("locked users should count"), - 2 - ); - - assert_eq!( - repository - .upsert_oauth_provider_config_guarded( - &UpsertOAuthProviderConfigRecord { - is_enabled: false, - ..sample_upsert("linuxdo") - }, - false, - false, - 0, - ) - .await - .expect("guarded disable should resolve"), - UpsertOAuthProviderConfigOutcome::DisableRequiresConfirmation { affected_count: 1 } - ); - let forced = repository - .upsert_oauth_provider_config_guarded( - &UpsertOAuthProviderConfigRecord { - is_enabled: false, - ..sample_upsert("linuxdo") - }, - false, - true, - 0, - ) - .await - .expect("forced provider disable should succeed"); - assert!(matches!( - forced, - UpsertOAuthProviderConfigOutcome::Upserted(provider) if !provider.is_enabled - )); - - assert!(!repository - .delete_oauth_provider_config_if_unlinked("linuxdo", false) - .await - .expect("linked provider deletion should resolve")); - assert!(repository - .get_oauth_provider_config("linuxdo") - .await - .expect("provider should fetch") - .is_some()); - - sqlx::query("DELETE FROM user_oauth_links WHERE provider_type = 'linuxdo'") - .execute(&pool) - .await - .expect("links should delete"); - assert!(repository - .delete_oauth_provider_config_if_unlinked("linuxdo", false) - .await - .expect("unlinked provider should delete")); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/pool.rs b/crates/aether-data/adapters/sqlite/src/pool.rs deleted file mode 100644 index 2e0c64128..000000000 --- a/crates/aether-data/adapters/sqlite/src/pool.rs +++ /dev/null @@ -1,159 +0,0 @@ -use std::path::PathBuf; -use std::str::FromStr; -use std::time::Duration; - -use crate::{DataLayerError, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; -use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; -use sqlx::SqlitePool as SqlxSqlitePool; - -pub type SqlitePool = SqlxSqlitePool; -pub type SqlitePoolConfig = SqlDatabaseConfig; - -#[derive(Debug, Clone)] -pub struct SqlitePoolFactory { - config: SqlitePoolConfig, -} - -impl SqlitePoolFactory { - pub fn new(config: SqlitePoolConfig) -> Result { - if config.driver != DatabaseDriver::Sqlite { - return Err(DataLayerError::InvalidConfiguration(format!( - "sqlite pool requires sqlite driver, got {}", - config.driver - ))); - } - config.validate()?; - Ok(Self { config }) - } - - pub fn config(&self) -> &SqlitePoolConfig { - &self.config - } - - pub fn connect_options(&self) -> Result { - ensure_sqlite_parent_dir(self.config.url.trim())?; - let is_memory = is_sqlite_memory_url(self.config.url.trim()); - SqliteConnectOptions::from_str(self.config.url.trim()) - .map(|options| { - let options = options - .create_if_missing(true) - .foreign_keys(true) - .statement_cache_capacity(self.config.pool.statement_cache_capacity); - if is_memory { - options - } else { - options.journal_mode(SqliteJournalMode::Wal) - } - }) - .map_err(|err| { - DataLayerError::InvalidConfiguration(format!("invalid sqlite database url: {err}")) - }) - } - - pub fn connect_lazy(&self) -> Result { - let SqlPoolConfig { - min_connections, - max_connections, - acquire_timeout_ms, - idle_timeout_ms, - max_lifetime_ms, - .. - } = self.config.pool; - - Ok(SqlitePoolOptions::new() - .min_connections(min_connections) - .max_connections(max_connections) - .acquire_timeout(Duration::from_millis(acquire_timeout_ms)) - .idle_timeout(Duration::from_millis(idle_timeout_ms)) - .max_lifetime(Duration::from_millis(max_lifetime_ms)) - .connect_lazy_with(self.connect_options()?)) - } -} - -fn is_sqlite_memory_url(url: &str) -> bool { - matches!(url.trim(), "sqlite::memory:" | "sqlite://:memory:") -} - -fn sqlite_file_path_from_url(url: &str) -> Option { - let url = url.trim(); - if is_sqlite_memory_url(url) { - return None; - } - let path = url - .strip_prefix("sqlite://") - .or_else(|| url.strip_prefix("sqlite:"))?; - if path.is_empty() { - return None; - } - Some(PathBuf::from(path)) -} - -fn ensure_sqlite_parent_dir(url: &str) -> Result<(), DataLayerError> { - let Some(path) = sqlite_file_path_from_url(url) else { - return Ok(()); - }; - let Some(parent) = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - else { - return Ok(()); - }; - std::fs::create_dir_all(parent).map_err(|err| { - DataLayerError::InvalidConfiguration(format!( - "failed to create sqlite database parent directory '{}': {err}", - parent.display() - )) - }) -} - -#[cfg(test)] -mod tests { - use super::SqlitePoolFactory; - use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; - use std::path::PathBuf; - - #[tokio::test] - async fn factory_builds_lazy_pool_from_valid_config() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite://./data/aether.db".to_string(), - pool: SqlPoolConfig { - min_connections: 1, - max_connections: 4, - acquire_timeout_ms: 1_000, - idle_timeout_ms: 5_000, - max_lifetime_ms: 30_000, - statement_cache_capacity: 64, - require_ssl: false, - }, - }; - - let factory = SqlitePoolFactory::new(config).expect("factory should build"); - let _pool = factory.connect_lazy().expect("lazy pool should build"); - } - - #[tokio::test] - async fn factory_creates_parent_directory_for_file_database() { - let db_path = unique_temp_db_path(); - let parent = db_path.parent().expect("temp db path should have parent"); - let _ = std::fs::remove_dir_all(parent); - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: format!("sqlite://{}", db_path.display()), - pool: SqlPoolConfig::default(), - }; - - let factory = SqlitePoolFactory::new(config).expect("factory should build"); - let _pool = factory.connect_lazy().expect("lazy pool should build"); - - assert!(parent.exists()); - let _ = std::fs::remove_dir_all(parent); - } - - fn unique_temp_db_path() -> PathBuf { - std::env::temp_dir() - .join(format!("aether-sqlite-{}", uuid::Uuid::new_v4())) - .join("nested") - .join("aether.db") - } -} diff --git a/crates/aether-data/adapters/sqlite/src/pool_scores.rs b/crates/aether-data/adapters/sqlite/src/pool_scores.rs deleted file mode 100644 index 1ac1124c0..000000000 --- a/crates/aether-data/adapters/sqlite/src/pool_scores.rs +++ /dev/null @@ -1,901 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::pool_scores::*; -use aether_data_query::{push_eq, push_in, push_limit, push_limit_offset, WhereClause}; - -use crate::error::SqlResultExt; -use crate::{DataLayerError, SqlitePool}; - -const SCORE_COLUMNS: &str = r#" -SELECT - id, - pool_kind, - pool_id, - member_kind, - member_id, - capability, - scope_kind, - scope_id, - score, - hard_state, - score_version, - score_reason, - last_ranked_at, - last_scheduled_at, - last_success_at, - last_failure_at, - failure_count, - last_probe_attempt_at, - last_probe_success_at, - last_probe_failure_at, - probe_failure_count, - probe_status, - updated_at -FROM pool_member_scores -"#; - -const UPSERT_PRESERVING_NULLABLE_TIMESTAMPS_SQL: &str = r#" -INSERT INTO pool_member_scores ( - id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id, - score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at, - last_success_at, last_failure_at, failure_count, last_probe_attempt_at, - last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at -) VALUES ( - ?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,? -) -ON CONFLICT(id) DO UPDATE SET - pool_kind = excluded.pool_kind, - pool_id = excluded.pool_id, - member_kind = excluded.member_kind, - member_id = excluded.member_id, - capability = excluded.capability, - scope_kind = excluded.scope_kind, - scope_id = excluded.scope_id, - score = excluded.score, - hard_state = excluded.hard_state, - score_version = excluded.score_version, - score_reason = excluded.score_reason, - last_ranked_at = excluded.last_ranked_at, - last_scheduled_at = COALESCE(excluded.last_scheduled_at, pool_member_scores.last_scheduled_at), - last_success_at = COALESCE(excluded.last_success_at, pool_member_scores.last_success_at), - last_failure_at = COALESCE(excluded.last_failure_at, pool_member_scores.last_failure_at), - failure_count = excluded.failure_count, - last_probe_attempt_at = COALESCE(excluded.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at), - last_probe_success_at = COALESCE(excluded.last_probe_success_at, pool_member_scores.last_probe_success_at), - last_probe_failure_at = COALESCE(excluded.last_probe_failure_at, pool_member_scores.last_probe_failure_at), - probe_failure_count = excluded.probe_failure_count, - probe_status = excluded.probe_status, - updated_at = excluded.updated_at -"#; - -const UPSERT_OAUTH_RECOVERY_SQL: &str = r#" -INSERT INTO pool_member_scores ( - id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id, - score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at, - last_success_at, last_failure_at, failure_count, last_probe_attempt_at, - last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at -) VALUES ( - ?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,? -) -ON CONFLICT(id) DO UPDATE SET - pool_kind = excluded.pool_kind, - pool_id = excluded.pool_id, - member_kind = excluded.member_kind, - member_id = excluded.member_id, - capability = excluded.capability, - scope_kind = excluded.scope_kind, - scope_id = excluded.scope_id, - score = CASE WHEN pool_member_scores.updated_at <= excluded.updated_at THEN excluded.score ELSE pool_member_scores.score END, - hard_state = CASE WHEN pool_member_scores.updated_at <= excluded.updated_at THEN excluded.hard_state ELSE pool_member_scores.hard_state END, - score_version = CASE WHEN pool_member_scores.updated_at <= excluded.updated_at THEN excluded.score_version ELSE pool_member_scores.score_version END, - score_reason = CASE WHEN pool_member_scores.updated_at <= excluded.updated_at THEN excluded.score_reason ELSE pool_member_scores.score_reason END, - last_ranked_at = CASE WHEN pool_member_scores.updated_at <= excluded.updated_at THEN excluded.last_ranked_at ELSE pool_member_scores.last_ranked_at END, - last_failure_at = CASE - WHEN pool_member_scores.last_failure_at IS NULL OR pool_member_scores.last_failure_at <= excluded.updated_at - THEN excluded.last_failure_at ELSE pool_member_scores.last_failure_at END, - failure_count = CASE - WHEN pool_member_scores.last_failure_at IS NULL OR pool_member_scores.last_failure_at <= excluded.updated_at - THEN excluded.failure_count ELSE pool_member_scores.failure_count END, - last_probe_failure_at = CASE - WHEN pool_member_scores.last_probe_failure_at IS NULL OR pool_member_scores.last_probe_failure_at <= excluded.updated_at - THEN excluded.last_probe_failure_at ELSE pool_member_scores.last_probe_failure_at END, - probe_failure_count = CASE - WHEN pool_member_scores.last_probe_failure_at IS NULL OR pool_member_scores.last_probe_failure_at <= excluded.updated_at - THEN excluded.probe_failure_count ELSE pool_member_scores.probe_failure_count END, - probe_status = CASE - WHEN (pool_member_scores.last_probe_attempt_at IS NOT NULL AND pool_member_scores.last_probe_attempt_at > excluded.updated_at) - OR (pool_member_scores.last_probe_success_at IS NOT NULL AND pool_member_scores.last_probe_success_at > excluded.updated_at) - OR (pool_member_scores.last_probe_failure_at IS NOT NULL AND pool_member_scores.last_probe_failure_at > excluded.updated_at) - THEN pool_member_scores.probe_status ELSE excluded.probe_status END, - updated_at = MAX(pool_member_scores.updated_at, excluded.updated_at) -"#; - -fn pool_member_score_upsert_sql(mode: PoolMemberScoreUpsertMode) -> &'static str { - match mode { - PoolMemberScoreUpsertMode::PreserveExistingNullableTimestamps => { - UPSERT_PRESERVING_NULLABLE_TIMESTAMPS_SQL - } - PoolMemberScoreUpsertMode::OAuthRecovery => UPSERT_OAUTH_RECOVERY_SQL, - } -} - -#[derive(Debug, Clone)] -pub struct SqlitePoolMemberScoreRepository { - pool: SqlitePool, -} - -impl SqlitePoolMemberScoreRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn find_scores_by_identity( - &self, - identity: &PoolMemberIdentity, - scope: Option<&PoolScoreScope>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "pool_kind", - identity.pool_kind.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "pool_id", - identity.pool_id.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "member_kind", - identity.member_kind.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "member_id", - identity.member_id.clone(), - ); - if let Some(scope) = scope { - push_eq( - &mut builder, - &mut where_clause, - "capability", - scope.capability.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "scope_kind", - scope.scope_kind.clone(), - ); - if let Some(scope_id) = &scope.scope_id { - push_eq( - &mut builder, - &mut where_clause, - "scope_id", - scope_id.clone(), - ); - } else { - where_clause.push_next(&mut builder); - builder.push("scope_id IS NULL"); - } - } - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } -} - -#[async_trait] -impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository { - async fn list_ranked_pool_members( - &self, - query: &ListRankedPoolMembersQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "pool_kind", - query.pool_kind.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "pool_id", - query.pool_id.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "capability", - query.capability.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "scope_kind", - query.scope_kind.clone(), - ); - if let Some(scope_id) = &query.scope_id { - push_eq( - &mut builder, - &mut where_clause, - "scope_id", - scope_id.clone(), - ); - } else { - where_clause.push_next(&mut builder); - builder.push("scope_id IS NULL"); - } - if !query.hard_states.is_empty() { - let states = query - .hard_states - .iter() - .map(|state| state.as_database()) - .collect::>(); - push_in(&mut builder, &mut where_clause, "hard_state", &states); - } - if let Some(statuses) = &query.probe_statuses { - if !statuses.is_empty() { - let statuses = statuses - .iter() - .map(|status| status.as_database()) - .collect::>(); - push_in(&mut builder, &mut where_clause, "probe_status", &statuses); - } - } - builder.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC"); - push_limit_offset( - &mut builder, - i64_from_usize(query.limit.max(1), "pool score limit")?, - i64_from_usize(query.offset, "pool score offset")?, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn list_pool_member_scores( - &self, - query: &ListPoolMemberScoresQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "pool_kind", - query.pool_kind.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "pool_id", - query.pool_id.clone(), - ); - if let Some(capability) = &query.capability { - push_eq( - &mut builder, - &mut where_clause, - "capability", - capability.clone(), - ); - } - if let Some(scope_kind) = &query.scope_kind { - push_eq( - &mut builder, - &mut where_clause, - "scope_kind", - scope_kind.clone(), - ); - } - if let Some(scope_id) = &query.scope_id { - push_eq( - &mut builder, - &mut where_clause, - "scope_id", - scope_id.clone(), - ); - } - if !query.hard_states.is_empty() { - let states = query - .hard_states - .iter() - .map(|state| state.as_database()) - .collect::>(); - push_in(&mut builder, &mut where_clause, "hard_state", &states); - } - if let Some(statuses) = &query.probe_statuses { - if !statuses.is_empty() { - let statuses = statuses - .iter() - .map(|status| status.as_database()) - .collect::>(); - push_in(&mut builder, &mut where_clause, "probe_status", &statuses); - } - } - builder.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC"); - push_limit_offset( - &mut builder, - i64_from_usize(query.limit.max(1), "pool score limit")?, - i64_from_usize(query.offset, "pool score offset")?, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn list_pool_member_probe_candidates( - &self, - query: &ListPoolMemberProbeCandidatesQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "pool_kind", - query.pool_kind.clone(), - ); - push_eq( - &mut builder, - &mut where_clause, - "pool_id", - query.pool_id.clone(), - ); - if let Some(capability) = &query.capability { - push_eq( - &mut builder, - &mut where_clause, - "capability", - capability.clone(), - ); - } - where_clause.push_next(&mut builder); - builder - .push("hard_state IN ('available','unknown','cooldown','quota_exhausted')") - .push(" AND (probe_status IN ('never','failed','stale')") - .push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ") - .push_bind(i64_from_u64( - query.stale_before_unix_secs, - "pool probe stale_before_unix_secs", - )?) - .push("))") - .push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ") - .push_bind(i64_from_u64( - query.stale_before_unix_secs, - "pool probe stale_before_unix_secs", - )?) - .push(")))") - .push( - r#" - ORDER BY - CASE - WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0 - WHEN hard_state = 'quota_exhausted' THEN 1 - WHEN hard_state = 'unknown' THEN 2 - WHEN probe_status = 'stale' THEN 3 - ELSE 4 - END ASC, - probe_failure_count DESC, - COALESCE(last_probe_success_at, 0) ASC, - COALESCE(last_scheduled_at, 0) DESC, - member_id ASC -"#, - ); - push_limit( - &mut builder, - i64_from_usize(query.limit.max(1), "pool probe candidate limit")?, - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } - - async fn get_pool_member_scores_by_ids( - &self, - query: &GetPoolMemberScoresByIdsQuery, - ) -> Result, DataLayerError> { - if query.ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_in(&mut builder, &mut where_clause, "id", &query.ids); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_score_row).collect() - } -} - -#[async_trait] -impl PoolMemberScoreWriteRepository for SqlitePoolMemberScoreRepository { - async fn upsert_pool_member_score_with_mode( - &self, - score: UpsertPoolMemberScore, - mode: PoolMemberScoreUpsertMode, - ) -> Result { - score.validate()?; - let stored = score.into_stored(); - let score_reason = serde_json::to_string(&stored.score_reason) - .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; - sqlx::query(pool_member_score_upsert_sql(mode)) - .bind(stored.id.as_str()) - .bind(stored.pool_kind.as_str()) - .bind(stored.pool_id.as_str()) - .bind(stored.member_kind.as_str()) - .bind(stored.member_id.as_str()) - .bind(stored.capability.as_str()) - .bind(stored.scope_kind.as_str()) - .bind(stored.scope_id.as_deref()) - .bind(stored.score) - .bind(stored.hard_state.as_database()) - .bind(i64_from_u64(stored.score_version, "pool score version")?) - .bind(score_reason) - .bind(i64_opt_from_u64( - stored.last_ranked_at, - "pool score last_ranked_at", - )?) - .bind(i64_opt_from_u64( - stored.last_scheduled_at, - "pool score last_scheduled_at", - )?) - .bind(i64_opt_from_u64( - stored.last_success_at, - "pool score last_success_at", - )?) - .bind(i64_opt_from_u64( - stored.last_failure_at, - "pool score last_failure_at", - )?) - .bind(i64_from_u64( - stored.failure_count, - "pool score failure_count", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_attempt_at, - "pool score last_probe_attempt_at", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_success_at, - "pool score last_probe_success_at", - )?) - .bind(i64_opt_from_u64( - stored.last_probe_failure_at, - "pool score last_probe_failure_at", - )?) - .bind(i64_from_u64( - stored.probe_failure_count, - "pool score probe_failure_count", - )?) - .bind(stored.probe_status.as_database()) - .bind(i64_from_u64(stored.updated_at, "pool score updated_at")?) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(stored) - } - - async fn mark_pool_member_probe_in_progress( - &self, - attempt: PoolMemberProbeAttempt, - ) -> Result { - let rows = self - .find_scores_by_identity(&attempt.identity, attempt.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_probe_attempt_at = Some(attempt.attempted_at); - row.probe_status = PoolMemberProbeStatus::InProgress; - row.score_reason = - merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone()); - row.updated_at = attempt.attempted_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn record_pool_member_probe_result( - &self, - result: PoolMemberProbeResult, - ) -> Result { - let rows = self - .find_scores_by_identity(&result.identity, result.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_probe_attempt_at = Some(result.attempted_at); - row.probe_status = result.probe_status; - if result.succeeded { - row.last_probe_success_at = Some(result.attempted_at); - row.probe_failure_count = 0; - } else { - row.last_probe_failure_at = Some(result.attempted_at); - row.probe_failure_count = row.probe_failure_count.saturating_add(1); - } - if let Some(hard_state) = result.hard_state { - row.hard_state = hard_state; - } - row.score_reason = - merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone()); - row.updated_at = result.attempted_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn record_pool_member_schedule_feedback( - &self, - feedback: PoolMemberScheduleFeedback, - ) -> Result { - let rows = self - .find_scores_by_identity(&feedback.identity, feedback.scope.as_ref()) - .await?; - let count = rows.len(); - for mut row in rows { - row.last_scheduled_at = Some(feedback.scheduled_at); - match feedback.succeeded { - Some(true) => row.last_success_at = Some(feedback.scheduled_at), - Some(false) => { - row.last_failure_at = Some(feedback.scheduled_at); - row.failure_count = row.failure_count.saturating_add(1); - } - None => {} - } - if let Some(hard_state) = feedback.hard_state { - row.hard_state = hard_state; - } - row.score = score_with_delta(row.score, feedback.score_delta); - row.score_reason = - merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone()); - row.updated_at = feedback.scheduled_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn mark_pool_member_hard_state( - &self, - identity: &PoolMemberIdentity, - scope: Option<&PoolScoreScope>, - hard_state: PoolMemberHardState, - updated_at: u64, - ) -> Result { - let rows = self.find_scores_by_identity(identity, scope).await?; - let count = rows.len(); - for mut row in rows { - row.hard_state = hard_state; - row.updated_at = updated_at; - self.upsert_pool_member_score(upsert_from_stored(row)) - .await?; - } - Ok(count) - } - - async fn delete_pool_member_scores_for_member( - &self, - identity: &PoolMemberIdentity, - ) -> Result { - let result = sqlx::query( - r#" -DELETE FROM pool_member_scores -WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ? -"#, - ) - .bind(identity.pool_kind.as_str()) - .bind(identity.pool_id.as_str()) - .bind(identity.member_kind.as_str()) - .bind(identity.member_id.as_str()) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } -} - -fn map_score_row(row: &SqliteRow) -> Result { - let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?; - Ok(StoredPoolMemberScore { - id: row.try_get("id").map_sql_err()?, - pool_kind: row.try_get("pool_kind").map_sql_err()?, - pool_id: row.try_get("pool_id").map_sql_err()?, - member_kind: row.try_get("member_kind").map_sql_err()?, - member_id: row.try_get("member_id").map_sql_err()?, - capability: row.try_get("capability").map_sql_err()?, - scope_kind: row.try_get("scope_kind").map_sql_err()?, - scope_id: row.try_get("scope_id").map_sql_err()?, - score: row.try_get("score").map_sql_err()?, - hard_state: PoolMemberHardState::from_database( - row.try_get::("hard_state") - .map_sql_err()? - .as_str(), - )?, - score_version: u64_from_i64( - row.try_get("score_version").map_sql_err()?, - "pool_member_scores.score_version", - )?, - score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null), - last_ranked_at: u64_opt_from_i64( - row.try_get("last_ranked_at").map_sql_err()?, - "pool_member_scores.last_ranked_at", - )?, - last_scheduled_at: u64_opt_from_i64( - row.try_get("last_scheduled_at").map_sql_err()?, - "pool_member_scores.last_scheduled_at", - )?, - last_success_at: u64_opt_from_i64( - row.try_get("last_success_at").map_sql_err()?, - "pool_member_scores.last_success_at", - )?, - last_failure_at: u64_opt_from_i64( - row.try_get("last_failure_at").map_sql_err()?, - "pool_member_scores.last_failure_at", - )?, - failure_count: u64_from_i64( - row.try_get("failure_count").map_sql_err()?, - "pool_member_scores.failure_count", - )?, - last_probe_attempt_at: u64_opt_from_i64( - row.try_get("last_probe_attempt_at").map_sql_err()?, - "pool_member_scores.last_probe_attempt_at", - )?, - last_probe_success_at: u64_opt_from_i64( - row.try_get("last_probe_success_at").map_sql_err()?, - "pool_member_scores.last_probe_success_at", - )?, - last_probe_failure_at: u64_opt_from_i64( - row.try_get("last_probe_failure_at").map_sql_err()?, - "pool_member_scores.last_probe_failure_at", - )?, - probe_failure_count: u64_from_i64( - row.try_get("probe_failure_count").map_sql_err()?, - "pool_member_scores.probe_failure_count", - )?, - probe_status: PoolMemberProbeStatus::from_database( - row.try_get::("probe_status") - .map_sql_err()? - .as_str(), - )?, - updated_at: u64_from_i64( - row.try_get("updated_at").map_sql_err()?, - "pool_member_scores.updated_at", - )?, - }) -} - -fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore { - UpsertPoolMemberScore { - id: score.id, - identity: PoolMemberIdentity { - pool_kind: score.pool_kind, - pool_id: score.pool_id, - member_kind: score.member_kind, - member_id: score.member_id, - }, - scope: PoolScoreScope { - capability: score.capability, - scope_kind: score.scope_kind, - scope_id: score.scope_id, - }, - score: score.score, - hard_state: score.hard_state, - score_version: score.score_version, - score_reason: score.score_reason, - last_ranked_at: score.last_ranked_at, - last_scheduled_at: score.last_scheduled_at, - last_success_at: score.last_success_at, - last_failure_at: score.last_failure_at, - failure_count: score.failure_count, - last_probe_attempt_at: score.last_probe_attempt_at, - last_probe_success_at: score.last_probe_success_at, - last_probe_failure_at: score.last_probe_failure_at, - probe_failure_count: score.probe_failure_count, - probe_status: score.probe_status, - updated_at: score.updated_at, - } -} - -fn i64_from_usize(value: usize, field: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::run_migrations; - - fn score(timestamp_base: Option) -> UpsertPoolMemberScore { - UpsertPoolMemberScore { - id: "score-1".to_string(), - identity: PoolMemberIdentity::provider_api_key("provider-1", "key-1"), - scope: PoolScoreScope { - capability: POOL_SCORE_CAPABILITY_ACCOUNT.to_string(), - scope_kind: POOL_SCORE_SCOPE_KIND_ACCOUNT.to_string(), - scope_id: None, - }, - score: 0.75, - hard_state: PoolMemberHardState::Available, - score_version: 1, - score_reason: serde_json::json!({}), - last_ranked_at: Some(20), - last_scheduled_at: timestamp_base, - last_success_at: timestamp_base.map(|value| value + 1), - last_failure_at: timestamp_base.map(|value| value + 2), - failure_count: 0, - last_probe_attempt_at: timestamp_base.map(|value| value + 3), - last_probe_success_at: timestamp_base.map(|value| value + 4), - last_probe_failure_at: timestamp_base.map(|value| value + 5), - probe_failure_count: 0, - probe_status: PoolMemberProbeStatus::Never, - updated_at: 20, - } - } - - async fn load_score(repository: &SqlitePoolMemberScoreRepository) -> StoredPoolMemberScore { - repository - .get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery { - ids: vec!["score-1".to_string()], - }) - .await - .expect("pool score should load") - .pop() - .expect("pool score should exist") - } - - async fn repository() -> SqlitePoolMemberScoreRepository { - 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"); - SqlitePoolMemberScoreRepository::new(pool) - } - - #[tokio::test] - async fn sqlite_ordinary_upsert_preserves_nullable_timestamps() { - let repository = repository().await; - - repository - .upsert_pool_member_score(score(Some(100))) - .await - .expect("initial pool score should insert"); - repository - .upsert_pool_member_score(score(None)) - .await - .expect("ordinary upsert should succeed"); - let preserved = load_score(&repository).await; - assert_eq!(preserved.last_scheduled_at, Some(100)); - assert_eq!(preserved.last_success_at, Some(101)); - assert_eq!(preserved.last_failure_at, Some(102)); - assert_eq!(preserved.last_probe_attempt_at, Some(103)); - assert_eq!(preserved.last_probe_success_at, Some(104)); - assert_eq!(preserved.last_probe_failure_at, Some(105)); - } - - #[tokio::test] - async fn sqlite_oauth_recovery_clears_old_failures_and_preserves_success_history() { - let repository = repository().await; - let mut invalid = score(None); - invalid.score = 0.2; - invalid.hard_state = PoolMemberHardState::AuthInvalid; - invalid.score_reason = serde_json::json!({"state": "invalid"}); - invalid.last_ranked_at = Some(90); - invalid.last_scheduled_at = Some(80); - invalid.last_success_at = Some(81); - invalid.last_failure_at = Some(90); - invalid.failure_count = 9; - invalid.last_probe_attempt_at = Some(82); - invalid.last_probe_success_at = Some(83); - invalid.last_probe_failure_at = Some(91); - invalid.probe_failure_count = 4; - invalid.probe_status = PoolMemberProbeStatus::Failed; - invalid.updated_at = 91; - repository - .upsert_pool_member_score(invalid) - .await - .expect("invalid score should insert"); - let mut recovery = score(None); - recovery.score = 0.9; - recovery.score_reason = serde_json::json!({"state": "recovered"}); - recovery.last_ranked_at = Some(100); - recovery.updated_at = 100; - - repository - .upsert_pool_member_score_with_mode(recovery, PoolMemberScoreUpsertMode::OAuthRecovery) - .await - .expect("OAuth recovery should succeed"); - let recovered = load_score(&repository).await; - assert_eq!(recovered.score, 0.9); - assert_eq!(recovered.hard_state, PoolMemberHardState::Available); - assert_eq!( - recovered.score_reason, - serde_json::json!({"state": "recovered"}) - ); - assert_eq!(recovered.last_ranked_at, Some(100)); - assert_eq!(recovered.last_scheduled_at, Some(80)); - assert_eq!(recovered.last_success_at, Some(81)); - assert_eq!(recovered.last_failure_at, None); - assert_eq!(recovered.failure_count, 0); - assert_eq!(recovered.last_probe_attempt_at, Some(82)); - assert_eq!(recovered.last_probe_success_at, Some(83)); - assert_eq!(recovered.last_probe_failure_at, None); - assert_eq!(recovered.probe_failure_count, 0); - assert_eq!(recovered.probe_status, PoolMemberProbeStatus::Never); - assert_eq!(recovered.updated_at, 100); - } - - #[tokio::test] - async fn sqlite_oauth_recovery_preserves_newer_feedback() { - let repository = repository().await; - let mut current = score(None); - current.score = 0.3; - current.hard_state = PoolMemberHardState::AuthInvalid; - current.score_version = 7; - current.score_reason = serde_json::json!({"state": "newer_failure"}); - current.last_ranked_at = Some(120); - current.last_scheduled_at = Some(120); - current.last_success_at = Some(80); - current.last_failure_at = Some(120); - current.failure_count = 3; - current.last_probe_attempt_at = Some(130); - current.last_probe_success_at = Some(85); - current.last_probe_failure_at = Some(130); - current.probe_failure_count = 2; - current.probe_status = PoolMemberProbeStatus::Failed; - current.updated_at = 130; - repository - .upsert_pool_member_score(current.clone()) - .await - .expect("newer score should insert"); - let mut stale_recovery = score(None); - stale_recovery.score = 1.0; - stale_recovery.score_version = 8; - stale_recovery.score_reason = serde_json::json!({"state": "recovered"}); - stale_recovery.last_ranked_at = Some(100); - stale_recovery.updated_at = 100; - - repository - .upsert_pool_member_score_with_mode( - stale_recovery, - PoolMemberScoreUpsertMode::OAuthRecovery, - ) - .await - .expect("stale OAuth recovery should succeed"); - let preserved = load_score(&repository).await; - let current = current.into_stored(); - assert_eq!(preserved.score, current.score); - assert_eq!(preserved.hard_state, current.hard_state); - assert_eq!(preserved.score_version, current.score_version); - assert_eq!(preserved.score_reason, current.score_reason); - assert_eq!(preserved.last_ranked_at, current.last_ranked_at); - assert_eq!(preserved.last_scheduled_at, current.last_scheduled_at); - assert_eq!(preserved.last_success_at, current.last_success_at); - assert_eq!(preserved.last_failure_at, current.last_failure_at); - assert_eq!(preserved.failure_count, current.failure_count); - assert_eq!( - preserved.last_probe_attempt_at, - current.last_probe_attempt_at - ); - assert_eq!( - preserved.last_probe_success_at, - current.last_probe_success_at - ); - assert_eq!( - preserved.last_probe_failure_at, - current.last_probe_failure_at - ); - assert_eq!(preserved.probe_failure_count, current.probe_failure_count); - assert_eq!(preserved.probe_status, current.probe_status); - assert_eq!(preserved.updated_at, current.updated_at); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs b/crates/aether-data/adapters/sqlite/src/provider_catalog.rs deleted file mode 100644 index 11be74256..000000000 --- a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs +++ /dev/null @@ -1,5131 +0,0 @@ -use std::collections::BTreeMap; - -use async_trait::async_trait; -use sqlx::{ - query::Query, - sqlite::{SqliteArguments, SqliteRow}, - QueryBuilder, Row, Sqlite, -}; - -use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, - ProviderCatalogKeyCredentialsCasUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - ProviderCatalogProviderConfigCasUpdate, ProviderCatalogProxyCasUpdate, - ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, - StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, - StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, SqlitePool}; -use aether_data_query::{ - push_ci_contains_any, push_eq, push_in, push_limit_offset, push_optional_eq, SqlDialect, - WhereClause, -}; - -const LIST_PROVIDERS_BY_IDS_PREFIX: &str = r#" -SELECT - id, - name, - description, - website, - provider_type, - billing_type, - CAST(monthly_quota_usd AS REAL) AS monthly_quota_usd, - CAST(monthly_used_usd AS REAL) AS 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 -WHERE id IN ( -"#; - -const LIST_ENDPOINTS_BY_IDS_PREFIX: &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 -WHERE id IN ( -"#; - -const LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX: &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 -WHERE provider_id IN ( -"#; - -const LIST_KEYS_BY_IDS_PREFIX: &str = r#" -SELECT - id, - provider_id, - name, - auth_type, - capabilities, - is_active, - api_formats, - auth_type_by_format, - allow_auth_channel_mismatch_formats, - COALESCE(api_key, encrypted_key) AS api_key, - auth_config, - note, - internal_priority, - rate_multipliers, - global_priority_by_format, - allowed_models, - expires_at AS expires_at_unix_secs, - cache_ttl_minutes, - max_probe_interval_minutes, - proxy, - fingerprint, - rpm_limit, - concurrent_limit, - learned_rpm_limit, - concurrent_429_count, - rpm_429_count, - last_429_at AS last_429_at_unix_secs, - last_429_type, - adjustment_history, - utilization_samples, - last_probe_increase_at AS last_probe_increase_at_unix_secs, - last_rpm_peak, - request_count, - total_tokens, - CAST(total_cost_usd AS REAL) AS total_cost_usd, - success_count, - error_count, - total_response_time_ms, - last_used_at AS last_used_at_unix_secs, - auto_fetch_models, - last_models_fetch_at AS last_models_fetch_at_unix_secs, - last_models_fetch_error, - locked_models, - model_include_patterns, - model_exclude_patterns, - upstream_metadata, - oauth_invalid_at AS oauth_invalid_at_unix_secs, - oauth_invalid_reason, - status_snapshot, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs, - health_by_format, - circuit_breaker_by_format -FROM provider_api_keys -WHERE id IN ( -"#; - -const LIST_KEYS_BY_PROVIDER_IDS_PREFIX: &str = r#" -SELECT - id, - provider_id, - name, - auth_type, - capabilities, - is_active, - api_formats, - auth_type_by_format, - allow_auth_channel_mismatch_formats, - COALESCE(api_key, encrypted_key) AS api_key, - auth_config, - note, - internal_priority, - rate_multipliers, - global_priority_by_format, - allowed_models, - expires_at AS expires_at_unix_secs, - cache_ttl_minutes, - max_probe_interval_minutes, - proxy, - fingerprint, - rpm_limit, - concurrent_limit, - learned_rpm_limit, - concurrent_429_count, - rpm_429_count, - last_429_at AS last_429_at_unix_secs, - last_429_type, - adjustment_history, - utilization_samples, - last_probe_increase_at AS last_probe_increase_at_unix_secs, - last_rpm_peak, - request_count, - total_tokens, - CAST(total_cost_usd AS REAL) AS total_cost_usd, - success_count, - error_count, - total_response_time_ms, - last_used_at AS last_used_at_unix_secs, - auto_fetch_models, - last_models_fetch_at AS last_models_fetch_at_unix_secs, - last_models_fetch_error, - locked_models, - model_include_patterns, - model_exclude_patterns, - upstream_metadata, - oauth_invalid_at AS oauth_invalid_at_unix_secs, - oauth_invalid_reason, - status_snapshot, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs, - health_by_format, - circuit_breaker_by_format -FROM provider_api_keys -WHERE provider_id IN ( -"#; - -const LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &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, - 0 AS total_tokens, - 0.0 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 -WHERE provider_id IN ( -"#; - -const LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX: &str = r#" -SELECT - provider_id, - COUNT(*) AS total_keys, - SUM(CASE WHEN is_active THEN 1 ELSE 0 END) AS active_keys -FROM provider_api_keys -WHERE provider_id IN ( -"#; - -const LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &str = r#" -SELECT - id, - provider_id, - is_active, - upstream_metadata -FROM provider_api_keys -WHERE provider_id IN ( -"#; - -#[derive(Debug, Clone)] -pub struct SqliteProviderCatalogReadRepository { - pool: SqlitePool, -} - -impl SqliteProviderCatalogReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn list_providers_by_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - let rows = build_list_query( - LIST_PROVIDERS_BY_IDS_PREFIX, - provider_ids, - " ORDER BY name ASC", - ) - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_provider_row).collect() - } - - pub async fn list_providers( - &self, - active_only: bool, - ) -> Result, DataLayerError> { - let mut builder = - QueryBuilder::::new(select_prefix_for_in(LIST_PROVIDERS_BY_IDS_PREFIX)); - 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_list_query( - LIST_ENDPOINTS_BY_IDS_PREFIX, - 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() - } - - 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_list_query( - LIST_ENDPOINTS_BY_PROVIDER_IDS_PREFIX, - 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_list_query( - LIST_KEYS_BY_IDS_PREFIX, - key_ids, - " ORDER BY name ASC, id ASC", - ) - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_key_row).collect() - } - - pub async fn list_keys_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - let rows = build_list_query( - LIST_KEYS_BY_PROVIDER_IDS_PREFIX, - 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_list_query( - LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX, - 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_list_query( - LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX, - 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 { - if query.provider_id.trim().is_empty() { - return Err(DataLayerError::InvalidInput( - "provider catalog provider_id is empty".to_string(), - )); - } - - 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 order_by = match query.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" - } - }; - - let mut count_builder = - QueryBuilder::::new("SELECT COUNT(*) AS total FROM provider_api_keys"); - let mut count_where = WhereClause::new(); - apply_key_page_filters(&mut count_builder, &mut count_where, query); - let total = count_builder - .build_query_scalar::() - .fetch_one(&self.pool) - .await - .map_sql_err()? - .max(0) as usize; - - let mut list_builder = - QueryBuilder::::new(select_prefix_for_in(LIST_KEYS_BY_IDS_PREFIX)); - let mut list_where = WhereClause::new(); - apply_key_page_filters(&mut list_builder, &mut list_where, query); - list_builder.push(" ORDER BY ").push(order_by); - push_limit_offset(&mut list_builder, limit, offset); - let rows = list_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - 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_list_query( - LIST_KEY_STATS_BY_PROVIDER_IDS_PREFIX, - provider_ids, - "\nGROUP BY provider_id\nORDER 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( - &self, - provider: &StoredProviderCatalogProvider, - shift_existing_priorities_from: Option, - ) -> Result { - validate_provider(provider)?; - let now = current_unix_secs(); - let created_at = provider.created_at_unix_ms.unwrap_or(now) as i64; - let updated_at = provider.updated_at_unix_secs.unwrap_or(now) as i64; - let mut tx = self.pool.begin().await.map_sql_err()?; - - if let Some(target_priority) = shift_existing_priorities_from { - sqlx::query( - r#" -UPDATE providers -SET provider_priority = provider_priority + 1 -WHERE provider_priority >= ? -"#, - ) - .bind(target_priority) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query( - r#" -INSERT INTO providers ( - id, name, description, website, provider_type, billing_type, - monthly_quota_usd, monthly_used_usd, quota_reset_day, - quota_last_reset_at, quota_expires_at, 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, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&provider.id) - .bind(&provider.name) - .bind(&provider.description) - .bind(&provider.website) - .bind(&provider.provider_type) - .bind( - provider - .billing_type - .clone() - .unwrap_or_else(|| "pay_as_you_go".to_string()), - ) - .bind(provider.monthly_quota_usd) - .bind(provider.monthly_used_usd) - .bind(optional_i64_from_u64( - provider.quota_reset_day, - "providers.quota_reset_day", - )?) - .bind(optional_i64_from_u64( - provider.quota_last_reset_at_unix_secs, - "providers.quota_last_reset_at", - )?) - .bind(optional_i64_from_u64( - provider.quota_expires_at_unix_secs, - "providers.quota_expires_at", - )?) - .bind(provider.provider_priority) - .bind(provider.is_active) - .bind(provider.keep_priority_on_conversion) - .bind(provider.enable_format_conversion) - .bind(provider.concurrent_limit) - .bind(provider.max_retries) - .bind(optional_json_to_string(&provider.proxy, "providers.proxy")?) - .bind(provider.request_timeout_secs) - .bind(provider.stream_first_byte_timeout_secs) - .bind(optional_json_to_string( - &provider.config, - "providers.config", - )?) - .bind(created_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - self.reload_provider(&provider.id, "created").await - } - - pub async fn update_provider( - &self, - provider: &StoredProviderCatalogProvider, - ) -> Result { - validate_provider(provider)?; - let updated_at = provider - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET - name = ?, - description = ?, - website = ?, - provider_type = ?, - billing_type = ?, - monthly_quota_usd = ?, - monthly_used_usd = ?, - quota_reset_day = ?, - quota_last_reset_at = ?, - quota_expires_at = ?, - provider_priority = ?, - is_active = ?, - keep_priority_on_conversion = ?, - enable_format_conversion = ?, - concurrent_limit = ?, - max_retries = ?, - proxy = ?, - request_timeout = ?, - stream_first_byte_timeout = ?, - config = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&provider.name) - .bind(&provider.description) - .bind(&provider.website) - .bind(&provider.provider_type) - .bind( - provider - .billing_type - .clone() - .unwrap_or_else(|| "pay_as_you_go".to_string()), - ) - .bind(provider.monthly_quota_usd) - .bind(provider.monthly_used_usd) - .bind(optional_i64_from_u64( - provider.quota_reset_day, - "providers.quota_reset_day", - )?) - .bind(optional_i64_from_u64( - provider.quota_last_reset_at_unix_secs, - "providers.quota_last_reset_at", - )?) - .bind(optional_i64_from_u64( - provider.quota_expires_at_unix_secs, - "providers.quota_expires_at", - )?) - .bind(provider.provider_priority) - .bind(provider.is_active) - .bind(provider.keep_priority_on_conversion) - .bind(provider.enable_format_conversion) - .bind(provider.concurrent_limit) - .bind(provider.max_retries) - .bind(optional_json_to_string(&provider.proxy, "providers.proxy")?) - .bind(provider.request_timeout_secs) - .bind(provider.stream_first_byte_timeout_secs) - .bind(optional_json_to_string( - &provider.config, - "providers.config", - )?) - .bind(updated_at) - .bind(&provider.id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog provider {} not found", - provider.id - ))); - } - self.reload_provider(&provider.id, "updated").await - } - - pub async fn compare_and_swap_provider_config( - &self, - update: &ProviderCatalogProviderConfigCasUpdate, - ) -> Result { - validate_non_empty(&update.provider_id, "provider catalog provider_id")?; - let expected_config = - optional_json_to_string(&update.expected_config, "providers.expected_config")?; - let config = optional_json_to_string(&update.config, "providers.config")?; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET config = ?, updated_at = ? -WHERE id = ? - AND config IS ? -"#, - ) - .bind(config) - .bind(current_unix_secs() as i64) - .bind(&update.provider_id) - .bind(expected_config) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - pub async fn compare_and_swap_provider_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog provider_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM providers WHERE id = ?", - "UPDATE providers SET proxy = ?, updated_at = ? WHERE id = ? AND proxy IS ?", - update, - "providers.proxy", - ) - .await - } - - pub async fn delete_provider(&self, provider_id: &str) -> Result { - validate_non_empty(provider_id, "provider catalog provider_id")?; - let rows_affected = sqlx::query("DELETE FROM providers WHERE id = ?") - .bind(provider_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn cleanup_deleted_provider_refs( - &self, - provider_id: &str, - provider_deleted: bool, - endpoint_ids: &[String], - key_ids: &[String], - ) -> Result<(), DataLayerError> { - validate_non_empty(provider_id, "provider catalog provider_id")?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - if provider_deleted { - sqlx::query( - "UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id = ?", - ) - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE video_tasks SET provider_id = NULL WHERE provider_id = ?") - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM request_candidates WHERE provider_id = ?") - .bind(provider_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - for endpoint_id in endpoint_ids { - sqlx::query("UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id = ?") - .bind(endpoint_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM request_candidates WHERE endpoint_id = ?") - .bind(endpoint_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - for key_id in key_ids { - sqlx::query("DELETE FROM gemini_file_mappings WHERE key_id = ?") - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE video_tasks SET key_id = NULL WHERE key_id = ?") - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - tx.commit().await.map_sql_err()?; - Ok(()) - } - - pub async fn create_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - validate_endpoint(endpoint)?; - let now = current_unix_secs(); - sqlx::query( - r#" -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, api_family, endpoint_kind, - is_active, health_score, header_rules, body_rules, max_retries, - custom_path, config, format_acceptance_config, proxy, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&endpoint.id) - .bind(&endpoint.provider_id) - .bind(&endpoint.api_format) - .bind(&endpoint.base_url) - .bind(&endpoint.api_format) - .bind(&endpoint.api_family) - .bind(&endpoint.endpoint_kind) - .bind(endpoint.is_active) - .bind(endpoint.health_score) - .bind(optional_json_to_string( - &endpoint.header_rules, - "provider_endpoints.header_rules", - )?) - .bind(optional_json_to_string( - &endpoint.body_rules, - "provider_endpoints.body_rules", - )?) - .bind(endpoint.max_retries) - .bind(&endpoint.custom_path) - .bind(optional_json_to_string( - &endpoint.config, - "provider_endpoints.config", - )?) - .bind(optional_json_to_string( - &endpoint.format_acceptance_config, - "provider_endpoints.format_acceptance_config", - )?) - .bind(optional_json_to_string( - &endpoint.proxy, - "provider_endpoints.proxy", - )?) - .bind(endpoint.created_at_unix_ms.unwrap_or(now) as i64) - .bind(endpoint.updated_at_unix_secs.unwrap_or(now) as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.reload_endpoint(&endpoint.id, "created").await - } - - pub async fn update_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - validate_endpoint(endpoint)?; - let updated_at = endpoint - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64; - let rows_affected = sqlx::query( - r#" -UPDATE provider_endpoints -SET - provider_id = ?, - name = ?, - base_url = ?, - api_format = ?, - api_family = ?, - endpoint_kind = ?, - is_active = ?, - health_score = ?, - header_rules = ?, - body_rules = ?, - max_retries = ?, - custom_path = ?, - config = ?, - format_acceptance_config = ?, - proxy = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&endpoint.provider_id) - .bind(&endpoint.api_format) - .bind(&endpoint.base_url) - .bind(&endpoint.api_format) - .bind(&endpoint.api_family) - .bind(&endpoint.endpoint_kind) - .bind(endpoint.is_active) - .bind(endpoint.health_score) - .bind(optional_json_to_string( - &endpoint.header_rules, - "provider_endpoints.header_rules", - )?) - .bind(optional_json_to_string( - &endpoint.body_rules, - "provider_endpoints.body_rules", - )?) - .bind(endpoint.max_retries) - .bind(&endpoint.custom_path) - .bind(optional_json_to_string( - &endpoint.config, - "provider_endpoints.config", - )?) - .bind(optional_json_to_string( - &endpoint.format_acceptance_config, - "provider_endpoints.format_acceptance_config", - )?) - .bind(optional_json_to_string( - &endpoint.proxy, - "provider_endpoints.proxy", - )?) - .bind(updated_at) - .bind(&endpoint.id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog endpoint {} not found", - endpoint.id - ))); - } - self.reload_endpoint(&endpoint.id, "updated").await - } - - pub async fn compare_and_swap_endpoint_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog endpoint_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM provider_endpoints WHERE id = ?", - "UPDATE provider_endpoints SET proxy = ?, updated_at = ? WHERE id = ? AND proxy IS ?", - update, - "provider_endpoints.proxy", - ) - .await - } - - pub async fn delete_endpoint(&self, endpoint_id: &str) -> Result { - validate_non_empty(endpoint_id, "provider catalog endpoint_id")?; - let rows_affected = sqlx::query("DELETE FROM provider_endpoints WHERE id = ?") - .bind(endpoint_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn create_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - validate_key(key)?; - let now = current_unix_secs(); - sqlx::query(key_insert_sql()) - .bind(&key.id) - .bind(&key.provider_id) - .bind(&key.name) - .bind(&key.encrypted_api_key) - .bind(&key.auth_type) - .bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .bind(key.is_active) - .bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .bind(&key.encrypted_auth_config) - .bind(&key.note) - .bind(key.internal_priority) - .bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .bind(key.cache_ttl_minutes) - .bind(key.max_probe_interval_minutes) - .bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .bind(optional_i64_from_u32(key.rpm_limit)) - .bind(key.concurrent_limit) - .bind(optional_i64_from_u32(key.learned_rpm_limit)) - .bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0)) - .bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0)) - .bind(optional_i64_from_u64( - key.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .bind(&key.last_429_type) - .bind(optional_json_to_string( - &key.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .bind(optional_json_to_string( - &key.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .bind(optional_i64_from_u64( - key.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .bind(optional_i64_from_u32(key.last_rpm_peak)) - .bind(optional_i64_from_u32(key.request_count).unwrap_or(0)) - .bind(i64::try_from(key.total_tokens).map_err(|_| { - DataLayerError::InvalidInput(format!( - "provider catalog key.total_tokens exceeds i64: {}", - key.total_tokens - )) - })?) - .bind(key.total_cost_usd) - .bind(optional_i64_from_u32(key.success_count).unwrap_or(0)) - .bind(optional_i64_from_u32(key.error_count).unwrap_or(0)) - .bind( - optional_i64_from_u64( - key.total_response_time_ms, - "provider_api_keys.total_response_time_ms", - )? - .unwrap_or(0), - ) - .bind(optional_i64_from_u64( - key.last_used_at_unix_secs, - "provider_api_keys.last_used_at", - )?) - .bind(key.auto_fetch_models) - .bind(optional_i64_from_u64( - key.last_models_fetch_at_unix_secs, - "provider_api_keys.last_models_fetch_at", - )?) - .bind(&key.last_models_fetch_error) - .bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .bind(optional_json_to_string( - &key.upstream_metadata, - "provider_api_keys.upstream_metadata", - )?) - .bind(optional_i64_from_u64( - key.oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .bind(&key.oauth_invalid_reason) - .bind(optional_json_to_string( - &key.status_snapshot, - "provider_api_keys.status_snapshot", - )?) - .bind(optional_json_to_string( - &key.health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &key.circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(key.created_at_unix_ms.unwrap_or(now) as i64) - .bind(key.updated_at_unix_secs.unwrap_or(now) as i64) - .execute(&self.pool) - .await - .map_sql_err()?; - - self.reload_key(&key.id, "created").await - } - - pub async fn update_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - validate_key(key)?; - let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let rows_affected = key_update_query(key, updated_at)? - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog key {} not found", - key.id - ))); - } - self.reload_key(&key.id, "updated").await - } - - pub async fn compare_and_swap_key_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - validate_non_empty(&update.record_id, "provider catalog key_id")?; - compare_and_swap_proxy_json( - &self.pool, - "SELECT proxy FROM provider_api_keys WHERE id = ?", - "UPDATE provider_api_keys SET proxy = ?, updated_at = ? WHERE id = ? AND proxy IS ?", - update, - "provider_api_keys.proxy", - ) - .await - } - - pub async fn compare_and_swap_key_credentials( - &self, - update: &ProviderCatalogKeyCredentialsCasUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.expected_provider_id, - "provider catalog expected provider_id", - )?; - let rows_affected = sqlx::query( - "UPDATE provider_api_keys SET api_key = ?, encrypted_key = NULL, auth_config = ? WHERE id = ? AND provider_id = ? AND COALESCE(api_key, encrypted_key) IS ? AND auth_config IS ?", - ) - .bind(update.encrypted_api_key.as_deref()) - .bind(update.encrypted_auth_config.as_deref()) - .bind(&update.key_id) - .bind(&update.expected_provider_id) - .bind(update.expected_encrypted_api_key.as_deref()) - .bind(update.expected_encrypted_auth_config.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) - } - - pub async fn compare_and_update_key_admin_state( - &self, - update: &ProviderCatalogKeyAdminCasUpdate, - ) -> Result { - validate_admin_key_cas_update(update)?; - let key = &update.key; - let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let rotation_json = update - .codex_rotation - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys Codex rotation is not serializable: {err}" - )) - })?; - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET provider_id = "); - push_admin_key_assignments(&mut builder, key, updated_at)?; - if update.reset_oauth_runtime { - builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0"); - } - if let Some(rotation_json) = rotation_json.as_deref() { - builder - .push(", upstream_metadata = json_set(COALESCE(upstream_metadata, '{}'), '$.codex', json(") - .push_bind(rotation_json) - .push("))"); - } - if update.codex_rotation.is_some() || update.reset_oauth_runtime { - builder.push(", status_snapshot = "); - match (update.codex_rotation.is_some(), update.reset_oauth_runtime) { - (true, true) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.quota', json('null'), '$.oauth', json('null'))"), - (true, false) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.quota', json('null'))"), - (false, true) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.oauth', json('null'))"), - (false, false) => unreachable!(), - }; - } - builder - .push(" WHERE id = ") - .push_bind(&key.id) - .push(" AND api_key IS ") - .push_bind(update.expected_credential.encrypted_api_key.as_deref()) - .push(" AND auth_config IS ") - .push_bind(update.expected_encrypted_auth_config.as_deref()) - .push(" AND auth_type = ") - .push_bind(&update.expected_credential.auth_type) - .push(" AND provider_id = ") - .push_bind(&update.expected_credential.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE providers.id = provider_api_keys.provider_id AND providers.provider_type = ", - ) - .push_bind(&update.expected_credential.provider_type) - .push(")"); - if update.codex_rotation.is_some() { - builder - .push(" AND CASE WHEN upstream_metadata IS NULL THEN 1 WHEN json_valid(upstream_metadata) THEN json_type(upstream_metadata) = 'object' ELSE 0 END") - .push(" AND NOT (api_key IS ") - .push_bind(key.encrypted_api_key.as_deref()) - .push(" AND auth_config IS ") - .push_bind(key.encrypted_auth_config.as_deref()) - .push(" AND auth_type = ") - .push_bind(&key.auth_type) - .push(" AND provider_id = ") - .push_bind(&key.provider_id) - .push(")"); - } - if update.codex_rotation.is_some() || update.reset_oauth_runtime { - builder.push(" AND CASE WHEN status_snapshot IS NULL THEN 1 WHEN json_valid(status_snapshot) THEN json_type(status_snapshot) = 'object' ELSE 0 END"); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_keys( - &self, - keys: &[StoredProviderCatalogKey], - ) -> Result, DataLayerError> { - if keys.is_empty() { - return Ok(Vec::new()); - } - for key in keys { - validate_key(key)?; - } - - let updated_at = current_unix_secs() as i64; - let mut transaction = self.pool.begin().await.map_sql_err()?; - for key in keys { - let key_updated_at = key.updated_at_unix_secs.unwrap_or(updated_at as u64) as i64; - let rows_affected = key_update_query(key, key_updated_at)? - .execute(&mut *transaction) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - return Err(DataLayerError::UnexpectedValue(format!( - "provider catalog key {} not found", - key.id - ))); - } - } - transaction.commit().await.map_sql_err()?; - let key_ids = keys.iter().map(|key| key.id.clone()).collect::>(); - let mut reloaded = self - .list_keys_by_ids(&key_ids) - .await? - .into_iter() - .map(|key| (key.id.clone(), key)) - .collect::>(); - keys.iter() - .map(|key| { - reloaded.remove(&key.id).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "updated provider catalog key {} could not be reloaded", - key.id - )) - }) - }) - .collect() - } - - pub async fn delete_key(&self, key_id: &str) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query("DELETE FROM provider_api_keys WHERE id = ?") - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_delete_key_oauth_credential( - &self, - delete: &ProviderCatalogKeyOAuthCredentialCasDelete, - ) -> Result { - validate_non_empty(&delete.key_id, "provider catalog key_id")?; - let expected = &delete.expected_credential; - if expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - || delete - .expected_upstream_metadata_namespace - .as_ref() - .is_some_and(|expected| expected.namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth credential CAS delete contains empty fields".to_string(), - )); - } - let mut builder = QueryBuilder::::new("DELETE FROM provider_api_keys WHERE id = "); - builder - .push_bind(&delete.key_id) - .push(" AND auth_config IS ") - .push_bind(delete.expected_encrypted_auth_config.as_deref()) - .push(" AND api_key IS ") - .push_bind(expected.encrypted_api_key.as_deref()) - .push(" AND auth_type = ") - .push_bind(&expected.auth_type) - .push(" AND provider_id = ") - .push_bind(&expected.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE \ - providers.id = provider_api_keys.provider_id \ - AND providers.provider_type = ", - ) - .push_bind(&expected.provider_type) - .push(")"); - if let Some(expected) = delete.expected_upstream_metadata_namespace.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(&expected.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); - if let Some(expected_value) = expected.expected_value.as_ref() { - let expected_value = serde_json::to_string(expected_value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path.clone()) - .push(") = json_type(json(") - .push_bind(expected_value.clone()) - .push(")) AND json_extract(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS json_extract(json_object('value', json(") - .push_bind(expected_value) - .push(")), '$.value')"); - } else { - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS NULL"); - } - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_upstream_metadata( - &self, - key_id: &str, - upstream_metadata: Option<&serde_json::Value>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_json_ref_to_string( - upstream_metadata, - "provider_api_keys.upstream_metadata", - )?) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn upsert_key_upstream_metadata_namespace( - &self, - key_id: &str, - namespace: &str, - value: &serde_json::Value, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - validate_non_empty(namespace, "provider catalog upstream metadata namespace")?; - let value_json = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })?; - let namespace_path = format!( - "$.{}", - serde_json::to_string(namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = json_set( - COALESCE(NULLIF(upstream_metadata, ''), '{}'), - ?, json(?) - ), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(namespace_path) - .bind(value_json) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_model_fetch_state( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: Option, - last_models_fetch_error: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_json_ref_to_string( - allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - last_models_fetch_at_unix_secs, - "provider_api_keys.last_models_fetch_at", - )?) - .bind(last_models_fetch_error) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_model_fetch_success( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: u64, - upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let allowed_models = - optional_json_ref_to_string(allowed_models, "provider_api_keys.allowed_models")?; - let namespace_updates = upstream_metadata_updates - .iter() - .map(|update| { - validate_non_empty( - &update.namespace, - "provider catalog upstream metadata namespace", - )?; - let path = format!( - "$.{}", - serde_json::to_string(&update.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(&update.value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })?; - Ok((path, value)) - }) - .collect::, DataLayerError>>()?; - let updated_at = updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; - let mut tx = self.pool.begin().await.map_sql_err()?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = NULL, updated_at = ? -WHERE id = ? -"#, - ) - .bind(allowed_models) - .bind(optional_i64_from_u64( - Some(last_models_fetch_at_unix_secs), - "provider_api_keys.last_models_fetch_at", - )?) - .bind(updated_at) - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - for (path, value) in namespace_updates { - sqlx::query( - r#" -UPDATE provider_api_keys -SET upstream_metadata = json_set( - COALESCE(NULLIF(upstream_metadata, ''), '{}'), - ?, json(?) - ) -WHERE id = ? -"#, - ) - .bind(path) - .bind(value) - .bind(key_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - pub async fn clear_key_oauth_invalid_marker( - &self, - key_id: &str, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET oauth_invalid_at = NULL, oauth_invalid_reason = NULL, updated_at = ? -WHERE id = ? -"#, - ) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_oauth_runtime_state( - &self, - key_id: &str, - oauth_invalid_at_unix_secs: Option, - oauth_invalid_reason: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET oauth_invalid_at = ?, oauth_invalid_reason = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(optional_i64_from_u64( - oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .bind(oauth_invalid_reason) - .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_oauth_runtime_state( - &self, - update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.encrypted_auth_config, - "provider catalog OAuth auth_config", - )?; - if update - .encrypted_api_key_update - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth api_key update must not be empty".to_string(), - )); - } - if update.expected_credential.as_ref().is_some_and(|expected| { - expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - }) { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth credential fence must not contain empty fields".to_string(), - )); - } - if update - .expected_upstream_metadata_namespace - .as_ref() - .is_some_and(|expected| expected.namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace must not be empty".to_string(), - )); - } - if update - .upstream_metadata_namespace_to_remove - .as_deref() - .is_some_and(|namespace| namespace.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace to remove must not be empty" - .to_string(), - )); - } - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - )); - } - if update - .upstream_metadata_patch - .as_ref() - .is_some_and(|patch| !patch.is_object()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog upstream metadata patch must be an object".to_string(), - )); - } - if update - .upstream_metadata_namespace_to_remove - .as_ref() - .is_some_and(|namespace| { - update - .upstream_metadata_patch - .as_ref() - .and_then(serde_json::Value::as_object) - .is_some_and(|patch| patch.contains_key(namespace)) - }) - { - return Err(DataLayerError::InvalidInput( - "provider catalog OAuth runtime metadata namespace cannot be patched and removed in the same update" - .to_string(), - )); - } - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET oauth_invalid_at = "); - builder - .push_bind(optional_i64_from_u64( - update.oauth_invalid_at_unix_secs, - "provider_api_keys.oauth_invalid_at", - )?) - .push(", oauth_invalid_reason = ") - .push_bind(update.oauth_invalid_reason.as_deref()) - .push(", auth_config = ") - .push_bind(&update.encrypted_auth_config); - if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_deref() { - builder.push(", api_key = ").push_bind(encrypted_api_key); - } - if let Some(expires_at_unix_secs) = update.expires_at_unix_secs_update { - builder - .push(", expires_at = ") - .push_bind(optional_i64_from_u64( - expires_at_unix_secs, - "provider_api_keys.expires_at", - )?); - } - if update.upstream_metadata_patch.is_some() - || update.upstream_metadata_namespace_to_remove.is_some() - { - builder.push(", upstream_metadata = "); - if update.upstream_metadata_namespace_to_remove.is_some() { - builder.push("json_remove("); - } - if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { - push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; - } else { - builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - } - if let Some(namespace) = update.upstream_metadata_namespace_to_remove.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - builder.push(", ").push_bind(namespace_path).push(")"); - } - } - builder.push(", status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - if update.reset_error_count { - builder.push(", error_count = 0"); - } - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND auth_config IS ") - .push_bind(update.expected_encrypted_auth_config.as_deref()); - if let Some(expected) = update.expected_credential.as_ref() { - builder - .push(" AND api_key IS ") - .push_bind(expected.encrypted_api_key.as_deref()) - .push(" AND auth_type = ") - .push_bind(&expected.auth_type) - .push(" AND provider_id = ") - .push_bind(&expected.provider_id) - .push( - " AND EXISTS (SELECT 1 FROM providers WHERE \ - providers.id = provider_api_keys.provider_id \ - AND providers.provider_type = ", - ) - .push_bind(&expected.provider_type) - .push(")"); - } - if update.expected_upstream_metadata_namespace.is_some() - || update.upstream_metadata_patch.is_some() - || update.upstream_metadata_namespace_to_remove.is_some() - { - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); - } - if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { - let namespace_path = format!( - "$.{}", - serde_json::to_string(&expected.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - if let Some(expected_value) = expected.expected_value.as_ref() { - let expected_value = serde_json::to_string(expected_value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path.clone()) - .push(") = json_type(json(") - .push_bind(expected_value.clone()) - .push(")) AND json_extract(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS json_extract(json_object('value', json(") - .push_bind(expected_value) - .push(")), '$.value')"); - } else { - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS NULL"); - } - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_health_state( - &self, - key_id: &str, - is_active: bool, - health_by_format: Option<&serde_json::Value>, - circuit_breaker_by_format: Option<&serde_json::Value>, - ) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET is_active = ?, health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ? -WHERE id = ? -"#, - ) - .bind(is_active) - .bind(optional_json_ref_to_string( - health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_ref_to_string( - circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn reset_key_error_count(&self, key_id: &str) -> Result { - validate_non_empty(key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET error_count = 0, updated_at = ? -WHERE id = ? -"#, - ) - .bind(current_unix_secs() as i64) - .bind(key_id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_adaptive_state( - &self, - update: &ProviderCatalogKeyAdaptiveStateUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - let status_snapshot_patch = adaptive_status_snapshot_patch(&update.status_snapshot_patch)?; - let expected = update.expected.canonicalized(); - let next = update.next.canonicalized(); - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET learned_rpm_limit = "); - builder - .push_bind(optional_i64_from_u32(next.learned_rpm_limit)) - .push(", rpm_429_count = ") - .push_bind(optional_i64_from_u32(next.rpm_429_count)) - .push(", last_429_at = ") - .push_bind(optional_i64_from_u64( - next.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .push(", last_429_type = ") - .push_bind(&next.last_429_type) - .push(", adjustment_history = ") - .push_bind(optional_json_to_string( - &next.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .push(", utilization_samples = ") - .push_bind(optional_json_to_string( - &next.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .push(", last_probe_increase_at = ") - .push_bind(optional_i64_from_u64( - next.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .push(", last_rpm_peak = ") - .push_bind(optional_i64_from_u32(next.last_rpm_peak)) - .push(", concurrent_429_count = ") - .push_bind(optional_i64_from_u32(next.concurrent_429_count)) - .push(", status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND learned_rpm_limit IS ") - .push_bind(optional_i64_from_u32(expected.learned_rpm_limit)) - .push(" AND rpm_429_count IS ") - .push_bind(optional_i64_from_u32(expected.rpm_429_count)) - .push(" AND last_429_at IS ") - .push_bind(optional_i64_from_u64( - expected.last_429_at_unix_secs, - "provider_api_keys.last_429_at", - )?) - .push(" AND last_429_type IS ") - .push_bind(&expected.last_429_type) - .push(" AND json(adjustment_history) IS json(") - .push_bind(optional_json_to_string( - &expected.adjustment_history, - "provider_api_keys.adjustment_history", - )?) - .push(")") - .push(" AND json(utilization_samples) IS json(") - .push_bind(optional_json_to_string( - &expected.utilization_samples, - "provider_api_keys.utilization_samples", - )?) - .push(")") - .push(" AND last_probe_increase_at IS ") - .push_bind(optional_i64_from_u64( - expected.last_probe_increase_at_unix_secs, - "provider_api_keys.last_probe_increase_at", - )?) - .push(" AND last_rpm_peak IS ") - .push_bind(optional_i64_from_u32(expected.last_rpm_peak)) - .push(" AND concurrent_429_count IS ") - .push_bind(optional_i64_from_u32(expected.concurrent_429_count)); - if let Some(expected_encrypted_auth_config) = - update.expected_encrypted_auth_config.as_deref() - { - builder - .push(" AND auth_config IS ") - .push_bind(expected_encrypted_auth_config); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_runtime_metadata( - &self, - update: &ProviderCatalogKeyRuntimeMetadataUpdate, - ) -> Result { - validate_runtime_metadata_update(update)?; - let namespace_path = format!( - "$.{}", - serde_json::to_string(&update.namespace).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata namespace is not serializable: {err}" - )) - })? - ); - let metadata_value = - serde_json::to_string(&update.upstream_metadata_value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata value is not serializable: {err}" - )) - })?; - let expected_metadata_value = update - .expected_upstream_metadata_value - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata expected value is not serializable: {err}" - )) - })?; - let mut builder = QueryBuilder::::new( - "UPDATE provider_api_keys SET upstream_metadata = json_set(\ - COALESCE(NULLIF(upstream_metadata, ''), '{}'), ", - ); - builder - .push_bind(namespace_path.clone()) - .push(", json(") - .push_bind(metadata_value) - .push(")), status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id) - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); - if let Some(expected_metadata_value) = expected_metadata_value { - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path.clone()) - .push(") = json_type(json(") - .push_bind(expected_metadata_value.clone()) - .push(")) AND json_extract(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS json_extract(json_object('value', json(") - .push_bind(expected_metadata_value) - .push(")), '$.value')"); - } else { - builder - .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") - .push_bind(namespace_path) - .push(") IS NULL"); - } - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn update_key_status_snapshot( - &self, - update: &ProviderCatalogKeyStatusSnapshotUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - )); - } - let mut builder = - QueryBuilder::::new("UPDATE provider_api_keys SET status_snapshot = "); - push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; - builder - .push(", updated_at = ") - .push_bind( - update - .updated_at_unix_secs - .unwrap_or_else(current_unix_secs) as i64, - ) - .push(" WHERE id = ") - .push_bind(&update.key_id); - let rows_affected = builder - .build() - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - pub async fn compare_and_update_key_health_state( - &self, - update: &ProviderCatalogKeyHealthStateUpdate, - ) -> Result { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - let rows_affected = sqlx::query( - r#" -UPDATE provider_api_keys -SET health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ? -WHERE id = ? - AND json(health_by_format) IS json(?) - AND json(circuit_breaker_by_format) IS json(?) - AND (? IS NULL OR auth_config IS ?) -"#, - ) - .bind(optional_json_to_string( - &update.health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &update.circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(current_unix_secs() as i64) - .bind(&update.key_id) - .bind(optional_json_to_string( - &update.expected_health_by_format, - "provider_api_keys.health_by_format", - )?) - .bind(optional_json_to_string( - &update.expected_circuit_breaker_by_format, - "provider_api_keys.circuit_breaker_by_format", - )?) - .bind(update.expected_encrypted_auth_config.as_deref()) - .bind(update.expected_encrypted_auth_config.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected > 0) - } - - async fn reload_provider( - &self, - provider_id: &str, - operation: &str, - ) -> Result { - self.list_providers_by_ids(&[provider_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog provider {provider_id} could not be reloaded" - )) - }) - } - - async fn reload_endpoint( - &self, - endpoint_id: &str, - operation: &str, - ) -> Result { - self.list_endpoints_by_ids(&[endpoint_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog endpoint {endpoint_id} could not be reloaded" - )) - }) - } - - async fn reload_key( - &self, - key_id: &str, - operation: &str, - ) -> Result { - self.list_keys_by_ids(&[key_id.to_string()]) - .await? - .into_iter() - .next() - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{operation} provider catalog key {key_id} could not be reloaded" - )) - }) - } -} - -#[async_trait] -impl ProviderCatalogReadRepository for SqliteProviderCatalogReadRepository { - async fn list_providers( - &self, - active_only: bool, - ) -> Result, DataLayerError> { - Self::list_providers(self, active_only).await - } - - async fn list_providers_by_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_providers_by_ids(self, provider_ids).await - } - - async fn list_endpoints_by_ids( - &self, - endpoint_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_endpoints_by_ids(self, endpoint_ids).await - } - - async fn list_endpoints_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_endpoints_by_provider_ids(self, provider_ids).await - } - - async fn list_keys_by_ids( - &self, - key_ids: &[String], - ) -> Result, DataLayerError> { - 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(self, provider_ids).await - } - - async fn list_key_summaries_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - 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> { - Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await - } - - async fn list_keys_page( - &self, - query: &ProviderCatalogKeyListQuery, - ) -> Result { - Self::list_keys_page(self, query).await - } - - async fn list_key_stats_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - Self::list_key_stats_by_provider_ids(self, provider_ids).await - } -} - -#[async_trait] -impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository { - async fn create_provider( - &self, - provider: &StoredProviderCatalogProvider, - shift_existing_priorities_from: Option, - ) -> Result { - Self::create_provider(self, provider, shift_existing_priorities_from).await - } - - async fn update_provider( - &self, - provider: &StoredProviderCatalogProvider, - ) -> Result { - Self::update_provider(self, provider).await - } - - async fn compare_and_swap_provider_config( - &self, - update: &ProviderCatalogProviderConfigCasUpdate, - ) -> Result { - Self::compare_and_swap_provider_config(self, update).await - } - - async fn compare_and_swap_provider_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_provider_proxy(self, update).await - } - - async fn delete_provider(&self, provider_id: &str) -> Result { - Self::delete_provider(self, provider_id).await - } - - async fn cleanup_deleted_provider_refs( - &self, - provider_id: &str, - provider_deleted: bool, - endpoint_ids: &[String], - key_ids: &[String], - ) -> Result<(), DataLayerError> { - Self::cleanup_deleted_provider_refs( - self, - provider_id, - provider_deleted, - endpoint_ids, - key_ids, - ) - .await - } - - async fn create_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - Self::create_endpoint(self, endpoint).await - } - - async fn update_endpoint( - &self, - endpoint: &StoredProviderCatalogEndpoint, - ) -> Result { - Self::update_endpoint(self, endpoint).await - } - - async fn compare_and_swap_endpoint_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_endpoint_proxy(self, update).await - } - - async fn delete_endpoint(&self, endpoint_id: &str) -> Result { - Self::delete_endpoint(self, endpoint_id).await - } - - async fn create_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - Self::create_key(self, key).await - } - - async fn update_key( - &self, - key: &StoredProviderCatalogKey, - ) -> Result { - Self::update_key(self, key).await - } - - async fn compare_and_swap_key_proxy( - &self, - update: &ProviderCatalogProxyCasUpdate, - ) -> Result { - Self::compare_and_swap_key_proxy(self, update).await - } - - async fn compare_and_swap_key_credentials( - &self, - update: &ProviderCatalogKeyCredentialsCasUpdate, - ) -> Result { - Self::compare_and_swap_key_credentials(self, update).await - } - - async fn compare_and_update_key_admin_state( - &self, - update: &ProviderCatalogKeyAdminCasUpdate, - ) -> Result { - Self::compare_and_update_key_admin_state(self, update).await - } - - async fn update_keys( - &self, - keys: &[StoredProviderCatalogKey], - ) -> Result, DataLayerError> { - Self::update_keys(self, keys).await - } - - async fn update_key_upstream_metadata( - &self, - key_id: &str, - upstream_metadata: Option<&serde_json::Value>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_upstream_metadata(self, key_id, upstream_metadata, updated_at_unix_secs) - .await - } - - async fn upsert_key_upstream_metadata_namespace( - &self, - key_id: &str, - namespace: &str, - value: &serde_json::Value, - updated_at_unix_secs: Option, - ) -> Result { - Self::upsert_key_upstream_metadata_namespace( - self, - key_id, - namespace, - value, - updated_at_unix_secs, - ) - .await - } - - async fn update_key_model_fetch_state( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: Option, - last_models_fetch_error: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_model_fetch_state( - self, - key_id, - allowed_models, - last_models_fetch_at_unix_secs, - last_models_fetch_error, - updated_at_unix_secs, - ) - .await - } - - async fn update_key_model_fetch_success( - &self, - key_id: &str, - allowed_models: Option<&serde_json::Value>, - last_models_fetch_at_unix_secs: u64, - upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_model_fetch_success( - self, - key_id, - allowed_models, - last_models_fetch_at_unix_secs, - upstream_metadata_updates, - updated_at_unix_secs, - ) - .await - } - - async fn delete_key(&self, key_id: &str) -> Result { - Self::delete_key(self, key_id).await - } - - async fn compare_and_delete_key_oauth_credential( - &self, - delete: &ProviderCatalogKeyOAuthCredentialCasDelete, - ) -> Result { - Self::compare_and_delete_key_oauth_credential(self, delete).await - } - - async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result { - Self::clear_key_oauth_invalid_marker(self, key_id).await - } - - async fn update_key_oauth_runtime_state( - &self, - key_id: &str, - oauth_invalid_at_unix_secs: Option, - oauth_invalid_reason: Option<&str>, - updated_at_unix_secs: Option, - ) -> Result { - Self::update_key_oauth_runtime_state( - self, - key_id, - oauth_invalid_at_unix_secs, - oauth_invalid_reason, - updated_at_unix_secs, - ) - .await - } - - async fn compare_and_update_key_oauth_runtime_state( - &self, - update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ) -> Result { - Self::compare_and_update_key_oauth_runtime_state(self, update).await - } - - async fn update_key_health_state( - &self, - key_id: &str, - is_active: bool, - health_by_format: Option<&serde_json::Value>, - circuit_breaker_by_format: Option<&serde_json::Value>, - ) -> Result { - Self::update_key_health_state( - self, - key_id, - is_active, - health_by_format, - circuit_breaker_by_format, - ) - .await - } - - async fn reset_key_error_count(&self, key_id: &str) -> Result { - Self::reset_key_error_count(self, key_id).await - } - - async fn compare_and_update_key_adaptive_state( - &self, - update: &ProviderCatalogKeyAdaptiveStateUpdate, - ) -> Result { - Self::compare_and_update_key_adaptive_state(self, update).await - } - - async fn update_key_runtime_metadata( - &self, - update: &ProviderCatalogKeyRuntimeMetadataUpdate, - ) -> Result { - Self::update_key_runtime_metadata(self, update).await - } - - async fn update_key_status_snapshot( - &self, - update: &ProviderCatalogKeyStatusSnapshotUpdate, - ) -> Result { - Self::update_key_status_snapshot(self, update).await - } - - async fn compare_and_update_key_health_state( - &self, - update: &ProviderCatalogKeyHealthStateUpdate, - ) -> Result { - Self::compare_and_update_key_health_state(self, update).await - } -} - -fn build_list_query<'a>( - prefix: &'static str, - ids: &'a [String], - suffix: &'static str, -) -> QueryBuilder<'a, Sqlite> { - let mut builder = QueryBuilder::::new(select_prefix_for_in(prefix)); - let mut where_clause = WhereClause::new(); - push_in( - &mut builder, - &mut where_clause, - in_column_for_prefix(prefix), - ids, - ); - builder.push(suffix); - builder -} - -fn select_prefix_for_in(prefix: &'static str) -> &'static str { - prefix - .rsplit_once("\nWHERE ") - .map(|(select_prefix, _)| select_prefix) - .expect("provider catalog IN query prefix must contain WHERE") -} - -fn in_column_for_prefix(prefix: &'static str) -> &'static str { - prefix - .rsplit_once("\nWHERE ") - .and_then(|(_, predicate)| predicate.trim().strip_suffix("IN (")) - .map(str::trim) - .expect("provider catalog IN query prefix must end with IN (") -} - -fn apply_key_page_filters<'a>( - builder: &mut QueryBuilder<'a, Sqlite>, - 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::Sqlite, - &["name", "id"], - search, - ); - } - push_optional_eq(builder, where_clause, "is_active", query.is_active); -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn validate_non_empty(value: &str, field_name: &str) -> Result<(), DataLayerError> { - if value.trim().is_empty() { - return Err(DataLayerError::InvalidInput(format!( - "{field_name} is empty" - ))); - } - Ok(()) -} - -fn adaptive_status_snapshot_patch( - patch: &serde_json::Value, -) -> Result { - const OWNED_FIELDS: [&str; 6] = [ - "observation_count", - "header_observation_count", - "latest_upstream_limit", - "learning_confidence", - "enforcement_active", - "known_boundary", - ]; - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog adaptive status snapshot patch must be an object".to_string(), - ) - })?; - Ok(serde_json::Value::Object( - OWNED_FIELDS - .into_iter() - .filter_map(|field| { - object - .get(field) - .cloned() - .map(|value| (field.to_string(), value)) - }) - .collect(), - )) -} - -fn validate_runtime_metadata_update( - update: &ProviderCatalogKeyRuntimeMetadataUpdate, -) -> Result<(), DataLayerError> { - validate_non_empty(&update.key_id, "provider catalog key_id")?; - validate_non_empty( - &update.namespace, - "provider catalog runtime metadata namespace", - )?; - if !update.status_snapshot_patch.is_object() { - return Err(DataLayerError::InvalidInput( - "provider catalog runtime status snapshot patch must be an object".to_string(), - )); - } - Ok(()) -} - -fn push_status_snapshot_shallow_patch<'args>( - builder: &mut QueryBuilder<'args, Sqlite>, - patch: &serde_json::Value, -) -> Result<(), DataLayerError> { - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog status snapshot patch must be an object".to_string(), - ) - })?; - if object.is_empty() { - builder.push("COALESCE(NULLIF(status_snapshot, ''), '{}')"); - return Ok(()); - } - - builder.push("json_set(COALESCE(NULLIF(status_snapshot, ''), '{}')"); - for (field, value) in object { - let path = format!( - "$.{}", - serde_json::to_string(field).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.status_snapshot field is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.status_snapshot value is not serializable: {err}" - )) - })?; - builder.push(", ").push_bind(path).push(", json("); - builder.push_bind(value).push(")"); - } - builder.push(")"); - Ok(()) -} - -fn push_upstream_metadata_shallow_patch<'args>( - builder: &mut QueryBuilder<'args, Sqlite>, - patch: &serde_json::Value, -) -> Result<(), DataLayerError> { - let object = patch.as_object().ok_or_else(|| { - DataLayerError::InvalidInput( - "provider catalog upstream metadata patch must be an object".to_string(), - ) - })?; - if object.is_empty() { - builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - return Ok(()); - } - - builder.push("json_set(COALESCE(NULLIF(upstream_metadata, ''), '{}')"); - for (field, value) in object { - let path = format!( - "$.{}", - serde_json::to_string(field).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata field is not serializable: {err}" - )) - })? - ); - let value = serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "provider_api_keys.upstream_metadata value is not serializable: {err}" - )) - })?; - builder.push(", ").push_bind(path).push(", json("); - builder.push_bind(value).push(")"); - } - builder.push(")"); - Ok(()) -} - -fn validate_provider(provider: &StoredProviderCatalogProvider) -> Result<(), DataLayerError> { - validate_non_empty(&provider.id, "provider catalog provider.id")?; - validate_non_empty(&provider.name, "provider catalog provider.name")?; - validate_non_empty( - &provider.provider_type, - "provider catalog provider.provider_type", - )?; - if provider - .billing_type - .as_deref() - .map(str::trim) - .is_some_and(str::is_empty) - { - return Err(DataLayerError::InvalidInput( - "provider catalog provider.billing_type is empty".to_string(), - )); - } - Ok(()) -} - -fn validate_endpoint(endpoint: &StoredProviderCatalogEndpoint) -> Result<(), DataLayerError> { - validate_non_empty(&endpoint.id, "provider catalog endpoint.id")?; - validate_non_empty( - &endpoint.provider_id, - "provider catalog endpoint.provider_id", - )?; - validate_non_empty(&endpoint.api_format, "provider catalog endpoint.api_format")?; - validate_non_empty(&endpoint.base_url, "provider catalog endpoint.base_url")?; - Ok(()) -} - -fn validate_key(key: &StoredProviderCatalogKey) -> Result<(), DataLayerError> { - validate_non_empty(&key.id, "provider catalog key.id")?; - validate_non_empty(&key.provider_id, "provider catalog key.provider_id")?; - validate_non_empty(&key.name, "provider catalog key.name")?; - validate_non_empty(&key.auth_type, "provider catalog key.auth_type")?; - Ok(()) -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}")) - }) - }) - .transpose() -} - -fn optional_i64_from_u32(value: Option) -> Option { - value.map(i64::from) -} - -fn optional_json_ref_to_string( - value: Option<&serde_json::Value>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - optional_json_ref_to_string(value.as_ref(), field_name) -} - -async fn compare_and_swap_proxy_json( - pool: &SqlitePool, - select_sql: &'static str, - update_sql: &'static str, - update: &ProviderCatalogProxyCasUpdate, - field_name: &'static str, -) -> Result { - // SQLite stores these JSON fields as TEXT. Legacy Python rows commonly include - // insignificant whitespace, so comparing a serde_json re-serialization directly would - // make lazy credential migration conflict forever. Compare semantic JSON first and use - // the exact observed bytes as the atomic write fence. - // Outer None means the row does not exist; inner None is an existing SQL NULL proxy. - let observed_raw: Option> = sqlx::query_scalar::<_, Option>(select_sql) - .bind(&update.record_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - let Some(observed_raw) = observed_raw else { - return Ok(false); - }; - let observed = optional_json_from_string(observed_raw.clone(), field_name)?; - if observed != update.expected_proxy { - return Ok(false); - } - - let replacement = optional_json_to_string(&update.proxy, field_name)?; - let rows_affected = sqlx::query(update_sql) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(&update.record_id) - .bind(observed_raw) - .execute(pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(rows_affected == 1) -} - -fn key_insert_sql() -> &'static str { - r#" -INSERT INTO provider_api_keys ( - id, provider_id, name, api_key, auth_type, capabilities, is_active, - api_formats, auth_type_by_format, allow_auth_channel_mismatch_formats, - auth_config, note, internal_priority, rate_multipliers, - global_priority_by_format, allowed_models, expires_at, cache_ttl_minutes, - max_probe_interval_minutes, proxy, fingerprint, rpm_limit, concurrent_limit, - learned_rpm_limit, concurrent_429_count, rpm_429_count, last_429_at, - last_429_type, adjustment_history, utilization_samples, - last_probe_increase_at, last_rpm_peak, request_count, total_tokens, - total_cost_usd, success_count, error_count, total_response_time_ms, - last_used_at, auto_fetch_models, last_models_fetch_at, - last_models_fetch_error, locked_models, model_include_patterns, - model_exclude_patterns, upstream_metadata, oauth_invalid_at, - oauth_invalid_reason, status_snapshot, health_by_format, - circuit_breaker_by_format, created_at, updated_at -) -VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ? -) -"# -} - -fn key_update_sql() -> &'static str { - r#" -UPDATE provider_api_keys -SET - provider_id = ?, - name = ?, - api_key = ?, - auth_type = ?, - capabilities = ?, - is_active = ?, - api_formats = ?, - auth_type_by_format = ?, - allow_auth_channel_mismatch_formats = ?, - auth_config = ?, - note = ?, - internal_priority = ?, - rate_multipliers = ?, - global_priority_by_format = ?, - allowed_models = ?, - expires_at = ?, - cache_ttl_minutes = ?, - max_probe_interval_minutes = ?, - proxy = ?, - fingerprint = ?, - rpm_limit = ?, - concurrent_limit = ?, - auto_fetch_models = ?, - locked_models = ?, - model_include_patterns = ?, - model_exclude_patterns = ?, - updated_at = ? -WHERE id = ? - AND provider_id = ? - AND auth_type = ? - AND api_key IS ? - AND auth_config IS ? -"# -} - -fn key_update_query( - key: &StoredProviderCatalogKey, - updated_at: i64, -) -> Result>, DataLayerError> { - Ok(sqlx::query(key_update_sql()) - .bind(&key.provider_id) - .bind(&key.name) - .bind(&key.encrypted_api_key) - .bind(&key.auth_type) - .bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .bind(key.is_active) - .bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .bind(&key.encrypted_auth_config) - .bind(&key.note) - .bind(key.internal_priority) - .bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .bind(key.cache_ttl_minutes) - .bind(key.max_probe_interval_minutes) - .bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .bind(optional_i64_from_u32(key.rpm_limit)) - .bind(key.concurrent_limit) - .bind(key.auto_fetch_models) - .bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .bind(updated_at) - .bind(&key.id) - .bind(&key.provider_id) - .bind(&key.auth_type) - .bind(&key.encrypted_api_key) - .bind(&key.encrypted_auth_config)) -} - -fn push_admin_key_assignments<'args>( - builder: &mut QueryBuilder<'args, Sqlite>, - key: &'args StoredProviderCatalogKey, - updated_at: i64, -) -> Result<(), DataLayerError> { - builder - .push_bind(&key.provider_id) - .push(", name = ") - .push_bind(&key.name) - .push(", api_key = ") - .push_bind(&key.encrypted_api_key) - .push(", auth_type = ") - .push_bind(&key.auth_type) - .push(", capabilities = ") - .push_bind(optional_json_to_string( - &key.capabilities, - "provider_api_keys.capabilities", - )?) - .push(", is_active = ") - .push_bind(key.is_active) - .push(", api_formats = ") - .push_bind(optional_json_to_string( - &key.api_formats, - "provider_api_keys.api_formats", - )?) - .push(", auth_type_by_format = ") - .push_bind(optional_json_to_string( - &key.auth_type_by_format, - "provider_api_keys.auth_type_by_format", - )?) - .push(", allow_auth_channel_mismatch_formats = ") - .push_bind(optional_json_to_string( - &key.allow_auth_channel_mismatch_formats, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?) - .push(", auth_config = ") - .push_bind(&key.encrypted_auth_config) - .push(", note = ") - .push_bind(&key.note) - .push(", internal_priority = ") - .push_bind(key.internal_priority) - .push(", rate_multipliers = ") - .push_bind(optional_json_to_string( - &key.rate_multipliers, - "provider_api_keys.rate_multipliers", - )?) - .push(", global_priority_by_format = ") - .push_bind(optional_json_to_string( - &key.global_priority_by_format, - "provider_api_keys.global_priority_by_format", - )?) - .push(", allowed_models = ") - .push_bind(optional_json_to_string( - &key.allowed_models, - "provider_api_keys.allowed_models", - )?) - .push(", expires_at = ") - .push_bind(optional_i64_from_u64( - key.expires_at_unix_secs, - "provider_api_keys.expires_at", - )?) - .push(", cache_ttl_minutes = ") - .push_bind(key.cache_ttl_minutes) - .push(", max_probe_interval_minutes = ") - .push_bind(key.max_probe_interval_minutes) - .push(", proxy = ") - .push_bind(optional_json_to_string( - &key.proxy, - "provider_api_keys.proxy", - )?) - .push(", fingerprint = ") - .push_bind(optional_json_to_string( - &key.fingerprint, - "provider_api_keys.fingerprint", - )?) - .push(", rpm_limit = ") - .push_bind(optional_i64_from_u32(key.rpm_limit)) - .push(", concurrent_limit = ") - .push_bind(key.concurrent_limit) - .push(", auto_fetch_models = ") - .push_bind(key.auto_fetch_models) - .push(", locked_models = ") - .push_bind(optional_json_to_string( - &key.locked_models, - "provider_api_keys.locked_models", - )?) - .push(", model_include_patterns = ") - .push_bind(optional_json_to_string( - &key.model_include_patterns, - "provider_api_keys.model_include_patterns", - )?) - .push(", model_exclude_patterns = ") - .push_bind(optional_json_to_string( - &key.model_exclude_patterns, - "provider_api_keys.model_exclude_patterns", - )?) - .push(", updated_at = ") - .push_bind(updated_at); - Ok(()) -} - -fn validate_admin_key_cas_update( - update: &ProviderCatalogKeyAdminCasUpdate, -) -> Result<(), DataLayerError> { - validate_key(&update.key)?; - let expected = &update.expected_credential; - if expected.auth_type.trim().is_empty() - || expected.provider_id.trim().is_empty() - || expected.provider_type.trim().is_empty() - || expected - .encrypted_api_key - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - || update - .expected_encrypted_auth_config - .as_deref() - .is_some_and(|value| value.trim().is_empty()) - { - return Err(DataLayerError::InvalidInput( - "provider catalog admin credential fence contains empty fields".to_string(), - )); - } - let Some(rotation) = update.codex_rotation.as_ref() else { - return Ok(()); - }; - let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex") - && rotation.as_object().is_some_and(|object| { - object.len() == 1 - && object - .get("credential_generation") - .and_then(serde_json::Value::as_str) - .is_some_and(|generation| !generation.trim().is_empty()) - }); - if !valid_rotation { - return Err(DataLayerError::InvalidInput( - "provider catalog Codex rotation must contain only credential_generation".to_string(), - )); - } - Ok(()) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_u64(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn optional_u32(value: Option, field_name: &str) -> Result, DataLayerError> { - value - .map(|value| { - u32::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")) - }) - }) - .transpose() -} - -fn map_provider_row(row: &SqliteRow) -> Result { - Ok(StoredProviderCatalogProvider::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("website").map_sql_err()?, - row.try_get("provider_type").map_sql_err()?, - )? - .with_description(row.try_get("description").map_sql_err()?) - .with_billing_fields( - row.try_get("billing_type").map_sql_err()?, - sqlite_optional_real(row, "monthly_quota_usd")?, - sqlite_optional_real(row, "monthly_used_usd")?, - optional_u64( - row.try_get("quota_reset_day").map_sql_err()?, - "providers.quota_reset_day", - )?, - optional_u64( - row.try_get("quota_last_reset_at_unix_secs").map_sql_err()?, - "providers.quota_last_reset_at", - )?, - optional_u64( - row.try_get("quota_expires_at_unix_secs").map_sql_err()?, - "providers.quota_expires_at", - )?, - ) - .with_routing_fields(row.try_get("provider_priority").map_sql_err()?) - .with_transport_fields( - row.try_get("is_active").map_sql_err()?, - row.try_get("keep_priority_on_conversion").map_sql_err()?, - row.try_get("enable_format_conversion").map_sql_err()?, - row.try_get("concurrent_limit").map_sql_err()?, - row.try_get("max_retries").map_sql_err()?, - optional_json_from_string(row.try_get("proxy").map_sql_err()?, "providers.proxy")?, - row.try_get("request_timeout").map_sql_err()?, - row.try_get("stream_first_byte_timeout").map_sql_err()?, - optional_json_from_string(row.try_get("config").map_sql_err()?, "providers.config")?, - ) - .with_timestamps( - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "providers.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "providers.updated_at", - )?, - )) -} - -fn map_endpoint_row(row: &SqliteRow) -> Result { - StoredProviderCatalogEndpoint::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("api_format").map_sql_err()?, - row.try_get("api_family").map_sql_err()?, - row.try_get("endpoint_kind").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - )? - .with_timestamps( - optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "provider_endpoints.created_at", - )?, - optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "provider_endpoints.updated_at", - )?, - ) - .with_health_score(sqlite_optional_real(row, "health_score")?.unwrap_or(1.0)) - .with_transport_fields( - row.try_get("base_url").map_sql_err()?, - optional_json_from_string( - row.try_get("header_rules").map_sql_err()?, - "provider_endpoints.header_rules", - )?, - optional_json_from_string( - row.try_get("body_rules").map_sql_err()?, - "provider_endpoints.body_rules", - )?, - row.try_get("max_retries").map_sql_err()?, - row.try_get("custom_path").map_sql_err()?, - optional_json_from_string( - row.try_get("config").map_sql_err()?, - "provider_endpoints.config", - )?, - optional_json_from_string( - row.try_get("format_acceptance_config").map_sql_err()?, - "provider_endpoints.format_acceptance_config", - )?, - optional_json_from_string( - row.try_get("proxy").map_sql_err()?, - "provider_endpoints.proxy", - )?, - ) -} - -fn map_key_stats_row(row: &SqliteRow) -> 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: &SqliteRow, -) -> 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: &SqliteRow) -> Result { - let total_cost_usd = sqlite_optional_real(row, "total_cost_usd")?.unwrap_or(0.0); - if !total_cost_usd.is_finite() { - return Err(DataLayerError::UnexpectedValue( - "invalid provider_api_keys.total_cost_usd".to_string(), - )); - } - - StoredProviderCatalogKey::new( - row.try_get("id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("auth_type").map_sql_err()?, - optional_json_from_string( - row.try_get("capabilities").map_sql_err()?, - "provider_api_keys.capabilities", - )?, - row.try_get("is_active").map_sql_err()?, - )? - .with_transport_fields( - optional_json_from_string( - row.try_get("api_formats").map_sql_err()?, - "provider_api_keys.api_formats", - )?, - row.try_get::, _>("api_key").map_sql_err()?, - row.try_get("auth_config").map_sql_err()?, - optional_json_from_string( - row.try_get("rate_multipliers").map_sql_err()?, - "provider_api_keys.rate_multipliers", - )?, - optional_json_from_string( - row.try_get("global_priority_by_format").map_sql_err()?, - "provider_api_keys.global_priority_by_format", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "provider_api_keys.allowed_models", - )?, - optional_u64( - row.try_get("expires_at_unix_secs").map_sql_err()?, - "provider_api_keys.expires_at", - )?, - optional_json_from_string( - row.try_get("proxy").map_sql_err()?, - "provider_api_keys.proxy", - )?, - optional_json_from_string( - row.try_get("fingerprint").map_sql_err()?, - "provider_api_keys.fingerprint", - )?, - ) - .map(|key| { - let mut key = key - .with_rate_limit_fields( - optional_u32( - row.try_get("rpm_limit").map_sql_err()?, - "provider_api_keys.rpm_limit", - )?, - row.try_get("concurrent_limit").map_sql_err()?, - optional_u32( - row.try_get("learned_rpm_limit").map_sql_err()?, - "provider_api_keys.learned_rpm_limit", - )?, - optional_u32( - row.try_get("concurrent_429_count").map_sql_err()?, - "provider_api_keys.concurrent_429_count", - )?, - optional_u32( - row.try_get("rpm_429_count").map_sql_err()?, - "provider_api_keys.rpm_429_count", - )?, - optional_u64( - row.try_get("last_429_at_unix_secs").map_sql_err()?, - "provider_api_keys.last_429_at", - )?, - optional_json_from_string( - row.try_get("adjustment_history").map_sql_err()?, - "provider_api_keys.adjustment_history", - )?, - optional_u32( - row.try_get("request_count").map_sql_err()?, - "provider_api_keys.request_count", - )?, - optional_u32( - row.try_get("success_count").map_sql_err()?, - "provider_api_keys.success_count", - )?, - ) - .with_usage_fields( - optional_u32( - row.try_get("error_count").map_sql_err()?, - "provider_api_keys.error_count", - )?, - optional_u64( - row.try_get("total_response_time_ms").map_sql_err()?, - "provider_api_keys.total_response_time_ms", - )?, - ) - .with_usage_totals( - optional_u64( - row.try_get("total_tokens").map_sql_err()?, - "provider_api_keys.total_tokens", - )? - .unwrap_or(0), - total_cost_usd, - ) - .with_health_fields( - optional_json_from_string( - row.try_get("health_by_format").map_sql_err()?, - "provider_api_keys.health_by_format", - )?, - optional_json_from_string( - row.try_get("circuit_breaker_by_format").map_sql_err()?, - "provider_api_keys.circuit_breaker_by_format", - )?, - ); - key.note = row.try_get("note").map_sql_err()?; - let auth_type_by_format = optional_json_from_string( - row.try_get("auth_type_by_format").map_sql_err()?, - "provider_api_keys.auth_type_by_format", - )?; - let allow_auth_channel_mismatch_formats = optional_json_from_string( - row.try_get("allow_auth_channel_mismatch_formats") - .map_sql_err()?, - "provider_api_keys.allow_auth_channel_mismatch_formats", - )?; - key.internal_priority = row.try_get("internal_priority").map_sql_err()?; - key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").map_sql_err()?; - key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").map_sql_err()?; - key.last_429_type = row.try_get("last_429_type").map_sql_err()?; - key.utilization_samples = optional_json_from_string( - row.try_get("utilization_samples").map_sql_err()?, - "provider_api_keys.utilization_samples", - )?; - key.last_probe_increase_at_unix_secs = optional_u64( - row.try_get("last_probe_increase_at_unix_secs") - .map_sql_err()?, - "provider_api_keys.last_probe_increase_at", - )?; - key.last_rpm_peak = optional_u32( - row.try_get("last_rpm_peak").map_sql_err()?, - "provider_api_keys.last_rpm_peak", - )?; - key.last_used_at_unix_secs = optional_u64( - row.try_get("last_used_at_unix_secs").map_sql_err()?, - "provider_api_keys.last_used_at", - )?; - key.auto_fetch_models = row.try_get("auto_fetch_models").map_sql_err()?; - key.last_models_fetch_at_unix_secs = optional_u64( - row.try_get("last_models_fetch_at_unix_secs") - .map_sql_err()?, - "provider_api_keys.last_models_fetch_at", - )?; - key.last_models_fetch_error = row.try_get("last_models_fetch_error").map_sql_err()?; - key.locked_models = optional_json_from_string( - row.try_get("locked_models").map_sql_err()?, - "provider_api_keys.locked_models", - )?; - key.model_include_patterns = optional_json_from_string( - row.try_get("model_include_patterns").map_sql_err()?, - "provider_api_keys.model_include_patterns", - )?; - key.model_exclude_patterns = optional_json_from_string( - row.try_get("model_exclude_patterns").map_sql_err()?, - "provider_api_keys.model_exclude_patterns", - )?; - key.upstream_metadata = optional_json_from_string( - row.try_get("upstream_metadata").map_sql_err()?, - "provider_api_keys.upstream_metadata", - )?; - key.oauth_invalid_at_unix_secs = optional_u64( - row.try_get("oauth_invalid_at_unix_secs").map_sql_err()?, - "provider_api_keys.oauth_invalid_at", - )?; - key.oauth_invalid_reason = row.try_get("oauth_invalid_reason").map_sql_err()?; - key.status_snapshot = optional_json_from_string( - row.try_get("status_snapshot").map_sql_err()?, - "provider_api_keys.status_snapshot", - )?; - key.created_at_unix_ms = optional_u64( - row.try_get("created_at_unix_ms").map_sql_err()?, - "provider_api_keys.created_at", - )?; - key.updated_at_unix_secs = optional_u64( - row.try_get("updated_at_unix_secs").map_sql_err()?, - "provider_api_keys.updated_at", - )?; - key.with_auth_channel_policy_fields( - auth_type_by_format, - allow_auth_channel_mismatch_formats, - ) - })? -} - -#[cfg(test)] -mod tests { - use super::SqliteProviderCatalogReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, - ProviderCatalogProxyCasUpdate, ProviderCatalogUpstreamMetadataNamespaceExpectation, - ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, - StoredProviderCatalogKey, StoredProviderCatalogProvider, - }; - use serde_json::json; - - #[test] - fn credential_cas_migrates_legacy_encrypted_key_with_null_safe_fence() { - let source = include_str!("provider_catalog.rs"); - assert!(source.contains( - "UPDATE provider_api_keys SET api_key = ?, encrypted_key = NULL, auth_config = ? WHERE id = ? AND provider_id = ? AND COALESCE(api_key, encrypted_key) IS ? AND auth_config IS ?" - )); - } - - #[tokio::test] - async fn sqlite_proxy_cas_accepts_python_spaced_json_and_distinguishes_null_from_missing_rows() - { - 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 = SqliteProviderCatalogReadRepository::new(pool.clone()); - - let expected_proxy = json!({ - "url": "http://proxy.example.test:8080/", - "username": "alice" - }); - let replacement_proxy = json!({ - "url": "http://proxy.example.test:8080/", - "username": "aether-runtime-secret:v1:sealed" - }); - let mut spaced_provider = StoredProviderCatalogProvider::new( - "proxy-cas-spaced-provider".to_string(), - "Proxy CAS Spaced Provider".to_string(), - None, - "custom".to_string(), - ) - .expect("provider should build"); - spaced_provider.proxy = Some(expected_proxy.clone()); - repository - .create_provider(&spaced_provider, None) - .await - .expect("provider should create"); - - let legacy_python_json = - r#"{"url": "http://proxy.example.test:8080/", "username": "alice"}"#; - sqlx::query("UPDATE providers SET proxy = ? WHERE id = ?") - .bind(legacy_python_json) - .bind(&spaced_provider.id) - .execute(&pool) - .await - .expect("legacy proxy JSON should be seeded"); - - let update = ProviderCatalogProxyCasUpdate { - record_id: spaced_provider.id.clone(), - expected_proxy: Some(expected_proxy), - proxy: Some(replacement_proxy.clone()), - }; - assert!(repository - .compare_and_swap_provider_proxy(&update) - .await - .expect("semantic proxy CAS should run")); - assert!(!repository - .compare_and_swap_provider_proxy(&update) - .await - .expect("stale semantic proxy CAS should run")); - let stored_raw: Option = - sqlx::query_scalar("SELECT proxy FROM providers WHERE id = ?") - .bind(&spaced_provider.id) - .fetch_one(&pool) - .await - .expect("updated proxy should load"); - assert_eq!( - stored_raw - .as_deref() - .map(serde_json::from_str::) - .transpose() - .expect("stored proxy should remain valid JSON"), - Some(replacement_proxy.clone()) - ); - assert_ne!(stored_raw.as_deref(), Some(legacy_python_json)); - - let endpoint = StoredProviderCatalogEndpoint::new( - "proxy-cas-spaced-endpoint".to_string(), - spaced_provider.id.clone(), - "openai:chat".to_string(), - Some("openai".to_string()), - Some("chat".to_string()), - true, - ) - .expect("endpoint should build") - .with_transport_fields( - "https://api.example.test/v1".to_string(), - None, - None, - None, - None, - None, - None, - Some(json!({ - "url": "http://proxy.example.test:8080/", - "username": "alice" - })), - ) - .expect("endpoint transport should build"); - repository - .create_endpoint(&endpoint) - .await - .expect("endpoint should create"); - sqlx::query("UPDATE provider_endpoints SET proxy = ? WHERE id = ?") - .bind(legacy_python_json) - .bind(&endpoint.id) - .execute(&pool) - .await - .expect("legacy endpoint proxy JSON should be seeded"); - let endpoint_update = ProviderCatalogProxyCasUpdate { - record_id: endpoint.id.clone(), - expected_proxy: Some(json!({ - "url": "http://proxy.example.test:8080/", - "username": "alice" - })), - proxy: Some(replacement_proxy.clone()), - }; - assert!(repository - .compare_and_swap_endpoint_proxy(&endpoint_update) - .await - .expect("semantic endpoint proxy CAS should run")); - assert!(!repository - .compare_and_swap_endpoint_proxy(&endpoint_update) - .await - .expect("stale endpoint proxy CAS should run")); - - let key = StoredProviderCatalogKey::new( - "proxy-cas-spaced-key".to_string(), - spaced_provider.id.clone(), - "Proxy CAS Key".to_string(), - "api_key".to_string(), - None, - true, - ) - .expect("key should build") - .with_transport_fields( - None, - None::, - None, - None, - None, - None, - None, - Some(json!({ - "url": "http://proxy.example.test:8080/", - "username": "alice" - })), - None, - ) - .expect("key transport should build"); - repository - .create_key(&key) - .await - .expect("key should create"); - sqlx::query("UPDATE provider_api_keys SET proxy = ? WHERE id = ?") - .bind(legacy_python_json) - .bind(&key.id) - .execute(&pool) - .await - .expect("legacy key proxy JSON should be seeded"); - let key_update = ProviderCatalogProxyCasUpdate { - record_id: key.id.clone(), - expected_proxy: Some(json!({ - "url": "http://proxy.example.test:8080/", - "username": "alice" - })), - proxy: Some(replacement_proxy), - }; - assert!(repository - .compare_and_swap_key_proxy(&key_update) - .await - .expect("semantic key proxy CAS should run")); - assert!(!repository - .compare_and_swap_key_proxy(&key_update) - .await - .expect("stale key proxy CAS should run")); - - let null_provider = StoredProviderCatalogProvider::new( - "proxy-cas-null-provider".to_string(), - "Proxy CAS Null Provider".to_string(), - None, - "custom".to_string(), - ) - .expect("provider should build"); - repository - .create_provider(&null_provider, None) - .await - .expect("null-proxy provider should create"); - assert!(repository - .compare_and_swap_provider_proxy(&ProviderCatalogProxyCasUpdate { - record_id: null_provider.id, - expected_proxy: None, - proxy: Some(json!({"url": "http://proxy.example.test:8081/"})), - }) - .await - .expect("existing NULL proxy should be distinguishable from a missing row")); - assert!(!repository - .compare_and_swap_provider_proxy(&ProviderCatalogProxyCasUpdate { - record_id: "proxy-cas-missing-provider".to_string(), - expected_proxy: None, - proxy: Some(json!({"url": "http://proxy.example.test:8082/"})), - }) - .await - .expect("missing-row proxy CAS should run")); - } - - #[tokio::test] - async fn sqlite_admin_credential_cas_rotates_codex_namespace_atomically() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "admin-cas-provider".to_string(), - "Admin CAS Provider".to_string(), - None, - "codex".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - - let mut original = StoredProviderCatalogKey::new( - "admin-cas-key".to_string(), - "admin-cas-provider".to_string(), - "Before".to_string(), - "oauth".to_string(), - None, - true, - ) - .expect("key should build"); - original.encrypted_api_key = Some("api-v1".to_string()); - original.encrypted_auth_config = Some("auth-v1".to_string()); - original.learned_rpm_limit = Some(17); - original.error_count = Some(9); - original.oauth_invalid_at_unix_secs = Some(123); - original.oauth_invalid_reason = Some("old credential invalid".to_string()); - original.upstream_metadata = Some(json!({ - "codex": {"credential_generation":"generation-v1","used_percent":90}, - "other": {"preserved":true} - })); - original.status_snapshot = Some(json!({ - "quota":{"used_ratio":0.9}, - "oauth":{"status":"valid"} - })); - repository - .create_key(&original) - .await - .expect("key should create"); - - let mut replacement = original.clone(); - replacement.name = "After".to_string(); - replacement.auth_type = "api_key".to_string(); - replacement.encrypted_api_key = Some("api-v2".to_string()); - replacement.encrypted_auth_config = None; - replacement.updated_at_unix_secs = Some(200); - let update = ProviderCatalogKeyAdminCasUpdate { - expected_encrypted_auth_config: Some("auth-v1".to_string()), - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: Some("api-v1".to_string()), - auth_type: "oauth".to_string(), - provider_id: "admin-cas-provider".to_string(), - provider_type: "codex".to_string(), - }, - key: replacement, - codex_rotation: Some(json!({"credential_generation":"generation-v2"})), - reset_oauth_runtime: true, - }; - assert!(repository - .compare_and_update_key_admin_state(&update) - .await - .expect("credential rotation should run")); - assert!(!repository - .compare_and_update_key_admin_state(&update) - .await - .expect("stale credential fence should be a CAS miss")); - - let stored = repository - .list_keys_by_ids(&["admin-cas-key".to_string()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!(stored.name, "After"); - assert_eq!(stored.auth_type, "api_key"); - assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v2")); - assert_eq!(stored.encrypted_auth_config, None); - assert_eq!(stored.learned_rpm_limit, Some(17)); - assert_eq!(stored.error_count, Some(0)); - assert_eq!(stored.oauth_invalid_at_unix_secs, None); - assert_eq!(stored.oauth_invalid_reason, None); - assert_eq!( - stored - .upstream_metadata - .as_ref() - .and_then(|metadata| metadata.get("codex")), - Some(&json!({"credential_generation":"generation-v2"})) - ); - assert_eq!( - stored - .upstream_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/other/preserved")), - Some(&json!(true)) - ); - assert_eq!( - stored - .status_snapshot - .as_ref() - .and_then(|status| status.get("quota")), - Some(&serde_json::Value::Null) - ); - assert_eq!( - stored - .status_snapshot - .as_ref() - .and_then(|status| status.get("oauth")), - Some(&serde_json::Value::Null) - ); - } - - #[tokio::test] - async fn sqlite_admin_credential_rotation_rejects_non_object_runtime_roots() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "admin-invalid-provider".to_string(), - "Admin Invalid Provider".to_string(), - None, - "codex".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - let mut original = StoredProviderCatalogKey::new( - "admin-invalid-key".to_string(), - "admin-invalid-provider".to_string(), - "Before".to_string(), - "oauth".to_string(), - None, - true, - ) - .expect("key should build"); - original.encrypted_api_key = Some("api-v1".to_string()); - original.encrypted_auth_config = Some("auth-v1".to_string()); - original.upstream_metadata = Some(json!(null)); - original.status_snapshot = Some(json!({"quota":{"used_ratio":0.9}})); - repository - .create_key(&original) - .await - .expect("key should create"); - let mut replacement = original.clone(); - replacement.encrypted_api_key = Some("api-v2".to_string()); - let update = ProviderCatalogKeyAdminCasUpdate { - expected_encrypted_auth_config: Some("auth-v1".to_string()), - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: Some("api-v1".to_string()), - auth_type: "oauth".to_string(), - provider_id: "admin-invalid-provider".to_string(), - provider_type: "codex".to_string(), - }, - key: replacement, - codex_rotation: Some(json!({"credential_generation":"generation-v2"})), - reset_oauth_runtime: true, - }; - assert!(!repository - .compare_and_update_key_admin_state(&update) - .await - .expect("invalid metadata root should be a CAS miss")); - let stored = repository - .list_keys_by_ids(&["admin-invalid-key".to_string()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v1")); - assert_eq!(stored.upstream_metadata, Some(json!(null))); - } - - #[tokio::test] - async fn sqlite_ordinary_admin_update_cannot_replace_credentials() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "ordinary-cas-provider".to_string(), - "Ordinary CAS Provider".to_string(), - None, - "codex".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - let mut original = StoredProviderCatalogKey::new( - "ordinary-cas-key".to_string(), - "ordinary-cas-provider".to_string(), - "Before".to_string(), - "oauth".to_string(), - None, - true, - ) - .expect("key should build"); - original.encrypted_api_key = Some("api-v1".to_string()); - original.encrypted_auth_config = Some("auth-v1".to_string()); - repository - .create_key(&original) - .await - .expect("key should create"); - - let mut forbidden = original.clone(); - forbidden.name = "Must Not Persist".to_string(); - forbidden.encrypted_api_key = Some("api-v2".to_string()); - assert!(repository.update_key(&forbidden).await.is_err()); - - let stored = repository - .list_keys_by_ids(&["ordinary-cas-key".to_string()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!(stored.name, "Before"); - assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v1")); - assert_eq!(stored.encrypted_auth_config.as_deref(), Some("auth-v1")); - } - - #[tokio::test] - async fn sqlite_repository_reads_provider_catalog_contract_views() { - 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_rows(&pool).await; - - let repository = SqliteProviderCatalogReadRepository::new(pool); - let providers = repository - .list_providers(true) - .await - .expect("providers should list"); - assert_eq!(providers.len(), 1); - assert_eq!(providers[0].provider_priority, 10); - assert_eq!(providers[0].monthly_quota_usd, Some(0.0)); - assert_eq!(providers[0].monthly_used_usd, Some(0.0)); - - let endpoints = repository - .list_endpoints_by_provider_ids(&["provider-1".to_string()]) - .await - .expect("endpoints should list"); - assert_eq!(endpoints.len(), 1); - assert_eq!(endpoints[0].health_score, 0.95); - - let keys = repository - .list_keys_by_provider_ids(&["provider-1".to_string()]) - .await - .expect("keys should list"); - assert_eq!(keys.len(), 1); - assert_eq!(keys[0].total_tokens, 1234); - assert_eq!(keys[0].total_response_time_ms, Some(u32::MAX as u64 + 1)); - assert_eq!(keys[0].concurrent_limit, Some(3)); - - let page = repository - .list_keys_page(&ProviderCatalogKeyListQuery { - provider_id: "provider-1".to_string(), - search: Some("default".to_string()), - is_active: Some(true), - offset: 0, - limit: 10, - order: ProviderCatalogKeyListOrder::Name, - }) - .await - .expect("key page should load"); - assert_eq!(page.total, 1); - - let stats = repository - .list_key_stats_by_provider_ids(&["provider-1".to_string()]) - .await - .expect("stats should load"); - assert_eq!(stats[0].active_keys, 1); - } - - #[tokio::test] - async fn sqlite_runtime_key_mutations_are_field_scoped_and_compare_and_swap() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "runtime-provider".to_string(), - "Runtime Provider".to_string(), - None, - "custom".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - let mut key = StoredProviderCatalogKey::new( - "runtime-key".to_string(), - "runtime-provider".to_string(), - "Runtime Key".to_string(), - "api_key".to_string(), - None, - false, - ) - .expect("key should build") - .with_rate_limit_fields( - None, - None, - Some(10), - None, - Some(1), - Some(100), - Some(json!([{"limit":10}])), - None, - None, - ) - .with_health_fields( - Some(json!({"openai:chat":{"consecutive_failures":1}})), - None, - ); - key.upstream_metadata = Some(json!({ - "codex": {"remaining": 5}, - "grok": {"remaining": 7} - })); - key.status_snapshot = Some(json!({ - "quota": {"remaining": 5, "window": "day"}, - "oauth": {"invalid": false}, - "observation_count": 1, - "known_boundary": "old" - })); - key.encrypted_auth_config = Some("auth-current".to_string()); - let mut stale_admin_key = key.clone(); - repository - .create_key(&key) - .await - .expect("key should create"); - - let health_update = ProviderCatalogKeyHealthStateUpdate { - key_id: "runtime-key".to_string(), - expected_encrypted_auth_config: None, - expected_health_by_format: key.health_by_format.clone(), - expected_circuit_breaker_by_format: None, - health_by_format: Some(json!({"openai:chat":{"consecutive_failures":2}})), - circuit_breaker_by_format: None, - }; - assert!(repository - .compare_and_update_key_health_state(&health_update) - .await - .expect("health CAS should succeed")); - assert!(!repository - .compare_and_update_key_health_state(&health_update) - .await - .expect("stale health CAS should conflict")); - - let adaptive_current = repository - .list_keys_by_ids(&["runtime-key".to_string()]) - .await - .expect("key should reload before adaptive CAS") - .pop() - .expect("key should exist"); - let expected = ProviderCatalogKeyAdaptiveState::from(&adaptive_current); - let mut next = expected.clone(); - next.learned_rpm_limit = Some(8); - next.rpm_429_count = Some(2); - let adaptive_update = ProviderCatalogKeyAdaptiveStateUpdate { - key_id: "runtime-key".to_string(), - expected_encrypted_auth_config: Some("auth-current".to_string()), - expected: expected.clone(), - next, - status_snapshot_patch: json!({ - "observation_count": 2, - "learning_confidence": 0.5, - "known_boundary": null, - "quota": {"remaining": 0} - }), - updated_at_unix_secs: Some(200), - }; - let stale_generation_update = ProviderCatalogKeyAdaptiveStateUpdate { - expected_encrypted_auth_config: Some("auth-stale".to_string()), - ..adaptive_update.clone() - }; - assert!(!repository - .compare_and_update_key_adaptive_state(&stale_generation_update) - .await - .expect("stale auth generation should conflict")); - assert!(repository - .compare_and_update_key_adaptive_state(&adaptive_update) - .await - .expect("adaptive CAS should succeed")); - assert!(!repository - .compare_and_update_key_adaptive_state(&ProviderCatalogKeyAdaptiveStateUpdate { - key_id: "runtime-key".to_string(), - expected_encrypted_auth_config: Some("auth-current".to_string()), - expected: expected.clone(), - next: expected, - status_snapshot_patch: json!({}), - updated_at_unix_secs: Some(201), - }) - .await - .expect("stale adaptive CAS should conflict")); - - assert!(repository - .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { - key_id: "runtime-key".to_string(), - namespace: "codex".to_string(), - expected_upstream_metadata_value: Some(json!({"remaining":5})), - upstream_metadata_value: json!({"remaining":3}), - status_snapshot_patch: json!({"quota":{"remaining":3}}), - updated_at_unix_secs: Some(300), - }) - .await - .expect("runtime metadata should update")); - - stale_admin_key.name = "Admin Renamed".to_string(); - stale_admin_key.is_active = true; - repository - .update_key(&stale_admin_key) - .await - .expect("stale admin update should preserve runtime fields"); - - let stored = repository - .list_keys_by_ids(&["runtime-key".to_string()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!(stored.name, "Admin Renamed"); - assert!(stored.is_active); - assert_eq!(stored.learned_rpm_limit, Some(8)); - assert_eq!(stored.rpm_429_count, Some(2)); - assert_eq!( - stored.health_by_format, - Some(json!({"openai:chat":{"consecutive_failures":2}})) - ); - assert_eq!( - stored.upstream_metadata.as_ref().unwrap()["codex"], - json!({"remaining":3}) - ); - assert_eq!( - stored.upstream_metadata.as_ref().unwrap()["grok"], - json!({"remaining":7}) - ); - let status = stored.status_snapshot.expect("status should exist"); - assert_eq!(status["quota"], json!({"remaining":3})); - assert!(status["quota"].get("window").is_none()); - assert_eq!(status["oauth"], json!({"invalid":false})); - assert_eq!(status["observation_count"], json!(2)); - assert_eq!(status["learning_confidence"], json!(0.5)); - assert!(status - .as_object() - .expect("status should be an object") - .contains_key("known_boundary")); - assert_eq!(status["known_boundary"], serde_json::Value::Null); - } - - #[tokio::test] - async fn sqlite_runtime_metadata_cas_rejects_non_object_metadata_roots() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "invalid-root-provider".to_string(), - "Invalid Root Provider".to_string(), - None, - "codex".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - - for (index, invalid_root) in [json!(null), json!([]), json!("invalid"), json!(1)] - .into_iter() - .enumerate() - { - let key_id = format!("invalid-root-{index}"); - let mut key = StoredProviderCatalogKey::new( - key_id.clone(), - "invalid-root-provider".to_string(), - format!("Invalid Root {index}"), - "oauth".to_string(), - None, - true, - ) - .expect("key should build"); - key.upstream_metadata = Some(invalid_root.clone()); - key.status_snapshot = Some(json!({"quota":{"remaining":9}})); - key.updated_at_unix_secs = Some(10); - repository - .create_key(&key) - .await - .expect("key should create"); - - assert!(!repository - .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { - key_id: key_id.clone(), - namespace: "codex".to_string(), - expected_upstream_metadata_value: None, - upstream_metadata_value: json!({"remaining":1}), - status_snapshot_patch: json!({"quota":{"remaining":1}}), - updated_at_unix_secs: Some(20), - }) - .await - .expect("non-object root should be a CAS miss")); - assert!(!repository - .compare_and_update_key_oauth_runtime_state( - &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: key_id.clone(), - expected_encrypted_auth_config: None, - expected_credential: None, - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: None, - }, - ), - encrypted_auth_config: "next-auth".to_string(), - encrypted_api_key_update: None, - expires_at_unix_secs_update: None, - oauth_invalid_at_unix_secs: None, - oauth_invalid_reason: None, - upstream_metadata_patch: Some(json!({"codex":{"remaining":1}})), - upstream_metadata_namespace_to_remove: None, - status_snapshot_patch: json!({"quota":{"remaining":1}}), - reset_error_count: false, - updated_at_unix_secs: Some(20), - }, - ) - .await - .expect("OAuth metadata update should be a CAS miss")); - assert!(!repository - .compare_and_delete_key_oauth_credential( - &ProviderCatalogKeyOAuthCredentialCasDelete { - key_id: key_id.clone(), - expected_encrypted_auth_config: None, - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: None, - auth_type: "oauth".to_string(), - provider_id: "invalid-root-provider".to_string(), - provider_type: "codex".to_string(), - }, - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: None, - }, - ), - }, - ) - .await - .expect("OAuth credential delete should be a CAS miss")); - - let stored = repository - .list_keys_by_ids(&[key_id]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!(stored.upstream_metadata, Some(invalid_root)); - assert_eq!( - stored.status_snapshot, - Some(json!({"quota":{"remaining":9}})) - ); - assert_eq!(stored.encrypted_auth_config, None); - assert_eq!(stored.updated_at_unix_secs, Some(10)); - } - } - - #[tokio::test] - async fn sqlite_oauth_runtime_cas_fences_auth_config_and_preserves_admin_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"); - let repository = SqliteProviderCatalogReadRepository::new(pool); - repository - .create_provider( - &StoredProviderCatalogProvider::new( - "oauth-cas-provider".to_string(), - "OAuth CAS Provider".to_string(), - None, - "codex".to_string(), - ) - .expect("provider should build"), - None, - ) - .await - .expect("provider should create"); - - let mut key = StoredProviderCatalogKey::new( - "oauth-cas-key".to_string(), - "oauth-cas-provider".to_string(), - "Admin Managed Name".to_string(), - "oauth".to_string(), - None, - false, - ) - .expect("key should build") - .with_transport_fields( - Some(json!(["openai:responses"])), - Some("encrypted-api-key".to_string()), - Some("encrypted-auth-v1".to_string()), - None, - Some(json!({"openai:responses": 17})), - Some(json!(["gpt-5"])), - Some(4_102_444_800), - None, - None, - ) - .expect("key transport should build"); - key.note = Some("admin note".to_string()); - key.internal_priority = 23; - key.status_snapshot = Some(json!({ - "oauth": {"invalid": true, "source": "old-task"}, - "quota": {"remaining": 7}, - "admin": {"label": "keep"} - })); - key.upstream_metadata = Some(json!({ - "codex": {"remaining": 5}, - "admin": {"keep": true} - })); - repository - .create_key(&key) - .await - .expect("key should create"); - - let update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: key.id.clone(), - expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()), - expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: Some("encrypted-api-key".to_string()), - auth_type: "oauth".to_string(), - provider_id: "oauth-cas-provider".to_string(), - provider_type: "codex".to_string(), - }), - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: Some(json!({"remaining": 5})), - }, - ), - encrypted_auth_config: "encrypted-auth-v2".to_string(), - encrypted_api_key_update: Some("encrypted-api-v2".to_string()), - expires_at_unix_secs_update: Some(Some(4_102_555_900)), - oauth_invalid_at_unix_secs: None, - oauth_invalid_reason: None, - upstream_metadata_patch: Some(json!({"codex": {"remaining": 3}})), - upstream_metadata_namespace_to_remove: None, - status_snapshot_patch: json!({ - "oauth": {"invalid": false, "task_id": "task-v2"}, - "runtime": {"generation": 2} - }), - reset_error_count: false, - updated_at_unix_secs: Some(200), - }; - assert!(repository - .compare_and_update_key_oauth_runtime_state(&update) - .await - .expect("matching OAuth runtime CAS should succeed")); - - let stored = repository - .list_keys_by_ids(&[key.id.clone()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!( - stored.encrypted_auth_config.as_deref(), - Some("encrypted-auth-v2") - ); - assert_eq!(stored.oauth_invalid_at_unix_secs, None); - assert_eq!(stored.oauth_invalid_reason, None); - assert_eq!(stored.name, "Admin Managed Name"); - assert!(!stored.is_active); - assert_eq!(stored.note.as_deref(), Some("admin note")); - assert_eq!(stored.internal_priority, 23); - assert_eq!( - stored.global_priority_by_format, - Some(json!({"openai:responses": 17})) - ); - assert_eq!(stored.allowed_models, Some(json!(["gpt-5"]))); - assert_eq!( - stored.encrypted_api_key.as_deref(), - Some("encrypted-api-v2") - ); - assert_eq!(stored.expires_at_unix_secs, Some(4_102_555_900)); - assert_eq!( - stored.upstream_metadata.as_ref().unwrap()["codex"]["remaining"], - 3 - ); - let status = stored.status_snapshot.expect("status should exist"); - assert_eq!( - status["oauth"], - json!({"invalid": false, "task_id": "task-v2"}) - ); - assert!(status["oauth"].get("source").is_none()); - assert_eq!(status["quota"], json!({"remaining": 7})); - assert_eq!(status["admin"], json!({"label": "keep"})); - assert_eq!(status["runtime"], json!({"generation": 2})); - - let stale_metadata_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - expected_encrypted_auth_config: Some("encrypted-auth-v2".to_string()), - expected_credential: None, - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: Some(json!({"remaining": 5})), - }, - ), - encrypted_auth_config: "encrypted-auth-v3".to_string(), - encrypted_api_key_update: None, - expires_at_unix_secs_update: None, - upstream_metadata_patch: Some(json!({"codex": {"remaining": 0}})), - status_snapshot_patch: json!({"quota": {"remaining": 0}}), - updated_at_unix_secs: Some(201), - ..update.clone() - }; - assert!(!repository - .compare_and_update_key_oauth_runtime_state(&stale_metadata_update) - .await - .expect("stale metadata namespace should conflict")); - let stored_after_metadata_conflict = repository - .list_keys_by_ids(&[key.id.clone()]) - .await - .expect("key should reload after metadata conflict") - .pop() - .expect("key should exist"); - assert_eq!( - stored_after_metadata_conflict - .encrypted_auth_config - .as_deref(), - Some("encrypted-auth-v2") - ); - assert_eq!( - stored_after_metadata_conflict - .upstream_metadata - .as_ref() - .unwrap()["codex"], - json!({"remaining": 3}) - ); - assert_eq!( - stored_after_metadata_conflict - .status_snapshot - .as_ref() - .unwrap()["quota"], - json!({"remaining": 7}) - ); - - let mut absent_key = StoredProviderCatalogKey::new( - "oauth-cas-absent-key".to_string(), - "oauth-cas-provider".to_string(), - "Absent Namespace".to_string(), - "oauth".to_string(), - None, - true, - ) - .expect("absent namespace key should build") - .with_transport_fields( - None, - Some("encrypted-api-absent".to_string()), - Some("encrypted-auth-absent".to_string()), - None, - None, - None, - None, - None, - None, - ) - .expect("absent namespace key transport should build"); - absent_key.upstream_metadata = Some(json!({"admin": {"keep": true}})); - repository - .create_key(&absent_key) - .await - .expect("absent namespace key should create"); - let absent_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: absent_key.id.clone(), - expected_encrypted_auth_config: Some("encrypted-auth-absent".to_string()), - expected_credential: None, - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: None, - }, - ), - encrypted_auth_config: "encrypted-auth-absent".to_string(), - encrypted_api_key_update: None, - expires_at_unix_secs_update: None, - oauth_invalid_at_unix_secs: None, - oauth_invalid_reason: None, - upstream_metadata_patch: Some(json!({"codex": {"remaining": 1}})), - upstream_metadata_namespace_to_remove: None, - status_snapshot_patch: json!({}), - reset_error_count: false, - updated_at_unix_secs: Some(202), - }; - assert!(repository - .compare_and_update_key_oauth_runtime_state(&absent_update) - .await - .expect("absent metadata namespace should match")); - let stored_absent = repository - .list_keys_by_ids(&[absent_key.id]) - .await - .expect("absent namespace key should reload") - .pop() - .expect("absent namespace key should exist"); - assert_eq!( - stored_absent.upstream_metadata.as_ref().unwrap()["codex"], - json!({"remaining": 1}) - ); - assert_eq!( - stored_absent.upstream_metadata.as_ref().unwrap()["admin"], - json!({"keep": true}) - ); - - let remove_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: Some(json!({"remaining": 1})), - }, - ), - upstream_metadata_patch: Some(json!({"runtime": {"generation": 2}})), - upstream_metadata_namespace_to_remove: Some("codex".to_string()), - status_snapshot_patch: json!({"quota": null}), - updated_at_unix_secs: Some(203), - ..absent_update - }; - assert!(repository - .compare_and_update_key_oauth_runtime_state(&remove_update) - .await - .expect("matching metadata namespace removal should succeed")); - assert!(!repository - .compare_and_update_key_oauth_runtime_state(&remove_update) - .await - .expect("stale metadata namespace removal should be a CAS miss")); - let stored_after_remove = repository - .list_keys_by_ids(&[stored_absent.id]) - .await - .expect("removed namespace key should reload") - .pop() - .expect("removed namespace key should exist"); - let metadata = stored_after_remove - .upstream_metadata - .as_ref() - .and_then(serde_json::Value::as_object) - .expect("metadata should remain an object"); - assert!(!metadata.contains_key("codex")); - assert_eq!(metadata["admin"], json!({"keep": true})); - assert_eq!(metadata["runtime"], json!({"generation": 2})); - assert_eq!( - stored_after_remove.status_snapshot.as_ref().unwrap()["quota"], - serde_json::Value::Null - ); - - let stale_api_key_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - expected_encrypted_auth_config: Some("encrypted-auth-v2".to_string()), - expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: Some("encrypted-api-key".to_string()), - auth_type: "oauth".to_string(), - provider_id: "oauth-cas-provider".to_string(), - provider_type: "codex".to_string(), - }), - encrypted_auth_config: "encrypted-auth-v3".to_string(), - status_snapshot_patch: json!({}), - updated_at_unix_secs: Some(201), - ..update.clone() - }; - assert!(!repository - .compare_and_update_key_oauth_runtime_state(&stale_api_key_update) - .await - .expect("stale API key generation should conflict")); - - let stale_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()), - encrypted_auth_config: "encrypted-auth-v3".to_string(), - status_snapshot_patch: json!({"quota": {"remaining": 0}}), - updated_at_unix_secs: Some(201), - ..update - }; - assert!(!repository - .compare_and_update_key_oauth_runtime_state(&stale_update) - .await - .expect("stale OAuth runtime CAS should conflict")); - - let stored_after_stale = repository - .list_keys_by_ids(&[key.id]) - .await - .expect("key should reload after stale CAS") - .pop() - .expect("key should exist"); - assert_eq!( - stored_after_stale.encrypted_auth_config.as_deref(), - Some("encrypted-auth-v2") - ); - assert_eq!( - stored_after_stale - .status_snapshot - .as_ref() - .expect("status should remain")["quota"], - json!({"remaining": 7}) - ); - assert_eq!( - stored_after_stale.upstream_metadata.as_ref().unwrap()["codex"]["remaining"], - 3 - ); - - let stale_delete = ProviderCatalogKeyOAuthCredentialCasDelete { - key_id: stored_after_stale.id.clone(), - expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()), - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: Some("encrypted-api-key".to_string()), - auth_type: "oauth".to_string(), - provider_id: "oauth-cas-provider".to_string(), - provider_type: "codex".to_string(), - }, - expected_upstream_metadata_namespace: None, - }; - assert!(!repository - .compare_and_delete_key_oauth_credential(&stale_delete) - .await - .expect("stale credential delete should conflict")); - - let stale_namespace_delete = ProviderCatalogKeyOAuthCredentialCasDelete { - key_id: stored_after_stale.id.clone(), - expected_encrypted_auth_config: stored_after_stale.encrypted_auth_config.clone(), - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: stored_after_stale.encrypted_api_key.clone(), - auth_type: stored_after_stale.auth_type.clone(), - provider_id: stored_after_stale.provider_id.clone(), - provider_type: "codex".to_string(), - }, - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: Some(json!({"remaining": 3})), - }, - ), - }; - let current_codex = json!({"remaining": 2, "request": "newer"}); - assert!(repository - .upsert_key_upstream_metadata_namespace( - &stored_after_stale.id, - "codex", - ¤t_codex, - Some(202), - ) - .await - .expect("newer codex namespace should persist")); - assert!(!repository - .compare_and_delete_key_oauth_credential(&stale_namespace_delete) - .await - .expect("stale namespace delete should conflict")); - let current_delete = ProviderCatalogKeyOAuthCredentialCasDelete { - expected_upstream_metadata_namespace: Some( - ProviderCatalogUpstreamMetadataNamespaceExpectation { - namespace: "codex".to_string(), - expected_value: Some(current_codex), - }, - ), - ..stale_namespace_delete - }; - assert!(repository - .compare_and_delete_key_oauth_credential(¤t_delete) - .await - .expect("current credential generation should delete")); - assert!(repository - .list_keys_by_ids(&[stored_after_stale.id]) - .await - .expect("deleted key lookup should succeed") - .is_empty()); - } - - #[tokio::test] - async fn sqlite_repository_writes_provider_catalog_contract_views() { - 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 = SqliteProviderCatalogReadRepository::new(pool); - let provider = StoredProviderCatalogProvider::new( - "provider-write-1".to_string(), - "Provider Write".to_string(), - Some("https://write.example.com".to_string()), - "custom".to_string(), - ) - .expect("provider should build") - .with_description(Some("write provider".to_string())) - .with_billing_fields( - Some("pay_as_you_go".to_string()), - Some(100.0), - Some(7.5), - Some(1), - Some(1_710_000_000), - None, - ) - .with_routing_fields(20) - .with_transport_fields( - true, - true, - true, - Some(4), - Some(2), - Some(json!({"http":"proxy"})), - Some(30.0), - Some(2.5), - Some(json!({"region":"us"})), - ); - let created_provider = repository - .create_provider(&provider, None) - .await - .expect("provider should create"); - assert_eq!(created_provider.provider_priority, 20); - assert_eq!(created_provider.proxy, Some(json!({"http":"proxy"}))); - - let mut updated_provider = created_provider.clone(); - updated_provider.description = Some("updated provider".to_string()); - updated_provider.provider_priority = 30; - updated_provider.is_active = false; - let updated_provider = repository - .update_provider(&updated_provider) - .await - .expect("provider should update"); - assert_eq!( - updated_provider.description, - Some("updated provider".to_string()) - ); - assert!(!updated_provider.is_active); - - let endpoint = StoredProviderCatalogEndpoint::new( - "endpoint-write-1".to_string(), - "provider-write-1".to_string(), - "openai:chat".to_string(), - Some("openai".to_string()), - Some("chat".to_string()), - true, - ) - .expect("endpoint should build") - .with_health_score(0.88) - .with_transport_fields( - "https://write.example.com/v1".to_string(), - Some(json!({"Authorization":"Bearer"})), - Some(json!({"model":"gpt"})), - Some(3), - Some("/chat/completions".to_string()), - Some(json!({"timeout":30})), - Some(json!({"accept":["openai:chat"]})), - Some(json!({"https":"proxy"})), - ) - .expect("endpoint transport should build"); - let created_endpoint = repository - .create_endpoint(&endpoint) - .await - .expect("endpoint should create"); - assert_eq!(created_endpoint.health_score, 0.88); - - let mut updated_endpoint = created_endpoint.clone(); - updated_endpoint.health_score = 0.5; - updated_endpoint.is_active = false; - let updated_endpoint = repository - .update_endpoint(&updated_endpoint) - .await - .expect("endpoint should update"); - assert_eq!(updated_endpoint.health_score, 0.5); - assert!(!updated_endpoint.is_active); - - let mut key = StoredProviderCatalogKey::new( - "key-write-1".to_string(), - "provider-write-1".to_string(), - "Default Key".to_string(), - "api_key".to_string(), - Some(json!({"cache_1h":true})), - true, - ) - .expect("key should build") - .with_transport_fields( - Some(json!(["openai:chat"])), - Some("enc-key".to_string()), - Some("enc-auth".to_string()), - Some(json!({"openai:chat":1.0})), - Some(json!({"openai:chat":10})), - Some(json!(["gpt-4.1"])), - Some(1_730_000_000), - Some(json!({"http":"proxy"})), - Some(json!({"fp":"abc"})), - ) - .expect("key transport should build") - .with_rate_limit_fields( - Some(120), - Some(3), - Some(110), - Some(1), - Some(2), - Some(1_720_000_000), - Some(json!([{"rpm":110}])), - Some(10), - Some(9), - ) - .with_usage_fields(Some(1), Some(u32::MAX as u64 + 42)) - .with_usage_totals(1234, 1.5) - .with_health_fields( - Some(json!({"openai:chat":{"score":1}})), - Some(json!({"openai:chat":{"open":false}})), - ); - key.last_models_fetch_at_unix_secs = Some(1_730_000_100); - key.last_models_fetch_error = Some("stale models fetch error".to_string()); - let created_key = repository - .create_key(&key) - .await - .expect("key should create"); - assert_eq!(created_key.concurrent_limit, Some(3)); - assert_eq!(created_key.total_tokens, 1234); - assert_eq!( - created_key.total_response_time_ms, - Some(u32::MAX as u64 + 42) - ); - assert_eq!( - created_key.last_models_fetch_error.as_deref(), - Some("stale models fetch error") - ); - let mut second_key = key.clone(); - second_key.id = "key-write-2".to_string(); - second_key.name = "Secondary Key".to_string(); - let created_second_key = repository - .create_key(&second_key) - .await - .expect("second key should create"); - - let mut updated_key = created_key.clone(); - updated_key.name = "Updated Key".to_string(); - updated_key.is_active = false; - updated_key.upstream_metadata = Some(json!({"models":["gpt-4.1"]})); - updated_key.last_models_fetch_at_unix_secs = Some(1_730_000_200); - updated_key.last_models_fetch_error = None; - let updated_key = repository - .update_key(&updated_key) - .await - .expect("key should update"); - assert_eq!(updated_key.name, "Updated Key"); - assert!(!updated_key.is_active); - assert_eq!( - updated_key.last_models_fetch_at_unix_secs, - Some(1_730_000_100) - ); - assert_eq!( - updated_key.last_models_fetch_error.as_deref(), - Some("stale models fetch error") - ); - assert_eq!(updated_key.upstream_metadata, created_key.upstream_metadata); - - let mut batch_first = updated_key.clone(); - batch_first.auto_fetch_models = true; - batch_first.allowed_models = Some(json!(["gpt-4.1", "gpt-4.1-mini"])); - batch_first.locked_models = Some(json!(["gpt-4.1"])); - batch_first.model_include_patterns = Some(json!(["gpt-*"])); - batch_first.model_exclude_patterns = Some(json!(["*-preview"])); - let mut batch_second = created_second_key; - batch_second.auto_fetch_models = true; - batch_second.allowed_models = batch_first.allowed_models.clone(); - batch_second.locked_models = batch_first.locked_models.clone(); - batch_second.model_include_patterns = batch_first.model_include_patterns.clone(); - batch_second.model_exclude_patterns = batch_first.model_exclude_patterns.clone(); - - let batch_updated = repository - .update_keys(&[batch_first, batch_second]) - .await - .expect("keys should update in one transaction"); - assert_eq!(batch_updated.len(), 2); - assert!(batch_updated.iter().all(|key| key.auto_fetch_models)); - assert!(batch_updated - .iter() - .all(|key| key.locked_models == Some(json!(["gpt-4.1"])))); - assert!(batch_updated - .iter() - .all(|key| key.model_include_patterns == Some(json!(["gpt-*"])))); - assert!(batch_updated - .iter() - .all(|key| key.model_exclude_patterns == Some(json!(["*-preview"])))); - - let mut valid_change = batch_updated - .iter() - .find(|key| key.id == "key-write-1") - .expect("first key should be returned") - .clone(); - valid_change.name = "Must Roll Back".to_string(); - let mut missing_change = valid_change.clone(); - missing_change.id = "missing-key".to_string(); - assert!(repository - .update_keys(&[valid_change, missing_change]) - .await - .is_err()); - let rolled_back = repository - .list_keys_by_ids(&["key-write-1".to_string()]) - .await - .expect("first key should reload") - .pop() - .expect("first key should exist"); - assert_eq!(rolled_back.name, "Updated Key"); - - assert!(repository - .update_key_upstream_metadata( - "key-write-1", - Some(&json!({ - "codex": { - "quota_by_model": { - "gpt-5.6-sol": {"remaining_fraction": 0.75} - } - }, - "codex_models": {"cards": {"old": {"slug": "old"}}} - })), - Some(1_740_000_000), - ) - .await - .expect("upstream metadata should update")); - assert!(repository - .update_key_model_fetch_success( - "key-write-1", - Some(&json!(["gpt-5.6-sol"])), - 1_740_000_002, - &[ProviderCatalogUpstreamMetadataNamespaceUpdate { - namespace: "codex_models".to_string(), - value: json!({ - "cards": { - "gpt-5.6-sol": { - "slug": "gpt-5.6-sol", - "use_responses_lite": true - } - } - }), - }], - Some(1_740_000_002), - ) - .await - .expect("model fetch success should update atomically")); - assert!(repository - .update_key_health_state( - "key-write-1", - true, - Some(&json!({"openai:chat":{"score":0.9}})), - None, - ) - .await - .expect("health state should update")); - assert!(repository - .clear_key_oauth_invalid_marker("key-write-1") - .await - .expect("oauth invalid marker should clear")); - - let reloaded_key = repository - .list_keys_by_ids(&["key-write-1".to_string()]) - .await - .expect("key should reload") - .pop() - .expect("key should exist"); - assert_eq!( - reloaded_key.upstream_metadata, - Some(json!({ - "codex": { - "quota_by_model": { - "gpt-5.6-sol": {"remaining_fraction": 0.75} - } - }, - "codex_models": { - "cards": { - "gpt-5.6-sol": { - "slug": "gpt-5.6-sol", - "use_responses_lite": true - } - } - } - })) - ); - assert_eq!(reloaded_key.allowed_models, Some(json!(["gpt-5.6-sol"]))); - assert_eq!( - reloaded_key.last_models_fetch_at_unix_secs, - Some(1_740_000_002) - ); - assert!(reloaded_key.is_active); - - assert!(repository - .delete_key("key-write-1") - .await - .expect("key should delete")); - assert!(repository - .delete_key("key-write-2") - .await - .expect("second key should delete")); - assert!(repository - .delete_endpoint("endpoint-write-1") - .await - .expect("endpoint should delete")); - assert!(repository - .delete_provider("provider-write-1") - .await - .expect("provider should delete")); - } - - async fn seed_rows(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers ( - id, name, description, website, provider_type, provider_priority, - monthly_quota_usd, monthly_used_usd, - is_active, keep_priority_on_conversion, enable_format_conversion, - config, created_at, updated_at -) VALUES ( - 'provider-1', 'Provider One', 'test provider', 'https://example.com', - 'custom', 10, 0, 0, 1, 1, 1, '{"region":"us"}', 1, 2 -) -"#, - ) - .execute(pool) - .await - .expect("provider should seed"); - sqlx::query( - r#" -INSERT INTO provider_endpoints ( - id, provider_id, name, base_url, api_format, api_family, endpoint_kind, - is_active, health_score, header_rules, created_at, updated_at -) VALUES ( - 'endpoint-1', 'provider-1', 'primary', 'https://api.example.com', - 'openai:chat', 'openai', 'chat', 1, 0.95, '{"Authorization":"Bearer"}', 3, 4 -) -"#, - ) - .execute(pool) - .await - .expect("endpoint should seed"); - sqlx::query( - r#" -INSERT INTO provider_api_keys ( - id, provider_id, name, api_key, auth_type, capabilities, is_active, - api_formats, auth_type_by_format, internal_priority, rpm_limit, - concurrent_limit, request_count, total_tokens, total_cost_usd, - success_count, error_count, total_response_time_ms, health_by_format, - created_at, updated_at -) VALUES ( - 'key-1', 'provider-1', 'default', 'enc-key', 'api_key', - '{"cache_1h":true}', 1, '["openai:chat"]', '{"openai:chat":"api_key"}', - 5, 120, 3, 10, 1234, 1.5, 9, 1, 4294967296, '{"openai:chat":{"score":1}}', - 5, 6 -) -"#, - ) - .execute(pool) - .await - .expect("key should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/proxy_nodes.rs b/crates/aether-data/adapters/sqlite/src/proxy_nodes.rs deleted file mode 100644 index ee4d29d59..000000000 --- a/crates/aether-data/adapters/sqlite/src/proxy_nodes.rs +++ /dev/null @@ -1,2931 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::proxy_nodes::{ - bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample, - merge_proxy_metadata_for_registration, normalize_heartbeat_proxy_metadata, - normalize_proxy_metadata, proxy_metadata_has_explicit_tunnel_security, - reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, - ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, - ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation, - ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, - ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent, - StoredProxyNodeMetricsBucket, TunnelErrorEventRecord, TunnelMetricsSample, - PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR, -}; -use aether_data_contracts::DataLayerError; -use aether_data_query::{push_eq, push_limit, WhereClause}; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const PROXY_NODE_REGISTRATION_CAS_RETRIES: usize = 8; - -fn log_reported_tunnel_error_event( - node_id: &str, - event: &TunnelErrorEventRecord, - received_at_unix_secs: u64, -) { - tracing::warn!( - event_name = "proxy_tunnel_error_reported", - source = "heartbeat", - node_id = %node_id, - category = %event.category, - message = %event.message, - severity = ?event.severity, - component = ?event.component, - summary = ?event.summary, - operator_action = ?event.operator_action, - error_reported_at_unix_secs = event.timestamp_unix_secs, - error_reported_at_unix_ms = ?event.timestamp_unix_ms, - report_received_at_unix_secs = received_at_unix_secs, - "proxy reported tunnel error via heartbeat" - ); -} - -#[derive(Debug, Clone)] -pub struct SqliteProxyNodeReadRepository { - pool: SqlitePool, -} - -impl SqliteProxyNodeReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn write_node( - &self, - node: &StoredProxyNode, - update_existing: bool, - ) -> Result<(), DataLayerError> { - let now = current_unix_secs(); - let upsert_sql = r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, region, status, registered_by, last_heartbeat_at, - heartbeat_interval, active_connections, total_requests, avg_latency_ms, - is_manual, proxy_url, proxy_username, proxy_password, created_at, - updated_at, remote_config, config_version, hardware_info, - estimated_max_concurrency, tunnel_mode, tunnel_connected, tunnel_connected_at, - failed_requests, dns_failures, stream_errors, proxy_metadata -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(id) DO UPDATE SET - name = excluded.name, - ip = excluded.ip, - port = excluded.port, - region = excluded.region, - status = excluded.status, - registered_by = excluded.registered_by, - last_heartbeat_at = excluded.last_heartbeat_at, - heartbeat_interval = excluded.heartbeat_interval, - active_connections = excluded.active_connections, - total_requests = excluded.total_requests, - avg_latency_ms = excluded.avg_latency_ms, - is_manual = excluded.is_manual, - proxy_url = excluded.proxy_url, - proxy_username = excluded.proxy_username, - proxy_password = excluded.proxy_password, - updated_at = excluded.updated_at, - remote_config = excluded.remote_config, - config_version = excluded.config_version, - hardware_info = excluded.hardware_info, - estimated_max_concurrency = excluded.estimated_max_concurrency, - tunnel_mode = excluded.tunnel_mode, - tunnel_connected = excluded.tunnel_connected, - tunnel_connected_at = excluded.tunnel_connected_at, - failed_requests = excluded.failed_requests, - dns_failures = excluded.dns_failures, - stream_errors = excluded.stream_errors, - proxy_metadata = excluded.proxy_metadata -"#; - let sql = if update_existing { - upsert_sql - } else { - upsert_sql - .split_once("\nON CONFLICT(id) DO UPDATE SET") - .map(|(insert_sql, _)| insert_sql) - .expect("proxy node upsert SQL should contain its conflict clause") - }; - sqlx::query(sql) - .bind(&node.id) - .bind(&node.tunnel_generation) - .bind(&node.name) - .bind(&node.ip) - .bind(node.port) - .bind(&node.region) - .bind(&node.status) - .bind(&node.registered_by) - .bind(optional_i64_from_u64( - node.last_heartbeat_at_unix_secs, - "proxy_nodes.last_heartbeat_at", - )?) - .bind(node.heartbeat_interval) - .bind(node.active_connections) - .bind(node.total_requests) - .bind(node.avg_latency_ms) - .bind(node.is_manual) - .bind(&node.proxy_url) - .bind(&node.proxy_username) - .bind(&node.proxy_password) - .bind(node.created_at_unix_ms.unwrap_or(now) as i64) - .bind(node.updated_at_unix_secs.unwrap_or(now) as i64) - .bind(optional_json_to_string( - &node.remote_config, - "proxy_nodes.remote_config", - )?) - .bind(node.config_version) - .bind(optional_json_to_string( - &node.hardware_info, - "proxy_nodes.hardware_info", - )?) - .bind(node.estimated_max_concurrency) - .bind(node.tunnel_mode) - .bind(node.tunnel_connected) - .bind(optional_i64_from_u64( - node.tunnel_connected_at_unix_secs, - "proxy_nodes.tunnel_connected_at", - )?) - .bind(node.failed_requests) - .bind(node.dns_failures) - .bind(node.stream_errors) - .bind(optional_json_to_string( - &node.proxy_metadata, - "proxy_nodes.proxy_metadata", - )?) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - async fn insert_node(&self, node: &StoredProxyNode) -> Result<(), DataLayerError> { - self.write_node(node, false).await - } - - async fn update_existing_registration_if_unchanged( - &self, - mutation: &ProxyNodeRegistrationMutation, - existing: &StoredProxyNode, - replacement_proxy_metadata: Option<&serde_json::Value>, - now: u64, - ) -> Result { - let hardware_info = - optional_json_to_string(&mutation.hardware_info, "proxy_nodes.hardware_info")?; - let replacement_proxy_metadata = optional_json_to_string( - &replacement_proxy_metadata.cloned(), - "proxy_nodes.proxy_metadata", - )?; - let expected_proxy_metadata = - optional_json_to_string(&existing.proxy_metadata, "proxy_nodes.proxy_metadata")?; - let result = sqlx::query(UPDATE_PROXY_NODE_REGISTRATION_SQL) - .bind(&mutation.name) - .bind(&mutation.ip) - .bind(mutation.port) - .bind(mutation.region.as_deref()) - .bind(mutation.registered_by.as_deref()) - .bind(now as i64) - .bind(mutation.heartbeat_interval) - .bind(mutation.active_connections) - .bind(mutation.total_requests) - .bind(mutation.avg_latency_ms) - .bind(hardware_info) - .bind(mutation.estimated_max_concurrency) - .bind(mutation.tunnel_mode) - .bind(replacement_proxy_metadata) - .bind(now as i64) - .bind(&existing.id) - .bind(&existing.tunnel_generation) - .bind(&existing.ip) - .bind(existing.port) - .bind(expected_proxy_metadata.as_deref()) - .bind(expected_proxy_metadata.as_deref()) - .bind(expected_proxy_metadata.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn find_duplicate_proxy_node( - &self, - ip: &str, - port: i32, - excluding_node_id: Option<&str>, - ) -> Result, DataLayerError> { - let row = if let Some(excluding_node_id) = excluding_node_id { - sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE ip = ? AND port = ? AND id <> ? LIMIT 1" - )) - .bind(ip) - .bind(port) - .bind(excluding_node_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE ip = ? AND port = ? LIMIT 1" - )) - .bind(ip) - .bind(port) - .fetch_optional(&self.pool) - .await - .map_sql_err()? - }; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn find_registered_proxy_node_by_endpoint( - &self, - ip: &str, - port: i32, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{PROXY_NODE_COLUMNS} WHERE ip = ? AND port = ? AND is_manual = 0 ORDER BY created_at ASC, id ASC LIMIT 1" - )) - .bind(ip) - .bind(port) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn insert_event( - &self, - node_id: &str, - expected_tunnel_generation: Option<&str>, - event_type: &str, - detail: Option<&str>, - event_metadata: Option<&serde_json::Value>, - created_at_unix_secs: Option, - ) -> Result<(), DataLayerError> { - sqlx::query( - r#" -INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at) -SELECT id, ?, ?, ?, ? -FROM proxy_nodes -WHERE id = ? AND (? IS NULL OR tunnel_generation = ?) -"#, - ) - .bind(event_type) - .bind(detail) - .bind(optional_json_to_string( - &event_metadata.cloned(), - "proxy_node_events.event_metadata", - )?) - .bind(created_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) - .bind(node_id) - .bind(expected_tunnel_generation) - .bind(expected_tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - async fn upsert_metrics_bucket( - &self, - table: &str, - node_id: &str, - expected_tunnel_generation: Option<&str>, - bucket_start: u64, - sample: &TunnelMetricsSample, - ) -> Result<(), DataLayerError> { - sqlx::query(&format!( - r#" -INSERT INTO {table} ( - node_id, - bucket_start_unix_secs, - samples, - uptime_samples, - active_connections_sum, - active_connections_max, - heartbeat_rtt_ms_sum, - heartbeat_rtt_ms_max, - connect_errors_delta, - disconnects_delta, - error_events_delta, - ws_in_bytes_delta, - ws_out_bytes_delta, - ws_in_frames_delta, - ws_out_frames_delta -) -SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? -FROM proxy_nodes -WHERE id = ? AND (? IS NULL OR tunnel_generation = ?) -ON CONFLICT(node_id, bucket_start_unix_secs) DO UPDATE SET - samples = {table}.samples + excluded.samples, - uptime_samples = {table}.uptime_samples + excluded.uptime_samples, - active_connections_sum = {table}.active_connections_sum + excluded.active_connections_sum, - active_connections_max = MAX({table}.active_connections_max, excluded.active_connections_max), - heartbeat_rtt_ms_sum = {table}.heartbeat_rtt_ms_sum + excluded.heartbeat_rtt_ms_sum, - heartbeat_rtt_ms_max = MAX({table}.heartbeat_rtt_ms_max, excluded.heartbeat_rtt_ms_max), - connect_errors_delta = {table}.connect_errors_delta + excluded.connect_errors_delta, - disconnects_delta = {table}.disconnects_delta + excluded.disconnects_delta, - error_events_delta = {table}.error_events_delta + excluded.error_events_delta, - ws_in_bytes_delta = {table}.ws_in_bytes_delta + excluded.ws_in_bytes_delta, - ws_out_bytes_delta = {table}.ws_out_bytes_delta + excluded.ws_out_bytes_delta, - ws_in_frames_delta = {table}.ws_in_frames_delta + excluded.ws_in_frames_delta, - ws_out_frames_delta = {table}.ws_out_frames_delta + excluded.ws_out_frames_delta -"# - )) - .bind(node_id) - .bind(i64::try_from(bucket_start).unwrap_or(i64::MAX)) - .bind(sample.samples) - .bind(sample.uptime_samples) - .bind(sample.active_connections_sum) - .bind(sample.active_connections_max) - .bind(sample.heartbeat_rtt_ms_sum) - .bind(sample.heartbeat_rtt_ms_max) - .bind(sample.connect_errors_delta) - .bind(sample.disconnects_delta) - .bind(sample.error_events_delta) - .bind(sample.ws_in_bytes_delta) - .bind(sample.ws_out_bytes_delta) - .bind(sample.ws_in_frames_delta) - .bind(sample.ws_out_frames_delta) - .bind(node_id) - .bind(expected_tunnel_generation) - .bind(expected_tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(()) - } - - fn normalize_remote_config( - mutation: &ProxyNodeRemoteConfigMutation, - existing: Option<&serde_json::Value>, - ) -> Option { - let mut config = match existing { - Some(serde_json::Value::Object(map)) => map.clone(), - _ => serde_json::Map::new(), - }; - - if let Some(node_name) = mutation.node_name.as_ref() { - config.insert( - "node_name".to_string(), - serde_json::Value::String(node_name.clone()), - ); - } - if let Some(allowed_ports) = mutation.allowed_ports.as_ref() { - config.insert( - "allowed_ports".to_string(), - serde_json::json!(allowed_ports), - ); - } - if let Some(log_level) = mutation.log_level.as_ref() { - config.insert( - "log_level".to_string(), - serde_json::Value::String(log_level.clone()), - ); - } - if let Some(heartbeat_interval) = mutation.heartbeat_interval { - config.insert( - "heartbeat_interval".to_string(), - serde_json::json!(heartbeat_interval), - ); - } - if let Some(scheduling_state) = mutation.scheduling_state.as_ref() { - match scheduling_state { - Some(state) => { - config.insert( - "scheduling_state".to_string(), - serde_json::Value::String(state.clone()), - ); - } - None => { - config.remove("scheduling_state"); - } - } - } - if let Some(upgrade_to) = mutation.upgrade_to.as_ref() { - match upgrade_to { - Some(version) => { - config.insert( - "upgrade_to".to_string(), - serde_json::Value::String(version.clone()), - ); - } - None => { - config.remove("upgrade_to"); - } - } - } - - (!config.is_empty()).then_some(serde_json::Value::Object(config)) - } -} - -const PROXY_NODE_COLUMNS: &str = r#" -SELECT - id, - tunnel_generation, - name, - ip, - port, - region, - is_manual, - proxy_url, - proxy_username, - proxy_password, - status, - registered_by, - last_heartbeat_at AS last_heartbeat_at_unix_secs, - heartbeat_interval, - active_connections, - total_requests, - avg_latency_ms, - failed_requests, - dns_failures, - stream_errors, - proxy_metadata, - hardware_info, - estimated_max_concurrency, - tunnel_mode, - tunnel_connected, - tunnel_connected_at AS tunnel_connected_at_unix_secs, - remote_config, - config_version, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs -FROM proxy_nodes -"#; - -const PROXY_NODE_EVENT_COLUMNS: &str = r#" -SELECT - id, - node_id, - event_type, - detail, - event_metadata, - created_at AS created_at_unix_ms -FROM proxy_node_events -"#; - -const APPLY_HEARTBEAT_SQL: &str = r#" -UPDATE proxy_nodes -SET last_heartbeat_at = ?, - tunnel_connected_at = CASE - WHEN status <> 'online' OR tunnel_connected = 0 THEN ? - ELSE tunnel_connected_at - END, - updated_at = CASE - WHEN status <> 'online' OR tunnel_connected = 0 THEN ? - ELSE updated_at - END, - status = 'online', - tunnel_connected = 1, - heartbeat_interval = COALESCE(?, heartbeat_interval), - active_connections = COALESCE(?, active_connections), - avg_latency_ms = COALESCE(?, avg_latency_ms), - total_requests = total_requests + MAX(COALESCE(?, 0), 0), - failed_requests = failed_requests + MAX(COALESCE(?, 0), 0), - dns_failures = dns_failures + MAX(COALESCE(?, 0), 0), - stream_errors = stream_errors + MAX(COALESCE(?, 0), 0) -WHERE id = ? - AND tunnel_mode = 1 - AND tunnel_generation = ? -"#; - -const CAS_HEARTBEAT_PROXY_METADATA_SQL: &str = r#" -UPDATE proxy_nodes -SET proxy_metadata = ?, updated_at = ? -WHERE id = ? AND tunnel_generation = ? - AND ( - (proxy_metadata IS NULL AND ? IS NULL) - OR ( - proxy_metadata IS NOT NULL AND ? IS NOT NULL - AND json(proxy_metadata) = json(?) - ) - ) -"#; - -const UPDATE_TUNNEL_STATUS_SQL: &str = r#" -UPDATE proxy_nodes -SET tunnel_connected = ?, - active_connections = CASE WHEN ? THEN active_connections ELSE 0 END, - tunnel_connected_at = ?, - status = CASE WHEN ? THEN 'online' ELSE 'offline' END, - updated_at = ? -WHERE id = ? - AND tunnel_generation = ? - AND (tunnel_connected_at IS NULL OR tunnel_connected_at <= ?) -"#; - -const UPDATE_MANUAL_PROXY_NODE_SQL: &str = r#" -UPDATE proxy_nodes -SET name = COALESCE(?, name), - ip = COALESCE(?, ip), - port = COALESCE(?, port), - region = COALESCE(?, region), - proxy_url = COALESCE(?, proxy_url), - proxy_username = COALESCE(?, proxy_username), - proxy_password = COALESCE(?, proxy_password), - updated_at = ? -WHERE id = ? AND is_manual = 1 - AND tunnel_generation = ? -"#; - -const UPDATE_PROXY_NODE_REGISTRATION_SQL: &str = r#" -UPDATE proxy_nodes -SET name = ?, ip = ?, port = ?, region = ?, registered_by = ?, - last_heartbeat_at = ?, heartbeat_interval = ?, - active_connections = COALESCE(?, active_connections), - total_requests = COALESCE(?, total_requests), - avg_latency_ms = COALESCE(?, avg_latency_ms), - hardware_info = COALESCE(?, hardware_info), - estimated_max_concurrency = COALESCE(?, estimated_max_concurrency), - tunnel_mode = ?, proxy_metadata = COALESCE(?, proxy_metadata), updated_at = ? -WHERE id = ? AND tunnel_generation = ? - AND is_manual = 0 AND ip = ? AND port = ? - AND ( - (proxy_metadata IS NULL AND ? IS NULL) - OR ( - proxy_metadata IS NOT NULL AND ? IS NOT NULL - AND json(proxy_metadata) = json(?) - ) - ) -"#; - -const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#" -UPDATE proxy_nodes -SET name = COALESCE(?, name), remote_config = ?, - config_version = config_version + 1, updated_at = ? -WHERE id = ? AND tunnel_generation = ? AND config_version = ? - AND is_manual = 0 -"#; - -const RECORD_PROXY_NODE_TRAFFIC_SQL: &str = 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 = ? AND is_manual = 1 - AND tunnel_generation = ? -"#; - -const INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL: &str = r#" -UPDATE proxy_nodes -SET total_requests = total_requests + MAX(?, 0), - failed_requests = failed_requests + MAX(?, 0), - avg_latency_ms = COALESCE(?, avg_latency_ms), - updated_at = ? -WHERE id = ? AND is_manual = 1 - AND tunnel_generation = ? -"#; - -const UNREGISTER_PROXY_NODE_SQL: &str = r#" -UPDATE proxy_nodes -SET status = 'offline', tunnel_connected = 0, active_connections = 0, - tunnel_connected_at = ?, updated_at = ? -WHERE id = ? - AND tunnel_generation = ? -"#; - -// Retire counter rows after the parent delete commits. Keeping this statement -// outside the parent-locking transaction avoids the outbox -> proxy_nodes -// versus proxy_nodes -> outbox lock inversion with the counter flusher. -const RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL: &str = r#" -DELETE FROM usage_counter_deltas -WHERE kind = 'proxy_node' - AND target_id = ? - AND target_tunnel_generation = ? - AND processed_at IS NULL -"#; - -#[async_trait] -impl ProxyNodeReadRepository for SqliteProxyNodeReadRepository { - async fn list_proxy_nodes(&self) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(PROXY_NODE_COLUMNS); - builder.push(" ORDER BY name ASC, id ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_proxy_node_row).collect() - } - - async fn find_proxy_node( - &self, - node_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(PROXY_NODE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq(&mut builder, &mut where_clause, "id", node_id.to_string()); - push_limit(&mut builder, 1); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_proxy_node_row).transpose() - } - - async fn list_proxy_node_events( - &self, - node_id: &str, - limit: usize, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(PROXY_NODE_EVENT_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "node_id", - node_id.to_string(), - ); - builder.push(" ORDER BY created_at DESC, id DESC"); - push_limit(&mut builder, i64::try_from(limit).unwrap_or(i64::MAX)); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_proxy_node_event_row).collect() - } - - async fn list_proxy_node_events_filtered( - &self, - node_id: &str, - query: &ProxyNodeEventQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(PROXY_NODE_EVENT_COLUMNS); - let mut where_clause = WhereClause::new(); - push_eq( - &mut builder, - &mut where_clause, - "node_id", - node_id.to_string(), - ); - if let Some(from_unix_secs) = query.from_unix_secs { - where_clause.push_next(&mut builder); - builder - .push("created_at >= ") - .push_bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX)); - } - if let Some(to_unix_secs) = query.to_unix_secs { - where_clause.push_next(&mut builder); - builder - .push("created_at <= ") - .push_bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX)); - } - if let Some(event_type) = query.event_type.as_deref() { - where_clause.push_next(&mut builder); - builder - .push("LOWER(event_type) = LOWER(") - .push_bind(event_type.to_string()) - .push(")"); - } - builder.push(" ORDER BY created_at DESC, id DESC"); - push_limit(&mut builder, i64::try_from(query.limit).unwrap_or(i64::MAX)); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_proxy_node_event_row).collect() - } - - async fn list_proxy_node_metrics( - &self, - node_id: &str, - step: ProxyNodeMetricsStep, - from_unix_secs: u64, - to_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let table = match step { - ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m", - ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h", - }; - let rows = sqlx::query(&format!( - r#" -SELECT - node_id, - bucket_start_unix_secs, - samples, - uptime_samples, - active_connections_sum, - active_connections_max, - heartbeat_rtt_ms_sum, - heartbeat_rtt_ms_max, - connect_errors_delta, - disconnects_delta, - error_events_delta, - ws_in_bytes_delta, - ws_out_bytes_delta, - ws_in_frames_delta, - ws_out_frames_delta -FROM {table} -WHERE node_id = ? - AND bucket_start_unix_secs >= ? - AND bucket_start_unix_secs <= ? -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"# - )) - .bind(node_id) - .bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_node_metric_row).collect() - } - - async fn list_proxy_fleet_metrics( - &self, - step: ProxyNodeMetricsStep, - from_unix_secs: u64, - to_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - let table = match step { - ProxyNodeMetricsStep::OneMinute => "proxy_node_metrics_1m", - ProxyNodeMetricsStep::OneHour => "proxy_node_metrics_1h", - }; - let rows = sqlx::query(&format!( - r#" -SELECT - bucket_start_unix_secs, - SUM(samples) AS samples, - SUM(uptime_samples) AS uptime_samples, - SUM(active_connections_sum) AS active_connections_sum, - MAX(active_connections_max) AS active_connections_max, - SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum, - MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max, - SUM(connect_errors_delta) AS connect_errors_delta, - SUM(disconnects_delta) AS disconnects_delta, - SUM(error_events_delta) AS error_events_delta, - SUM(ws_in_bytes_delta) AS ws_in_bytes_delta, - SUM(ws_out_bytes_delta) AS ws_out_bytes_delta, - SUM(ws_in_frames_delta) AS ws_in_frames_delta, - SUM(ws_out_frames_delta) AS ws_out_frames_delta -FROM {table} -WHERE bucket_start_unix_secs >= ? - AND bucket_start_unix_secs <= ? -GROUP BY bucket_start_unix_secs -ORDER BY bucket_start_unix_secs ASC -LIMIT ? -"# - )) - .bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX)) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_proxy_fleet_metric_row).collect() - } -} - -#[async_trait] -impl ProxyNodeWriteRepository for SqliteProxyNodeReadRepository { - async fn reset_stale_tunnel_statuses(&self) -> Result { - let now = current_unix_secs() as i64; - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET tunnel_connected = 0, - status = 'offline', - active_connections = 0, - tunnel_connected_at = ?, - updated_at = ? -WHERE is_manual = 0 - AND tunnel_connected = 1 -"#, - ) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn compare_and_set_proxy_password( - &self, - node_id: &str, - expected: &str, - replacement: &str, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET proxy_password = ?, updated_at = ? -WHERE id = ? AND proxy_password = ? -"#, - ) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn compare_and_set_proxy_metadata( - &self, - node_id: &str, - expected: &serde_json::Value, - replacement: &serde_json::Value, - ) -> Result { - let expected = serde_json::to_string(expected).map_err(|err| { - DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}")) - })?; - let replacement = serde_json::to_string(replacement).map_err(|err| { - DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}")) - })?; - let result = sqlx::query( - r#" -UPDATE proxy_nodes -SET proxy_metadata = ?, updated_at = ? -WHERE id = ? - AND proxy_metadata IS NOT NULL - AND json(proxy_metadata) = json(?) -"#, - ) - .bind(replacement) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(expected) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn create_manual_node( - &self, - mutation: &ProxyNodeManualCreateMutation, - ) -> Result { - if let Some(existing) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - return Err(duplicate_proxy_node_error(&existing)); - } - - let node_id = requested_proxy_node_id(mutation.node_id.as_deref())? - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - if let Some(existing) = self.find_proxy_node(&node_id).await? { - return Err(proxy_node_id_in_use_error(&existing)); - } - let now = Some(current_unix_secs()); - let node = StoredProxyNode::new( - node_id, - mutation.name.clone(), - mutation.ip.clone(), - mutation.port, - true, - "online".to_string(), - 0, - 0, - 0, - 0, - 0, - 0, - false, - false, - 0, - )? - .with_manual_proxy_fields( - Some(mutation.proxy_url.clone()), - mutation.proxy_username.clone(), - mutation.proxy_password.clone(), - ) - .with_runtime_fields( - mutation.region.clone(), - mutation.registered_by.clone(), - None, - None, - None, - None, - None, - None, - None, - now, - now, - ); - - if let Err(error) = self.insert_node(&node).await { - if let Some(duplicate) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - return Err(duplicate_proxy_node_error(&duplicate)); - } - if let Some(owner) = self.find_proxy_node(&node.id).await? { - return Err(proxy_node_id_in_use_error(&owner)); - } - return Err(error); - } - Ok(node) - } - - async fn update_manual_node( - &self, - mutation: &ProxyNodeManualUpdateMutation, - ) -> Result, DataLayerError> { - let Some(existing) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if !existing.is_manual { - return Err(DataLayerError::InvalidInput( - "只能编辑手动添加的代理节点".to_string(), - )); - } - - let next_ip = mutation.ip.as_deref().unwrap_or(existing.ip.as_str()); - let next_port = mutation.port.unwrap_or(existing.port); - if let Some(existing) = self - .find_duplicate_proxy_node(next_ip, next_port, Some(&mutation.node_id)) - .await? - { - return Err(duplicate_proxy_node_error(&existing)); - } - - let result = sqlx::query(UPDATE_MANUAL_PROXY_NODE_SQL) - .bind(mutation.name.as_deref()) - .bind(mutation.ip.as_deref()) - .bind(mutation.port) - .bind(mutation.region.as_deref()) - .bind(mutation.proxy_url.as_deref()) - .bind(mutation.proxy_username.as_deref()) - .bind(mutation.proxy_password.as_deref()) - .bind(current_unix_secs() as i64) - .bind(&mutation.node_id) - .bind(&existing.tunnel_generation) - .execute(&self.pool) - .await; - let result = match result { - Ok(result) => result, - Err(error) => { - if let Some(duplicate) = self - .find_duplicate_proxy_node(next_ip, next_port, Some(&mutation.node_id)) - .await? - { - return Err(duplicate_proxy_node_error(&duplicate)); - } - return Err(DataLayerError::sql(error)); - } - }; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_proxy_node(&mutation.node_id).await - } - - async fn register_node( - &self, - mutation: &ProxyNodeRegistrationMutation, - ) -> Result { - let requested_id = requested_proxy_node_id(mutation.node_id.as_deref())?; - let normalized_proxy_metadata = normalize_proxy_metadata( - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ); - let rotates_tunnel_security = - proxy_metadata_has_explicit_tunnel_security(normalized_proxy_metadata.as_ref()); - - let Some(initial_existing) = self - .find_registered_proxy_node_by_endpoint(&mutation.ip, mutation.port) - .await? - else { - let node_id = requested_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - if let Some(existing) = self.find_proxy_node(&node_id).await? { - return Err(proxy_node_id_in_use_error(&existing)); - } - let now = Some(current_unix_secs()); - let node = StoredProxyNode::new( - node_id, - mutation.name.clone(), - mutation.ip.clone(), - mutation.port, - false, - "offline".to_string(), - mutation.heartbeat_interval, - mutation.active_connections.unwrap_or(0), - mutation.total_requests.unwrap_or(0), - 0, - 0, - 0, - mutation.tunnel_mode, - false, - 0, - )? - .with_runtime_fields( - mutation.region.clone(), - mutation.registered_by.clone(), - now, - mutation.avg_latency_ms, - merge_proxy_metadata_for_registration(None, normalized_proxy_metadata.clone()), - mutation.hardware_info.clone(), - mutation.estimated_max_concurrency, - None, - None, - now, - now, - ); - if let Err(error) = self.insert_node(&node).await { - if let Some(winner) = self - .find_duplicate_proxy_node(&mutation.ip, mutation.port, None) - .await? - { - if winner.is_manual { - return Err(duplicate_proxy_node_error(&winner)); - } - if let Some(requested_id) = requested_id.as_deref() { - if requested_id != winner.id { - return Err(proxy_node_registration_identity_error( - requested_id, - &winner.id, - )); - } - } - return Ok(winner); - } - if let Some(owner) = self.find_proxy_node(&node.id).await? { - return Err(proxy_node_id_in_use_error(&owner)); - } - return Err(error); - } - return Ok(node); - }; - - if let Some(requested_id) = requested_id.as_deref() { - if requested_id != initial_existing.id { - return Err(proxy_node_registration_identity_error( - requested_id, - &initial_existing.id, - )); - } - } - - let pinned_id = initial_existing.id.clone(); - let pinned_generation = initial_existing.tunnel_generation.clone(); - let mut existing = initial_existing; - for attempt in 0..PROXY_NODE_REGISTRATION_CAS_RETRIES { - if attempt != 0 { - existing = self - .find_registered_proxy_node_by_endpoint(&mutation.ip, mutation.port) - .await? - .ok_or_else(proxy_node_registration_changed_error)?; - } - if existing.id != pinned_id || existing.tunnel_generation != pinned_generation { - return Err(proxy_node_registration_changed_error()); - } - - let replacement_proxy_metadata = merge_proxy_metadata_for_registration( - existing.proxy_metadata.as_ref(), - normalized_proxy_metadata.clone(), - ); - let now = current_unix_secs(); - if self - .update_existing_registration_if_unchanged( - mutation, - &existing, - replacement_proxy_metadata.as_ref(), - now, - ) - .await? - { - return self - .find_proxy_node(&pinned_id) - .await? - .filter(|current| current.tunnel_generation == pinned_generation) - .ok_or_else(proxy_node_registration_changed_error); - } - if rotates_tunnel_security { - return Err(DataLayerError::UnexpectedValue( - "proxy node changed during explicit tunnel security rotation".to_string(), - )); - } - } - - Err(DataLayerError::UnexpectedValue( - "proxy node registration changed during every CAS retry".to_string(), - )) - } - - async fn apply_heartbeat( - &self, - mutation: &ProxyNodeHeartbeatMutation, - ) -> Result, DataLayerError> { - let Some(existing) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != existing.tunnel_generation) - { - return Ok(None); - } - if !existing.tunnel_mode { - return Err(DataLayerError::InvalidInput( - "non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode" - .to_string(), - )); - } - - let tunnel_generation = existing.tunnel_generation.clone(); - let now_unix_secs = current_unix_secs(); - let now = i64::try_from(now_unix_secs).unwrap_or(i64::MAX); - let has_proxy_metadata_update = normalize_heartbeat_proxy_metadata( - None, - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ) - .is_some(); - - let result = sqlx::query(APPLY_HEARTBEAT_SQL) - .bind(now) - .bind(now) - .bind(now) - .bind(mutation.heartbeat_interval) - .bind(mutation.active_connections) - .bind(mutation.avg_latency_ms) - .bind(mutation.total_requests_delta) - .bind(mutation.failed_requests_delta) - .bind(mutation.dns_failures_delta) - .bind(mutation.stream_errors_delta) - .bind(&mutation.node_id) - .bind(&tunnel_generation) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - - let mut updated = None; - let mut tunnel_metrics_sample = None; - if has_proxy_metadata_update { - for _ in 0..8 { - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != tunnel_generation { - return Ok(None); - } - let Some(replacement) = normalize_heartbeat_proxy_metadata( - current.proxy_metadata.as_ref(), - mutation.proxy_metadata.as_ref(), - mutation.proxy_version.as_deref(), - ) else { - break; - }; - if current.proxy_metadata.as_ref() == Some(&replacement) { - tunnel_metrics_sample = build_tunnel_metrics_sample( - current.proxy_metadata.as_ref(), - Some(&replacement), - current.active_connections, - current.tunnel_connected, - ); - updated = Some(current); - break; - } - - let expected = - optional_json_to_string(¤t.proxy_metadata, "proxy_nodes.proxy_metadata")?; - let replacement_json = serde_json::to_string(&replacement).map_err(|error| { - DataLayerError::UnexpectedValue(format!( - "proxy_nodes.proxy_metadata contains unserializable JSON: {error}" - )) - })?; - let result = sqlx::query(CAS_HEARTBEAT_PROXY_METADATA_SQL) - .bind(replacement_json) - .bind(now) - .bind(&mutation.node_id) - .bind(&tunnel_generation) - .bind(expected.as_deref()) - .bind(expected.as_deref()) - .bind(expected.as_deref()) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - continue; - } - let Some(after_cas) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if after_cas.tunnel_generation != tunnel_generation { - return Ok(None); - } - tunnel_metrics_sample = build_tunnel_metrics_sample( - current.proxy_metadata.as_ref(), - after_cas.proxy_metadata.as_ref(), - after_cas.active_connections, - after_cas.tunnel_connected, - ); - updated = Some(after_cas); - break; - } - } - let updated = if let Some(updated) = updated { - updated - } else { - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != tunnel_generation { - return Ok(None); - } - current - }; - - if let Some(sample) = tunnel_metrics_sample.as_ref() { - self.upsert_metrics_bucket( - "proxy_node_metrics_1m", - &updated.id, - Some(tunnel_generation.as_str()), - bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute), - sample, - ) - .await?; - self.upsert_metrics_bucket( - "proxy_node_metrics_1h", - &updated.id, - Some(tunnel_generation.as_str()), - bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour), - sample, - ) - .await?; - - for error in &sample.recent_error_events { - log_reported_tunnel_error_event(&updated.id, error, now_unix_secs); - let detail = build_tunnel_error_event_detail(error); - let event_metadata = serde_json::json!({ - "source": "heartbeat", - "category": error.category, - "message": error.message, - "severity": error.severity.as_deref(), - "component": error.component.as_deref(), - "summary": error.summary.as_deref(), - "operator_action": error.operator_action.as_deref(), - "timestamp_unix_secs": error.timestamp_unix_secs, - "timestamp_unix_ms": error.timestamp_unix_ms, - }); - self.insert_event( - &updated.id, - Some(tunnel_generation.as_str()), - PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR, - Some(detail.as_str()), - Some(&event_metadata), - Some(if error.timestamp_unix_secs == 0 { - now_unix_secs - } else { - error.timestamp_unix_secs - }), - ) - .await?; - } - } - - if reconcile_remote_config_after_heartbeat( - updated.remote_config.as_ref(), - mutation.proxy_version.as_deref(), - ) != updated.remote_config - { - return self - .update_remote_config(&ProxyNodeRemoteConfigMutation { - node_id: mutation.node_id.clone(), - expected_tunnel_generation: Some(tunnel_generation), - node_name: None, - allowed_ports: None, - log_level: None, - heartbeat_interval: None, - scheduling_state: None, - upgrade_to: Some(None), - }) - .await; - } - - Ok(Some(updated)) - } - - async fn record_traffic( - &self, - mutation: &ProxyNodeTrafficMutation, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let lock = sqlx::query("UPDATE proxy_nodes SET id = id WHERE id = ?") - .bind(&mutation.node_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if lock.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let row = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(&mutation.node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - let Some(expected_generation) = mutation.expected_tunnel_generation.as_deref() else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - if expected_generation != generation { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query(RECORD_PROXY_NODE_TRAFFIC_SQL) - .bind(mutation.total_requests_delta) - .bind(mutation.failed_requests_delta) - .bind(mutation.dns_failures_delta) - .bind(mutation.stream_errors_delta) - .bind(current_unix_secs() as i64) - .bind(&mutation.node_id) - .bind(expected_generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - let applied = result.rows_affected() > 0; - tx.commit().await.map_sql_err()?; - Ok(applied) - } - - async fn update_tunnel_status( - &self, - mutation: &ProxyNodeTunnelStatusMutation, - ) -> Result, DataLayerError> { - let Some(node) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != node.tunnel_generation) - { - return Ok(None); - } - - let event_time = mutation - .observed_at_unix_secs - .unwrap_or_else(current_unix_secs); - let event_type = if mutation.connected { - "connected" - } else { - "disconnected" - }; - let event_detail = mutation.detail.clone().unwrap_or_else(|| { - format!( - "[tunnel_node_status] conn_count={}", - i32::max(mutation.conn_count, 0) - ) - }); - - let event_time_i64 = i64::try_from(event_time).unwrap_or(i64::MAX); - let result = sqlx::query(UPDATE_TUNNEL_STATUS_SQL) - .bind(mutation.connected) - .bind(mutation.connected) - .bind(event_time_i64) - .bind(mutation.connected) - .bind(event_time_i64) - .bind(&mutation.node_id) - .bind(&node.tunnel_generation) - .bind(event_time_i64) - .execute(&self.pool) - .await - .map_sql_err()?; - let Some(current) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if current.tunnel_generation != node.tunnel_generation { - return Ok(None); - } - let stale = result.rows_affected() == 0 - && current - .tunnel_connected_at_unix_secs - .is_some_and(|last_transition| event_time < last_transition); - let persisted_detail = if stale { - format!("[stale_ignored] {event_detail}") - } else { - event_detail - }; - self.insert_event( - &mutation.node_id, - Some(node.tunnel_generation.as_str()), - event_type, - Some(&persisted_detail), - None, - Some(if stale { - current_unix_secs() - } else { - event_time - }), - ) - .await?; - Ok(Some(current)) - } - - async fn unregister_node( - &self, - node_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let lock = sqlx::query("UPDATE proxy_nodes SET id = id WHERE id = ?") - .bind(node_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if lock.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - let row = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - let now = current_unix_secs() as i64; - sqlx::query(UNREGISTER_PROXY_NODE_SQL) - .bind(now) - .bind(now) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err() - .and_then(|row| map_proxy_node_row(&row)); - let updated = updated?; - tx.commit().await.map_sql_err()?; - Ok(Some(updated)) - } - - async fn delete_node(&self, node_id: &str) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - - // SQLite has no SELECT ... FOR UPDATE. A matched no-op UPDATE acquires the - // connection's write lock before we read the generation, so a concurrent - // unregister/re-register cannot interleave with the cleanup below. - let lock = sqlx::query("UPDATE proxy_nodes SET id = id WHERE id = ?") - .bind(node_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if lock.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - - let row = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let existing = map_proxy_node_row(&row)?; - let generation = existing.tunnel_generation.as_str(); - - // Child tables do not carry the generation themselves. Keep the parent - // identity predicate on every delete so this remains correct if schema - // constraints differ between installations. - sqlx::query( - "DELETE FROM proxy_node_events WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes WHERE id = ? AND tunnel_generation = ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "DELETE FROM proxy_node_metrics_1m WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes WHERE id = ? AND tunnel_generation = ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "DELETE FROM proxy_node_metrics_1h WHERE node_id = ? AND EXISTS (SELECT 1 FROM proxy_nodes WHERE id = ? AND tunnel_generation = ?)", - ) - .bind(node_id) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let deleted = sqlx::query("DELETE FROM proxy_nodes WHERE id = ? AND tunnel_generation = ?") - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - if deleted.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - tx.commit().await.map_sql_err()?; - // The parent lock is released before this cleanup. If a flusher already - // claimed one of the rows it can finish (and fail the generation - // predicate) without forming a lock cycle with the delete transaction. - if let Err(error) = sqlx::query(RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL) - .bind(node_id) - .bind(generation) - .execute(&self.pool) - .await - .map_sql_err() - { - tracing::warn!( - node_id = %node_id, - tunnel_generation = %generation, - error = ?error, - "failed to retire deleted proxy node counter rows" - ); - } - Ok(Some(existing)) - } - - async fn update_remote_config( - &self, - mutation: &ProxyNodeRemoteConfigMutation, - ) -> Result, DataLayerError> { - for _ in 0..8 { - let Some(node) = self.find_proxy_node(&mutation.node_id).await? else { - return Ok(None); - }; - if mutation - .expected_tunnel_generation - .as_deref() - .is_some_and(|expected| expected != node.tunnel_generation) - { - return Ok(None); - } - if node.is_manual { - return Err(DataLayerError::InvalidInput( - "手动节点不支持远程配置下发".to_string(), - )); - } - - let remote_config = - Self::normalize_remote_config(mutation, node.remote_config.as_ref()); - let remote_config = - optional_json_to_string(&remote_config, "proxy_nodes.remote_config")?; - let now = current_unix_secs() as i64; - let result = sqlx::query(UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL) - .bind(mutation.node_name.as_deref()) - .bind(remote_config) - .bind(now) - .bind(&mutation.node_id) - .bind(&node.tunnel_generation) - .bind(node.config_version) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - continue; - } - - let current = self.find_proxy_node(&mutation.node_id).await?; - return Ok( - current.filter(|current| current.tunnel_generation == node.tunnel_generation) - ); - } - - Err(DataLayerError::UnexpectedValue( - "proxy node remote config changed during every CAS retry".to_string(), - )) - } - - async fn increment_manual_node_requests( - &self, - node_id: &str, - total_delta: i64, - failed_delta: i64, - latency_ms: Option, - ) -> Result<(), DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let lock = sqlx::query("UPDATE proxy_nodes SET id = id WHERE id = ?") - .bind(node_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if lock.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(()); - } - let row = sqlx::query(&format!("{PROXY_NODE_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(node_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let generation = map_proxy_node_row(&row)?.tunnel_generation; - sqlx::query(INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL) - .bind(total_delta) - .bind(failed_delta) - .bind(latency_ms.map(|value| value as f64)) - .bind(current_unix_secs() as i64) - .bind(node_id) - .bind(generation) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(()) - } - - async fn cleanup_proxy_node_metrics( - &self, - retain_1m_from_unix_secs: u64, - retain_1h_from_unix_secs: u64, - delete_limit: usize, - ) -> Result { - let delete_limit_i64 = i64::try_from(delete_limit.max(1)).unwrap_or(i64::MAX); - let deleted_1m = sqlx::query( - r#" -DELETE FROM proxy_node_metrics_1m -WHERE (node_id, bucket_start_unix_secs) IN ( - SELECT node_id, bucket_start_unix_secs - FROM proxy_node_metrics_1m - WHERE bucket_start_unix_secs < ? - ORDER BY bucket_start_unix_secs ASC - LIMIT ? -) -"#, - ) - .bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX)) - .bind(delete_limit_i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() as usize; - - let deleted_1h = sqlx::query( - r#" -DELETE FROM proxy_node_metrics_1h -WHERE (node_id, bucket_start_unix_secs) IN ( - SELECT node_id, bucket_start_unix_secs - FROM proxy_node_metrics_1h - WHERE bucket_start_unix_secs < ? - ORDER BY bucket_start_unix_secs ASC - LIMIT ? -) -"#, - ) - .bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX)) - .bind(delete_limit_i64) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() as usize; - - Ok(ProxyNodeMetricsCleanupSummary { - deleted_1m_rows: deleted_1m, - deleted_1h_rows: deleted_1h, - }) - } -} - -fn optional_unix_secs(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - -fn optional_i64_from_u64( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}")) - }) - }) - .transpose() -} - -fn optional_json_to_string( - value: &Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains unserializable JSON: {err}" - )) - }) - }) - .transpose() -} - -fn duplicate_proxy_node_error(node: &StoredProxyNode) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "已存在相同地址的代理节点: {} ({}:{})", - node.name, node.ip, node.port - )) -} - -fn requested_proxy_node_id(value: Option<&str>) -> Result, DataLayerError> { - let Some(value) = value else { - return Ok(None); - }; - if value.is_empty() || value.trim() != value { - return Err(DataLayerError::InvalidInput( - "proxy node id must be non-empty and unpadded".to_string(), - )); - } - Ok(Some(value.to_string())) -} - -fn proxy_node_registration_identity_error(requested_id: &str, existing_id: &str) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "proxy node registration identity changed: requested {requested_id}, existing {existing_id}" - )) -} - -fn proxy_node_registration_changed_error() -> DataLayerError { - DataLayerError::UnexpectedValue( - "registered proxy node identity changed during registration".to_string(), - ) -} - -fn proxy_node_id_in_use_error(node: &StoredProxyNode) -> DataLayerError { - DataLayerError::InvalidInput(format!( - "proxy node id is already in use: {} ({}:{})", - node.id, node.ip, node.port - )) -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn map_proxy_node_row(row: &SqliteRow) -> Result { - let tunnel_generation: String = row.try_get("tunnel_generation").map_sql_err()?; - if tunnel_generation.trim().is_empty() { - return Err(DataLayerError::UnexpectedValue( - "proxy_nodes.tunnel_generation must not be empty".to_string(), - )); - } - Ok(StoredProxyNode::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("ip").map_sql_err()?, - row.try_get("port").map_sql_err()?, - row.try_get("is_manual").map_sql_err()?, - row.try_get("status").map_sql_err()?, - row.try_get("heartbeat_interval").map_sql_err()?, - row.try_get("active_connections").map_sql_err()?, - row.try_get("total_requests").map_sql_err()?, - row.try_get("failed_requests").map_sql_err()?, - row.try_get("dns_failures").map_sql_err()?, - row.try_get("stream_errors").map_sql_err()?, - row.try_get("tunnel_mode").map_sql_err()?, - row.try_get("tunnel_connected").map_sql_err()?, - row.try_get("config_version").map_sql_err()?, - )? - .with_tunnel_generation(tunnel_generation) - .with_manual_proxy_fields( - row.try_get("proxy_url").map_sql_err()?, - row.try_get("proxy_username").map_sql_err()?, - row.try_get("proxy_password").map_sql_err()?, - ) - .with_runtime_fields( - row.try_get("region").map_sql_err()?, - row.try_get("registered_by").map_sql_err()?, - optional_unix_secs(row.try_get("last_heartbeat_at_unix_secs").map_sql_err()?), - row.try_get("avg_latency_ms").map_sql_err()?, - optional_json_from_string( - row.try_get("proxy_metadata").map_sql_err()?, - "proxy_nodes.proxy_metadata", - )?, - optional_json_from_string( - row.try_get("hardware_info").map_sql_err()?, - "proxy_nodes.hardware_info", - )?, - row.try_get("estimated_max_concurrency").map_sql_err()?, - optional_unix_secs(row.try_get("tunnel_connected_at_unix_secs").map_sql_err()?), - optional_json_from_string( - row.try_get("remote_config").map_sql_err()?, - "proxy_nodes.remote_config", - )?, - optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - optional_unix_secs(row.try_get("updated_at_unix_secs").map_sql_err()?), - )) -} - -fn map_proxy_node_event_row(row: &SqliteRow) -> Result { - Ok(StoredProxyNodeEvent { - id: row.try_get("id").map_sql_err()?, - node_id: row.try_get("node_id").map_sql_err()?, - event_type: row.try_get("event_type").map_sql_err()?, - detail: row.try_get("detail").map_sql_err()?, - event_metadata: optional_json_from_string( - row.try_get("event_metadata").map_sql_err()?, - "proxy_node_events.event_metadata", - )?, - created_at_unix_ms: optional_unix_secs(row.try_get("created_at_unix_ms").map_sql_err()?), - }) -} - -fn map_proxy_node_metric_row( - row: &SqliteRow, -) -> Result { - Ok(StoredProxyNodeMetricsBucket { - node_id: row.try_get("node_id").map_sql_err()?, - bucket_start_unix_secs: optional_unix_secs( - row.try_get("bucket_start_unix_secs").map_sql_err()?, - ) - .unwrap_or_default(), - samples: row.try_get("samples").map_sql_err()?, - uptime_samples: row.try_get("uptime_samples").map_sql_err()?, - active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?, - active_connections_max: row.try_get("active_connections_max").map_sql_err()?, - heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?, - heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?, - connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?, - disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?, - error_events_delta: row.try_get("error_events_delta").map_sql_err()?, - ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?, - ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?, - ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?, - ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?, - }) -} - -fn map_proxy_fleet_metric_row( - row: &SqliteRow, -) -> Result { - Ok(StoredProxyFleetMetricsBucket { - bucket_start_unix_secs: optional_unix_secs( - row.try_get("bucket_start_unix_secs").map_sql_err()?, - ) - .unwrap_or_default(), - samples: row.try_get("samples").map_sql_err()?, - uptime_samples: row.try_get("uptime_samples").map_sql_err()?, - active_connections_sum: row.try_get("active_connections_sum").map_sql_err()?, - active_connections_max: row.try_get("active_connections_max").map_sql_err()?, - heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_sql_err()?, - heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_sql_err()?, - connect_errors_delta: row.try_get("connect_errors_delta").map_sql_err()?, - disconnects_delta: row.try_get("disconnects_delta").map_sql_err()?, - error_events_delta: row.try_get("error_events_delta").map_sql_err()?, - ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_sql_err()?, - ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_sql_err()?, - ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_sql_err()?, - ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_sql_err()?, - }) -} - -#[cfg(test)] -mod tests { - use std::{sync::Arc, time::Duration}; - - use super::SqliteProxyNodeReadRepository; - use crate::run_migrations; - use aether_data_contracts::repository::proxy_nodes::{ - merge_proxy_metadata_for_registration, normalize_proxy_metadata, ProxyNodeEventQuery, - ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, - ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation, - ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, - ProxyNodeWriteRepository, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR, - }; - use serde_json::json; - use tokio::sync::Barrier; - - #[tokio::test] - async fn sqlite_proxy_node_mutations_remain_atomic_under_concurrency() { - let database_path = std::env::temp_dir().join(format!( - "aether-proxy-node-atomicity-{}.sqlite", - uuid::Uuid::new_v4().simple() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .busy_timeout(Duration::from_secs(10)); - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(12) - .connect_with(options) - .await - .expect("concurrent sqlite pool should connect"); - run_migrations(&pool) - .await - .expect("sqlite migrations should run"); - let repository = SqliteProxyNodeReadRepository::new(pool.clone()); - - let manual = repository - .create_manual_node(&ProxyNodeManualCreateMutation { - node_id: Some("atomic-manual".to_string()), - name: "atomic-manual".to_string(), - ip: "127.0.0.40".to_string(), - port: 7040, - region: None, - proxy_url: "http://127.0.0.40:7040".to_string(), - proxy_username: None, - proxy_password: None, - registered_by: None, - }) - .await - .expect("manual node should create"); - let tunnel = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some("atomic-tunnel".to_string()), - name: "atomic-tunnel".to_string(), - ip: "127.0.0.41".to_string(), - port: 7041, - region: None, - heartbeat_interval: 30, - active_connections: Some(0), - total_requests: Some(0), - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: None, - proxy_version: None, - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("tunnel node should register"); - - const WORKERS: usize = 24; - let barrier = Arc::new(Barrier::new(WORKERS * 3)); - let mut tasks = Vec::with_capacity(WORKERS * 3); - for _ in 0..WORKERS { - let traffic_repository = repository.clone(); - let traffic_barrier = Arc::clone(&barrier); - let node_id = manual.id.clone(); - let tunnel_generation = manual.tunnel_generation.clone(); - tasks.push(tokio::spawn(async move { - traffic_barrier.wait().await; - traffic_repository - .record_traffic(&ProxyNodeTrafficMutation { - node_id, - expected_tunnel_generation: Some(tunnel_generation), - total_requests_delta: 1, - failed_requests_delta: 1, - dns_failures_delta: 1, - stream_errors_delta: 1, - }) - .await - .expect("traffic delta should apply") - })); - - let request_repository = repository.clone(); - let request_barrier = Arc::clone(&barrier); - let node_id = manual.id.clone(); - tasks.push(tokio::spawn(async move { - request_barrier.wait().await; - request_repository - .increment_manual_node_requests(&node_id, 1, 1, Some(7)) - .await - .expect("manual request delta should apply"); - true - })); - - let heartbeat_repository = repository.clone(); - let heartbeat_barrier = Arc::clone(&barrier); - let node_id = tunnel.id.clone(); - let tunnel_generation = tunnel.tunnel_generation.clone(); - tasks.push(tokio::spawn(async move { - heartbeat_barrier.wait().await; - heartbeat_repository - .apply_heartbeat(&ProxyNodeHeartbeatMutation { - node_id, - expected_tunnel_generation: Some(tunnel_generation), - heartbeat_interval: None, - active_connections: None, - total_requests_delta: Some(1), - avg_latency_ms: None, - failed_requests_delta: Some(1), - dns_failures_delta: Some(1), - stream_errors_delta: Some(1), - proxy_metadata: None, - proxy_version: None, - }) - .await - .expect("heartbeat delta should apply") - .is_some() - })); - } - for task in tasks { - assert!(task.await.expect("mutation task should join")); - } - - let manual = repository - .find_proxy_node(&manual.id) - .await - .expect("manual node should read") - .expect("manual node should exist"); - assert_eq!(manual.total_requests, (WORKERS * 2) as i64); - assert_eq!(manual.failed_requests, (WORKERS * 2) as i64); - assert_eq!(manual.dns_failures, WORKERS as i64); - assert_eq!(manual.stream_errors, WORKERS as i64); - assert_eq!(manual.avg_latency_ms, Some(7.0)); - - let tunnel = repository - .find_proxy_node(&tunnel.id) - .await - .expect("tunnel node should read") - .expect("tunnel node should exist"); - assert_eq!(tunnel.total_requests, WORKERS as i64); - assert_eq!(tunnel.failed_requests, WORKERS as i64); - assert_eq!(tunnel.dns_failures, WORKERS as i64); - assert_eq!(tunnel.stream_errors, WORKERS as i64); - - sqlx::query("UPDATE proxy_nodes SET remote_config = ?, config_version = 9 WHERE id = ?") - .bind(r#"{"admin_revision":9}"#) - .bind(&manual.id) - .execute(&pool) - .await - .expect("manual sentinel fields should seed"); - repository - .update_manual_node(&ProxyNodeManualUpdateMutation { - node_id: manual.id.clone(), - name: Some("atomic-manual-renamed".to_string()), - ip: None, - port: None, - region: None, - proxy_url: None, - proxy_username: None, - proxy_password: None, - }) - .await - .expect("manual node should update") - .expect("manual node should remain"); - let after_manual_update = repository - .find_proxy_node(&manual.id) - .await - .expect("manual node should reload") - .expect("manual node should remain"); - assert_eq!(after_manual_update.total_requests, (WORKERS * 2) as i64); - assert_eq!(after_manual_update.config_version, 9); - assert_eq!( - after_manual_update.remote_config, - Some(json!({"admin_revision": 9})) - ); - - let base_transition = super::current_unix_secs().saturating_add(10_000); - let barrier = Arc::new(Barrier::new(WORKERS)); - let mut status_tasks = Vec::with_capacity(WORKERS); - for offset in 0..WORKERS { - let repository = repository.clone(); - let barrier = Arc::clone(&barrier); - let node_id = tunnel.id.clone(); - let tunnel_generation = tunnel.tunnel_generation.clone(); - status_tasks.push(tokio::spawn(async move { - barrier.wait().await; - repository - .update_tunnel_status(&ProxyNodeTunnelStatusMutation { - node_id, - expected_tunnel_generation: Some(tunnel_generation), - connected: offset % 2 == 0, - conn_count: 0, - detail: None, - observed_at_unix_secs: Some(base_transition + offset as u64), - }) - .await - .expect("status mutation should execute") - .is_some() - })); - } - for task in status_tasks { - assert!(task.await.expect("status task should join")); - } - let tunnel = repository - .find_proxy_node(&tunnel.id) - .await - .expect("tunnel node should reload") - .expect("tunnel node should remain"); - assert_eq!( - tunnel.tunnel_connected_at_unix_secs, - Some(base_transition + (WORKERS - 1) as u64) - ); - - let config_barrier = Arc::new(Barrier::new(2)); - let allowed_ports_update = { - let repository = repository.clone(); - let barrier = Arc::clone(&config_barrier); - let node_id = tunnel.id.clone(); - let tunnel_generation = tunnel.tunnel_generation.clone(); - tokio::spawn(async move { - barrier.wait().await; - repository - .update_remote_config(&ProxyNodeRemoteConfigMutation { - node_id, - expected_tunnel_generation: Some(tunnel_generation), - node_name: None, - allowed_ports: Some(vec![443, 8443]), - log_level: None, - heartbeat_interval: None, - scheduling_state: None, - upgrade_to: None, - }) - .await - }) - }; - let log_level_update = { - let repository = repository.clone(); - let barrier = Arc::clone(&config_barrier); - let node_id = tunnel.id.clone(); - let tunnel_generation = tunnel.tunnel_generation.clone(); - tokio::spawn(async move { - barrier.wait().await; - repository - .update_remote_config(&ProxyNodeRemoteConfigMutation { - node_id, - expected_tunnel_generation: Some(tunnel_generation), - node_name: None, - allowed_ports: None, - log_level: Some("trace".to_string()), - heartbeat_interval: None, - scheduling_state: None, - upgrade_to: None, - }) - .await - }) - }; - assert!(allowed_ports_update - .await - .expect("allowed ports task should join") - .expect("allowed ports update should execute") - .is_some()); - assert!(log_level_update - .await - .expect("log level task should join") - .expect("log level update should execute") - .is_some()); - let tunnel = repository - .find_proxy_node(&tunnel.id) - .await - .expect("tunnel config should reload") - .expect("tunnel node should remain"); - assert_eq!(tunnel.config_version, 2); - assert_eq!( - tunnel - .remote_config - .as_ref() - .and_then(|value| value.get("allowed_ports")), - Some(&json!([443, 8443])) - ); - assert_eq!( - tunnel - .remote_config - .as_ref() - .and_then(|value| value.get("log_level")), - Some(&json!("trace")) - ); - - let registration_barrier = Arc::new(Barrier::new(2)); - let mut registrations = Vec::new(); - for node_id in ["endpoint-race-a", "endpoint-race-b"] { - let repository = repository.clone(); - let barrier = Arc::clone(®istration_barrier); - registrations.push(tokio::spawn(async move { - barrier.wait().await; - repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some(node_id.to_string()), - name: node_id.to_string(), - ip: "127.0.0.42".to_string(), - port: 7042, - region: None, - heartbeat_interval: 30, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: None, - proxy_version: None, - registered_by: None, - tunnel_mode: true, - }) - .await - })); - } - let mut success_count = 0; - let mut identity_error_count = 0; - for task in registrations { - match task.await.expect("registration task should join") { - Ok(_) => success_count += 1, - Err(error) if error.to_string().contains("identity changed") => { - identity_error_count += 1 - } - Err(error) => panic!("unexpected concurrent registration error: {error}"), - } - } - assert_eq!(success_count, 1); - assert_eq!(identity_error_count, 1); - let endpoint_rows: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM proxy_nodes WHERE ip = ? AND port = ?") - .bind("127.0.0.42") - .bind(7042_i32) - .fetch_one(&pool) - .await - .expect("endpoint identity count should read"); - assert_eq!(endpoint_rows, 1); - - drop(repository); - drop(pool); - let _ = std::fs::remove_file(database_path); - } - - #[tokio::test] - async fn sqlite_repository_reads_proxy_nodes_and_events() { - 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"); - sqlx::query( - r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, status, heartbeat_interval, active_connections, - total_requests, failed_requests, dns_failures, stream_errors, - tunnel_mode, tunnel_connected, config_version, proxy_metadata, - hardware_info, remote_config, created_at, updated_at -) VALUES ( - 'node-1', 'test-generation-node-1', 'Node 1', '127.0.0.1', 8080, 'online', 30, 1, - 10, 2, 1, 0, 1, 1, 3, '{"version":"1.0.0"}', - '{"cpu":"m1"}', '{"log_level":"debug"}', 1, 2 -) -"#, - ) - .execute(&pool) - .await - .expect("proxy node should seed"); - sqlx::query( - r#" -INSERT INTO proxy_node_events (node_id, event_type, detail, created_at) -VALUES ('node-1', 'registered', 'ok', 3) -"#, - ) - .execute(&pool) - .await - .expect("proxy node event should seed"); - - let repository = SqliteProxyNodeReadRepository::new(pool); - let nodes = repository - .list_proxy_nodes() - .await - .expect("proxy nodes should list"); - assert_eq!(nodes.len(), 1); - assert_eq!(nodes[0].id, "node-1"); - assert_eq!( - nodes[0].proxy_metadata, - Some(serde_json::json!({"version": "1.0.0"})) - ); - - let node = repository - .find_proxy_node("node-1") - .await - .expect("proxy node should load") - .expect("proxy node should exist"); - assert_eq!(node.config_version, 3); - - let events = repository - .list_proxy_node_events("node-1", 10) - .await - .expect("events should list"); - assert_eq!(events.len(), 1); - assert_eq!(events[0].event_type, "registered"); - } - - #[tokio::test] - async fn sqlite_repository_writes_proxy_node_contract_views() { - 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 = SqliteProxyNodeReadRepository::new(pool); - let manual = repository - .create_manual_node(&ProxyNodeManualCreateMutation { - node_id: Some("manual-1-fixed-id".to_string()), - name: "manual-1".to_string(), - ip: "127.0.0.2".to_string(), - port: 8081, - region: Some("local".to_string()), - proxy_url: "http://127.0.0.2:8081".to_string(), - proxy_username: Some("user".to_string()), - proxy_password: Some("pass".to_string()), - registered_by: Some("admin".to_string()), - }) - .await - .expect("manual node should create"); - assert!(manual.is_manual); - assert_eq!(manual.id, "manual-1-fixed-id"); - assert_eq!(manual.status, "online"); - assert!(repository - .create_manual_node(&ProxyNodeManualCreateMutation { - node_id: Some("manual-1-fixed-id".to_string()), - name: "manual-id-collision".to_string(), - ip: "127.0.0.3".to_string(), - port: 8082, - region: None, - proxy_url: "http://127.0.0.3:8082".to_string(), - proxy_username: Some("attacker".to_string()), - proxy_password: Some("replacement-pass".to_string()), - registered_by: None, - }) - .await - .is_err()); - let after_collision = repository - .find_proxy_node("manual-1-fixed-id") - .await - .expect("manual node should reload") - .expect("manual node should remain"); - assert_eq!(after_collision.proxy_password.as_deref(), Some("pass")); - - let manual = repository - .update_manual_node(&ProxyNodeManualUpdateMutation { - node_id: manual.id.clone(), - name: Some("manual-updated".to_string()), - ip: None, - port: None, - region: Some("edge".to_string()), - proxy_url: None, - proxy_username: None, - proxy_password: None, - }) - .await - .expect("manual node should update") - .expect("manual node should exist"); - assert_eq!(manual.name, "manual-updated"); - assert_eq!(manual.region, Some("edge".to_string())); - - assert!(repository - .record_traffic(&ProxyNodeTrafficMutation { - node_id: manual.id.clone(), - expected_tunnel_generation: Some(manual.tunnel_generation.clone()), - total_requests_delta: 5, - failed_requests_delta: 1, - dns_failures_delta: 1, - stream_errors_delta: 0, - }) - .await - .expect("manual traffic should record")); - repository - .increment_manual_node_requests(&manual.id, 3, 1, Some(42)) - .await - .expect("manual request counters should increment"); - let manual = repository - .find_proxy_node(&manual.id) - .await - .expect("manual node should reload") - .expect("manual node should exist"); - assert_eq!(manual.total_requests, 8); - assert_eq!(manual.failed_requests, 2); - assert_eq!(manual.avg_latency_ms, Some(42.0)); - - let registered = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some("tunnel-1-fixed-id".to_string()), - name: "tunnel-1".to_string(), - ip: "10.0.0.1".to_string(), - port: 7000, - region: Some("us".to_string()), - heartbeat_interval: 30, - active_connections: Some(2), - total_requests: Some(10), - avg_latency_ms: Some(12.5), - hardware_info: Some(json!({"cpu":"m1"})), - estimated_max_concurrency: Some(100), - proxy_metadata: Some(json!({ - "arch":"arm64", - "tunnel_security": { - "mode": "non_tls_required", - "encryption_key_encrypted": "aether-proxy-node-secret-v2:aether-runtime-secret-v1:trusted-registration-key" - } - })), - proxy_version: Some("1.0.0".to_string()), - registered_by: Some("proxy".to_string()), - tunnel_mode: true, - }) - .await - .expect("tunnel node should register"); - assert!(!registered.is_manual); - assert_eq!(registered.id, "tunnel-1-fixed-id"); - assert!(registered.tunnel_mode); - assert_eq!( - registered - .proxy_metadata - .as_ref() - .and_then(|value| value.get("version")) - .and_then(serde_json::Value::as_str), - Some("1.0.0") - ); - - let configured = repository - .update_remote_config(&ProxyNodeRemoteConfigMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - node_name: Some("tunnel-renamed".to_string()), - allowed_ports: Some(vec![443, 8443]), - log_level: Some("debug".to_string()), - heartbeat_interval: Some(45), - scheduling_state: Some(Some("draining".to_string())), - upgrade_to: Some(Some("tunnel-v2.0.0".to_string())), - }) - .await - .expect("remote config should update") - .expect("tunnel node should exist"); - assert_eq!(configured.name, "tunnel-renamed"); - assert_eq!(configured.config_version, 1); - - let heartbeat = repository - .apply_heartbeat(&ProxyNodeHeartbeatMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - heartbeat_interval: Some(45), - active_connections: Some(4), - total_requests_delta: Some(6), - avg_latency_ms: Some(10.0), - failed_requests_delta: Some(1), - dns_failures_delta: Some(0), - stream_errors_delta: Some(2), - proxy_metadata: Some(json!({ - "arch":"arm64", - "tunnel_security": { - "mode": "disabled", - "encryption_key": "heartbeat-attacker-controlled" - } - })), - proxy_version: Some("2.0.0".to_string()), - }) - .await - .expect("heartbeat should apply") - .expect("tunnel node should exist"); - assert_eq!(heartbeat.status, "online"); - assert!(heartbeat.tunnel_connected); - assert_eq!(heartbeat.active_connections, 4); - assert_eq!(heartbeat.total_requests, 16); - assert_eq!(heartbeat.config_version, 2); - assert_eq!( - heartbeat - .proxy_metadata - .as_ref() - .and_then(|value| value.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:trusted-registration-key") - ); - assert!(heartbeat - .remote_config - .as_ref() - .and_then(|value| value.get("upgrade_to")) - .is_none()); - - let stale = repository - .update_tunnel_status(&ProxyNodeTunnelStatusMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - connected: false, - conn_count: 0, - detail: None, - observed_at_unix_secs: Some(1), - }) - .await - .expect("stale tunnel status should be recorded") - .expect("tunnel node should exist"); - assert_eq!(stale.status, "online"); - - let disconnected = repository - .update_tunnel_status(&ProxyNodeTunnelStatusMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - connected: false, - conn_count: 0, - detail: Some("closed".to_string()), - observed_at_unix_secs: Some(1_800_000_000), - }) - .await - .expect("tunnel status should update") - .expect("tunnel node should exist"); - assert_eq!(disconnected.status, "offline"); - assert!(!disconnected.tunnel_connected); - - let events = repository - .list_proxy_node_events(®istered.id, 10) - .await - .expect("events should list"); - assert_eq!(events.len(), 2); - assert_eq!(events[0].detail.as_deref(), Some("closed")); - assert_eq!( - events[1].detail.as_deref(), - Some("[stale_ignored] [tunnel_node_status] conn_count=0") - ); - - assert_eq!( - repository - .reset_stale_tunnel_statuses() - .await - .expect("stale tunnels should reset"), - 0 - ); - assert!(repository - .unregister_node(®istered.id) - .await - .expect("node should unregister") - .is_some()); - assert!(repository - .delete_node(®istered.id) - .await - .expect("node should delete") - .is_some()); - assert!(repository - .delete_node(&manual.id) - .await - .expect("manual node should delete") - .is_some()); - } - - #[tokio::test] - async fn sqlite_repository_aggregates_proxy_node_metrics_and_filters_events() { - 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 = SqliteProxyNodeReadRepository::new(pool); - let registered = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: None, - name: "tunnel-1".to_string(), - ip: "10.0.0.1".to_string(), - port: 7000, - region: None, - heartbeat_interval: 30, - active_connections: Some(0), - total_requests: Some(0), - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: None, - proxy_version: Some("1.0.0".to_string()), - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("node should register"); - let now = super::current_unix_secs(); - repository - .apply_heartbeat(&ProxyNodeHeartbeatMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - heartbeat_interval: Some(30), - active_connections: Some(0), - total_requests_delta: None, - avg_latency_ms: None, - failed_requests_delta: None, - dns_failures_delta: None, - stream_errors_delta: None, - proxy_metadata: Some(json!({ - "tunnel_metrics": { - "connect_errors": 0, - "disconnects": 0, - "error_events_total": 0, - "ws_in_bytes": 0, - "ws_out_bytes": 0, - "ws_in_frames": 0, - "ws_out_frames": 0, - "heartbeat_rtt_last_ms": 0 - } - })), - proxy_version: Some("1.0.0".to_string()), - }) - .await - .expect("baseline heartbeat should apply") - .expect("node should exist"); - repository - .apply_heartbeat(&ProxyNodeHeartbeatMutation { - node_id: registered.id.clone(), - expected_tunnel_generation: None, - heartbeat_interval: Some(30), - active_connections: Some(5), - total_requests_delta: None, - avg_latency_ms: None, - failed_requests_delta: None, - dns_failures_delta: None, - stream_errors_delta: None, - proxy_metadata: Some(json!({ - "tunnel_metrics": { - "connect_errors": 4, - "disconnects": 1, - "error_events_total": 1, - "ws_in_bytes": 100, - "ws_out_bytes": 200, - "ws_in_frames": 3, - "ws_out_frames": 6, - "heartbeat_rtt_last_ms": 33 - }, - "recent_tunnel_errors": [{ - "timestamp_unix_secs": now, - "category": "tcp_connect_timeout", - "message": "timeout" - }] - })), - proxy_version: Some("1.0.0".to_string()), - }) - .await - .expect("heartbeat should apply") - .expect("node should exist"); - - let metrics = repository - .list_proxy_node_metrics( - ®istered.id, - ProxyNodeMetricsStep::OneMinute, - now.saturating_sub(120), - now.saturating_add(120), - 10, - ) - .await - .expect("metrics should list"); - assert!(!metrics.is_empty()); - assert_eq!(metrics.iter().map(|bucket| bucket.samples).sum::(), 2); - assert_eq!( - metrics - .iter() - .map(|bucket| bucket.uptime_samples) - .sum::(), - 2 - ); - assert_eq!( - metrics - .iter() - .map(|bucket| bucket.active_connections_max) - .max() - .unwrap_or_default(), - 5 - ); - assert_eq!( - metrics - .iter() - .map(|bucket| bucket.heartbeat_rtt_ms_sum) - .sum::(), - 33 - ); - assert_eq!( - metrics - .iter() - .map(|bucket| bucket.connect_errors_delta) - .sum::(), - 4 - ); - assert_eq!( - metrics - .iter() - .map(|bucket| bucket.ws_out_frames_delta) - .sum::(), - 6 - ); - - let fleet = repository - .list_proxy_fleet_metrics( - ProxyNodeMetricsStep::OneMinute, - now.saturating_sub(120), - now.saturating_add(120), - 10, - ) - .await - .expect("fleet metrics should list"); - assert!(!fleet.is_empty()); - assert_eq!(fleet.iter().map(|bucket| bucket.samples).sum::(), 2); - assert_eq!( - fleet - .iter() - .map(|bucket| bucket.error_events_delta) - .sum::(), - 1 - ); - - let events = repository - .list_proxy_node_events_filtered( - ®istered.id, - &ProxyNodeEventQuery { - limit: 10, - from_unix_secs: Some(now.saturating_sub(120)), - to_unix_secs: Some(now.saturating_add(120)), - event_type: Some(PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR.to_string()), - }, - ) - .await - .expect("events should list"); - assert_eq!(events.len(), 1); - assert_eq!(events[0].event_type, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR); - assert_eq!( - events[0] - .event_metadata - .as_ref() - .and_then(|value| value.get("category")) - .and_then(serde_json::Value::as_str), - Some("tcp_connect_timeout") - ); - - let cleanup = repository - .cleanup_proxy_node_metrics(now.saturating_add(1), now.saturating_add(1), 10) - .await - .expect("cleanup should run"); - assert_eq!(cleanup.deleted_1m_rows, 1); - assert_eq!(cleanup.deleted_1h_rows, 1); - } - - #[tokio::test] - async fn sqlite_registration_preserves_omitted_security_and_allows_rotation() { - 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 = SqliteProxyNodeReadRepository::new(pool); - - let first = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some("registration-security-node".to_string()), - name: "registration-security-node".to_string(), - ip: "127.0.0.70".to_string(), - port: 7070, - region: None, - heartbeat_interval: 30, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({ - "version": "1.0.0", - "tunnel_security": { - "mode": "non_tls_required", - "encryption_key_encrypted": "aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-old" - } - })), - proxy_version: None, - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("first registration should succeed"); - - let refreshed = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some(first.id.clone()), - name: "registration-security-node-refreshed".to_string(), - ip: "127.0.0.70".to_string(), - port: 7070, - region: None, - heartbeat_interval: 45, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({"runtime": "refreshed"})), - proxy_version: Some("2.0.0".to_string()), - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("metadata-only re-registration should succeed"); - assert_eq!(refreshed.id, first.id); - assert_eq!( - refreshed - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-old") - ); - assert_eq!( - refreshed - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.get("runtime")), - Some(&json!("refreshed")) - ); - - let rotated = repository - .register_node(&ProxyNodeRegistrationMutation { - node_id: Some(first.id.clone()), - name: "registration-security-node-rotated".to_string(), - ip: "127.0.0.70".to_string(), - port: 7070, - region: None, - heartbeat_interval: 45, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({ - "tunnel_security": { - "mode": "non_tls_required", - "encryption_key_encrypted": "aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new" - } - })), - proxy_version: Some("2.1.0".to_string()), - registered_by: None, - tunnel_mode: true, - }) - .await - .expect("explicit security rotation should succeed"); - assert_eq!( - rotated - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new") - ); - - let stale_refresh = ProxyNodeRegistrationMutation { - node_id: Some(first.id.clone()), - name: "registration-security-node-stale".to_string(), - ip: "127.0.0.70".to_string(), - port: 7070, - region: None, - heartbeat_interval: 45, - active_connections: None, - total_requests: None, - avg_latency_ms: None, - hardware_info: None, - estimated_max_concurrency: None, - proxy_metadata: Some(json!({"runtime": "stale-writer"})), - proxy_version: Some("2.2.0".to_string()), - registered_by: None, - tunnel_mode: true, - }; - let stale_replacement = merge_proxy_metadata_for_registration( - refreshed.proxy_metadata.as_ref(), - normalize_proxy_metadata( - stale_refresh.proxy_metadata.as_ref(), - stale_refresh.proxy_version.as_deref(), - ), - ); - assert!(!repository - .update_existing_registration_if_unchanged( - &stale_refresh, - &refreshed, - stale_replacement.as_ref(), - super::current_unix_secs(), - ) - .await - .expect("stale registration CAS should execute")); - - let committed_refresh = repository - .register_node(&ProxyNodeRegistrationMutation { - name: "registration-security-node-committed".to_string(), - proxy_metadata: Some(json!({"runtime": "committed-after-rotation"})), - ..stale_refresh - }) - .await - .expect("metadata refresh should retry from current security state"); - assert_eq!( - committed_refresh - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key_encrypted")) - .and_then(serde_json::Value::as_str), - Some("aether-proxy-node-secret-v2:aether-runtime-secret-v1:sealed-new") - ); - assert_eq!( - committed_refresh - .proxy_metadata - .as_ref() - .and_then(|metadata| metadata.get("runtime")), - Some(&json!("committed-after-rotation")) - ); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/quota.rs b/crates/aether-data/adapters/sqlite/src/quota.rs deleted file mode 100644 index 5140f03e1..000000000 --- a/crates/aether-data/adapters/sqlite/src/quota.rs +++ /dev/null @@ -1,217 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, Row, Sqlite}; - -use aether_data_contracts::repository::quota::{ - ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot, -}; -use aether_data_query::{DialectSql, SelectColumn, SelectQuery, SqlDialect}; - -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, sqlite_real, DataLayerError, SqlitePool}; - -fn quota_snapshot_select() -> SelectQuery<'static> { - SelectQuery::new("providers").select_columns([ - SelectColumn::expr("id").alias("provider_id"), - SelectColumn::expr( - DialectSql::common("billing_type").with_postgres("CAST(billing_type AS TEXT)"), - ) - .alias("billing_type"), - SelectColumn::expr(DialectSql::dialect( - "CAST(monthly_quota_usd AS DOUBLE PRECISION)", - "CAST(monthly_quota_usd AS REAL)", - )) - .alias("monthly_quota_usd"), - SelectColumn::expr(DialectSql::dialect( - "CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE PRECISION)", - "CAST(COALESCE(monthly_used_usd, 0) AS REAL)", - )) - .alias("monthly_used_usd"), - SelectColumn::expr("quota_reset_day"), - SelectColumn::expr(DialectSql::dialect( - "CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT)", - "quota_last_reset_at", - )) - .alias("quota_last_reset_at_unix_secs"), - SelectColumn::expr(DialectSql::dialect( - "CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT)", - "quota_expires_at", - )) - .alias("quota_expires_at_unix_secs"), - SelectColumn::expr("is_active"), - ]) -} - -#[derive(Debug, Clone)] -pub struct SqliteProviderQuotaRepository { - pool: SqlitePool, -} - -impl SqliteProviderQuotaRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl ProviderQuotaReadRepository for SqliteProviderQuotaRepository { - async fn find_by_provider_id( - &self, - provider_id: &str, - ) -> Result, DataLayerError> { - let mut statement = quota_snapshot_select().statement::(SqlDialect::Sqlite); - statement.where_eq("id", provider_id.to_string()).limit(1); - let row = statement - .finish() - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_row).transpose() - } - - async fn find_by_provider_ids( - &self, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut statement = quota_snapshot_select().statement::(SqlDialect::Sqlite); - statement - .where_in("id", provider_ids) - .order_by_sql("id ASC"); - let rows = statement - .finish() - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_row).collect() - } -} - -#[async_trait] -impl ProviderQuotaWriteRepository for SqliteProviderQuotaRepository { - async fn reset_due(&self, now_unix_secs: u64) -> Result { - let now = i64::try_from(now_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("provider quota reset timestamp overflow".to_string()) - })?; - let rows_affected = sqlx::query( - r#" -UPDATE providers -SET monthly_used_usd = 0.0, - quota_last_reset_at = ?, - updated_at = ? -WHERE billing_type = 'monthly_quota' - AND is_active = 1 - AND ( - quota_last_reset_at IS NULL - OR (? - quota_last_reset_at) >= (quota_reset_day * 86400) - ) -"#, - ) - .bind(now) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - Ok(usize::try_from(rows_affected).unwrap_or_default()) - } -} - -fn map_row(row: &SqliteRow) -> Result { - StoredProviderQuotaSnapshot::new( - row.try_get("provider_id").map_sql_err()?, - row.try_get("billing_type").map_sql_err()?, - sqlite_optional_real(row, "monthly_quota_usd")?, - sqlite_real(row, "monthly_used_usd")?, - row.try_get("quota_reset_day").map_sql_err()?, - row.try_get("quota_last_reset_at_unix_secs").map_sql_err()?, - row.try_get("quota_expires_at_unix_secs").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - ) -} - -#[cfg(test)] -mod tests { - use super::SqliteProviderQuotaRepository; - use aether_data_contracts::repository::quota::{ - ProviderQuotaReadRepository, ProviderQuotaWriteRepository, - }; - - use crate::run_migrations; - - #[tokio::test] - async fn sqlite_repository_reads_and_resets_provider_quotas() { - 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_quotas(&pool).await; - - let repository = SqliteProviderQuotaRepository::new(pool); - let quota = repository - .find_by_provider_id("provider-1") - .await - .expect("quota should load") - .expect("quota should exist"); - assert_eq!(quota.monthly_used_usd, 5.0); - - let quota = repository - .find_by_provider_id("provider-null-used") - .await - .expect("quota with null usage should load") - .expect("quota with null usage should exist"); - assert_eq!(quota.monthly_used_usd, 0.0); - - let quotas = repository - .find_by_provider_ids(&["provider-2".to_string(), "provider-1".to_string()]) - .await - .expect("quotas should load"); - assert_eq!( - quotas - .iter() - .map(|quota| quota.provider_id.as_str()) - .collect::>(), - vec!["provider-1", "provider-2"] - ); - - let reset = repository - .reset_due(1_000 + 7 * 24 * 60 * 60) - .await - .expect("quota reset should run"); - assert_eq!(reset, 1); - let quota = repository - .find_by_provider_id("provider-1") - .await - .expect("quota should reload") - .expect("quota should exist"); - assert_eq!(quota.monthly_used_usd, 0.0); - assert_eq!(quota.quota_last_reset_at_unix_secs, Some(605_800)); - } - - async fn seed_provider_quotas(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO providers ( - id, name, provider_type, billing_type, monthly_quota_usd, monthly_used_usd, - quota_reset_day, quota_last_reset_at, is_active, created_at, updated_at -) -VALUES - ('provider-1', 'Provider One', 'openai', 'monthly_quota', 20.0, 5.0, 7, 1000, 1, 1, 1), - ('provider-2', 'Provider Two', 'openai', 'payg', NULL, 1.5, NULL, NULL, 1, 1, 1), - ('provider-null-used', 'Provider Null Used', 'openai', 'payg', NULL, NULL, NULL, NULL, 1, 1, 1) -"#, - ) - .execute(pool) - .await - .expect("providers should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/routing_profiles.rs b/crates/aether-data/adapters/sqlite/src/routing_profiles.rs deleted file mode 100644 index e1fa9778b..000000000 --- a/crates/aether-data/adapters/sqlite/src/routing_profiles.rs +++ /dev/null @@ -1,856 +0,0 @@ -use async_trait::async_trait; -use serde_json::Value; -use sqlx::{sqlite::SqliteRow, Row}; - -use aether_data_contracts::repository::routing_profiles::*; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::pool::SqlitePool; - -const ROUTING_GROUP_SELECT: &str = r#" -SELECT - id, - name, - description, - enabled, - is_system_default, - sort_order, - config_json, - version, - created_at, - updated_at, - published_at -FROM routing_groups -"#; - -const ROUTING_GROUP_BINDING_SELECT: &str = r#" -SELECT - id, - group_id, - subject_type, - subject_id, - is_default, - allow_explicit_select, - created_at, - updated_at -FROM routing_group_bindings -"#; - -const ROUTING_GROUP_VERSION_SELECT: &str = r#" -SELECT - id, - group_id, - version, - config_json, - created_at, - created_by -FROM routing_group_versions -"#; - -#[derive(Debug, Clone)] -pub struct SqliteRoutingGroupRepository { - pool: SqlitePool, -} - -impl SqliteRoutingGroupRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl RoutingGroupReadRepository for SqliteRoutingGroupRepository { - async fn list_routing_groups(&self) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} ORDER BY enabled DESC, sort_order ASC, name ASC, id ASC" - )) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_group_row).collect() - } - - async fn find_routing_group( - &self, - lookup: RoutingGroupLookupKey<'_>, - ) -> Result, DataLayerError> { - let row = match lookup { - RoutingGroupLookupKey::Id(id) => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1" - )) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - RoutingGroupLookupKey::Name(name) => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE name = ? LIMIT 1" - )) - .bind(name) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!( - "{ROUTING_GROUP_SELECT} WHERE is_system_default = 1 AND enabled = 1 ORDER BY updated_at DESC, id ASC LIMIT 1" - )) - .fetch_optional(&self.pool) - .await - .map_sql_err()?, - }; - row.as_ref().map(map_group_row).transpose() - } - - async fn list_routing_group_bindings( - &self, - query: &RoutingGroupBindingQuery, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - r#" -{ROUTING_GROUP_BINDING_SELECT} -WHERE (? IS NULL OR group_id = ?) - AND (? IS NULL OR subject_type = ?) - AND (? IS NULL OR subject_id = ?) -ORDER BY created_at ASC, id ASC -"# - )) - .bind(query.group_id.as_deref()) - .bind(query.group_id.as_deref()) - .bind(query.subject_type.map(binding_subject_to_database)) - .bind(query.subject_type.map(binding_subject_to_database)) - .bind(query.subject_id.as_deref()) - .bind(query.subject_id.as_deref()) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_binding_row).collect() - } - - async fn has_any_routing_group_binding(&self) -> Result { - let row = sqlx::query("SELECT 1 FROM routing_group_bindings LIMIT 1") - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(row.is_some()) - } - - async fn list_routing_group_versions( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let rows = sqlx::query(&format!( - "{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = ? ORDER BY version DESC, created_at DESC, id ASC" - )) - .bind(group_id) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_version_row).collect() - } -} - -#[async_trait] -impl RoutingGroupWriteRepository for SqliteRoutingGroupRepository { - async fn create_routing_group( - &self, - record: CreateRoutingGroupRecord, - ) -> Result { - let group = StoredRoutingGroup::new(record)?; - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("UPDATE routing_groups SET is_system_default = is_system_default WHERE 0") - .execute(&mut *tx) - .await - .map_sql_err()?; - if group.is_system_default { - sqlx::query( - "UPDATE routing_groups SET is_system_default = 0 WHERE is_system_default = 1", - ) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -INSERT INTO routing_groups ( - id, name, description, enabled, is_system_default, sort_order, config_json, - version, created_at, updated_at, published_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&group.id) - .bind(&group.name) - .bind(&group.description) - .bind(group.enabled) - .bind(group.is_system_default) - .bind(group.sort_order) - .bind(json_to_string( - &group.config_json, - "routing_groups.config_json", - )?) - .bind(group.version) - .bind(group.created_at) - .bind(group.updated_at) - .bind(group.published_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(group) - } - - async fn update_routing_group( - &self, - id: &str, - patch: UpdateRoutingGroupRecord, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("UPDATE routing_groups SET is_system_default = is_system_default WHERE 0") - .execute(&mut *tx) - .await - .map_sql_err()?; - let row = sqlx::query(&format!("{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1")) - .bind(id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(mut group) = row.as_ref().map(map_group_row).transpose()? else { - return Ok(None); - }; - apply_group_patch(&mut group, patch)?; - if group.is_system_default { - sqlx::query( - "UPDATE routing_groups SET is_system_default = 0 WHERE is_system_default = 1 AND id <> ?", - ) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -UPDATE routing_groups -SET name = ?, - description = ?, - enabled = ?, - is_system_default = ?, - sort_order = ?, - config_json = ?, - version = ?, - updated_at = ?, - published_at = ? -WHERE id = ? -"#, - ) - .bind(&group.name) - .bind(&group.description) - .bind(group.enabled) - .bind(group.is_system_default) - .bind(group.sort_order) - .bind(json_to_string( - &group.config_json, - "routing_groups.config_json", - )?) - .bind(group.version) - .bind(group.updated_at) - .bind(group.published_at) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(Some(group)) - } - - async fn delete_routing_group(&self, id: &str) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM routing_group_versions WHERE group_id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - tx.commit().await.map_sql_err()?; - Ok(rows_affected > 0) - } - - async fn create_routing_group_binding( - &self, - record: CreateRoutingGroupBindingRecord, - ) -> Result { - let binding = StoredRoutingGroupBinding::new(record)?; - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("UPDATE routing_group_bindings SET is_default = is_default WHERE 0") - .execute(&mut *tx) - .await - .map_sql_err()?; - if binding.is_default { - sqlx::query( - r#" -UPDATE routing_group_bindings -SET is_default = 0 -WHERE is_default = 1 AND subject_type = ? AND subject_id = ? -"#, - ) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -INSERT INTO routing_group_bindings ( - id, group_id, subject_type, subject_id, is_default, - allow_explicit_select, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&binding.id) - .bind(&binding.group_id) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(binding.is_default) - .bind(binding.allow_explicit_select) - .bind(binding.created_at) - .bind(binding.updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(binding) - } - - async fn delete_routing_group_binding(&self, id: &str) -> Result { - Ok( - sqlx::query("DELETE FROM routing_group_bindings WHERE id = ?") - .bind(id) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected() - > 0, - ) - } - - async fn update_routing_group_binding( - &self, - id: &str, - patch: UpdateRoutingGroupBindingRecord, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - sqlx::query("UPDATE routing_group_bindings SET is_default = is_default WHERE 0") - .execute(&mut *tx) - .await - .map_sql_err()?; - let row = sqlx::query(&format!( - "{ROUTING_GROUP_BINDING_SELECT} WHERE id = ? LIMIT 1" - )) - .bind(id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(mut binding) = row.as_ref().map(map_binding_row).transpose()? else { - return Ok(None); - }; - apply_binding_patch(&mut binding, patch)?; - if binding.is_default { - sqlx::query( - r#" -UPDATE routing_group_bindings -SET is_default = 0 -WHERE is_default = 1 - AND subject_type = ? - AND subject_id = ? - AND id <> ? -"#, - ) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query( - r#" -UPDATE routing_group_bindings -SET group_id = ?, - subject_type = ?, - subject_id = ?, - is_default = ?, - allow_explicit_select = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&binding.group_id) - .bind(binding_subject_to_database(binding.subject_type)) - .bind(&binding.subject_id) - .bind(binding.is_default) - .bind(binding.allow_explicit_select) - .bind(binding.updated_at) - .bind(id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(Some(binding)) - } - - async fn create_routing_group_version( - &self, - record: CreateRoutingGroupVersionRecord, - ) -> Result { - let version = StoredRoutingGroupVersion::new(record)?; - sqlx::query( - r#" -INSERT INTO routing_group_versions ( - id, group_id, version, config_json, created_at, created_by -) -VALUES (?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&version.id) - .bind(&version.group_id) - .bind(version.version) - .bind(json_to_string( - &version.config_json, - "routing_group_versions.config_json", - )?) - .bind(version.created_at) - .bind(&version.created_by) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(version) - } -} - -fn map_group_row(row: &SqliteRow) -> Result { - Ok(StoredRoutingGroup { - id: row.try_get("id").map_sql_err()?, - name: row.try_get("name").map_sql_err()?, - description: row.try_get("description").map_sql_err()?, - enabled: row.try_get("enabled").map_sql_err()?, - is_system_default: row.try_get("is_system_default").map_sql_err()?, - sort_order: row.try_get("sort_order").map_sql_err()?, - config_json: json_from_string( - row.try_get("config_json").map_sql_err()?, - "routing_groups.config_json", - )?, - version: row.try_get("version").map_sql_err()?, - created_at: row.try_get("created_at").map_sql_err()?, - updated_at: row.try_get("updated_at").map_sql_err()?, - published_at: row.try_get("published_at").map_sql_err()?, - }) -} - -fn map_binding_row(row: &SqliteRow) -> Result { - Ok(StoredRoutingGroupBinding { - id: row.try_get("id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - subject_type: binding_subject_from_database(row.try_get("subject_type").map_sql_err()?)?, - subject_id: row.try_get("subject_id").map_sql_err()?, - is_default: row.try_get("is_default").map_sql_err()?, - allow_explicit_select: row.try_get("allow_explicit_select").map_sql_err()?, - created_at: row.try_get("created_at").map_sql_err()?, - updated_at: row.try_get("updated_at").map_sql_err()?, - }) -} - -fn map_version_row(row: &SqliteRow) -> Result { - Ok(StoredRoutingGroupVersion { - id: row.try_get("id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - version: row.try_get("version").map_sql_err()?, - config_json: json_from_string( - row.try_get("config_json").map_sql_err()?, - "routing_group_versions.config_json", - )?, - created_at: row.try_get("created_at").map_sql_err()?, - created_by: row.try_get("created_by").map_sql_err()?, - }) -} - -fn json_to_string(value: &Value, field_name: &str) -> Result { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains unserializable JSON: {err}")) - }) -} - -fn json_from_string(value: String, field_name: &str) -> Result { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}")) - }) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - use crate::run_migrations as run_sqlite_migrations; - - #[tokio::test] - async fn sqlite_routing_group_repository_round_trips() { - 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 repository = SqliteRoutingGroupRepository::new(pool); - repository - .create_routing_group(CreateRoutingGroupRecord { - id: "routing-group-1".to_string(), - name: "default".to_string(), - description: Some("initial".to_string()), - enabled: true, - is_system_default: true, - sort_order: 0, - config_json: json!({"allowed_models": ["gpt-*"]}), - version: 1, - created_at: 10, - updated_at: 10, - published_at: None, - }) - .await - .expect("group should create"); - - let system_default = repository - .find_routing_group(RoutingGroupLookupKey::SystemDefault) - .await - .expect("group lookup should succeed") - .expect("system default should exist"); - assert_eq!(system_default.id, "routing-group-1"); - - repository - .update_routing_group( - "routing-group-1", - UpdateRoutingGroupRecord { - description: Some(None), - version: Some(2), - updated_at: 20, - published_at: Some(Some(20)), - ..UpdateRoutingGroupRecord::default() - }, - ) - .await - .expect("group should update"); - - let binding = repository - .create_routing_group_binding(CreateRoutingGroupBindingRecord { - id: "binding-1".to_string(), - group_id: "routing-group-1".to_string(), - subject_type: RoutingGroupBindingSubject::ApiKey, - subject_id: "api-key-1".to_string(), - is_default: true, - allow_explicit_select: true, - created_at: 10, - updated_at: 10, - }) - .await - .expect("binding should create"); - - assert_eq!(binding.subject_type, RoutingGroupBindingSubject::ApiKey); - assert_eq!( - repository - .list_routing_group_bindings(&RoutingGroupBindingQuery { - group_id: Some("routing-group-1".to_string()), - subject_type: Some(RoutingGroupBindingSubject::ApiKey), - subject_id: Some("api-key-1".to_string()), - }) - .await - .expect("bindings should list") - .len(), - 1 - ); - - repository - .create_routing_group_version(CreateRoutingGroupVersionRecord { - id: "version-1".to_string(), - group_id: "routing-group-1".to_string(), - version: 2, - config_json: json!({"allowed_models": ["gpt-*"]}), - created_at: 20, - created_by: Some("admin".to_string()), - }) - .await - .expect("version should create"); - - assert_eq!( - repository - .list_routing_group_versions("routing-group-1") - .await - .expect("versions should list") - .len(), - 1 - ); - } - - #[tokio::test] - async fn sqlite_keeps_system_and_subject_defaults_unique() { - 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 repository = SqliteRoutingGroupRepository::new(pool); - - for (id, is_system_default) in [("group-1", true), ("group-2", true), ("group-3", false)] { - repository - .create_routing_group(group_record(id, is_system_default)) - .await - .expect("group should create"); - } - assert_eq!(system_default_ids(&repository).await, vec!["group-2"]); - - repository - .update_routing_group( - "group-1", - UpdateRoutingGroupRecord { - is_system_default: Some(true), - updated_at: 2, - ..UpdateRoutingGroupRecord::default() - }, - ) - .await - .expect("group should update"); - assert_eq!(system_default_ids(&repository).await, vec!["group-1"]); - - repository - .create_routing_group_binding(binding_record("binding-1", "group-1", "subject-1", true)) - .await - .expect("binding should create"); - repository - .create_routing_group_binding(binding_record("binding-2", "group-2", "subject-1", true)) - .await - .expect("binding should create"); - repository - .create_routing_group_binding(binding_record("binding-3", "group-3", "subject-2", true)) - .await - .expect("binding should create"); - - assert_eq!( - default_binding_ids(&repository, "subject-1").await, - vec!["binding-2"] - ); - assert_eq!( - default_binding_ids(&repository, "subject-2").await, - vec!["binding-3"] - ); - - repository - .update_routing_group_binding( - "binding-1", - UpdateRoutingGroupBindingRecord { - is_default: Some(true), - updated_at: 2, - ..UpdateRoutingGroupBindingRecord::default() - }, - ) - .await - .expect("binding should update"); - assert_eq!( - default_binding_ids(&repository, "subject-1").await, - vec!["binding-1"] - ); - assert_eq!( - default_binding_ids(&repository, "subject-2").await, - vec!["binding-3"] - ); - - repository - .update_routing_group_binding( - "binding-3", - UpdateRoutingGroupBindingRecord { - subject_id: Some("subject-1".to_string()), - updated_at: 3, - ..UpdateRoutingGroupBindingRecord::default() - }, - ) - .await - .expect("binding should move"); - assert_eq!( - default_binding_ids(&repository, "subject-1").await, - vec!["binding-3"] - ); - assert!(default_binding_ids(&repository, "subject-2") - .await - .is_empty()); - } - - #[tokio::test] - async fn sqlite_repair_migration_resolves_existing_duplicate_defaults() { - 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 repository = SqliteRoutingGroupRepository::new(pool.clone()); - sqlx::raw_sql( - r#" -DROP INDEX routing_groups_one_system_default_key; -DROP INDEX routing_group_bindings_subject_default_key; -"#, - ) - .execute(&pool) - .await - .expect("unique indexes should be removable to simulate a pre-repair database"); - - for id in ["group-1", "group-2", "group-3"] { - repository - .create_routing_group(group_record(id, false)) - .await - .expect("group should create"); - } - sqlx::query( - r#" -UPDATE routing_groups -SET is_system_default = 1, - enabled = CASE id WHEN 'group-3' THEN 0 ELSE 1 END, - updated_at = CASE id - WHEN 'group-1' THEN 1 - WHEN 'group-2' THEN 2 - ELSE 3 - END -"#, - ) - .execute(&pool) - .await - .expect("duplicate system defaults should seed"); - - for (id, subject_id) in [ - ("binding-3", "subject-1"), - ("binding-2", "subject-1"), - ("binding-1", "subject-1"), - ("binding-4", "subject-2"), - ] { - repository - .create_routing_group_binding(binding_record(id, "group-1", subject_id, false)) - .await - .expect("binding should create"); - } - sqlx::query( - r#" -UPDATE routing_group_bindings -SET is_default = 1, - created_at = CASE id WHEN 'binding-3' THEN 2 ELSE 1 END -"#, - ) - .execute(&pool) - .await - .expect("duplicate binding defaults should seed"); - - let repair_migration = - include_str!("../migrations/20260727000000_repair_routing_default_uniqueness.sql"); - for _ in 0..2 { - sqlx::raw_sql(repair_migration) - .execute(&pool) - .await - .expect("repair migration should be idempotent"); - } - - assert_eq!(system_default_ids(&repository).await, vec!["group-2"]); - assert_eq!( - default_binding_ids(&repository, "subject-1").await, - vec!["binding-1"] - ); - assert_eq!( - default_binding_ids(&repository, "subject-2").await, - vec!["binding-4"] - ); - - sqlx::query("UPDATE routing_groups SET is_system_default = 1 WHERE id = 'group-3'") - .execute(&pool) - .await - .expect_err("database should reject a second system default"); - sqlx::query("UPDATE routing_group_bindings SET is_default = 1 WHERE id = 'binding-2'") - .execute(&pool) - .await - .expect_err("database should reject a second default for the same subject"); - } - - fn group_record(id: &str, is_system_default: bool) -> CreateRoutingGroupRecord { - CreateRoutingGroupRecord { - id: id.to_string(), - name: id.to_string(), - description: None, - enabled: true, - is_system_default, - sort_order: 0, - config_json: json!({}), - version: 1, - created_at: 1, - updated_at: 1, - published_at: None, - } - } - - fn binding_record( - id: &str, - group_id: &str, - subject_id: &str, - is_default: bool, - ) -> CreateRoutingGroupBindingRecord { - CreateRoutingGroupBindingRecord { - id: id.to_string(), - group_id: group_id.to_string(), - subject_type: RoutingGroupBindingSubject::ApiKey, - subject_id: subject_id.to_string(), - is_default, - allow_explicit_select: true, - created_at: 1, - updated_at: 1, - } - } - - async fn system_default_ids(repository: &SqliteRoutingGroupRepository) -> Vec { - repository - .list_routing_groups() - .await - .expect("groups should list") - .into_iter() - .filter(|group| group.is_system_default) - .map(|group| group.id) - .collect() - } - - async fn default_binding_ids( - repository: &SqliteRoutingGroupRepository, - subject_id: &str, - ) -> Vec { - repository - .list_routing_group_bindings(&RoutingGroupBindingQuery { - group_id: None, - subject_type: Some(RoutingGroupBindingSubject::ApiKey), - subject_id: Some(subject_id.to_string()), - }) - .await - .expect("bindings should list") - .into_iter() - .filter(|binding| binding.is_default) - .map(|binding| binding.id) - .collect() - } -} diff --git a/crates/aether-data/adapters/sqlite/src/settlement.rs b/crates/aether-data/adapters/sqlite/src/settlement.rs deleted file mode 100644 index 6be6bfa54..000000000 --- a/crates/aether-data/adapters/sqlite/src/settlement.rs +++ /dev/null @@ -1,2156 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, Row}; - -use aether_data_contracts::repository::settlement::{ - finite_wallet_available_usd, plan_finite_wallet_debit, settlement_billable_cost_usd, - settlement_billing_status_for_usage_status, validate_wallet_settlement_values, - ReconcileUsagePolicyCostInput, ReleaseUsagePolicyRequestAdmissionInput, - ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput, - ReserveUsagePolicyRequestOutcome, SettlementWriteRepository, StoredUsagePolicyCostReservation, - StoredUsagePolicyRequestAdmission, StoredUsageSettlement, UsagePolicyCostReservationState, - UsagePolicyRequestAdmissionState, UsageSettlementInput, SETTLEMENT_EPSILON_USD, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, sqlite_real, SqlitePool}; - -const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#" -SELECT - usage_record.request_id, - COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id, - COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status) AS billing_status, - COALESCE( - usage_settlement_snapshots.wallet_balance_before, - usage_record.wallet_balance_before - ) AS wallet_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_balance_after, - usage_record.wallet_balance_after - ) AS wallet_balance_after, - COALESCE( - usage_settlement_snapshots.wallet_recharge_balance_before, - usage_record.wallet_recharge_balance_before - ) AS wallet_recharge_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_recharge_balance_after, - usage_record.wallet_recharge_balance_after - ) AS wallet_recharge_balance_after, - COALESCE( - usage_settlement_snapshots.wallet_gift_balance_before, - usage_record.wallet_gift_balance_before - ) AS wallet_gift_balance_before, - COALESCE( - usage_settlement_snapshots.wallet_gift_balance_after, - usage_record.wallet_gift_balance_after - ) AS wallet_gift_balance_after, - CAST(usage_settlement_snapshots.provider_monthly_used_usd AS REAL) AS provider_monthly_used_usd, - usage_record.provider_id, - COALESCE(usage_settlement_snapshots.finalized_at, usage_record.finalized_at) AS finalized_at_unix_secs -FROM "usage" AS usage_record -LEFT JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = usage_record.request_id -WHERE usage_record.request_id = ? -"#; - -const FINALIZE_USAGE_BILLING_SQL: &str = r#" -UPDATE "usage" -SET - billing_status = ?, - finalized_at = COALESCE(finalized_at, ?) -WHERE request_id = ? -"#; - -const UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL: &str = r#" -INSERT INTO usage_settlement_snapshots ( - request_id, - billing_status, - wallet_id, - wallet_balance_before, - wallet_balance_after, - wallet_recharge_balance_before, - wallet_recharge_balance_after, - wallet_gift_balance_before, - wallet_gift_balance_after, - provider_monthly_used_usd, - finalized_at, - created_at, - updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT (request_id) -DO UPDATE SET - billing_status = excluded.billing_status, - wallet_id = COALESCE(excluded.wallet_id, usage_settlement_snapshots.wallet_id), - wallet_balance_before = COALESCE( - excluded.wallet_balance_before, - usage_settlement_snapshots.wallet_balance_before - ), - wallet_balance_after = COALESCE( - excluded.wallet_balance_after, - usage_settlement_snapshots.wallet_balance_after - ), - wallet_recharge_balance_before = COALESCE( - excluded.wallet_recharge_balance_before, - usage_settlement_snapshots.wallet_recharge_balance_before - ), - wallet_recharge_balance_after = COALESCE( - excluded.wallet_recharge_balance_after, - usage_settlement_snapshots.wallet_recharge_balance_after - ), - wallet_gift_balance_before = COALESCE( - excluded.wallet_gift_balance_before, - usage_settlement_snapshots.wallet_gift_balance_before - ), - wallet_gift_balance_after = COALESCE( - excluded.wallet_gift_balance_after, - usage_settlement_snapshots.wallet_gift_balance_after - ), - provider_monthly_used_usd = COALESCE( - excluded.provider_monthly_used_usd, - usage_settlement_snapshots.provider_monthly_used_usd - ), - finalized_at = COALESCE(excluded.finalized_at, usage_settlement_snapshots.finalized_at), - 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, -} - -impl SqliteSettlementRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -fn usage_policy_cost_i64(value: u64, field: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds the integer range"))) -} - -fn usage_policy_cost_u64(value: i64, field: &str) -> Result { - u64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("{field} must not be negative"))) -} - -fn usage_policy_request_admission_from_sqlite_row( - row: &SqliteRow, -) -> Result { - let state: String = row.try_get("state").map_sql_err()?; - Ok(StoredUsagePolicyRequestAdmission { - request_id: row.try_get("request_id").map_sql_err()?, - subject_id: row.try_get("subject_id").map_sql_err()?, - event_token: row.try_get("event_token").map_sql_err()?, - admitted_at_unix_secs: usage_policy_cost_u64( - row.try_get("admitted_at_unix_secs").map_sql_err()?, - "usage policy request admitted_at", - )?, - retain_until_unix_secs: usage_policy_cost_u64( - row.try_get("retain_until_unix_secs").map_sql_err()?, - "usage policy request retain_until", - )?, - state: UsagePolicyRequestAdmissionState::parse(&state).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "unknown usage policy request admission state {state}" - )) - })?, - released_at_unix_secs: row - .try_get::, _>("released_at_unix_secs") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy request released_at")) - .transpose()?, - }) -} - -const FIND_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL: &str = r#" -SELECT request_id, subject_id, event_token, - admitted_at AS admitted_at_unix_secs, - retain_until AS retain_until_unix_secs, - state, released_at AS released_at_unix_secs -FROM usage_request_admissions -WHERE event_token = ? -"#; - -const INSERT_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL: &str = r#" -INSERT INTO usage_request_admissions ( - request_id, subject_id, event_token, admitted_at, retain_until, - state, released_at, created_at -) VALUES (?, ?, ?, ?, ?, 'active', NULL, ?) -ON CONFLICT(event_token) DO NOTHING -"#; - -fn usage_policy_cost_reservation_from_sqlite_row( - row: &SqliteRow, -) -> Result { - let state: String = row.try_get("state").map_sql_err()?; - Ok(StoredUsagePolicyCostReservation { - request_id: row.try_get("request_id").map_sql_err()?, - subject_id: row.try_get("subject_id").map_sql_err()?, - reservation_token: row.try_get("reservation_token").map_sql_err()?, - admitted_at_unix_secs: usage_policy_cost_u64( - row.try_get("admitted_at").map_sql_err()?, - "usage policy admitted_at", - )?, - reserved_cost_units: usage_policy_cost_u64( - row.try_get("reserved_cost_units").map_sql_err()?, - "usage policy reserved_cost_units", - )?, - actual_cost_units: row - .try_get::, _>("actual_cost_units") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy actual_cost_units")) - .transpose()?, - state: UsagePolicyCostReservationState::parse(&state).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "unknown usage policy reservation state {state}" - )) - })?, - reservation_expires_at_unix_secs: usage_policy_cost_u64( - row.try_get("reservation_expires_at").map_sql_err()?, - "usage policy reservation_expires_at", - )?, - retain_until_unix_secs: usage_policy_cost_u64( - row.try_get("retain_until").map_sql_err()?, - "usage policy retain_until", - )?, - finalized_at_unix_secs: row - .try_get::, _>("finalized_at") - .map_sql_err()? - .map(|value| usage_policy_cost_u64(value, "usage policy finalized_at")) - .transpose()?, - }) -} - -const FIND_USAGE_POLICY_COST_RESERVATION_SQLITE_SQL: &str = r#" -SELECT request_id, subject_id, reservation_token, admitted_at, - reserved_cost_units, actual_cost_units, state, - reservation_expires_at, retain_until, finalized_at -FROM usage_cost_reservations -WHERE reservation_token = ? -"#; - -async fn lock_usage_policy_subject_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - subject_id: &str, -) -> Result { - let result = sqlx::query("UPDATE users SET updated_at = updated_at WHERE id = ?") - .bind(subject_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) -} - -fn usage_policy_subject_missing() -> DataLayerError { - DataLayerError::InvalidInput("usage policy subject does not exist".to_string()) -} - -fn settlement_from_row(row: &SqliteRow) -> Result { - Ok(StoredUsageSettlement { - request_id: row.try_get("request_id").map_sql_err()?, - wallet_id: row.try_get("wallet_id").map_sql_err()?, - billing_status: row.try_get("billing_status").map_sql_err()?, - wallet_balance_before: sqlite_optional_real(row, "wallet_balance_before")?, - wallet_balance_after: sqlite_optional_real(row, "wallet_balance_after")?, - wallet_recharge_balance_before: sqlite_optional_real( - row, - "wallet_recharge_balance_before", - )?, - wallet_recharge_balance_after: sqlite_optional_real(row, "wallet_recharge_balance_after")?, - wallet_gift_balance_before: sqlite_optional_real(row, "wallet_gift_balance_before")?, - wallet_gift_balance_after: sqlite_optional_real(row, "wallet_gift_balance_after")?, - provider_monthly_used_usd: sqlite_optional_real(row, "provider_monthly_used_usd")?, - finalized_at_unix_secs: row - .try_get::, _>("finalized_at_unix_secs") - .map_sql_err()? - .map(|value| value as u64), - }) -} - -fn now_unix_secs() -> Result { - i64::try_from( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - ) - .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, - insufficient: bool, -} - -#[derive(Debug)] -struct DailyQuotaGrant { - entitlement_id: String, - daily_quota_usd: f64, - usage_date: String, - allow_wallet_overage: bool, -} - -fn daily_quota_usage_date( - reset_timezone: Option<&str>, - now: chrono::DateTime, -) -> Result { - let timezone = reset_timezone - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("Asia/Shanghai") - .parse::() - .map_err(|err| DataLayerError::InvalidInput(format!("invalid reset_timezone: {err}")))?; - Ok(now.with_timezone(&timezone).date_naive().to_string()) -} - -fn daily_quota_grants_from_entitlement( - entitlement_id: &str, - entitlements: &serde_json::Value, - current_allow_wallet_overage: Option, - now: chrono::DateTime, -) -> Result, DataLayerError> { - let mut grants = Vec::new(); - let Some(items) = entitlements.as_array() else { - return Ok(grants); - }; - for item in items { - if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") { - continue; - } - let daily_quota_usd = item - .get("daily_quota_usd") - .and_then(serde_json::Value::as_f64) - .unwrap_or(0.0); - if !daily_quota_usd.is_finite() || daily_quota_usd <= 0.0 { - continue; - } - grants.push(DailyQuotaGrant { - entitlement_id: entitlement_id.to_string(), - daily_quota_usd, - usage_date: daily_quota_usage_date( - item.get("reset_timezone") - .and_then(serde_json::Value::as_str), - now, - )?, - allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - }), - }); - } - Ok(grants) -} - -fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option { - entitlements.as_array()?.iter().find_map(|item| { - (item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota")) - .then(|| { - item.get("allow_wallet_overage") - .and_then(serde_json::Value::as_bool) - }) - .flatten() - }) -} - -async fn consume_daily_quota_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - user_id: &str, - request_id: &str, - total_cost_usd: f64, - wallet_available_usd: Option, - wallet_can_overdraft: bool, - now_unix_secs: i64, -) -> Result { - if !total_cost_usd.is_finite() || total_cost_usd < 0.0 { - return Err(DataLayerError::InvalidInput( - "daily quota settlement cost must be finite and non-negative".to_string(), - )); - } - if total_cost_usd == 0.0 { - return Ok(DailyQuotaDebitResult::default()); - } - let rows = sqlx::query( - r#" -SELECT - user_plan_entitlements.id, - user_plan_entitlements.entitlements_snapshot, - billing_plans.entitlements_json AS plan_entitlements_json -FROM user_plan_entitlements -JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id -WHERE user_plan_entitlements.user_id = ? - AND user_plan_entitlements.status = 'active' - AND user_plan_entitlements.starts_at <= ? - AND user_plan_entitlements.expires_at > ? -ORDER BY user_plan_entitlements.expires_at ASC, - user_plan_entitlements.created_at ASC, - user_plan_entitlements.id ASC -"#, - ) - .bind(user_id) - .bind(now_unix_secs) - .bind(now_unix_secs) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - let now = chrono::Utc::now(); - let mut grants = Vec::new(); - for row in rows { - let entitlement_id: String = row.try_get("id").map_sql_err()?; - let entitlements_raw: String = row.try_get("entitlements_snapshot").map_sql_err()?; - let entitlements = - serde_json::from_str::(&entitlements_raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "user_plan_entitlements.entitlements_snapshot invalid json: {err}" - )) - })?; - let plan_entitlements_raw: String = row.try_get("plan_entitlements_json").map_sql_err()?; - let plan_entitlements = serde_json::from_str::(&plan_entitlements_raw) - .map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "billing_plans.entitlements_json invalid json: {err}" - )) - })?; - grants.extend(daily_quota_grants_from_entitlement( - &entitlement_id, - &entitlements, - daily_quota_wallet_overage_policy(&plan_entitlements), - now, - )?); - } - if grants.is_empty() { - return Ok(DailyQuotaDebitResult::default()); - } - - let mut grants_with_remaining = Vec::new(); - let mut total_remaining = 0.0; - let mut allow_wallet_overage = true; - for grant in grants { - allow_wallet_overage &= grant.allow_wallet_overage; - let used = sqlx::query_scalar::<_, f64>( - r#" -SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) -FROM entitlement_usage_ledgers -WHERE user_entitlement_id = ? - AND usage_date = ? -"#, - ) - .bind(&grant.entitlement_id) - .bind(&grant.usage_date) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - if !used.is_finite() || used < 0.0 { - return Err(DataLayerError::UnexpectedValue( - "daily quota usage ledger total is invalid".to_string(), - )); - } - let remaining = (grant.daily_quota_usd - used).max(0.0); - total_remaining += remaining; - if !total_remaining.is_finite() { - return Err(DataLayerError::UnexpectedValue( - "daily quota remaining total overflowed".to_string(), - )); - } - grants_with_remaining.push((grant, remaining)); - } - let insufficient = (!allow_wallet_overage && total_remaining + 0.000_000_01 < total_cost_usd) - || (allow_wallet_overage - && !wallet_can_overdraft - && wallet_available_usd.is_some_and(|available| { - total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd - })); - - let mut remaining_cost = total_cost_usd; - let mut debited = 0.0; - for (grant, balance_before) in grants_with_remaining { - if remaining_cost <= 0.000_000_01 || balance_before <= 0.0 { - continue; - } - let amount = remaining_cost.min(balance_before); - let balance_after = balance_before - amount; - sqlx::query( - r#" -INSERT OR IGNORE INTO entitlement_usage_ledgers ( - id, user_entitlement_id, user_id, request_id, amount_usd, - balance_before, balance_after, usage_date, created_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&grant.entitlement_id) - .bind(user_id) - .bind(request_id) - .bind(amount) - .bind(balance_before) - .bind(balance_after) - .bind(&grant.usage_date) - .bind(now_unix_secs) - .execute(&mut **tx) - .await - .map_sql_err()?; - remaining_cost -= amount; - debited += amount; - } - Ok(DailyQuotaDebitResult { - debited_usd: debited, - insufficient, - }) -} - -#[async_trait] -impl SettlementWriteRepository for SqliteSettlementRepository { - async fn reserve_usage_policy_request( - &self, - input: ReserveUsagePolicyRequestInput, - ) -> Result { - input.validate()?; - let now = now_unix_secs()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - // This no-op update acquires SQLite's single writer slot before any admission reads. - if !lock_usage_policy_subject_sqlite(&mut tx, &input.subject_id).await? { - return Err(usage_policy_subject_missing()); - } - let existing_row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL) - .bind(&input.event_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if let Some(row) = existing_row.as_ref() { - let existing = usage_policy_request_admission_from_sqlite_row(row)?; - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Conflict); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy event_token must keep its original admitted_at".to_string(), - )); - } - sqlx::query( - "UPDATE usage_request_admissions SET retain_until = MAX(retain_until, ?) WHERE event_token = ?", - ) - .bind(usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy request retain_until", - )?) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - let outcome = match existing.state { - UsagePolicyRequestAdmissionState::Active => { - ReserveUsagePolicyRequestOutcome::Allowed - } - UsagePolicyRequestAdmissionState::Released => { - ReserveUsagePolicyRequestOutcome::AlreadyReleased - } - }; - tx.commit().await.map_sql_err()?; - return Ok(outcome); - } - - for (window_index, window) in input.windows.iter().enumerate() { - let used_requests = sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM usage_request_admissions -WHERE subject_id = ? - AND state = 'active' - AND admitted_at >= ? - AND admitted_at < ? - "#, - ) - .bind(&input.subject_id) - .bind(usage_policy_cost_i64( - window.starts_at_unix_secs, - "usage policy request window start", - )?) - .bind(usage_policy_cost_i64( - window.ends_at_unix_secs, - "usage policy request window end", - )?) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let used_requests = - usage_policy_cost_u64(used_requests, "usage policy request used_requests")?; - if used_requests >= window.limit_requests { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Rejected { - window_index, - limit_requests: window.limit_requests, - used_requests, - }); - } - } - - let insert_result = sqlx::query(INSERT_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL) - .bind(&input.request_id) - .bind(&input.subject_id) - .bind(&input.event_token) - .bind(usage_policy_cost_i64( - input.admitted_at_unix_secs, - "usage policy request admitted_at", - )?) - .bind(usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy request retain_until", - )?) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if insert_result.rows_affected() == 1 { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Allowed); - } - - // The writer lock above normally makes this branch unreachable for concurrent reserves, - // but classify the unique-token race explicitly so future lock changes cannot surface a - // raw SQLite constraint error or accidentally reactivate a released tombstone. - let row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL) - .bind(&input.event_token) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let existing = usage_policy_request_admission_from_sqlite_row(&row)?; - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyRequestOutcome::Conflict); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy event_token must keep its original admitted_at".to_string(), - )); - } - sqlx::query( - "UPDATE usage_request_admissions SET retain_until = MAX(retain_until, ?) WHERE event_token = ?", - ) - .bind(usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy request retain_until", - )?) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - let outcome = match existing.state { - UsagePolicyRequestAdmissionState::Active => ReserveUsagePolicyRequestOutcome::Allowed, - UsagePolicyRequestAdmissionState::Released => { - ReserveUsagePolicyRequestOutcome::AlreadyReleased - } - }; - tx.commit().await.map_sql_err()?; - Ok(outcome) - } - - async fn release_usage_policy_request_admission( - &self, - input: ReleaseUsagePolicyRequestAdmissionInput, - ) -> Result, DataLayerError> { - input.validate()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_sqlite(&mut tx, &input.subject_id).await? { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - let row = sqlx::query(FIND_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL) - .bind(&input.event_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - let mut admission = usage_policy_request_admission_from_sqlite_row(&row)?; - if admission.request_id != input.request_id || admission.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - if input.released_at_unix_secs < admission.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy released_at must not precede admitted_at".to_string(), - )); - } - if admission.state == UsagePolicyRequestAdmissionState::Active { - sqlx::query( - "UPDATE usage_request_admissions SET state = 'released', released_at = ? WHERE event_token = ? AND state = 'active'", - ) - .bind(usage_policy_cost_i64( - input.released_at_unix_secs, - "usage policy request released_at", - )?) - .bind(&input.event_token) - .execute(&mut *tx) - .await - .map_sql_err()?; - admission.state = UsagePolicyRequestAdmissionState::Released; - admission.released_at_unix_secs = Some(input.released_at_unix_secs); - } - tx.commit().await.map_sql_err()?; - Ok(Some(admission)) - } - - async fn cleanup_usage_policy_request_admissions( - &self, - now_unix_secs: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(0); - } - let now = usage_policy_cost_i64(now_unix_secs, "usage policy request cleanup timestamp")?; - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - let result = sqlx::query( - r#" -DELETE FROM usage_request_admissions -WHERE rowid IN ( - SELECT rowid - FROM usage_request_admissions - WHERE retain_until <= ? - ORDER BY retain_until, event_token - LIMIT ? -) - "#, - ) - .bind(now) - .bind(limit) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn reserve_usage_policy_cost( - &self, - input: ReserveUsagePolicyCostInput, - ) -> Result { - input.validate()?; - let now = now_unix_secs()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_sqlite(&mut tx, &input.subject_id).await? { - return Err(usage_policy_subject_missing()); - } - let existing_row = sqlx::query(FIND_USAGE_POLICY_COST_RESERVATION_SQLITE_SQL) - .bind(&input.reservation_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let existing = existing_row - .as_ref() - .map(usage_policy_cost_reservation_from_sqlite_row) - .transpose()?; - if let Some(existing) = existing.as_ref() { - if existing.request_id != input.request_id || existing.subject_id != input.subject_id { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyCostOutcome::Conflict); - } - if existing.state != UsagePolicyCostReservationState::Reserved { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyCostOutcome::AlreadyTerminal { - state: existing.state, - }); - } - if existing.admitted_at_unix_secs != input.admitted_at_unix_secs { - return Err(DataLayerError::InvalidInput( - "usage policy reservation_token must keep its original admitted_at".to_string(), - )); - } - } - - let previous_reserved_cost_units = existing - .as_ref() - .map(|reservation| reservation.reserved_cost_units) - .unwrap_or(0); - let target_reserved_cost_units = - previous_reserved_cost_units.max(input.reserved_cost_units); - for (window_index, window) in input.windows.iter().enumerate() { - let used_cost_units = sqlx::query_scalar::<_, i64>( - r#" -SELECT COALESCE(SUM( - CASE - WHEN state = 'finalized' THEN COALESCE(actual_cost_units, 0) - WHEN state = 'reserved' AND reservation_expires_at > ? THEN reserved_cost_units - ELSE 0 - END -), 0) -FROM usage_cost_reservations -WHERE subject_id = ? - AND admitted_at >= ? - AND admitted_at < ? - AND reservation_token <> ? - "#, - ) - .bind(usage_policy_cost_i64( - input.admitted_at_unix_secs, - "usage policy admitted_at", - )?) - .bind(&input.subject_id) - .bind(usage_policy_cost_i64( - window.starts_at_unix_secs, - "usage policy window start", - )?) - .bind(usage_policy_cost_i64( - window.ends_at_unix_secs, - "usage policy window end", - )?) - .bind(&input.reservation_token) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let used_cost_units = - usage_policy_cost_u64(used_cost_units, "usage policy used_cost_units")?; - if used_cost_units - .checked_add(target_reserved_cost_units) - .is_none_or(|total| total > window.limit_cost_units) - { - tx.commit().await.map_sql_err()?; - return Ok(ReserveUsagePolicyCostOutcome::Rejected { - window_index, - limit_cost_units: window.limit_cost_units, - used_cost_units, - }); - } - } - - sqlx::query( - r#" -INSERT INTO usage_cost_reservations ( - request_id, subject_id, reservation_token, admitted_at, - reserved_cost_units, actual_cost_units, - state, reservation_expires_at, retain_until, finalized_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, NULL, 'reserved', ?, ?, NULL, ?, ?) -ON CONFLICT (reservation_token) DO UPDATE SET - reserved_cost_units = MAX( - usage_cost_reservations.reserved_cost_units, - excluded.reserved_cost_units - ), - reservation_expires_at = MAX( - usage_cost_reservations.reservation_expires_at, - excluded.reservation_expires_at - ), - retain_until = MAX( - usage_cost_reservations.retain_until, - excluded.retain_until - ), - updated_at = excluded.updated_at - "#, - ) - .bind(&input.request_id) - .bind(&input.subject_id) - .bind(&input.reservation_token) - .bind(usage_policy_cost_i64( - input.admitted_at_unix_secs, - "usage policy admitted_at", - )?) - .bind(usage_policy_cost_i64( - target_reserved_cost_units, - "usage policy reserved_cost_units", - )?) - .bind(usage_policy_cost_i64( - input.reservation_expires_at_unix_secs, - "usage policy reservation_expires_at", - )?) - .bind(usage_policy_cost_i64( - input.retain_until_unix_secs, - "usage policy retain_until", - )?) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(ReserveUsagePolicyCostOutcome::Allowed { - reserved_cost_units: target_reserved_cost_units, - additional_reserved_cost_units: target_reserved_cost_units - .saturating_sub(previous_reserved_cost_units), - }) - } - - async fn reconcile_usage_policy_cost( - &self, - input: ReconcileUsagePolicyCostInput, - ) -> Result, DataLayerError> { - input.validate()?; - let now = now_unix_secs()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - if !lock_usage_policy_subject_sqlite(&mut tx, &input.subject_id).await? { - tx.commit().await.map_sql_err()?; - return Ok(None); - } - let row = sqlx::query(FIND_USAGE_POLICY_COST_RESERVATION_SQLITE_SQL) - .bind(&input.reservation_token) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - let mut reservation = usage_policy_cost_reservation_from_sqlite_row(&row)?; - if reservation.request_id != input.request_id || reservation.subject_id != input.subject_id - { - // The token selects the row; audit identity must still match before the reservation - // can be finalized. - tx.commit().await.map_sql_err()?; - return Ok(None); - } - if reservation.state == UsagePolicyCostReservationState::Reserved { - sqlx::query( - r#" -UPDATE usage_cost_reservations -SET state = ?, actual_cost_units = ?, finalized_at = ?, updated_at = ? -WHERE reservation_token = ? - AND request_id = ? - AND subject_id = ? - AND state = 'reserved' - "#, - ) - .bind(input.terminal_state.as_str()) - .bind(usage_policy_cost_i64( - input.actual_cost_units, - "usage policy actual_cost_units", - )?) - .bind(usage_policy_cost_i64( - input.finalized_at_unix_secs, - "usage policy finalized_at", - )?) - .bind(now) - .bind(&input.reservation_token) - .bind(&input.request_id) - .bind(&input.subject_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - reservation.state = input.terminal_state; - reservation.actual_cost_units = Some(input.actual_cost_units); - reservation.finalized_at_unix_secs = Some(input.finalized_at_unix_secs); - } - tx.commit().await.map_sql_err()?; - Ok(Some(reservation)) - } - - async fn cleanup_usage_policy_cost_reservations( - &self, - now_unix_secs: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(0); - } - let now = usage_policy_cost_i64(now_unix_secs, "usage policy cleanup timestamp")?; - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - let result = sqlx::query( - r#" -DELETE FROM usage_cost_reservations -WHERE rowid IN ( - SELECT rowid - FROM usage_cost_reservations - WHERE retain_until <= ? - ORDER BY retain_until, reservation_token - LIMIT ? -) - "#, - ) - .bind(now) - .bind(limit) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() as usize) - } - - async fn settle_usage( - &self, - input: UsageSettlementInput, - ) -> Result, DataLayerError> { - input.validate()?; - let finalized_at = i64::try_from( - input - .finalized_at_unix_secs - .unwrap_or(now_unix_secs()? as u64), - ) - .map_err(|_| DataLayerError::InvalidInput("finalized_at overflow".to_string()))?; - 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) - .await - .map_sql_err()?; - - let Some(usage_row) = row else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let current_billing_status: String = usage_row.try_get("billing_status").map_sql_err()?; - if matches!( - current_billing_status.as_str(), - "settled" | "void" | "insufficient_quota" - ) { - let settlement = settlement_from_row(&usage_row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - let mut final_billing_status = - settlement_billing_status_for_usage_status(&input.status).to_string(); - let mut settlement = StoredUsageSettlement { - request_id: input.request_id.clone(), - wallet_id: None, - billing_status: final_billing_status.clone(), - wallet_balance_before: None, - wallet_balance_after: None, - wallet_recharge_balance_before: None, - wallet_recharge_balance_after: None, - wallet_gift_balance_before: None, - wallet_gift_balance_after: None, - provider_monthly_used_usd: None, - finalized_at_unix_secs: Some(finalized_at as u64), - }; - - if final_billing_status == "settled" { - let api_key_id = input - .api_key_id - .as_deref() - .filter(|value| !value.is_empty()); - let api_key_is_standalone = if input.api_key_is_standalone { - true - } else if let Some(api_key_id) = api_key_id { - sqlx::query_scalar::<_, bool>( - r#" -SELECT is_standalone -FROM api_keys -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .unwrap_or(false) - } else { - false - }; - - let wallet_row = if let Some(api_key_id) = api_key_id { - sqlx::query( - r#" -SELECT id, balance, gift_balance, total_consumed, limit_mode -FROM wallets -WHERE api_key_id = ? -LIMIT 1 -"#, - ) - .bind(api_key_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - None - }; - - let wallet_row = if wallet_row.is_some() { - wallet_row - } else if !api_key_is_standalone { - if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) { - sqlx::query( - r#" -SELECT id, balance, gift_balance, total_consumed, limit_mode -FROM wallets -WHERE user_id = ? -LIMIT 1 -"#, - ) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - } else { - None - } - } else { - None - }; - - let wallet_can_overdraft = wallet_row.is_some(); - let wallet_available_usd = match wallet_row.as_ref() { - Some(row) => { - let recharge_balance = sqlite_real(row, "balance")?; - let gift_balance = sqlite_real(row, "gift_balance")?; - let total_consumed = sqlite_real(row, "total_consumed")?; - validate_wallet_settlement_values( - recharge_balance, - gift_balance, - total_consumed, - 0.0, - )?; - let limit_mode: String = row.try_get("limit_mode").map_sql_err()?; - if limit_mode.eq_ignore_ascii_case("unlimited") { - None - } else { - Some(finite_wallet_available_usd(recharge_balance, gift_balance)) - } - } - None => Some(0.0), - }; - if let Some(row) = wallet_row.as_ref() { - let wallet_id: String = row.try_get("id").map_sql_err()?; - let before_recharge = sqlite_real(row, "balance")?; - let before_gift = sqlite_real(row, "gift_balance")?; - let before_total = before_recharge + before_gift; - settlement.wallet_id = Some(wallet_id); - settlement.wallet_balance_before = Some(before_total); - settlement.wallet_balance_after = Some(before_total); - settlement.wallet_recharge_balance_before = Some(before_recharge); - settlement.wallet_recharge_balance_after = Some(before_recharge); - settlement.wallet_gift_balance_before = Some(before_gift); - settlement.wallet_gift_balance_after = Some(before_gift); - } - - let billable_cost_usd = settlement_billable_cost_usd(&input); - let wallet_debit_cost_usd = if !api_key_is_standalone { - if let Some(user_id) = input.user_id.as_deref().filter(|value| !value.is_empty()) { - let quota = consume_daily_quota_sqlite( - &mut tx, - user_id, - &input.request_id, - billable_cost_usd, - wallet_available_usd, - wallet_can_overdraft, - updated_at, - ) - .await?; - if quota.insufficient { - final_billing_status = "insufficient_quota".to_string(); - settlement.billing_status = final_billing_status.clone(); - 0.0 - } else { - (billable_cost_usd - quota.debited_usd).max(0.0) - } - } else { - billable_cost_usd - } - } else { - billable_cost_usd - }; - if final_billing_status != "settled" { - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - if wallet_debit_cost_usd > SETTLEMENT_EPSILON_USD { - if let Some(wallet_row) = wallet_row { - let wallet_id: String = wallet_row.try_get("id").map_sql_err()?; - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let total_consumed = sqlite_real(&wallet_row, "total_consumed")?; - let limit_mode: String = wallet_row.try_get("limit_mode").map_sql_err()?; - let before_total = before_recharge + before_gift; - let mut after_recharge = before_recharge; - let mut after_gift = before_gift; - if !limit_mode.eq_ignore_ascii_case("unlimited") { - let debit_plan = plan_finite_wallet_debit( - before_recharge, - before_gift, - wallet_debit_cost_usd, - ); - (after_recharge, after_gift) = - debit_plan.after_balances(before_recharge, before_gift); - } - let total_consumed_after = total_consumed + wallet_debit_cost_usd; - validate_wallet_settlement_values( - after_recharge, - after_gift, - total_consumed_after, - 0.0, - )?; - if final_billing_status == "settled" { - sqlx::query( - r#" -UPDATE wallets -SET - balance = ?, - gift_balance = ?, - total_consumed = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(total_consumed_after) - .bind(updated_at) - .bind(&wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - settlement.wallet_id = Some(wallet_id); - settlement.wallet_balance_before = Some(before_total); - settlement.wallet_balance_after = Some(after_recharge + after_gift); - settlement.wallet_recharge_balance_before = Some(before_recharge); - settlement.wallet_recharge_balance_after = Some(after_recharge); - settlement.wallet_gift_balance_before = Some(before_gift); - settlement.wallet_gift_balance_after = Some(after_gift); - } else { - final_billing_status = "insufficient_quota".to_string(); - settlement.billing_status = final_billing_status.clone(); - } - } - - if final_billing_status != "settled" { - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - return Ok(Some(settlement)); - } - - if let Some(provider_id) = input - .provider_id - .as_deref() - .filter(|value| !value.is_empty()) - { - enqueue_provider_monthly_usage_delta_sqlite( - &mut tx, - &input.request_id, - provider_id, - input.actual_total_cost_usd, - updated_at, - ) - .await?; - } - } - - sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL) - .bind(&settlement.request_id) - .bind(&settlement.billing_status) - .bind(settlement.wallet_id.as_deref()) - .bind(settlement.wallet_balance_before) - .bind(settlement.wallet_balance_after) - .bind(settlement.wallet_recharge_balance_before) - .bind(settlement.wallet_recharge_balance_after) - .bind(settlement.wallet_gift_balance_before) - .bind(settlement.wallet_gift_balance_after) - .bind(settlement.provider_monthly_used_usd) - .bind(settlement.finalized_at_unix_secs.map(|value| value as i64)) - .bind(updated_at) - .bind(updated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query(FINALIZE_USAGE_BILLING_SQL) - .bind(&final_billing_status) - .bind(finalized_at) - .bind(&input.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - tx.commit().await.map_sql_err()?; - Ok(Some(settlement)) - } -} - -#[cfg(test)] -mod tests { - use super::{SqliteSettlementRepository, INSERT_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL}; - use crate::{run_migrations, SqliteUserReadRepository}; - use aether_data_contracts::repository::settlement::{ - ReconcileUsagePolicyCostInput, ReleaseUsagePolicyRequestAdmissionInput, - ReserveUsagePolicyCostInput, ReserveUsagePolicyRequestInput, SettlementWriteRepository, - UsagePolicyCostReservationState, UsagePolicyCostWindow, UsagePolicyRequestWindow, - UsageSettlementInput, - }; - use aether_data_contracts::repository::users::UserReadRepository; - use sqlx::Row; - use std::time::Duration; - - #[tokio::test] - async fn sqlite_repository_settles_usage_once() { - 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_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(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), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "settled"); - assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-1")); - assert_eq!(settlement.wallet_balance_before, Some(12.0)); - 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, None); - - let wallet = sqlx::query( - "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'", - ) - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet.try_get::("balance").unwrap(), 4.0); - assert_eq!(wallet.try_get::("gift_balance").unwrap(), 2.0); - assert_eq!(wallet.try_get::("total_consumed").unwrap(), 6.0); - - let second = repository - .settle_usage(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(9_999), - }) - .await - .expect("second settlement should run") - .expect("usage should exist"); - assert_eq!(second.finalized_at_unix_secs, Some(1_234)); - - let provider_used: f64 = - sqlx::query_scalar("SELECT monthly_used_usd FROM providers WHERE id = '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 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] - async fn sqlite_settlement_rejects_corrupt_wallet_before_financial_mutation() { - 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_settlement_rows(&pool).await; - sqlx::query("UPDATE wallets SET balance = ? WHERE id = 'wallet-1'") - .bind(f64::INFINITY) - .execute(&pool) - .await - .expect("corrupt wallet fixture should update"); - - let result = SqliteSettlementRepository::new(pool.clone()) - .settle_usage(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), - }) - .await; - assert!(result.is_err()); - - let billing_status: String = - sqlx::query_scalar("SELECT billing_status FROM usage WHERE request_id = 'request-1'") - .fetch_one(&pool) - .await - .expect("usage should load"); - assert_eq!(billing_status, "pending"); - let settlement_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = 'request-1'", - ) - .fetch_one(&pool) - .await - .expect("settlement snapshots should count"); - assert_eq!(settlement_count, 0); - let delta_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = 'request-1'", - ) - .fetch_one(&pool) - .await - .expect("usage deltas should count"); - assert_eq!(delta_count, 0); - } - - #[tokio::test] - async fn request_admission_insert_defensively_preserves_the_existing_token() { - 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"); - sqlx::query( - r#" -INSERT INTO users (id, username, auth_source, created_at, updated_at) -VALUES - ('defensive-user-1', 'defensive-user-1', 'local', 1, 1), - ('defensive-user-2', 'defensive-user-2', 'local', 1, 1) - "#, - ) - .execute(&pool) - .await - .expect("usage policy subjects should insert"); - - let insert = |request_id: &'static str, subject_id: &'static str| { - sqlx::query(INSERT_USAGE_POLICY_REQUEST_ADMISSION_SQLITE_SQL) - .bind(request_id) - .bind(subject_id) - .bind("defensive-event-token") - .bind(100_i64) - .bind(200_i64) - .bind(100_i64) - }; - assert_eq!( - insert("defensive-request-1", "defensive-user-1") - .execute(&pool) - .await - .expect("initial admission should insert") - .rows_affected(), - 1 - ); - assert_eq!( - insert("defensive-request-2", "defensive-user-2") - .execute(&pool) - .await - .expect("duplicate token should be ignored") - .rows_affected(), - 0 - ); - let stored: (String, String, i64) = sqlx::query_as( - "SELECT request_id, subject_id, admitted_at FROM usage_request_admissions WHERE event_token = 'defensive-event-token'", - ) - .fetch_one(&pool) - .await - .expect("original admission should remain"); - assert_eq!( - stored, - ( - "defensive-request-1".to_string(), - "defensive-user-1".to_string(), - 100, - ) - ); - } - - #[tokio::test] - async fn deleting_user_cascades_usage_policy_ledgers_and_terminal_calls_are_noops() { - 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"); - sqlx::query( - r#" -INSERT INTO users (id, username, auth_source, created_at, updated_at) -VALUES ('usage-policy-delete-user', 'usage-policy-delete-user', 'local', 1, 1) - "#, - ) - .execute(&pool) - .await - .expect("usage policy user should insert"); - - let repository = SqliteSettlementRepository::new(pool.clone()); - repository - .reserve_usage_policy_request(ReserveUsagePolicyRequestInput { - request_id: "usage-policy-delete-request".to_string(), - subject_id: "usage-policy-delete-user".to_string(), - event_token: "usage-policy-delete-event".to_string(), - admitted_at_unix_secs: 100, - retain_until_unix_secs: 200, - windows: vec![UsagePolicyRequestWindow { - starts_at_unix_secs: 50, - ends_at_unix_secs: 200, - limit_requests: 10, - }], - }) - .await - .expect("request admission should reserve"); - repository - .reserve_usage_policy_cost(ReserveUsagePolicyCostInput { - request_id: "usage-policy-delete-request".to_string(), - subject_id: "usage-policy-delete-user".to_string(), - reservation_token: "usage-policy-delete-reservation".to_string(), - admitted_at_unix_secs: 100, - reserved_cost_units: 1, - reservation_expires_at_unix_secs: 150, - retain_until_unix_secs: 200, - windows: vec![UsagePolicyCostWindow { - window_id: "usage-policy-delete-window".to_string(), - starts_at_unix_secs: 50, - ends_at_unix_secs: 200, - limit_cost_units: 10, - }], - }) - .await - .expect("cost reservation should reserve"); - - assert!(SqliteUserReadRepository::new(pool.clone()) - .delete_local_auth_user("usage-policy-delete-user") - .await - .expect("user deletion should succeed")); - let ledger_count: i64 = sqlx::query_scalar( - r#" -SELECT - (SELECT COUNT(*) FROM usage_request_admissions) - + (SELECT COUNT(*) FROM usage_cost_reservations) - "#, - ) - .fetch_one(&pool) - .await - .expect("usage policy ledgers should count"); - assert_eq!(ledger_count, 0); - - assert!(repository - .release_usage_policy_request_admission(ReleaseUsagePolicyRequestAdmissionInput { - request_id: "usage-policy-delete-request".to_string(), - subject_id: "usage-policy-delete-user".to_string(), - event_token: "usage-policy-delete-event".to_string(), - released_at_unix_secs: 150, - },) - .await - .expect("post-delete release should be a no-op") - .is_none()); - assert!(repository - .reconcile_usage_policy_cost(ReconcileUsagePolicyCostInput { - request_id: "usage-policy-delete-request".to_string(), - subject_id: "usage-policy-delete-user".to_string(), - reservation_token: "usage-policy-delete-reservation".to_string(), - actual_cost_units: 1, - terminal_state: UsagePolicyCostReservationState::Finalized, - finalized_at_unix_secs: 150, - }) - .await - .expect("post-delete reconciliation should be a no-op") - .is_none()); - } - - #[tokio::test] - async fn sqlite_repository_voids_failed_usage_without_wallet_mutation() { - 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_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(UsageSettlementInput { - request_id: "request-2".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: "failed".to_string(), - billing_status: "pending".to_string(), - total_cost_usd: 3.0, - actual_total_cost_usd: 2.0, - finalized_at_unix_secs: Some(1_235), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "void"); - assert_eq!(settlement.wallet_id, None); - let wallet_total: f64 = - sqlx::query_scalar("SELECT balance + gift_balance FROM wallets WHERE id = 'wallet-1'") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet_total, 12.0); - } - - #[tokio::test] - async fn sqlite_repository_overdraws_finite_wallet_and_settles_usage() { - 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_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(UsageSettlementInput { - request_id: "request-overdraw".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: 15.0, - actual_total_cost_usd: 15.0, - finalized_at_unix_secs: Some(1_236), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "settled"); - assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-1")); - assert_eq!(settlement.wallet_balance_before, Some(12.0)); - 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, None); - - let wallet = sqlx::query( - "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'", - ) - .fetch_one(&pool) - .await - .expect("wallet should load"); - 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] - async fn sqlite_repository_records_wallet_for_quota_covered_user_usage() { - 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_quota_covered_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(UsageSettlementInput { - request_id: "request-quota-covered".to_string(), - user_id: Some("user-quota".to_string()), - api_key_id: Some("key-quota".to_string()), - api_key_is_standalone: false, - provider_id: None, - 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_260), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "settled"); - assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-quota")); - assert_eq!(settlement.wallet_balance_before, Some(0.0)); - assert_eq!(settlement.wallet_balance_after, Some(0.0)); - - let wallet_total: f64 = sqlx::query_scalar( - "SELECT balance + gift_balance FROM wallets WHERE id = 'wallet-quota'", - ) - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet_total, 0.0); - - let quota_used: f64 = sqlx::query_scalar( - "SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers WHERE request_id = 'request-quota-covered'", - ) - .fetch_one(&pool) - .await - .expect("quota ledger should load"); - assert_eq!(quota_used, 6.0); - } - - #[tokio::test] - async fn sqlite_repository_exhausts_strict_quota_after_actual_cost_overrun() { - 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_quota_covered_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(UsageSettlementInput { - request_id: "request-quota-overrun".to_string(), - user_id: Some("user-quota".to_string()), - api_key_id: Some("key-quota".to_string()), - api_key_is_standalone: false, - provider_id: None, - status: "completed".to_string(), - billing_status: "pending".to_string(), - total_cost_usd: 12.0, - actual_total_cost_usd: 12.0, - finalized_at_unix_secs: Some(1_261), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "insufficient_quota"); - let quota_used: f64 = sqlx::query_scalar( - "SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers WHERE request_id = 'request-quota-overrun'", - ) - .fetch_one(&pool) - .await - .expect("quota ledger should load"); - assert_eq!(quota_used, 10.0); - } - - #[tokio::test] - async fn sqlite_repository_uses_current_plan_wallet_overage_policy() { - 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_quota_covered_settlement_rows(&pool).await; - sqlx::query( - r#" -UPDATE wallets SET balance = 5.0 WHERE id = 'wallet-quota'; -UPDATE billing_plans -SET entitlements_json = '[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":true}]' -WHERE id = 'plan-quota'; -"#, - ) - .execute(&pool) - .await - .expect("plan overage policy should update"); - - let repository = SqliteSettlementRepository::new(pool.clone()); - let settlement = repository - .settle_usage(UsageSettlementInput { - request_id: "request-quota-overrun".to_string(), - user_id: Some("user-quota".to_string()), - api_key_id: Some("key-quota".to_string()), - api_key_is_standalone: false, - provider_id: None, - status: "completed".to_string(), - billing_status: "pending".to_string(), - total_cost_usd: 12.0, - actual_total_cost_usd: 12.0, - finalized_at_unix_secs: Some(1_261), - }) - .await - .expect("settlement should run") - .expect("usage should exist"); - - assert_eq!(settlement.billing_status, "settled"); - assert_eq!(settlement.wallet_balance_after, Some(3.0)); - let quota_used: f64 = sqlx::query_scalar( - "SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers WHERE request_id = 'request-quota-overrun'", - ) - .fetch_one(&pool) - .await - .expect("quota ledger should load"); - assert_eq!(quota_used, 10.0); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn sqlite_repository_exhausts_strict_quota_across_concurrent_requests() { - let database_path = std::env::temp_dir().join(format!( - "aether-quota-settlement-{}.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_quota_covered_settlement_rows(&pool).await; - - let repository = SqliteSettlementRepository::new(pool.clone()); - let input = |request_id: &str| UsageSettlementInput { - request_id: request_id.to_string(), - user_id: Some("user-quota".to_string()), - api_key_id: Some("key-quota".to_string()), - api_key_is_standalone: false, - provider_id: None, - status: "completed".to_string(), - billing_status: "pending".to_string(), - total_cost_usd: 6.0, - actual_total_cost_usd: 6.0, - finalized_at_unix_secs: Some(1_262), - }; - let (first, second) = tokio::join!( - repository.settle_usage(input("request-quota-race-1")), - repository.settle_usage(input("request-quota-race-2")), - ); - let first = first - .expect("first settlement should succeed") - .expect("first usage should exist"); - let second = second - .expect("second settlement should succeed") - .expect("second usage should exist"); - let mut statuses = [first.billing_status, second.billing_status]; - statuses.sort(); - assert_eq!(statuses, ["insufficient_quota", "settled"]); - - let quota_used: f64 = sqlx::query_scalar( - "SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers", - ) - .fetch_one(&pool) - .await - .expect("quota ledger should load"); - assert_eq!(quota_used, 10.0); - - 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(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#" -INSERT INTO providers ( - id, name, provider_type, monthly_used_usd, created_at, updated_at -) -VALUES ('provider-1', 'Provider One', 'openai', 5.0, 1, 1); - -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, created_at, updated_at -) -VALUES ('wallet-1', 'user-1', 10.0, 2.0, 'finite', 1, 1); - -INSERT INTO "usage" ( - request_id, user_id, provider_id, status, billing_status, total_cost_usd, actual_total_cost_usd -) -VALUES - ('request-1', 'user-1', 'provider-1', 'completed', 'pending', 3.0, 6.0), - ('request-2', 'user-1', 'provider-1', 'failed', 'pending', 3.0, 2.0), - ('request-overdraw', 'user-1', 'provider-1', 'completed', 'pending', 15.0, 15.0); -"#, - ) - .execute(pool) - .await - .expect("settlement rows should seed"); - } - - async fn seed_quota_covered_settlement_rows(pool: &sqlx::SqlitePool) { - sqlx::query( - r#" -INSERT INTO users ( - id, username, email, role, auth_source, password_hash, is_active, - is_deleted, created_at, updated_at -) VALUES ( - 'user-quota', 'quota-user', 'quota@example.com', 'user', 'local', - 'hash', 1, 0, 1, 1 -); - -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, created_at, updated_at -) VALUES ( - 'wallet-quota', 'user-quota', 0.0, 0.0, 'finite', 1, 1 -); - -INSERT INTO "usage" ( - request_id, user_id, api_key_id, status, billing_status, - total_cost_usd, actual_total_cost_usd -) VALUES - ( - 'request-quota-covered', 'user-quota', 'key-quota', 'completed', - 'pending', 3.0, 6.0 - ), - ( - 'request-quota-overrun', 'user-quota', 'key-quota', 'completed', - 'pending', 12.0, 12.0 - ), - ( - 'request-quota-race-1', 'user-quota', 'key-quota', 'completed', - 'pending', 6.0, 6.0 - ), - ( - 'request-quota-race-2', 'user-quota', 'key-quota', 'completed', - 'pending', 6.0, 6.0 - ); - -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, - duration_value, entitlements_json, created_at, updated_at -) VALUES ( - 'plan-quota', 'Quota Plan', 0.0, 'USD', 'month', 1, - '[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":false}]', - 1, 1 -); - -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-quota', 'order-quota', 'wallet-quota', 'user-quota', 0.0, 0.0, - 0.0, 'admin_manual', '{}', 'credited', 1 -); - -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) VALUES ( - 'entitlement-quota', 'user-quota', 'plan-quota', 'order-quota', - 'active', 1, 9999999999, - '[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":false}]', - 1, 1 -); -"#, - ) - .execute(pool) - .await - .expect("quota settlement rows should seed"); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/usage.rs b/crates/aether-data/adapters/sqlite/src/usage.rs deleted file mode 100644 index c1e5aeba5..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage.rs +++ /dev/null @@ -1,5260 +0,0 @@ -use std::collections::{BTreeMap, HashSet}; -use std::time::{SystemTime, UNIX_EPOCH}; - -use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; -use aether_data_contracts::repository::usage::UsageBodyField; -use async_trait::async_trait; -use flate2::read::GzDecoder; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, sqlite_real, SqlitePool}; -use aether_data_contracts::repository::usage::{ - read_decompressed_usage_json, sanitize_usage_capture_controls_for_persistence, - sanitize_usage_for_persistence, sanitize_usage_request_metadata, - usage_can_recover_terminal_failure, usage_error_category_for_status_code, - usage_lifecycle_update_allowed, usage_request_metadata_client_family, - PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary, - StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit, - StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, - StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow, - StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary, - StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount, - StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary, - StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance, - StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary, - StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, - StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UpsertUsageRecord, - UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, - UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, - UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, - 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, - "usage".request_id, - user_id, - api_key_id, - provider_name, - model, - target_model, - 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, - endpoint_kind, - endpoint_api_format, - provider_api_family, - provider_endpoint_kind, - 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, - COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0) - AS output_tokens, - total_tokens, - 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, - 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, - 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#" -INSERT INTO "usage" ( - request_id, - 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, - input_tokens, - 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, - status_code, - error_message, - error_category, - response_time_ms, - first_byte_time_ms, - status, - billing_status, - request_metadata, - candidate_id, - candidate_index, - key_name, - planner_kind, - route_family, - route_kind, - execution_path, - local_execution_runtime_miss_reason, - finalized_at, - created_at_unix_ms, - updated_at_unix_secs -) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ? -) -ON CONFLICT (request_id) DO UPDATE SET - user_id = excluded.user_id, - api_key_id = excluded.api_key_id, - provider_name = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_name ELSE excluded.provider_name END, - model = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".model ELSE excluded.model END, - target_model = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".target_model ELSE excluded.target_model END, - provider_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_id ELSE excluded.provider_id END, - provider_endpoint_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_endpoint_id ELSE excluded.provider_endpoint_id END, - provider_api_key_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_api_key_id ELSE excluded.provider_api_key_id END, - request_type = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".request_type ELSE excluded.request_type END, - api_format = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".api_format ELSE excluded.api_format END, - api_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".api_family ELSE excluded.api_family END, - endpoint_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".endpoint_kind ELSE excluded.endpoint_kind END, - endpoint_api_format = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".endpoint_api_format ELSE excluded.endpoint_api_format END, - provider_api_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_api_family ELSE excluded.provider_api_family END, - provider_endpoint_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_endpoint_kind ELSE excluded.provider_endpoint_kind END, - has_format_conversion = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".has_format_conversion ELSE excluded.has_format_conversion END, - is_stream = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".is_stream ELSE excluded.is_stream END, - upstream_is_stream = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".upstream_is_stream ELSE excluded.upstream_is_stream END, - input_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".input_tokens - ELSE excluded.input_tokens - END, - output_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_tokens - ELSE excluded.output_tokens - END, - total_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_tokens - ELSE excluded.total_tokens - END, - cache_creation_input_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_input_tokens - ELSE excluded.cache_creation_input_tokens - END, - cache_creation_ephemeral_5m_input_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_5m_input_tokens - ELSE excluded.cache_creation_ephemeral_5m_input_tokens - END, - cache_creation_ephemeral_1h_input_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_1h_input_tokens - ELSE excluded.cache_creation_ephemeral_1h_input_tokens - END, - cache_read_input_tokens = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_input_tokens - ELSE excluded.cache_read_input_tokens - END, - cache_creation_cost_usd = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_cost_usd - ELSE excluded.cache_creation_cost_usd - END, - cache_read_cost_usd = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_cost_usd - ELSE excluded.cache_read_cost_usd - END, - output_price_per_1m = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_price_per_1m - ELSE excluded.output_price_per_1m - END, - total_cost_usd = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_cost_usd - ELSE excluded.total_cost_usd - END, - actual_total_cost_usd = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".actual_total_cost_usd - ELSE excluded.actual_total_cost_usd - END, - status_code = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status_code - WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status_code - WHEN "usage".status = 'streaming' AND excluded.status = 'streaming' AND excluded.status_code IS NULL THEN "usage".status_code - ELSE excluded.status_code - END, - error_message = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".error_message - WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".error_message - ELSE excluded.error_message - END, - error_category = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".error_category - WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".error_category - ELSE excluded.error_category - END, - response_time_ms = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".response_time_ms - WHEN excluded.response_time_ms IS NULL OR excluded.response_time_ms = 0 THEN COALESCE("usage".response_time_ms, excluded.response_time_ms) - ELSE excluded.response_time_ms - END, - first_byte_time_ms = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".first_byte_time_ms - WHEN excluded.first_byte_time_ms IS NULL OR excluded.first_byte_time_ms = 0 THEN COALESCE("usage".first_byte_time_ms, excluded.first_byte_time_ms) - ELSE excluded.first_byte_time_ms - END, - status = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status - WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status - ELSE excluded.status - END, - billing_status = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".billing_status - ELSE excluded.billing_status - END, - request_metadata = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".request_metadata ELSE excluded.request_metadata END, - candidate_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".candidate_id WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.candidate_id ELSE COALESCE(excluded.candidate_id, "usage".candidate_id) END, - candidate_index = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".candidate_index WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.candidate_index ELSE COALESCE(excluded.candidate_index, "usage".candidate_index) END, - key_name = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".key_name WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.key_name ELSE COALESCE(excluded.key_name, "usage".key_name) END, - planner_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".planner_kind ELSE excluded.planner_kind END, - route_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".route_family ELSE excluded.route_family END, - route_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".route_kind ELSE excluded.route_kind END, - execution_path = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".execution_path ELSE excluded.execution_path END, - local_execution_runtime_miss_reason = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".local_execution_runtime_miss_reason ELSE excluded.local_execution_runtime_miss_reason END, - finalized_at = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".finalized_at - ELSE excluded.finalized_at - END, - updated_at_unix_secs = CASE - WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".updated_at_unix_secs - ELSE excluded.updated_at_unix_secs - END -"#; - -const 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 - AND excluded.updated_at_unix_secs >= COALESCE( - NULLIF("usage".updated_at_unix_secs, 0), - COALESCE("usage".created_at_unix_ms, 0) - ) -"#; - -const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#" -SELECT - "usage".request_id, - COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) AS billing_status -FROM "usage" -LEFT JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = "usage".request_id -WHERE "usage".status IN ('pending', 'streaming') - AND "usage".created_at_unix_ms < ? -ORDER BY "usage".created_at_unix_ms ASC, "usage".request_id ASC -LIMIT ? -"#; - -const SELECT_COMPLETED_REQUEST_CANDIDATES_SQL: &str = r#" -SELECT status, extra_data -FROM request_candidates -WHERE request_id = ? - AND status IN ('streaming', 'success') -"#; - -const SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED: &str = r#" -( - ( - provider_id IS NOT NULL - AND TRIM(provider_id) <> '' - AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending') - ) - OR ( - provider_name IS NOT NULL - AND TRIM(provider_name) <> '' - AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') - ) -) -"#; - -const SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR: &str = r#" -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 -"#; - -const SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR: &str = r#" -CASE - WHEN ( - LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'openai' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'openai:%' - ) - AND COALESCE(input_tokens, 0) > 0 - AND ( - COALESCE(cache_creation_input_tokens, 0) > 0 - OR COALESCE(cache_read_input_tokens, 0) > 0 - ) - THEN MAX( - COALESCE(input_tokens, 0) - - COALESCE(cache_creation_input_tokens, 0) - - COALESCE(cache_read_input_tokens, 0), - 0 - ) - WHEN ( - 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 -"#; - -const SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR: &str = r#" -CASE - WHEN ( - LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'openai' - OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'openai:%' - ) - THEN ( - CASE - WHEN COALESCE(input_tokens, 0) > 0 - AND ( - COALESCE(cache_creation_input_tokens, 0) > 0 - OR COALESCE(cache_read_input_tokens, 0) > 0 - ) - THEN MAX( - COALESCE(input_tokens, 0) - - COALESCE(cache_creation_input_tokens, 0) - - COALESCE(cache_read_input_tokens, 0), - 0 - ) - ELSE MAX(COALESCE(input_tokens, 0), 0) - END - ) + MAX(COALESCE(cache_creation_input_tokens, 0), 0) - + MAX(COALESCE(cache_read_input_tokens, 0), 0) - WHEN ( - 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:%' - ) - THEN ( - CASE - WHEN 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(cache_read_input_tokens, 0), 0) - ELSE MAX(COALESCE(input_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 MAX(COALESCE(cache_creation_input_tokens, 0), 0) - END - ) - + MAX(COALESCE(cache_read_input_tokens, 0), 0) -END -"#; - -const SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR: &str = r#" -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 ( - 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 -) -"#; - -const SQLITE_USAGE_SUCCESS_FLAG_EXPR: &str = r#" -CASE - WHEN status <> 'failed' - AND (status_code IS NULL OR status_code < 400) - AND error_message IS NULL - THEN 1 - ELSE 0 -END -"#; - -const SQLITE_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 SQLITE_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 -"#; - -const SQLITE_MONITORING_ERROR_PREDICATE: &str = r#" -( - LOWER(TRIM(COALESCE(status, ''))) IN ('failed', 'error') - OR (error_category IS NOT NULL AND TRIM(error_category) <> '') - OR ( - TRIM(COALESCE(status, '')) = '' - AND ( - COALESCE(status_code, 0) >= 400 - OR (error_message IS NOT NULL AND TRIM(error_message) <> '') - ) - ) -) -"#; - -const SQLITE_FINALIZED_USAGE_PREDICATE: &str = r#" -status NOT IN ('pending', 'streaming') -AND provider_name NOT IN ('unknown', 'pending') -"#; - -fn push_sqlite_usage_where(builder: &mut QueryBuilder<'_, Sqlite>, has_where: &mut bool) { - builder.push(if *has_where { " AND " } else { " WHERE " }); - *has_where = true; -} - -fn push_sqlite_usage_list_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - query: &UsageAuditListQuery, - has_where: &mut bool, -) { - if let Some(created_from_unix_secs) = query.created_from_unix_secs { - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms >= ") - .push_bind(created_from_unix_secs as i64); - } - if let Some(created_until_unix_secs) = query.created_until_unix_secs { - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms < ") - .push_bind(created_until_unix_secs as i64); - } - if let Some(user_id) = query.user_id.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder.push("user_id = ").push_bind(user_id.to_string()); - } - if let Some(provider_name) = query.provider_name.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder - .push("provider_name = ") - .push_bind(provider_name.to_string()); - } - if let Some(model) = query.model.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder.push("model = ").push_bind(model.to_string()); - } - if let Some(api_format) = query.api_format.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder - .push("api_format = ") - .push_bind(api_format.to_string()); - } - if let Some(client_family) = query.client_family.as_deref().map(str::trim) { - if !client_family.is_empty() { - push_sqlite_usage_where(builder, has_where); - builder - .push("LOWER(COALESCE(NULLIF(TRIM(CAST(json_extract(request_metadata, '$.client_session_affinity.client_family') AS TEXT)), ''), NULLIF(TRIM(CAST(json_extract(request_metadata, '$.client_family') AS TEXT)), ''))) = ") - .push_bind(client_family.to_ascii_lowercase()); - } - } - if query.exclude_unknown_model_or_provider { - push_sqlite_usage_where(builder, has_where); - builder.push( - "(LOWER(TRIM(COALESCE(model, ''))) NOT IN ('unknown', 'unknow') \ -AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow'))", - ); - } - if let Some(statuses) = query.statuses.as_deref() { - if !statuses.is_empty() { - push_sqlite_usage_where(builder, has_where); - builder.push("status IN ("); - let mut separated = builder.separated(", "); - for status in statuses { - separated.push_bind(status.to_string()); - } - separated.push_unseparated(")"); - } - } - push_sqlite_usage_excluded_status_codes(builder, has_where, &query.exclude_status_codes); - if let Some(is_stream) = query.is_stream { - push_sqlite_usage_where(builder, has_where); - builder - .push("is_stream = ") - .push_bind(if is_stream { 1_i64 } else { 0_i64 }); - } - if let Some(is_websocket) = query.is_websocket { - push_sqlite_usage_where(builder, has_where); - builder - .push("COALESCE(CAST(json_extract(request_metadata, '$.websocket_mode') AS INTEGER), 0) = ") - .push_bind(if is_websocket { 1_i64 } else { 0_i64 }); - } - if query.error_only { - push_sqlite_usage_where(builder, has_where); - builder.push( - "(status = 'failed' \ -OR COALESCE(status_code, 0) >= 400 \ -OR (error_message IS NOT NULL AND TRIM(error_message) <> ''))", - ); - } -} - -fn push_sqlite_usage_excluded_status_codes( - builder: &mut QueryBuilder<'_, Sqlite>, - has_where: &mut bool, - status_codes: &[u16], -) { - if status_codes.is_empty() { - return; - } - push_sqlite_usage_where(builder, has_where); - builder.push("(status_code IS NULL OR status_code NOT IN ("); - let mut separated = builder.separated(", "); - for status_code in status_codes { - separated.push_bind(i64::from(*status_code)); - } - separated.push_unseparated("))"); -} - -fn push_sqlite_usage_keyword_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - query: &UsageAuditKeywordSearchQuery, - has_where: &mut bool, -) { - push_sqlite_usage_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, - is_websocket: query.is_websocket, - 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_sqlite_usage_where(builder, has_where); - builder.push("("); - builder - .push("LOWER(COALESCE(model, '')) LIKE ") - .push_bind(pattern.clone()); - builder - .push(" OR LOWER(COALESCE(provider_name, '')) LIKE ") - .push_bind(pattern.clone()); - if query.auth_user_reader_available { - let matched_user_ids = query - .matched_user_ids_by_keyword - .get(index) - .cloned() - .unwrap_or_default(); - if !matched_user_ids.is_empty() { - builder.push(" OR user_id IN ("); - let mut separated = builder.separated(", "); - for user_id in matched_user_ids { - separated.push_bind(user_id); - } - separated.push_unseparated(")"); - } - } else { - builder - .push(" OR user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(username, '')) LIKE ") - .push_bind(pattern.clone()); - builder.push(")"); - } - if query.auth_api_key_reader_available { - let matched_api_key_ids = query - .matched_api_key_ids_by_keyword - .get(index) - .cloned() - .unwrap_or_default(); - if !matched_api_key_ids.is_empty() { - builder.push(" OR api_key_id IN ("); - let mut separated = builder.separated(", "); - for api_key_id in matched_api_key_ids { - separated.push_bind(api_key_id); - } - separated.push_unseparated(")"); - } - } else { - builder - .push(" OR api_key_id IN (SELECT id FROM api_keys WHERE LOWER(COALESCE(name, '')) LIKE ") - .push_bind(pattern); - builder.push(")"); - } - builder.push(")"); - } - - if let Some(username_keyword) = query - .username_keyword - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - push_sqlite_usage_where(builder, has_where); - if query.auth_user_reader_available { - if query.matched_user_ids_for_username.is_empty() { - builder.push("0 = 1"); - } else { - builder.push("user_id IN ("); - let mut separated = builder.separated(", "); - for user_id in &query.matched_user_ids_for_username { - separated.push_bind(user_id.clone()); - } - separated.push_unseparated(")"); - } - } else { - builder - .push("user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(username, '')) LIKE ") - .push_bind(format!("%{}%", username_keyword.to_ascii_lowercase())); - builder.push(")"); - } - } -} - -fn push_sqlite_usage_summary_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - query: &UsageAuditSummaryQuery, - has_where: &mut bool, -) { - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms >= ") - .push_bind(query.created_from_unix_secs as i64); - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms < ") - .push_bind(query.created_until_unix_secs as i64); - if let Some(user_id) = query.user_id.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder.push("user_id = ").push_bind(user_id.to_string()); - } - if let Some(provider_name) = query.provider_name.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder - .push("provider_name = ") - .push_bind(provider_name.to_string()); - } - if let Some(model) = query.model.as_deref() { - push_sqlite_usage_where(builder, has_where); - builder.push("model = ").push_bind(model.to_string()); - } -} - -fn push_sqlite_usage_order_limit_offset( - builder: &mut QueryBuilder<'_, Sqlite>, - newest_first: bool, - limit: Option, - offset: Option, -) { - if newest_first { - builder.push(" ORDER BY created_at_unix_ms DESC, id ASC"); - } else { - 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); - } - if let Some(offset) = offset { - builder.push(" OFFSET ").push_bind(offset as i64); - } -} - -fn sqlite_usage_aggregation_group_expr(group_by: UsageAuditAggregationGroupBy) -> &'static str { - match group_by { - UsageAuditAggregationGroupBy::Model => "COALESCE(NULLIF(model, ''), 'unknown')", - UsageAuditAggregationGroupBy::Provider => { - "CASE WHEN provider_id IS NOT NULL \ -AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending') \ -THEN TRIM(provider_id) ELSE TRIM(provider_name) END" - } - UsageAuditAggregationGroupBy::ApiFormat => "COALESCE(NULLIF(api_format, ''), 'unknown')", - UsageAuditAggregationGroupBy::User => "user_id", - } -} - -fn sqlite_usage_aggregation_secondary_expr(group_by: UsageAuditAggregationGroupBy) -> &'static str { - match group_by { - UsageAuditAggregationGroupBy::Provider => { - "CASE WHEN SUM(CASE WHEN provider_id IS NOT NULL \ -AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending') \ -THEN 1 ELSE 0 END) > 0 THEN 'provider_id' \ -WHEN SUM(CASE WHEN provider_name IS NOT NULL \ -AND TRIM(provider_name) <> '' \ -AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') \ -THEN 1 ELSE 0 END) > 0 THEN 'legacy_name' \ -ELSE NULL END" - } - _ => "NULL", - } -} - -fn sqlite_aggregate_u64(row: &SqliteRow, field: &str) -> Result { - Ok(row.try_get::(field).map_sql_err()?.max(0) as u64) -} - -fn sqlite_optional_u64(row: &SqliteRow, field: &str) -> Result, DataLayerError> { - Ok(row - .try_get::, _>(field) - .map_sql_err()? - .map(|value| value.max(0) as u64)) -} - -fn push_sqlite_usage_provider_performance_base_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - query: &UsageProviderPerformanceQuery, - has_where: &mut bool, -) { - push_sqlite_usage_range( - builder, - has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(builder, has_where); - builder.push( - "COALESCE(status, '') NOT IN ('pending', 'streaming') \ -AND provider_id IS NOT NULL AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'pending') \ -AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'pending')", - ); - push_sqlite_usage_provider_performance_filters(builder, query, has_where); -} - -fn decode_sqlite_usage_audit_summary_row( - row: &SqliteRow, -) -> Result { - Ok(StoredUsageAuditSummary { - total_requests: sqlite_aggregate_u64(row, "total_requests")?, - input_tokens: sqlite_aggregate_u64(row, "input_tokens")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - recorded_total_tokens: sqlite_aggregate_u64(row, "recorded_total_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(row, "cache_creation_tokens")?, - cache_creation_ephemeral_5m_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_5m_tokens", - )?, - cache_creation_ephemeral_1h_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_1h_tokens", - )?, - cache_read_tokens: sqlite_aggregate_u64(row, "cache_read_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?, - cache_creation_cost_usd: sqlite_real(row, "cache_creation_cost_usd")?, - cache_read_cost_usd: sqlite_real(row, "cache_read_cost_usd")?, - total_response_time_ms: sqlite_real(row, "total_response_time_ms")?, - error_requests: sqlite_aggregate_u64(row, "error_requests")?, - }) -} - -fn decode_sqlite_usage_aggregation_row( - row: &SqliteRow, -) -> Result { - Ok(StoredUsageAuditAggregation { - group_key: row.try_get::("group_key").map_sql_err()?, - display_name: row.try_get("display_name").map_sql_err()?, - secondary_name: row.try_get("secondary_name").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - effective_input_tokens: sqlite_aggregate_u64(row, "effective_input_tokens")?, - total_input_context: sqlite_aggregate_u64(row, "total_input_context")?, - cache_creation_tokens: sqlite_aggregate_u64(row, "cache_creation_tokens")?, - cache_creation_ephemeral_5m_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_5m_tokens", - )?, - cache_creation_ephemeral_1h_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_1h_tokens", - )?, - cache_read_tokens: sqlite_aggregate_u64(row, "cache_read_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?, - avg_response_time_ms: sqlite_optional_real(row, "avg_response_time_ms")?, - success_count: row - .try_get::, _>("success_count") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }) -} - -fn push_sqlite_usage_range( - builder: &mut QueryBuilder<'_, Sqlite>, - has_where: &mut bool, - created_from_unix_secs: u64, - created_until_unix_secs: u64, -) { - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms >= ") - .push_bind(created_from_unix_secs as i64); - push_sqlite_usage_where(builder, has_where); - builder - .push("created_at_unix_ms < ") - .push_bind(created_until_unix_secs as i64); -} - -fn push_sqlite_usage_optional_text_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - has_where: &mut bool, - column: &str, - value: Option<&str>, -) { - if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) { - push_sqlite_usage_where(builder, has_where); - builder - .push(column) - .push(" = ") - .push_bind(value.to_string()); - } -} - -fn push_sqlite_usage_bool_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - has_where: &mut bool, - column: &str, - value: Option, -) { - if let Some(value) = value { - push_sqlite_usage_where(builder, has_where); - builder - .push(column) - .push(" = ") - .push_bind(if value { 1_i64 } else { 0_i64 }); - } -} - -fn push_sqlite_usage_finalized_filter( - builder: &mut QueryBuilder<'_, Sqlite>, - has_where: &mut bool, -) { - push_sqlite_usage_where(builder, has_where); - builder.push(SQLITE_FINALIZED_USAGE_PREDICATE); -} - -fn push_sqlite_usage_provider_performance_filters( - builder: &mut QueryBuilder<'_, Sqlite>, - query: &UsageProviderPerformanceQuery, - has_where: &mut bool, -) { - push_sqlite_usage_optional_text_filter( - builder, - has_where, - "provider_id", - query.provider_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter(builder, has_where, "model", query.model.as_deref()); - push_sqlite_usage_optional_text_filter( - builder, - has_where, - "api_format", - query.api_format.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - builder, - has_where, - "endpoint_kind", - query.endpoint_kind.as_deref(), - ); - push_sqlite_usage_bool_filter(builder, has_where, "is_stream", query.is_stream); - push_sqlite_usage_bool_filter( - builder, - has_where, - "has_format_conversion", - query.has_format_conversion, - ); -} - -fn sqlite_usage_metadata_input_price_expr() -> &'static str { - r#" -COALESCE( - CAST(json_extract(request_metadata, '$.input_price_per_1m') AS REAL), - CAST(json_extract(request_metadata, '$.settlement_snapshot.input_price_per_1m') AS REAL), - CAST(json_extract(request_metadata, '$.billing_snapshot.input_price_per_1m') AS REAL), - 0 -) -"# -} - -fn sqlite_usage_bucket_expr( - granularity: UsageTimeSeriesGranularity, - tz_offset_minutes: i32, -) -> String { - let offset = i64::from(tz_offset_minutes) * 60; - match granularity { - UsageTimeSeriesGranularity::Hour => { - format!( - "strftime('%Y-%m-%dT%H:00:00+00:00', created_at_unix_ms + ({offset}), 'unixepoch')" - ) - } - UsageTimeSeriesGranularity::Day => { - format!("date(created_at_unix_ms + ({offset}), 'unixepoch')") - } - } -} - -fn sqlite_usage_local_date_expr(tz_offset_minutes: i32) -> String { - let offset = i64::from(tz_offset_minutes) * 60; - format!("date(created_at_unix_ms + ({offset}), 'unixepoch')") -} - -fn sqlite_usage_breakdown_group_expr(group_by: UsageBreakdownGroupBy) -> &'static str { - match group_by { - UsageBreakdownGroupBy::Model => "COALESCE(NULLIF(model, ''), 'unknown')", - UsageBreakdownGroupBy::Provider => "COALESCE(NULLIF(provider_name, ''), 'unknown')", - UsageBreakdownGroupBy::ApiFormat => "COALESCE(NULLIF(api_format, ''), 'unknown')", - } -} - -fn sqlite_usage_leaderboard_group_expr( - group_by: UsageLeaderboardGroupBy, -) -> (&'static str, &'static str, &'static str) { - match group_by { - UsageLeaderboardGroupBy::Model => { - ("model", "NULL", "model IS NOT NULL AND TRIM(model) <> ''") - } - UsageLeaderboardGroupBy::User => ( - "user_id", - "(SELECT users.username FROM users WHERE users.id = user_id)", - "user_id IS NOT NULL AND TRIM(user_id) <> ''", - ), - UsageLeaderboardGroupBy::ApiKey => ( - "api_key_id", - "(SELECT api_keys.name FROM api_keys WHERE api_keys.id = api_key_id)", - "api_key_id IS NOT NULL AND TRIM(api_key_id) <> ''", - ), - } -} - -fn sqlite_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_usage_json_value(bytes: &[u8]) -> Result { - let json_bytes = read_decompressed_usage_json(GzDecoder::new(bytes))?; - serde_json::from_slice(&json_bytes).map_err(|err| { - DataLayerError::UnexpectedValue(format!("failed to parse decompressed usage json: {err}")) - }) -} - -fn parse_usage_json_text(raw: &str) -> Result { - serde_json::from_str(raw).map_err(|err| { - DataLayerError::UnexpectedValue(format!("failed to parse usage json: {err}")) - }) -} - -#[derive(Debug, Clone)] -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 = sanitize_usage_for_persistence(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, -} - -#[derive(Debug, Clone, Copy, Default)] -struct SqliteProviderPerformancePercentiles { - p90_response_time_ms: Option, - p99_response_time_ms: Option, - p90_first_byte_time_ms: Option, - p99_first_byte_time_ms: Option, -} - -impl SqliteUsageReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn summarize_usage_daily_heatmap_raw_from_range( - &self, - created_from_unix_secs: u64, - created_until_unix_secs: u64, - user_id: Option<&str>, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - date(created_at_unix_ms, 'unixepoch') AS date, - COUNT(*) AS requests, - COALESCE(SUM( - MAX(COALESCE(input_tokens, 0), 0) - + MAX(COALESCE(output_tokens, 0), 0) - + {cache_creation_expr} - + MAX(COALESCE(cache_read_input_tokens, 0), 0) - ), 0) AS total_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0) - AS actual_total_cost_usd -FROM "usage" -"#, - cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR - )); - let mut has_where = false; - push_sqlite_usage_where(&mut builder, &mut has_where); - builder - .push("created_at_unix_ms >= ") - .push_bind(created_from_unix_secs as i64); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder - .push("created_at_unix_ms < ") - .push_bind(created_until_unix_secs as i64); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_optional_text_filter(&mut builder, &mut has_where, "user_id", user_id); - builder.push(" GROUP BY date ORDER BY date ASC"); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_sqlite_usage_daily_summary).collect() - } - - async fn summarize_usage_daily_heatmap_from_daily_aggregates( - &self, - created_from_unix_secs: u64, - created_until_unix_secs: u64, - user_id: Option<&str>, - ) -> Result, DataLayerError> { - let rows = if let Some(user_id) = user_id { - sqlx::query( - r#" -SELECT - date("date", 'unixepoch') AS date, - total_requests AS requests, - input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, - total_cost AS total_cost_usd, - total_cost AS actual_total_cost_usd -FROM stats_user_daily -WHERE user_id = ? - AND "date" >= ? - AND "date" < ? - AND total_requests > 0 -ORDER BY "date" ASC -"#, - ) - .bind(user_id) - .bind(created_from_unix_secs as i64) - .bind(created_until_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - date("date", 'unixepoch') AS date, - total_requests AS requests, - input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens, - total_cost AS total_cost_usd, - actual_total_cost AS actual_total_cost_usd -FROM stats_daily -WHERE "date" >= ? - AND "date" < ? - AND total_requests > 0 -ORDER BY "date" ASC -"#, - ) - .bind(created_from_unix_secs as i64) - .bind(created_until_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()? - }; - - rows.iter().map(map_sqlite_usage_daily_summary).collect() - } - - async fn summarize_provider_performance_percentiles( - &self, - query: &UsageProviderPerformanceQuery, - ) -> Result { - let mut builder = QueryBuilder::::new( - r#" -WITH filtered_usage AS ( - SELECT - MAX(COALESCE(response_time_ms, 0), 0) AS response_time_ms, - MAX(COALESCE(first_byte_time_ms, 0), 0) AS first_byte_time_ms, - response_time_ms IS NOT NULL AS has_response_time, - first_byte_time_ms IS NOT NULL AS has_first_byte_time, - CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 - ELSE 0 - END AS success_flag - FROM "usage" -"#, - ); - let mut has_where = false; - push_sqlite_usage_provider_performance_base_filters(&mut builder, query, &mut has_where); - builder.push( - r#" -), -response_ranked AS ( - SELECT - response_time_ms AS value, - ROW_NUMBER() OVER (ORDER BY response_time_ms) AS rn, - COUNT(response_time_ms) OVER () AS n - FROM filtered_usage - WHERE success_flag = 1 AND has_response_time -), -first_byte_ranked AS ( - SELECT - first_byte_time_ms AS value, - ROW_NUMBER() OVER (ORDER BY first_byte_time_ms) AS rn, - COUNT(first_byte_time_ms) OVER () AS n - FROM filtered_usage - WHERE success_flag = 1 AND has_first_byte_time -), -response_positions AS ( - SELECT - n, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM response_ranked - GROUP BY n -), -first_byte_positions AS ( - SELECT - n, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM first_byte_ranked - GROUP BY n -), -response_percentiles AS ( - SELECT - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_response_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_response_time_ms - FROM response_positions AS positions - JOIN response_ranked AS ranked ON ranked.n = positions.n - GROUP BY positions.n, positions.p90_pos, positions.p99_pos -), -first_byte_percentiles AS ( - SELECT - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_first_byte_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_first_byte_time_ms - FROM first_byte_positions AS positions - JOIN first_byte_ranked AS ranked ON ranked.n = positions.n - GROUP BY positions.n, positions.p90_pos, positions.p99_pos -) -SELECT - (SELECT p90_response_time_ms FROM response_percentiles) AS p90_response_time_ms, - (SELECT p99_response_time_ms FROM response_percentiles) AS p99_response_time_ms, - (SELECT p90_first_byte_time_ms FROM first_byte_percentiles) AS p90_first_byte_time_ms, - (SELECT p99_first_byte_time_ms FROM first_byte_percentiles) AS p99_first_byte_time_ms -"#, - ); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(SqliteProviderPerformancePercentiles { - p90_response_time_ms: sqlite_optional_u64(&row, "p90_response_time_ms")?, - p99_response_time_ms: sqlite_optional_u64(&row, "p99_response_time_ms")?, - p90_first_byte_time_ms: sqlite_optional_u64(&row, "p90_first_byte_time_ms")?, - p99_first_byte_time_ms: sqlite_optional_u64(&row, "p99_first_byte_time_ms")?, - }) - } - - async fn summarize_provider_performance_provider_percentiles( - &self, - query: &UsageProviderPerformanceQuery, - provider_ids: &[String], - ) -> Result, DataLayerError> { - if provider_ids.is_empty() { - return Ok(BTreeMap::new()); - } - - let mut builder = QueryBuilder::::new( - r#" -WITH filtered_usage AS ( - SELECT - provider_id, - MAX(COALESCE(response_time_ms, 0), 0) AS response_time_ms, - MAX(COALESCE(first_byte_time_ms, 0), 0) AS first_byte_time_ms, - response_time_ms IS NOT NULL AS has_response_time, - first_byte_time_ms IS NOT NULL AS has_first_byte_time, - CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 - ELSE 0 - END AS success_flag - FROM "usage" -"#, - ); - let mut has_where = false; - push_sqlite_usage_provider_performance_base_filters(&mut builder, query, &mut has_where); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("provider_id IN ("); - { - let mut separated = builder.separated(", "); - for provider_id in provider_ids { - separated.push_bind(provider_id.clone()); - } - } - builder.push( - r#") -), -response_ranked AS ( - SELECT - provider_id, - response_time_ms AS value, - ROW_NUMBER() OVER (PARTITION BY provider_id ORDER BY response_time_ms) AS rn, - COUNT(response_time_ms) OVER (PARTITION BY provider_id) AS n - FROM filtered_usage - WHERE success_flag = 1 AND has_response_time -), -first_byte_ranked AS ( - SELECT - provider_id, - first_byte_time_ms AS value, - ROW_NUMBER() OVER (PARTITION BY provider_id ORDER BY first_byte_time_ms) AS rn, - COUNT(first_byte_time_ms) OVER (PARTITION BY provider_id) AS n - FROM filtered_usage - WHERE success_flag = 1 AND has_first_byte_time -), -response_positions AS ( - SELECT - provider_id, - n, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM response_ranked - GROUP BY provider_id, n -), -first_byte_positions AS ( - SELECT - provider_id, - n, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM first_byte_ranked - GROUP BY provider_id, n -), -response_percentiles AS ( - SELECT - positions.provider_id, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_response_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_response_time_ms - FROM response_positions AS positions - JOIN response_ranked AS ranked - ON ranked.provider_id = positions.provider_id AND ranked.n = positions.n - GROUP BY positions.provider_id, positions.n, positions.p90_pos, positions.p99_pos -), -first_byte_percentiles AS ( - SELECT - positions.provider_id, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_first_byte_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_first_byte_time_ms - FROM first_byte_positions AS positions - JOIN first_byte_ranked AS ranked - ON ranked.provider_id = positions.provider_id AND ranked.n = positions.n - GROUP BY positions.provider_id, positions.n, positions.p90_pos, positions.p99_pos -), -provider_ids AS ( - SELECT DISTINCT provider_id FROM filtered_usage -) -SELECT - provider_ids.provider_id, - response_percentiles.p90_response_time_ms, - response_percentiles.p99_response_time_ms, - first_byte_percentiles.p90_first_byte_time_ms, - first_byte_percentiles.p99_first_byte_time_ms -FROM provider_ids -LEFT JOIN response_percentiles ON response_percentiles.provider_id = provider_ids.provider_id -LEFT JOIN first_byte_percentiles ON first_byte_percentiles.provider_id = provider_ids.provider_id -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - let mut values = BTreeMap::new(); - for row in rows { - values.insert( - row.try_get::("provider_id").map_sql_err()?, - SqliteProviderPerformancePercentiles { - p90_response_time_ms: sqlite_optional_u64(&row, "p90_response_time_ms")?, - p99_response_time_ms: sqlite_optional_u64(&row, "p99_response_time_ms")?, - p90_first_byte_time_ms: sqlite_optional_u64(&row, "p90_first_byte_time_ms")?, - p99_first_byte_time_ms: sqlite_optional_u64(&row, "p99_first_byte_time_ms")?, - }, - ); - } - Ok(values) - } - - async fn fetch_usage_items( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - 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 list_usage_audits( - &self, - query: &UsageAuditListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USAGE_COLUMNS); - let mut has_where = false; - push_sqlite_usage_list_filters(&mut builder, query, &mut has_where); - push_sqlite_usage_order_limit_offset( - &mut builder, - query.newest_first, - query.limit, - query.offset, - ); - - 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_usage_audits( - &self, - query: &UsageAuditListQuery, - ) -> Result { - let mut builder = QueryBuilder::::new(r#"SELECT COUNT(*) AS total FROM "usage""#); - let mut has_where = false; - push_sqlite_usage_list_filters(&mut builder, query, &mut has_where); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - 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_sqlite_usage_keyword_filters(&mut builder, query, &mut has_where); - push_sqlite_usage_order_limit_offset( - &mut builder, - query.newest_first, - query.limit, - query.offset, - ); - - 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_usage_audits_by_keyword_search( - &self, - query: &UsageAuditKeywordSearchQuery, - ) -> Result { - let mut builder = QueryBuilder::::new(r#"SELECT COUNT(*) AS total FROM "usage""#); - let mut has_where = false; - push_sqlite_usage_keyword_filters(&mut builder, query, &mut has_where); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - pub async fn summarize_usage_audits( - &self, - query: &UsageAuditSummaryQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageAuditSummary::default()); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -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({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, - COALESCE(SUM(MAX(COALESCE(cache_creation_ephemeral_1h_input_tokens, 0), 0)), 0) - AS cache_creation_ephemeral_1h_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0) - AS actual_total_cost_usd, - COALESCE(SUM(COALESCE(CAST(cache_creation_cost_usd AS REAL), 0)), 0) - AS cache_creation_cost_usd, - COALESCE(SUM(COALESCE(CAST(cache_read_cost_usd AS REAL), 0)), 0) - AS cache_read_cost_usd, - COALESCE(SUM(MAX(COALESCE(response_time_ms, 0), 0)), 0) AS total_response_time_ms, - COALESCE(SUM( - CASE - WHEN COALESCE(status_code, 0) >= 400 - OR (error_message IS NOT NULL AND TRIM(error_message) <> '') - THEN 1 ELSE 0 - 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, - 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); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - decode_sqlite_usage_audit_summary_row(&row) - } - - pub async fn aggregate_usage_audits( - &self, - query: &UsageAuditAggregationQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs || query.limit == 0 { - return Ok(Vec::new()); - } - - let group_expr = sqlite_usage_aggregation_group_expr(query.group_by); - let secondary_expr = sqlite_usage_aggregation_secondary_expr(query.group_by); - let display_expr = if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) { - "CASE WHEN provider_name IS NOT NULL AND TRIM(provider_name) <> '' AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') THEN TRIM(provider_name) ELSE NULL END" - } else { - "NULL" - }; - let avg_response_expr = if matches!( - query.group_by, - UsageAuditAggregationGroupBy::Provider | UsageAuditAggregationGroupBy::ApiFormat - ) { - "CASE WHEN COUNT(*) = 0 THEN 0 ELSE COALESCE(SUM(MAX(COALESCE(response_time_ms, 0), 0)), 0) * 1.0 / COUNT(*) END" - } else { - "NULL" - }; - let success_count_expr = if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) - { - "COALESCE(SUM(CASE WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled') AND (status_code IS NULL OR status_code < 400) THEN 1 ELSE 0 END), 0)" - } else { - "NULL" - }; - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {group_expr} AS group_key, - {display_expr} AS display_name, - {secondary_expr} AS secondary_name, - COUNT(*) AS request_count, - 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, - 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, - COALESCE(SUM(MAX(COALESCE(cache_creation_ephemeral_1h_input_tokens, 0), 0)), 0) - AS cache_creation_ephemeral_1h_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0) - AS actual_total_cost_usd, - {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, - total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR - )); - let mut has_where = false; - push_sqlite_usage_where(&mut builder, &mut has_where); - builder - .push("created_at_unix_ms >= ") - .push_bind(query.created_from_unix_secs as i64); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder - .push("created_at_unix_ms < ") - .push_bind(query.created_until_unix_secs as i64); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("status NOT IN ('pending', 'streaming')"); - if query.exclude_reserved_provider_labels { - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push(SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED); - } - if matches!(query.group_by, UsageAuditAggregationGroupBy::User) { - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("user_id IS NOT NULL AND TRIM(user_id) <> ''"); - } - builder - .push(" GROUP BY group_key") - .push(" ORDER BY request_count DESC, group_key ASC") - .push(" LIMIT ") - .push_bind(query.limit as i64); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(decode_sqlite_usage_aggregation_row) - .collect() - } - - pub async fn summarize_usage_totals_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let unique_user_ids = user_ids - .iter() - .cloned() - .collect::>(); - let mut totals = BTreeMap::::new(); - let mut aggregate_cutoffs = BTreeMap::::new(); - - let mut aggregate_builder = QueryBuilder::::new( - r#" -SELECT - user_id, - COALESCE(SUM(total_requests), 0) AS request_count, - COALESCE( - SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), - 0 - ) AS total_tokens, - MAX("date") AS latest_date -FROM stats_user_daily -WHERE user_id IN ( -"#, - ); - { - let mut separated = aggregate_builder.separated(", "); - for user_id in &unique_user_ids { - separated.push_bind(user_id.clone()); - } - } - aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC"); - - let aggregate_rows = aggregate_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - for row in aggregate_rows { - let user_id: String = row.try_get("user_id").map_sql_err()?; - let latest_date = row.try_get::("latest_date").map_sql_err()?.max(0) as u64; - aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400)); - totals.insert( - user_id.clone(), - StoredUsageUserTotals { - user_id, - request_count: row_u64(&row, "request_count")?, - total_tokens: row_u64(&row, "total_tokens")?, - }, - ); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - "usage".user_id, - COUNT(*) AS request_count, - 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 "); - } - let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default(); - builder - .push("SELECT ") - .push_bind(user_id.clone()) - .push(" AS user_id, ") - .push_bind(to_i64(cutoff, "usage aggregate cutoff")?) - .push(" AS cutoff_unix_secs"); - } - 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') -GROUP BY "usage".user_id -ORDER BY "usage".user_id ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - for row in rows { - let user_id: String = row.try_get("user_id").map_sql_err()?; - let entry = totals - .entry(user_id.clone()) - .or_insert_with(|| StoredUsageUserTotals { - user_id, - request_count: 0, - total_tokens: 0, - }); - entry.request_count = entry - .request_count - .saturating_add(row_u64(&row, "request_count")?); - entry.total_tokens = entry - .total_tokens - .saturating_add(row_u64(&row, "total_tokens")?); - } - Ok(totals.into_values().collect()) - } - - async fn summarize_dashboard_usage_from_daily_aggregates( - &self, - query: &UsageDashboardSummaryQuery, - ) -> Result, DataLayerError> { - let row = if let Some(user_id) = query.user_id.as_deref() { - sqlx::query( - r#" -SELECT - COALESCE(SUM(total_requests), 0) AS total_requests, - COALESCE(SUM(input_tokens), 0) AS input_tokens, - COALESCE(SUM(input_tokens), 0) AS effective_input_tokens, - COALESCE(SUM(output_tokens), 0) AS output_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, - COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context, - 0.0 AS cache_creation_cost_usd, - 0.0 AS cache_read_cost_usd, - COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS actual_total_cost_usd, - COALESCE(SUM(error_requests), 0) AS error_requests, - 0.0 AS response_time_sum_ms, - 0 AS response_time_samples -FROM stats_user_daily -WHERE user_id = ? - AND "date" >= ? - AND "date" < ? -"#, - ) - .bind(user_id) - .bind(query.created_from_unix_secs as i64) - .bind(query.created_until_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - COALESCE(SUM(total_requests), 0) AS total_requests, - COALESCE(SUM(input_tokens), 0) AS input_tokens, - COALESCE(SUM(input_tokens), 0) AS effective_input_tokens, - COALESCE(SUM(output_tokens), 0) AS output_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, - COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context, - 0.0 AS cache_creation_cost_usd, - 0.0 AS cache_read_cost_usd, - COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost AS REAL), 0)), 0) AS actual_total_cost_usd, - COALESCE(SUM(error_requests), 0) AS error_requests, - 0.0 AS response_time_sum_ms, - 0 AS response_time_samples -FROM stats_daily -WHERE "date" >= ? - AND "date" < ? -"#, - ) - .bind(query.created_from_unix_secs as i64) - .bind(query.created_until_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()? - }; - - let total_requests = sqlite_aggregate_u64(&row, "total_requests")?; - if total_requests == 0 { - return Ok(None); - } - - Ok(Some(StoredUsageDashboardSummary { - total_requests, - input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?, - effective_input_tokens: sqlite_aggregate_u64(&row, "effective_input_tokens")?, - output_tokens: sqlite_aggregate_u64(&row, "output_tokens")?, - total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?, - cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?, - total_input_context: sqlite_aggregate_u64(&row, "total_input_context")?, - cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?, - cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?, - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(&row, "actual_total_cost_usd")?, - error_requests: sqlite_aggregate_u64(&row, "error_requests")?, - response_time_sum_ms: sqlite_real(&row, "response_time_sum_ms")?, - response_time_samples: sqlite_aggregate_u64(&row, "response_time_samples")?, - })) - } - - async fn list_dashboard_daily_breakdown_from_daily_aggregates( - &self, - query: &UsageDashboardDailyBreakdownQuery, - ) -> Result, DataLayerError> { - let rows = if let Some(user_id) = query.user_id.as_deref() { - sqlx::query( - r#" -SELECT - date("date", 'unixepoch') AS date, - 'aggregate' AS model, - 'aggregate' AS provider, - COALESCE(SUM(total_requests), 0) AS requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd, - 0.0 AS response_time_sum_ms, - 0 AS response_time_samples -FROM stats_user_daily -WHERE user_id = ? - AND "date" >= ? - AND "date" < ? - AND total_requests > 0 -GROUP BY "date" -ORDER BY "date" ASC -"#, - ) - .bind(user_id) - .bind(query.created_from_unix_secs as i64) - .bind(query.created_until_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()? - } else { - sqlx::query( - r#" -SELECT - date("date", 'unixepoch') AS date, - 'aggregate' AS model, - 'aggregate' AS provider, - COALESCE(SUM(total_requests), 0) AS requests, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd, - 0.0 AS response_time_sum_ms, - 0 AS response_time_samples -FROM stats_daily -WHERE "date" >= ? - AND "date" < ? - AND total_requests > 0 -GROUP BY "date" -ORDER BY "date" ASC -"#, - ) - .bind(query.created_from_unix_secs as i64) - .bind(query.created_until_unix_secs as i64) - .fetch_all(&self.pool) - .await - .map_sql_err()? - }; - - rows.iter() - .map(|row| { - Ok(StoredUsageDashboardDailyBreakdownRow { - date: row.try_get("date").map_sql_err()?, - model: row.try_get("model").map_sql_err()?, - provider: row.try_get("provider").map_sql_err()?, - requests: sqlite_aggregate_u64(row, "requests")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - response_time_sum_ms: sqlite_real(row, "response_time_sum_ms")?, - response_time_samples: sqlite_aggregate_u64(row, "response_time_samples")?, - }) - }) - .collect() - } -} - -#[async_trait] -impl UsageReadRepository for SqliteUsageReadRepository { - 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() - } - - 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); - } - } - builder.push(") ORDER BY created_at_unix_ms DESC, id ASC"); - self.fetch_usage_items(builder).await - } - - 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), - } - } - - async fn resolve_body_ref( - &self, - body_ref: &str, - ) -> Result, DataLayerError> { - http_capture::resolve_body_ref(&self.pool, body_ref).await - } - - async fn list_usage_audits( - &self, - query: &UsageAuditListQuery, - ) -> Result, DataLayerError> { - Self::list_usage_audits(self, query).await - } - - async fn count_usage_audits(&self, query: &UsageAuditListQuery) -> Result { - Self::count_usage_audits(self, query).await - } - - async fn list_usage_audits_by_keyword_search( - &self, - query: &UsageAuditKeywordSearchQuery, - ) -> Result, DataLayerError> { - Self::list_usage_audits_by_keyword_search(self, query).await - } - - async fn count_usage_audits_by_keyword_search( - &self, - query: &UsageAuditKeywordSearchQuery, - ) -> Result { - Self::count_usage_audits_by_keyword_search(self, query).await - } - - async fn aggregate_usage_audits( - &self, - query: &UsageAuditAggregationQuery, - ) -> Result, DataLayerError> { - Self::aggregate_usage_audits(self, query).await - } - - async fn summarize_usage_audits( - &self, - query: &UsageAuditSummaryQuery, - ) -> Result { - Self::summarize_usage_audits(self, query).await - } - - async fn summarize_usage_totals_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - Self::summarize_usage_totals_by_user_ids(self, user_ids).await - } - - async fn summarize_usage_cache_hit_summary( - &self, - query: &UsageCacheHitSummaryQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageCacheHitSummary::default()); - } - let mut builder = QueryBuilder::::new( - r#" -SELECT - COUNT(*) AS total_requests, - COALESCE(SUM(CASE WHEN COALESCE(cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) - AS cache_hit_requests -FROM "usage" -"#, - ); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(StoredUsageCacheHitSummary { - total_requests: sqlite_aggregate_u64(&row, "total_requests")?, - cache_hit_requests: sqlite_aggregate_u64(&row, "cache_hit_requests")?, - }) - } - - async fn summarize_usage_settled_cost( - &self, - query: &UsageSettledCostSummaryQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageSettledCostSummary::default()); - } - let mut builder = QueryBuilder::::new( - r#" -SELECT - COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) AS total_cost_usd, - 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(cache_creation_input_tokens, 0), 0)), 0) AS cache_creation_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - MIN(finalized_at) AS first_finalized_at_unix_secs, - MAX(finalized_at) AS last_finalized_at_unix_secs -FROM "usage" -"#, - ); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "api_key_id", - query.api_key_id.as_deref(), - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("billing_status = 'settled' AND COALESCE(total_cost_usd, 0) > 0"); - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(StoredUsageSettledCostSummary { - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - total_requests: sqlite_aggregate_u64(&row, "total_requests")?, - input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?, - output_tokens: sqlite_aggregate_u64(&row, "output_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?, - cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?, - first_finalized_at_unix_secs: row - .try_get::, _>("first_finalized_at_unix_secs") - .map_sql_err()? - .map(|value| value.max(0) as u64), - last_finalized_at_unix_secs: row - .try_get::, _>("last_finalized_at_unix_secs") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }) - } - - async fn summarize_usage_cache_affinity_hit_summary( - &self, - query: &UsageCacheAffinityHitSummaryQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageCacheAffinityHitSummary::default()); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - COUNT(*) AS total_requests, - COALESCE(SUM(CASE WHEN COALESCE(cache_read_input_tokens, 0) > 0 THEN 1 ELSE 0 END), 0) - AS requests_with_cache_hit, - COALESCE(SUM(MAX(COALESCE(input_tokens, 0), 0)), 0) AS input_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM({cache_creation_expr}), 0) AS cache_creation_tokens, - COALESCE(SUM({total_input_context_expr}), 0) AS total_input_context, - COALESCE(SUM(COALESCE(CAST(cache_read_cost_usd AS REAL), 0)), 0) AS cache_read_cost_usd, - COALESCE(SUM(COALESCE(CAST(cache_creation_cost_usd AS REAL), 0)), 0) - AS cache_creation_cost_usd -FROM "usage" -"#, - cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR, - total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("status = 'completed'"); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "api_key_id", - query.api_key_id.as_deref(), - ); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(StoredUsageCacheAffinityHitSummary { - total_requests: sqlite_aggregate_u64(&row, "total_requests")?, - requests_with_cache_hit: sqlite_aggregate_u64(&row, "requests_with_cache_hit")?, - input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?, - cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?, - total_input_context: sqlite_aggregate_u64(&row, "total_input_context")?, - cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?, - cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?, - }) - } - - async fn list_usage_cache_affinity_intervals( - &self, - query: &UsageCacheAffinityIntervalQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let group_expr = match query.group_by { - UsageCacheAffinityIntervalGroupBy::User => "user_id", - UsageCacheAffinityIntervalGroupBy::ApiKey => "api_key_id", - }; - let mut builder = QueryBuilder::::new(format!( - r#" -WITH filtered_usage AS ( - SELECT - {group_expr} AS group_id, - NULL AS username, - model, - created_at_unix_ms, - id, - LAG(created_at_unix_ms) OVER ( - PARTITION BY {group_expr} - ORDER BY created_at_unix_ms ASC, id ASC - ) AS previous_created_at_unix_secs - FROM "usage" -"# - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "api_key_id", - query.api_key_id.as_deref(), - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("status = 'completed'"); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder - .push(group_expr) - .push(" IS NOT NULL AND TRIM(") - .push(group_expr) - .push(") <> ''"); - builder.push( - r#" -) -SELECT - group_id, - username, - model, - created_at_unix_ms AS created_at_unix_secs, - (created_at_unix_ms - previous_created_at_unix_secs) * 1.0 / 60.0 AS interval_minutes -FROM filtered_usage -WHERE previous_created_at_unix_secs IS NOT NULL -ORDER BY created_at_unix_ms ASC, id ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageCacheAffinityIntervalRow { - group_id: row.try_get("group_id").map_sql_err()?, - username: row.try_get("username").map_sql_err()?, - model: row.try_get("model").map_sql_err()?, - created_at_unix_secs: sqlite_aggregate_u64(row, "created_at_unix_secs")?, - interval_minutes: sqlite_real(row, "interval_minutes")?, - }) - }) - .collect() - } - - 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 - .summarize_dashboard_usage_from_daily_aggregates(query) - .await? - { - return Ok(summary); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - COUNT(*) AS total_requests, - COALESCE(SUM(MAX(COALESCE(input_tokens, 0), 0)), 0) AS input_tokens, - COALESCE(SUM({effective_input_expr}), 0) AS effective_input_tokens, - COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, - COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens, - COALESCE(SUM({cache_creation_expr}), 0) AS cache_creation_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM({total_input_context_expr}), 0) AS total_input_context, - COALESCE(SUM(COALESCE(CAST(cache_creation_cost_usd AS REAL), 0)), 0) - AS cache_creation_cost_usd, - COALESCE(SUM(COALESCE(CAST(cache_read_cost_usd AS REAL), 0)), 0) - AS cache_read_cost_usd, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0) - AS actual_total_cost_usd, - COALESCE(SUM(CASE WHEN COALESCE(status_code, 0) >= 400 OR status = 'failed' THEN 1 ELSE 0 END), 0) - AS error_requests, - 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 -"#, - effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR, - total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR, - cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR, - total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(StoredUsageDashboardSummary { - total_requests: sqlite_aggregate_u64(&row, "total_requests")?, - input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?, - effective_input_tokens: sqlite_aggregate_u64(&row, "effective_input_tokens")?, - output_tokens: sqlite_aggregate_u64(&row, "output_tokens")?, - total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?, - cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?, - total_input_context: sqlite_aggregate_u64(&row, "total_input_context")?, - cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?, - cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?, - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(&row, "actual_total_cost_usd")?, - error_requests: sqlite_aggregate_u64(&row, "error_requests")?, - response_time_sum_ms: sqlite_real(&row, "response_time_sum_ms")?, - response_time_samples: sqlite_aggregate_u64(&row, "response_time_samples")?, - }) - } - - async fn list_dashboard_daily_breakdown( - &self, - query: &UsageDashboardDailyBreakdownQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let aggregate_rows = self - .list_dashboard_daily_breakdown_from_daily_aggregates(query) - .await?; - if !aggregate_rows.is_empty() { - return Ok(aggregate_rows); - } - - let date_expr = sqlite_usage_local_date_expr(query.tz_offset_minutes); - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {date_expr} AS date, - model, - provider_name AS provider, - COUNT(*) AS requests, - 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( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - builder.push( - r#" -GROUP BY date, model, provider -ORDER BY date ASC, total_cost_usd DESC, model ASC, provider ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageDashboardDailyBreakdownRow { - date: row.try_get("date").map_sql_err()?, - model: row.try_get("model").map_sql_err()?, - provider: row.try_get("provider").map_sql_err()?, - requests: sqlite_aggregate_u64(row, "requests")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - response_time_sum_ms: sqlite_real(row, "response_time_sum_ms")?, - response_time_samples: sqlite_aggregate_u64(row, "response_time_samples")?, - }) - }) - .collect() - } - - async fn summarize_dashboard_provider_counts( - &self, - query: &UsageDashboardProviderCountsQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new( - r#" -SELECT - provider_name, - COUNT(*) AS request_count -FROM "usage" -"#, - ); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - builder.push( - r#" -GROUP BY provider_name -ORDER BY request_count DESC, provider_name ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageDashboardProviderCount { - provider_name: row.try_get("provider_name").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - }) - }) - .collect() - } - - async fn summarize_usage_breakdown( - &self, - query: &UsageBreakdownSummaryQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let group_expr = sqlite_usage_breakdown_group_expr(query.group_by); - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {group_expr} AS group_key, - COUNT(*) AS request_count, - COALESCE(SUM(MAX(COALESCE(input_tokens, 0), 0)), 0) AS input_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, - 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, - COALESCE(SUM(MAX(COALESCE(cache_creation_ephemeral_1h_input_tokens, 0), 0)), 0) - AS cache_creation_ephemeral_1h_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0) - AS actual_total_cost_usd, - COALESCE(SUM({success_flag_expr}), 0) AS success_count, - COALESCE(SUM(CASE WHEN {success_flag_expr} = 1 AND 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 {success_flag_expr} = 1 AND response_time_ms IS NOT NULL THEN 1 ELSE 0 END), 0) - AS response_time_samples, - COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL THEN MAX(COALESCE(response_time_ms, 0), 0) ELSE 0 END), 0) - AS overall_response_time_sum_ms, - 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, - total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "provider_name", - query.provider_name.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "model", - query.model.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "api_format", - query.api_format.as_deref(), - ); - push_sqlite_usage_excluded_status_codes( - &mut builder, - &mut has_where, - &query.exclude_status_codes, - ); - if matches!(query.group_by, UsageBreakdownGroupBy::ApiFormat) { - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("api_format IS NOT NULL"); - } - builder.push(" GROUP BY group_key ORDER BY request_count DESC, group_key ASC"); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageBreakdownSummaryRow { - group_key: row.try_get("group_key").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - input_tokens: sqlite_aggregate_u64(row, "input_tokens")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - effective_input_tokens: sqlite_aggregate_u64(row, "effective_input_tokens")?, - total_input_context: sqlite_aggregate_u64(row, "total_input_context")?, - cache_creation_tokens: sqlite_aggregate_u64(row, "cache_creation_tokens")?, - cache_creation_ephemeral_5m_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_5m_tokens", - )?, - cache_creation_ephemeral_1h_tokens: sqlite_aggregate_u64( - row, - "cache_creation_ephemeral_1h_tokens", - )?, - cache_read_tokens: sqlite_aggregate_u64(row, "cache_read_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?, - success_count: sqlite_aggregate_u64(row, "success_count")?, - response_time_sum_ms: sqlite_real(row, "response_time_sum_ms")?, - response_time_samples: sqlite_aggregate_u64(row, "response_time_samples")?, - overall_response_time_sum_ms: sqlite_real(row, "overall_response_time_sum_ms")?, - overall_response_time_samples: sqlite_aggregate_u64( - row, - "overall_response_time_samples", - )?, - }) - }) - .collect() - } - - async fn count_monitoring_usage_errors( - &self, - query: &UsageMonitoringErrorCountQuery, - ) -> Result { - let row = sqlx::query(&format!( - r#" -SELECT COUNT(*) AS total -FROM "usage" -WHERE created_at_unix_ms >= ? - AND created_at_unix_ms < ? - AND {error_predicate} -"#, - error_predicate = SQLITE_MONITORING_ERROR_PREDICATE - )) - .bind(query.created_from_unix_secs as i64) - .bind(query.created_until_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - sqlite_aggregate_u64(&row, "total") - } - - async fn list_monitoring_usage_errors( - &self, - query: &UsageMonitoringErrorListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USAGE_COLUMNS); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push(SQLITE_MONITORING_ERROR_PREDICATE); - builder.push(" ORDER BY created_at_unix_ms DESC, id ASC"); - if let Some(limit) = query.limit { - builder.push(" LIMIT ").push_bind(limit as i64); - } - self.fetch_usage_items(builder).await - } - - async fn summarize_usage_error_distribution( - &self, - query: &UsageErrorDistributionQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let date_expr = sqlite_usage_local_date_expr(query.tz_offset_minutes); - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {date_expr} AS date, - error_category, - COUNT(*) AS count -FROM "usage" -"# - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("error_category IS NOT NULL AND TRIM(error_category) <> ''"); - builder.push( - r#" -GROUP BY date, error_category -ORDER BY date ASC, count DESC, error_category ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageErrorDistributionRow { - date: row.try_get("date").map_sql_err()?, - error_category: row.try_get("error_category").map_sql_err()?, - count: sqlite_aggregate_u64(row, "count")?, - }) - }) - .collect() - } - - async fn summarize_usage_performance_percentiles( - &self, - query: &UsagePerformancePercentilesQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let date_expr = sqlite_usage_local_date_expr(query.tz_offset_minutes); - let mut builder = QueryBuilder::::new(format!( - r#" -WITH filtered_usage AS ( - SELECT - {date_expr} AS date, - response_time_ms, - first_byte_time_ms - FROM "usage" -"# - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push("status = 'completed'"); - builder.push( - r#" -), -response_ranked AS ( - SELECT - date, - response_time_ms AS value, - ROW_NUMBER() OVER (PARTITION BY date ORDER BY response_time_ms) AS rn, - COUNT(response_time_ms) OVER (PARTITION BY date) AS n - FROM filtered_usage - WHERE response_time_ms IS NOT NULL -), -first_byte_ranked AS ( - SELECT - date, - first_byte_time_ms AS value, - ROW_NUMBER() OVER (PARTITION BY date ORDER BY first_byte_time_ms) AS rn, - COUNT(first_byte_time_ms) OVER (PARTITION BY date) AS n - FROM filtered_usage - WHERE first_byte_time_ms IS NOT NULL -), -response_positions AS ( - SELECT - date, - n, - 0.5 * (n - 1) AS p50_pos, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM response_ranked - GROUP BY date, n -), -first_byte_positions AS ( - SELECT - date, - n, - 0.5 * (n - 1) AS p50_pos, - 0.9 * (n - 1) AS p90_pos, - 0.99 * (n - 1) AS p99_pos - FROM first_byte_ranked - GROUP BY date, n -), -response_percentiles AS ( - SELECT - positions.date, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p50_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p50_pos = CAST(positions.p50_pos AS INTEGER) - THEN CAST(positions.p50_pos AS INTEGER) + 1 - ELSE CAST(positions.p50_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p50_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p50_pos - CAST(positions.p50_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p50_response_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_response_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_response_time_ms - FROM response_positions AS positions - JOIN response_ranked AS ranked ON ranked.date = positions.date - GROUP BY positions.date, positions.n, positions.p50_pos, positions.p90_pos, positions.p99_pos -), -first_byte_percentiles AS ( - SELECT - positions.date, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p50_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p50_pos = CAST(positions.p50_pos AS INTEGER) - THEN CAST(positions.p50_pos AS INTEGER) + 1 - ELSE CAST(positions.p50_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p50_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p50_pos - CAST(positions.p50_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p50_first_byte_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p90_pos = CAST(positions.p90_pos AS INTEGER) - THEN CAST(positions.p90_pos AS INTEGER) + 1 - ELSE CAST(positions.p90_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p90_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p90_pos - CAST(positions.p90_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p90_first_byte_time_ms, - CASE WHEN positions.n >= 10 THEN CAST(( - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - + ( - MAX(CASE WHEN ranked.rn = CASE WHEN positions.p99_pos = CAST(positions.p99_pos AS INTEGER) - THEN CAST(positions.p99_pos AS INTEGER) + 1 - ELSE CAST(positions.p99_pos AS INTEGER) + 2 - END THEN ranked.value END) - - MAX(CASE WHEN ranked.rn = CAST(positions.p99_pos AS INTEGER) + 1 THEN ranked.value END) - ) * (positions.p99_pos - CAST(positions.p99_pos AS INTEGER)) - ) AS INTEGER) ELSE NULL END AS p99_first_byte_time_ms - FROM first_byte_positions AS positions - JOIN first_byte_ranked AS ranked ON ranked.date = positions.date - GROUP BY positions.date, positions.n, positions.p50_pos, positions.p90_pos, positions.p99_pos -), -dates AS ( - SELECT date FROM response_percentiles - UNION - SELECT date FROM first_byte_percentiles -) -SELECT - dates.date, - response_percentiles.p50_response_time_ms, - response_percentiles.p90_response_time_ms, - response_percentiles.p99_response_time_ms, - first_byte_percentiles.p50_first_byte_time_ms, - first_byte_percentiles.p90_first_byte_time_ms, - first_byte_percentiles.p99_first_byte_time_ms -FROM dates -LEFT JOIN response_percentiles ON response_percentiles.date = dates.date -LEFT JOIN first_byte_percentiles ON first_byte_percentiles.date = dates.date -ORDER BY dates.date ASC -"#, - ); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsagePerformancePercentilesRow { - date: row.try_get("date").map_sql_err()?, - p50_response_time_ms: row - .try_get::, _>("p50_response_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - p90_response_time_ms: row - .try_get::, _>("p90_response_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - p99_response_time_ms: row - .try_get::, _>("p99_response_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - p50_first_byte_time_ms: row - .try_get::, _>("p50_first_byte_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - p90_first_byte_time_ms: row - .try_get::, _>("p90_first_byte_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - p99_first_byte_time_ms: row - .try_get::, _>("p99_first_byte_time_ms") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }) - }) - .collect() - } - - async fn summarize_usage_provider_performance( - &self, - query: &UsageProviderPerformanceQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageProviderPerformance::default()); - } - - let summary_sql = format!( - r#" -SELECT - COUNT(*) AS request_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 ELSE 0 END), 0) AS success_count, - CASE - WHEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) > 0 - THEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(output_tokens, 0), 0) - ELSE 0 - END), 0) * 1000.0 / COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) - ELSE NULL - END AS avg_output_tps, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND first_byte_time_ms IS NOT NULL - THEN MAX(COALESCE(first_byte_time_ms, 0), 0) - ELSE NULL - END) AS avg_first_byte_time_ms, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND response_time_ms IS NOT NULL - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE NULL - END) AS avg_response_time_ms, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN 1 ELSE 0 END), 0) AS tps_sample_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND response_time_ms IS NOT NULL - THEN 1 ELSE 0 END), 0) AS response_time_sample_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND first_byte_time_ms IS NOT NULL - THEN 1 ELSE 0 END), 0) AS first_byte_sample_count, - COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL AND response_time_ms >= {slow_threshold} THEN 1 ELSE 0 END), 0) - AS slow_request_count -FROM "usage" -"#, - slow_threshold = query.slow_threshold_ms - ); - let mut summary_builder = QueryBuilder::::new(summary_sql); - let mut has_where = false; - push_sqlite_usage_range( - &mut summary_builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut summary_builder, &mut has_where); - summary_builder.push( - "COALESCE(status, '') NOT IN ('pending', 'streaming') \ -AND provider_id IS NOT NULL AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'pending') \ -AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'pending')", - ); - push_sqlite_usage_provider_performance_filters(&mut summary_builder, query, &mut has_where); - let summary_row = summary_builder - .build() - .fetch_one(&self.pool) - .await - .map_sql_err()?; - let summary_percentiles = self - .summarize_provider_performance_percentiles(query) - .await?; - let summary = StoredUsageProviderPerformanceSummary { - request_count: sqlite_aggregate_u64(&summary_row, "request_count")?, - success_count: sqlite_aggregate_u64(&summary_row, "success_count")?, - avg_output_tps: sqlite_optional_real(&summary_row, "avg_output_tps")?, - avg_first_byte_time_ms: sqlite_optional_real(&summary_row, "avg_first_byte_time_ms")?, - avg_response_time_ms: sqlite_optional_real(&summary_row, "avg_response_time_ms")?, - p90_response_time_ms: summary_percentiles.p90_response_time_ms, - p99_response_time_ms: summary_percentiles.p99_response_time_ms, - p90_first_byte_time_ms: summary_percentiles.p90_first_byte_time_ms, - p99_first_byte_time_ms: summary_percentiles.p99_first_byte_time_ms, - tps_sample_count: sqlite_aggregate_u64(&summary_row, "tps_sample_count")?, - response_time_sample_count: sqlite_aggregate_u64( - &summary_row, - "response_time_sample_count", - )?, - first_byte_sample_count: sqlite_aggregate_u64(&summary_row, "first_byte_sample_count")?, - slow_request_count: sqlite_aggregate_u64(&summary_row, "slow_request_count")?, - }; - - let provider_sql = format!( - r#" -SELECT - provider_id, - COALESCE(MAX(NULLIF(TRIM(provider_name), '')), provider_id) AS provider, - COUNT(*) AS request_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 ELSE 0 END), 0) AS success_count, - COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, - CASE - WHEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) > 0 - THEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(output_tokens, 0), 0) - ELSE 0 - END), 0) * 1000.0 / COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) - ELSE NULL - END AS avg_output_tps, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND first_byte_time_ms IS NOT NULL - THEN MAX(COALESCE(first_byte_time_ms, 0), 0) - ELSE NULL - END) AS avg_first_byte_time_ms, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND response_time_ms IS NOT NULL - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE NULL - END) AS avg_response_time_ms, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN 1 ELSE 0 END), 0) AS tps_sample_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND response_time_ms IS NOT NULL - THEN 1 ELSE 0 END), 0) AS response_time_sample_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND first_byte_time_ms IS NOT NULL - THEN 1 ELSE 0 END), 0) AS first_byte_sample_count, - COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL AND response_time_ms >= {slow_threshold} THEN 1 ELSE 0 END), 0) - AS slow_request_count -FROM "usage" -"#, - slow_threshold = query.slow_threshold_ms - ); - let mut provider_builder = QueryBuilder::::new(provider_sql); - let mut has_where = false; - push_sqlite_usage_range( - &mut provider_builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut provider_builder, &mut has_where); - provider_builder.push( - "COALESCE(status, '') NOT IN ('pending', 'streaming') \ -AND provider_id IS NOT NULL AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'pending') \ -AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'pending')", - ); - push_sqlite_usage_provider_performance_filters( - &mut provider_builder, - query, - &mut has_where, - ); - provider_builder - .push(" GROUP BY provider_id ORDER BY request_count DESC, provider_id ASC LIMIT "); - provider_builder.push_bind(query.limit.max(1) as i64); - let provider_rows = provider_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let provider_ids = provider_rows - .iter() - .map(|row| row.try_get::("provider_id").map_sql_err()) - .collect::, DataLayerError>>()?; - let provider_percentiles = self - .summarize_provider_performance_provider_percentiles(query, &provider_ids) - .await?; - let providers = provider_rows - .iter() - .map(|row| { - let provider_id = row.try_get::("provider_id").map_sql_err()?; - let percentiles = provider_percentiles - .get(&provider_id) - .copied() - .unwrap_or_default(); - Ok(StoredUsageProviderPerformanceProviderRow { - provider_id, - provider: row.try_get("provider").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - success_count: sqlite_aggregate_u64(row, "success_count")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - avg_output_tps: sqlite_optional_real(row, "avg_output_tps")?, - avg_first_byte_time_ms: sqlite_optional_real(row, "avg_first_byte_time_ms")?, - avg_response_time_ms: sqlite_optional_real(row, "avg_response_time_ms")?, - p90_response_time_ms: percentiles.p90_response_time_ms, - p99_response_time_ms: percentiles.p99_response_time_ms, - p90_first_byte_time_ms: percentiles.p90_first_byte_time_ms, - p99_first_byte_time_ms: percentiles.p99_first_byte_time_ms, - tps_sample_count: sqlite_aggregate_u64(row, "tps_sample_count")?, - response_time_sample_count: sqlite_aggregate_u64( - row, - "response_time_sample_count", - )?, - first_byte_sample_count: sqlite_aggregate_u64(row, "first_byte_sample_count")?, - slow_request_count: sqlite_aggregate_u64(row, "slow_request_count")?, - }) - }) - .collect::, DataLayerError>>()?; - let timeline = if !query.include_timeline || provider_ids.is_empty() { - Vec::new() - } else { - let bucket_expr = sqlite_usage_bucket_expr(query.granularity, query.tz_offset_minutes); - let mut timeline_builder = QueryBuilder::::new(format!( - r#" -SELECT - {bucket_expr} AS date, - provider_id, - COALESCE(MAX(NULLIF(TRIM(provider_name), '')), provider_id) AS provider, - COUNT(*) AS request_count, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 ELSE 0 END), 0) AS success_count, - COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens, - CASE - WHEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) > 0 - THEN COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(output_tokens, 0), 0) - ELSE 0 - END), 0) * 1000.0 / COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND COALESCE(response_time_ms, 0) > 0 - AND COALESCE(output_tokens, 0) > 0 - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE 0 - END), 0) - ELSE NULL - END AS avg_output_tps, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND first_byte_time_ms IS NOT NULL - THEN MAX(COALESCE(first_byte_time_ms, 0), 0) - ELSE NULL - END) AS avg_first_byte_time_ms, - AVG(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - AND response_time_ms IS NOT NULL - THEN MAX(COALESCE(response_time_ms, 0), 0) - ELSE NULL - END) AS avg_response_time_ms, - COALESCE(SUM(CASE WHEN response_time_ms IS NOT NULL AND response_time_ms >= {slow_threshold} THEN 1 ELSE 0 END), 0) - AS slow_request_count -FROM "usage" -"#, - slow_threshold = query.slow_threshold_ms - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut timeline_builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_where(&mut timeline_builder, &mut has_where); - timeline_builder.push( - "COALESCE(status, '') NOT IN ('pending', 'streaming') \ -AND provider_id IS NOT NULL AND TRIM(provider_id) <> '' \ -AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'pending') \ -AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'pending')", - ); - push_sqlite_usage_provider_performance_filters( - &mut timeline_builder, - query, - &mut has_where, - ); - push_sqlite_usage_where(&mut timeline_builder, &mut has_where); - timeline_builder.push("provider_id IN ("); - { - let mut separated = timeline_builder.separated(", "); - for provider_id in &provider_ids { - separated.push_bind(provider_id.clone()); - } - } - timeline_builder - .push(") GROUP BY date, provider_id ORDER BY date ASC, provider_id ASC"); - let rows = timeline_builder - .build() - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageProviderPerformanceTimelineRow { - date: row.try_get("date").map_sql_err()?, - provider_id: row.try_get("provider_id").map_sql_err()?, - provider: row.try_get("provider").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - success_count: sqlite_aggregate_u64(row, "success_count")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - avg_output_tps: sqlite_optional_real(row, "avg_output_tps")?, - avg_first_byte_time_ms: sqlite_optional_real( - row, - "avg_first_byte_time_ms", - )?, - avg_response_time_ms: sqlite_optional_real(row, "avg_response_time_ms")?, - slow_request_count: sqlite_aggregate_u64(row, "slow_request_count")?, - }) - }) - .collect::, DataLayerError>>()? - }; - - Ok(StoredUsageProviderPerformance { - summary, - providers, - timeline, - }) - } - - async fn summarize_usage_cost_savings( - &self, - query: &UsageCostSavingsSummaryQuery, - ) -> Result { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(StoredUsageCostSavingsSummary::default()); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM(COALESCE(CAST(cache_read_cost_usd AS REAL), 0)), 0) AS cache_read_cost_usd, - COALESCE(SUM(COALESCE(CAST(cache_creation_cost_usd AS REAL), 0)), 0) - AS cache_creation_cost_usd, - COALESCE(SUM( - {input_price_expr} - * MAX(COALESCE(cache_read_input_tokens, 0), 0) / 1000000.0 - ), 0) AS estimated_full_cost_usd -FROM "usage" -"#, - input_price_expr = sqlite_usage_metadata_input_price_expr() - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "provider_name", - query.provider_name.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "model", - query.model.as_deref(), - ); - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(StoredUsageCostSavingsSummary { - cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?, - cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?, - cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?, - estimated_full_cost_usd: sqlite_real(&row, "estimated_full_cost_usd")?, - }) - } - - async fn summarize_usage_time_series( - &self, - query: &UsageTimeSeriesQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let bucket_expr = sqlite_usage_bucket_expr(query.granularity, query.tz_offset_minutes); - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {bucket_expr} AS bucket_key, - 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(cache_creation_input_tokens, 0), 0)), 0) - AS cache_creation_tokens, - COALESCE(SUM(MAX(COALESCE(cache_read_input_tokens, 0), 0)), 0) AS cache_read_tokens, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd, - COALESCE(SUM(MAX(COALESCE(response_time_ms, 0), 0)), 0) AS total_response_time_ms -FROM "usage" -"# - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "provider_name", - query.provider_name.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "model", - query.model.as_deref(), - ); - builder.push(" GROUP BY bucket_key ORDER BY bucket_key ASC"); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageTimeSeriesBucket { - bucket_key: row.try_get("bucket_key").map_sql_err()?, - total_requests: sqlite_aggregate_u64(row, "total_requests")?, - input_tokens: sqlite_aggregate_u64(row, "input_tokens")?, - output_tokens: sqlite_aggregate_u64(row, "output_tokens")?, - cache_creation_tokens: sqlite_aggregate_u64(row, "cache_creation_tokens")?, - cache_read_tokens: sqlite_aggregate_u64(row, "cache_read_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - total_response_time_ms: sqlite_real(row, "total_response_time_ms")?, - }) - }) - .collect() - } - - async fn summarize_usage_leaderboard( - &self, - query: &UsageLeaderboardQuery, - ) -> Result, DataLayerError> { - if query.created_from_unix_secs >= query.created_until_unix_secs { - return Ok(Vec::new()); - } - - let (group_key_expr, legacy_name_expr, extra_filter) = - sqlite_usage_leaderboard_group_expr(query.group_by); - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - {group_key_expr} AS group_key, - MAX({legacy_name_expr}) AS legacy_name, - COUNT(*) AS request_count, - 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 - )); - let mut has_where = false; - push_sqlite_usage_range( - &mut builder, - &mut has_where, - query.created_from_unix_secs, - query.created_until_unix_secs, - ); - push_sqlite_usage_finalized_filter(&mut builder, &mut has_where); - push_sqlite_usage_where(&mut builder, &mut has_where); - builder.push(extra_filter); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "user_id", - query.user_id.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "provider_name", - query.provider_name.as_deref(), - ); - push_sqlite_usage_optional_text_filter( - &mut builder, - &mut has_where, - "model", - query.model.as_deref(), - ); - builder.push(" GROUP BY group_key ORDER BY group_key ASC"); - - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(StoredUsageLeaderboardSummary { - group_key: row.try_get("group_key").map_sql_err()?, - legacy_name: row.try_get("legacy_name").map_sql_err()?, - request_count: sqlite_aggregate_u64(row, "request_count")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - }) - }) - .collect() - } - - 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); - let mut has_where = false; - push_sqlite_usage_optional_text_filter(&mut builder, &mut has_where, "user_id", user_id); - builder.push(" ORDER BY created_at_unix_ms DESC, id ASC LIMIT "); - builder.push_bind(limit as i64); - self.fetch_usage_items(builder).await - } - - async fn summarize_total_tokens_by_api_key_ids( - &self, - api_key_ids: &[String], - ) -> Result, DataLayerError> { - if api_key_ids.is_empty() { - return Ok(BTreeMap::new()); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - "usage".api_key_id, - COALESCE(SUM({total_tokens_expr}), 0) AS total_tokens -FROM "usage" -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 \"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(); - for row in rows { - totals.insert( - row.try_get("api_key_id").map_sql_err()?, - sqlite_aggregate_u64(&row, "total_tokens")?, - ); - } - Ok(totals) - } - - async fn summarize_usage_by_provider_api_key_ids( - &self, - provider_api_key_ids: &[String], - ) -> Result, DataLayerError> { - if provider_api_key_ids.is_empty() { - return Ok(BTreeMap::new()); - } - - let mut builder = QueryBuilder::::new(format!( - r#" -SELECT - "usage".provider_api_key_id, - COUNT(*) AS request_count, - 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" -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 \"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(); - for row in rows { - let provider_api_key_id: String = row.try_get("provider_api_key_id").map_sql_err()?; - summaries.insert( - provider_api_key_id.clone(), - StoredProviderApiKeyUsageSummary { - provider_api_key_id, - request_count: sqlite_aggregate_u64(&row, "request_count")?, - total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?, - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - last_used_at_unix_secs: row - .try_get::, _>("last_used_at_unix_secs") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }, - ); - } - Ok(summaries) - } - - async fn summarize_usage_by_provider_api_key_windows( - &self, - requests: &[ProviderApiKeyWindowUsageRequest], - ) -> Result, DataLayerError> { - let mut summaries = Vec::with_capacity(requests.len()); - 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 row = sqlx::query(&format!( - r#" -SELECT - COUNT(*) AS request_count, - 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" -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) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - summaries.push(StoredProviderApiKeyWindowUsageSummary { - provider_api_key_id: provider_api_key_id.to_string(), - window_code: window_code.to_string(), - request_count: sqlite_aggregate_u64(&row, "request_count")?, - total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?, - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - }); - } - Ok(summaries) - } - - async fn summarize_provider_usage_since( - &self, - provider_id: &str, - since_unix_secs: u64, - ) -> Result { - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total_requests, - COALESCE(SUM(CASE - WHEN LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - THEN 1 ELSE 0 END), 0) AS successful_requests, - COALESCE(SUM(CASE - WHEN status NOT IN ('pending', 'streaming') - AND NOT ( - LOWER(COALESCE(status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled') - AND (status_code IS NULL OR status_code < 400) - ) - THEN 1 ELSE 0 END), 0) AS failed_requests, - COALESCE(AVG(CASE WHEN response_time_ms IS NOT NULL THEN MAX(COALESCE(response_time_ms, 0), 0) ELSE NULL END), 0) - AS avg_response_time_ms, - COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd -FROM "usage" -WHERE provider_id = ? - AND created_at_unix_ms >= ? -"#, - ) - .bind(provider_id) - .bind(since_unix_secs as i64) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - Ok(StoredProviderUsageSummary { - total_requests: sqlite_aggregate_u64(&row, "total_requests")?, - successful_requests: sqlite_aggregate_u64(&row, "successful_requests")?, - failed_requests: sqlite_aggregate_u64(&row, "failed_requests")?, - avg_response_time_ms: sqlite_real(&row, "avg_response_time_ms")?, - total_cost_usd: sqlite_real(&row, "total_cost_usd")?, - }) - } - - async fn summarize_usage_daily_heatmap( - &self, - query: &UsageDailyHeatmapQuery, - ) -> Result, DataLayerError> { - let created_until_unix_secs = usage_current_unix_secs().saturating_add(1); - let user_id = query.user_id.as_deref(); - let mut summaries = BTreeMap::::new(); - - for item in self - .summarize_usage_daily_heatmap_from_daily_aggregates( - query.created_from_unix_secs, - created_until_unix_secs, - user_id, - ) - .await? - { - summaries.insert(item.date.clone(), item); - } - for item in self - .summarize_usage_daily_heatmap_raw_from_range( - query.created_from_unix_secs, - created_until_unix_secs, - user_id, - ) - .await? - { - summaries.entry(item.date.clone()).or_insert(item); - } - - 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( - row: &SqliteRow, -) -> Result { - Ok(StoredUsageDailySummary { - date: row.try_get("date").map_sql_err()?, - requests: sqlite_aggregate_u64(row, "requests")?, - total_tokens: sqlite_aggregate_u64(row, "total_tokens")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?, - }) -} - -fn usage_current_unix_secs() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs()) - .unwrap_or_default() -} - -fn first_byte_transition_allowed( - existing: &StoredRequestUsageAudit, - incoming: &UpsertUsageRecord, -) -> bool { - existing.billing_status == "pending" - && matches!(existing.status.as_str(), "pending" | "streaming") - && existing.finalized_at_unix_secs.is_none() - && usage_lifecycle_update_allowed( - &existing.status, - &existing.billing_status, - existing.updated_at_unix_secs, - existing.finalized_at_unix_secs, - &incoming.status, - &incoming.billing_status, - incoming.updated_at_unix_secs, - incoming.finalized_at_unix_secs, - ) -} - -impl SqliteUsageWriteRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - 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), - } - } - - async fn upsert_in_tx( - tx: &mut sqlx::Transaction<'_, Sqlite>, - usage: UpsertUsageRecord, - ) -> Result<(), DataLayerError> { - usage.validate()?; - // Auxiliary tables may receive only clear tombstones, never request or response content. - let capture_usage = usage.clone(); - let mut usage = sanitize_usage_for_persistence(usage); - usage.validate()?; - let existing = counters::lock_and_load_usage(tx, &usage.request_id).await?; - if existing.as_ref().is_some_and(|existing| { - !usage_lifecycle_update_allowed( - &existing.status, - &existing.billing_status, - existing.updated_at_unix_secs, - existing.finalized_at_unix_secs, - &usage.status, - &usage.billing_status, - usage.updated_at_unix_secs, - usage.finalized_at_unix_secs, - ) - }) { - return Ok(()); - } - 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 mut capture_usage = sanitize_usage_capture_controls_for_persistence(capture_usage); - let prepared_capture = http_capture::prepare_usage_http_capture(&mut capture_usage)?; - 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 capture_usage, existing.as_ref()); - usage.request_metadata = - sanitize_usage_request_metadata(capture_usage.request_metadata.clone()); - } - let prepared_snapshots = capture_update_allowed - // The control projection preserves safe typed routing and allow-listed billing facts. - .then(|| snapshots::from_usage(&capture_usage)) - .transpose()?; - bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)? - .execute(&mut **tx) - .await - .map_sql_err()?; - 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(|existing| first_byte_transition_allowed(existing, &row.usage)) - }) - .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, &row.usage)) - { - 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 { - 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()) - }) - } - - 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(), - )); - } - } - - 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 { - sqlx::query( - r#" -UPDATE api_keys -SET total_requests = 0, - total_tokens = 0, - total_cost_usd = 0.0, - last_used_at = NULL -"#, - ) - .execute(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query(&format!( - r#" -SELECT - "usage".api_key_id, - COUNT(*) AS total_requests, - 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" -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()?; - - for row in &rows { - sqlx::query( - r#" -UPDATE api_keys -SET total_requests = ?, - total_tokens = ?, - total_cost_usd = ?, - last_used_at = ? -WHERE id = ? -"#, - ) - .bind(row.try_get::("total_requests").map_sql_err()?) - .bind(row.try_get::("total_tokens").map_sql_err()?) - .bind(sqlite_real(row, "total_cost_usd")?) - .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) - } - - async fn rebuild_provider_api_key_usage_stats(&self) -> Result { - sqlx::query( - r#" -UPDATE provider_api_keys -SET request_count = 0, - success_count = 0, - error_count = 0, - total_tokens = 0, - total_cost_usd = 0.0, - total_response_time_ms = 0, - last_used_at = NULL -"#, - ) - .execute(&self.pool) - .await - .map_sql_err()?; - - let rows = sqlx::query(&format!( - r#" -SELECT - "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 "usage".status IN ('pending', 'streaming') THEN 0 - ELSE {total_tokens_expr} - END), 0) AS total_tokens, - COALESCE(SUM(CASE - 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 "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("usage".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 "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, - total_tokens_expr = SQLITE_USAGE_CANONICAL_TOTAL_TOKENS_EXPR - )) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - - for row in &rows { - 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 = ? -"#, - ) - .bind(row.try_get::("request_count").map_sql_err()?) - .bind(row.try_get::("success_count").map_sql_err()?) - .bind(row.try_get::("error_count").map_sql_err()?) - .bind(row.try_get::("total_tokens").map_sql_err()?) - .bind(sqlite_real(row, "total_cost_usd")?) - .bind( - row.try_get::("total_response_time_ms") - .map_sql_err()?, - ) - .bind( - row.try_get::, _>("last_used_at") - .map_sql_err()?, - ) - .bind( - row.try_get::("provider_api_key_id") - .map_sql_err()?, - ) - .execute(&self.pool) - .await - .map_sql_err()?; - } - - Ok(rows.len() as u64) - } - - async fn cleanup_stale_pending_requests( - &self, - cutoff_unix_secs: u64, - now_unix_secs: u64, - _timeout_minutes: u64, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(PendingUsageCleanupSummary::default()); - } - - 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(|_| { - DataLayerError::InvalidInput(format!( - "invalid stale pending usage batch size: {batch_size}" - )) - })?; - - 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_secs, "stale pending usage cutoff")?) - .bind(to_i64(batch_size_u64, "stale pending usage batch size")?) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - - if stale_rows.is_empty() { - tx.rollback().await.map_sql_err()?; - break; - } - - let stale_rows = stale_rows - .iter() - .map(|row| { - Ok(StalePendingUsageRow { - request_id: row.try_get("request_id").map_sql_err()?, - billing_status: row.try_get("billing_status").map_sql_err()?, - }) - }) - .collect::, DataLayerError>>()?; - let completed_request_ids = - completed_request_ids_sqlite(&mut tx, stale_rows.iter().map(|row| &row.request_id)) - .await?; - - for row in stale_rows { - if completed_request_ids.contains(&row.request_id) { - sqlx::query( - r#" -UPDATE "usage" -SET status = 'completed', - status_code = 200, - error_message = NULL, - error_category = NULL -WHERE request_id = ? -"#, - ) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE request_candidates -SET status = 'success', - finished_at = ? -WHERE request_id = ? - AND status = 'streaming' -"#, - ) - .bind(to_i64(now_unix_ms, "request candidate finished_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - summary.recovered += 1; - continue; - } - - let candidate_info = - latest_failed_candidate_sqlite(&mut tx, &row.request_id).await?; - let status_code = resolve_stale_pending_status_code(candidate_info.as_ref()); - let error_category = usage_error_category_for_status_code(status_code); - let status_code_i64 = i64::from(status_code); - if row.billing_status == "pending" { - sqlx::query( - r#" -UPDATE "usage" -SET status = 'failed', - status_code = ?, - error_message = NULL, - error_category = ?, - billing_status = 'void', - finalized_at = ?, - total_cost_usd = 0.0, - actual_total_cost_usd = 0.0 -WHERE request_id = ? -"#, - ) - .bind(status_code_i64) - .bind(error_category) - .bind(to_i64(now_unix_secs, "usage finalized_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - upsert_void_usage_settlement_snapshot_sqlite( - &mut tx, - &row.request_id, - now_unix_secs, - ) - .await?; - } else { - sqlx::query( - r#" -UPDATE "usage" -SET status = 'failed', - status_code = ?, - error_message = NULL, - error_category = ? -WHERE request_id = ? -"#, - ) - .bind(status_code_i64) - .bind(error_category) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query( - r#" -UPDATE request_candidates -SET status = 'failed', - finished_at = ?, - error_type = 'internal', - error_message = NULL -WHERE request_id = ? - AND status IN ('pending', 'streaming') -"#, - ) - .bind(to_i64(now_unix_ms, "request candidate finished_at")?) - .bind(&row.request_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - summary.failed += 1; - } - - tx.commit().await.map_sql_err()?; - } - - 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 { - request_id: String, - billing_status: String, -} - -async fn completed_request_ids_sqlite<'a>( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - request_ids: impl Iterator, -) -> Result, DataLayerError> { - let mut completed = HashSet::new(); - for request_id in request_ids { - let rows = sqlx::query(SELECT_COMPLETED_REQUEST_CANDIDATES_SQL) - .bind(request_id) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - let mut is_completed = false; - for row in &rows { - if candidate_row_is_completed(row)? { - is_completed = true; - break; - } - } - if is_completed { - completed.insert(request_id.clone()); - } - } - Ok(completed) -} - -fn candidate_row_is_completed(row: &SqliteRow) -> Result { - let status: String = row.try_get("status").map_sql_err()?; - if status == "streaming" { - return Ok(true); - } - if status != "success" { - return Ok(false); - } - let Some(extra_data) = row - .try_get::, _>("extra_data") - .map_sql_err()? - else { - return Ok(false); - }; - let Ok(value) = serde_json::from_str::(&extra_data) else { - return Ok(false); - }; - Ok(value - .get("stream_completed") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false)) -} - -async fn upsert_void_usage_settlement_snapshot_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - request_id: &str, - now_unix_secs: u64, -) -> Result<(), DataLayerError> { - let now = to_i64(now_unix_secs, "usage settlement snapshot timestamp")?; - sqlx::query( - r#" -INSERT INTO usage_settlement_snapshots ( - request_id, - billing_status, - finalized_at, - created_at, - updated_at -) VALUES (?, 'void', ?, ?, ?) -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 -"#, - ) - .bind(request_id) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -struct FailedCandidateCleanupInfo { - status_code: Option, -} - -fn resolve_stale_pending_status_code(candidate: Option<&FailedCandidateCleanupInfo>) -> u16 { - candidate - .and_then(|info| info.status_code) - .unwrap_or(if candidate.is_some() { 502 } else { 504 }) -} - -async fn latest_failed_candidate_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - request_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT status_code -FROM request_candidates -WHERE request_id = ? - AND status IN ('failed', 'cancelled') -ORDER BY - COALESCE(finished_at, started_at, created_at) DESC, - retry_index DESC, - candidate_index DESC -LIMIT 1 -"#, - ) - .bind(request_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()?; - - let Some(row) = row else { - return Ok(None); - }; - let status_code = row - .try_get::, _>("status_code") - .map_sql_err()? - .and_then(|value| u16::try_from(value).ok()); - Ok(Some(FailedCandidateCleanupInfo { status_code })) -} - -fn bind_upsert<'q>( - mut query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>, - usage: &'q UpsertUsageRecord, -) -> Result>, DataLayerError> -{ - let input_tokens = usage.input_tokens.unwrap_or_default(); - let output_tokens = usage.output_tokens.unwrap_or_default(); - let cache_creation_tokens = usage - .cache_creation_input_tokens - .or_else(|| { - Some( - usage - .cache_creation_ephemeral_5m_input_tokens - .unwrap_or_default() - + usage - .cache_creation_ephemeral_1h_input_tokens - .unwrap_or_default(), - ) - }) - .unwrap_or_default(); - let cache_read_tokens = usage.cache_read_input_tokens.unwrap_or_default(); - let total_tokens = usage - .total_tokens - .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); - let request_metadata = usage - .request_metadata - .as_ref() - .map(serde_json::to_string) - .transpose() - .map_err(|err| DataLayerError::InvalidInput(err.to_string()))?; - - query = query - .bind(&usage.request_id) - .bind(&usage.request_id) - .bind(usage.user_id.as_deref()) - .bind(usage.api_key_id.as_deref()) - .bind(&usage.provider_name) - .bind(&usage.model) - .bind(usage.target_model.as_deref()) - .bind(usage.provider_id.as_deref()) - .bind(usage.provider_endpoint_id.as_deref()) - .bind(usage.provider_api_key_id.as_deref()) - .bind(usage.request_type.as_deref()) - .bind(usage.api_format.as_deref()) - .bind(usage.api_family.as_deref()) - .bind(usage.endpoint_kind.as_deref()) - .bind(usage.endpoint_api_format.as_deref()) - .bind(usage.provider_api_family.as_deref()) - .bind(usage.provider_endpoint_kind.as_deref()) - .bind(i64::from(usage.has_format_conversion.unwrap_or(false))) - .bind(i64::from(usage.is_stream.unwrap_or(false))) - .bind(i64::from(usage_upstream_is_stream(usage))) - .bind(to_i64(input_tokens, "input_tokens")?) - .bind(to_i64(output_tokens, "output_tokens")?) - .bind(to_i64(total_tokens, "total_tokens")?) - .bind(to_i64( - cache_creation_tokens, - "cache_creation_input_tokens", - )?) - .bind(to_i64( - usage - .cache_creation_ephemeral_5m_input_tokens - .unwrap_or_default(), - "cache_creation_ephemeral_5m_input_tokens", - )?) - .bind(to_i64( - usage - .cache_creation_ephemeral_1h_input_tokens - .unwrap_or_default(), - "cache_creation_ephemeral_1h_input_tokens", - )?) - .bind(to_i64(cache_read_tokens, "cache_read_input_tokens")?) - .bind(usage.cache_creation_cost_usd.unwrap_or_default()) - .bind(usage.cache_read_cost_usd.unwrap_or_default()) - .bind(usage.output_price_per_1m) - .bind(usage.total_cost_usd.unwrap_or_default()) - .bind(usage.actual_total_cost_usd.unwrap_or_default()) - .bind(usage.status_code.map(i64::from)) - .bind(usage.error_message.as_deref()) - .bind(usage.error_category.as_deref()) - .bind(usage.response_time_ms.map(|value| value as i64)) - .bind(usage.first_byte_time_ms.map(|value| value as i64)) - .bind(&usage.status) - .bind(&usage.billing_status) - .bind(request_metadata) - .bind(usage.candidate_id.as_deref()) - .bind(usage.candidate_index.map(|value| value as i64)) - .bind(usage.key_name.as_deref()) - .bind(usage.planner_kind.as_deref()) - .bind(usage.route_family.as_deref()) - .bind(usage.route_kind.as_deref()) - .bind(usage.execution_path.as_deref()) - .bind(usage.local_execution_runtime_miss_reason.as_deref()) - .bind(usage.finalized_at_unix_secs.map(|value| value as i64)) - .bind(to_i64(created_at, "created_at_unix_ms")?) - .bind(to_i64(usage.updated_at_unix_secs, "updated_at_unix_secs")?); - Ok(query) -} - -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()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("provider_name").map_sql_err()?, - row.try_get("model").map_sql_err()?, - row.try_get("target_model").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("provider_endpoint_id").map_sql_err()?, - row.try_get("provider_api_key_id").map_sql_err()?, - row.try_get("request_type").map_sql_err()?, - row.try_get("api_format").map_sql_err()?, - row.try_get("api_family").map_sql_err()?, - row.try_get("endpoint_kind").map_sql_err()?, - row.try_get("endpoint_api_format").map_sql_err()?, - row.try_get("provider_api_family").map_sql_err()?, - row.try_get("provider_endpoint_kind").map_sql_err()?, - row.try_get::("has_format_conversion") - .map_sql_err()? - != 0, - row.try_get::("is_stream").map_sql_err()? != 0, - row_i32(row, "input_tokens")?, - row_i32(row, "output_tokens")?, - row_i32(row, "total_tokens")?, - sqlite_real(row, "total_cost_usd")?, - sqlite_real(row, "actual_total_cost_usd")?, - row_optional_i32(row, "status_code")?, - row.try_get("error_message").map_sql_err()?, - row.try_get("error_category").map_sql_err()?, - row_optional_i32(row, "response_time_ms")?, - row_optional_i32(row, "first_byte_time_ms")?, - row.try_get("status").map_sql_err()?, - row.try_get("billing_status").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - row.try_get("finalized_at_unix_secs").map_sql_err()?, - )?; - audit.cache_creation_input_tokens = row_u64(row, "cache_creation_input_tokens")?; - audit.cache_creation_ephemeral_5m_input_tokens = - row_u64(row, "cache_creation_ephemeral_5m_input_tokens")?; - audit.cache_creation_ephemeral_1h_input_tokens = - row_u64(row, "cache_creation_ephemeral_1h_input_tokens")?; - audit.cache_read_input_tokens = row_u64(row, "cache_read_input_tokens")?; - audit.cache_creation_cost_usd = - sqlite_optional_real(row, "cache_creation_cost_usd")?.unwrap_or(0.0); - audit.cache_read_cost_usd = sqlite_optional_real(row, "cache_read_cost_usd")?.unwrap_or(0.0); - audit.output_price_per_1m = sqlite_optional_real(row, "output_price_per_1m")?; - audit.request_metadata = row - .try_get::, _>("request_metadata") - .map_sql_err()? - .map(|raw| serde_json::from_str(&raw)) - .transpose() - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; - audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref()) - .map(ToOwned::to_owned); - http_capture::hydrate_usage_row(row, &mut audit, resolve_legacy_compressed)?; - let upstream_is_stream = row - .try_get::, _>("upstream_is_stream") - .map_sql_err()? - .map(|value| value != 0); - merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream); - snapshots::hydrate_row(row, &mut audit)?; - Ok(audit) -} - -fn to_i64(value: u64, field: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow"))) -} - -fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool { - 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_else(|| usage.is_stream.unwrap_or(false)) -} - -fn merge_usage_stream_metadata(metadata: &mut Option, upstream: Option) { - let Some(upstream) = upstream else { - return; - }; - let value = metadata.get_or_insert_with(|| serde_json::json!({})); - let Some(object) = value.as_object_mut() else { - return; - }; - object - .entry(UPSTREAM_IS_STREAM_KEY) - .or_insert(serde_json::Value::Bool(upstream)); -} - -fn row_i32(row: &SqliteRow, field: &str) -> Result { - let value: i64 = row.try_get(field).map_sql_err()?; - i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) -} - -fn row_optional_i32(row: &SqliteRow, field: &str) -> Result, DataLayerError> { - row.try_get::, _>(field) - .map_sql_err()? - .map(|value| { - i32::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow"))) - }) - .transpose() -} - -fn row_u64(row: &SqliteRow, field: &str) -> Result { - let value: i64 = row.try_get(field).map_sql_err()?; - u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative"))) -} - -#[cfg(test)] -mod tests; diff --git a/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs b/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs deleted file mode 100644 index c81d4bc8d..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage/cleanup.rs +++ /dev/null @@ -1,841 +0,0 @@ -use aether_data_contracts::repository::usage::{ - UsageCleanupExecutionMode, UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, - UsageCleanupWindow, -}; -use chrono::{DateTime, Utc}; -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 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 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 - ) -) -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, 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, - } - } -} - -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 { - purge_legacy_body_refs(pool, window.detail_cutoff, detail_newer_than, batch_size).await? - } else { - 0 - }; - let body_externalized = if targets.detail_body { - cleanup_body_fields( - pool, - window.detail_cutoff, - detail_newer_than, - batch_size, - BodyCleanupKind::All, - ) - .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, - cost_reservations_deleted: 0, - request_admissions_deleted: 0, - }) -} - -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()?; - } - - 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 purge_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 purge 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 purged = 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(metadata) = legacy_body_ref_purge_plan(metadata.as_deref())? else { - continue; - }; - 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(); - purge_detached_body_capture(&mut tx, &row.request_id).await?; - if updated > 0 { - purged += 1; - } - } - tx.commit().await.map_sql_err()?; - total = total.saturating_add(purged); - if row_count < batch_size || purged == 0 { - break; - } - } - Ok(total) -} - -fn legacy_body_ref_purge_plan( - 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 removed = false; - for key in [ - "request_body_ref", - "provider_request_body_ref", - "response_body_ref", - "client_response_body_ref", - ] { - if object.remove(key).is_some() { - removed = true; - } - } - 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(metadata)) -} - -async fn purge_detached_body_capture( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - request_id: &str, -) -> Result<(), DataLayerError> { - sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") - .bind(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(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - delete_empty_http_audit(tx, request_id).await -} - -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 deleted file mode 100644 index b76b936d4..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage/counters.rs +++ /dev/null @@ -1,1603 +0,0 @@ -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, - target_tunnel_generation, - 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, - target_tunnel_generation: Option, - 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, - // Generation is part of the key: rows from an old incarnation must never - // be coalesced with rows for a newly registered node using the same id. - proxy_nodes: BTreeMap<(String, String), ProxyNodeCounterDelta>, - 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 Some(tunnel_generation) = row - .target_tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - else { - // Legacy rows without an identity fence are intentionally - // discarded when marked processed below. - continue; - }; - let aggregate_key = (row.target_id.clone(), tunnel_generation.clone()); - let entry = aggregates.proxy_nodes.entry(aggregate_key).or_insert( - ProxyNodeCounterDelta { - node_id: row.target_id.clone(), - expected_tunnel_generation: Some(tunnel_generation), - 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, tunnel_generation), delta) in &aggregates.proxy_nodes { - apply_proxy_node(&mut tx, target_id, tunnel_generation, 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 Some(expected_tunnel_generation) = delta - .expected_tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .filter(|value| value.len() <= 64) - .map(ToOwned::to_owned) - else { - // Never infer an incarnation from a bare node id. A missing fence can - // otherwise make a stale request update a node recreated under that id. - 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()?; - // Read the generation as a predicate without taking an exclusive parent - // lock. Flush claims outbox rows before touching proxy_nodes; keeping this - // read lock-free avoids the inverse parent->outbox lock order. The value is - // persisted in the outbox row and the flush UPDATE repeats the generation - // predicate, so a delete/re-register race can only retire the delta. - let tunnel_generation: Option = sqlx::query_scalar( - "SELECT tunnel_generation FROM proxy_nodes WHERE id = ? AND tunnel_generation = ? LIMIT 1", - ) - .bind(&node_id) - .bind(&expected_tunnel_generation) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(_tunnel_generation) = tunnel_generation - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - 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, - target_tunnel_generation: Some(&expected_tunnel_generation), - ..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, - target_tunnel_generation: Option<&'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, - target_tunnel_generation, - 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.target_tunnel_generation) - .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()?, - target_tunnel_generation: row.try_get("target_tunnel_generation").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, - tunnel_generation: &str, - delta: &ProxyNodeCounterDelta, -) -> Result<(), DataLayerError> { - if target_id.trim().is_empty() || tunnel_generation.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 = ? AND tunnel_generation = ? -"#, - ) - .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) - .bind(tunnel_generation) - .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 std::{sync::Arc, time::Duration}; - - use super::{ - cleanup_processed, enqueue_api_key_last_used, enqueue_management_token, enqueue_proxy_node, - flush, read_health, read_pending_health, - }; - use crate::proxy_nodes::SqliteProxyNodeReadRepository; - use aether_data_contracts::repository::proxy_nodes::ProxyNodeWriteRepository; - use aether_data_contracts::repository::usage::{ - ApiKeyLastUsedDelta, ManagementTokenCounterDelta, ProxyNodeCounterDelta, - }; - use tokio::sync::Barrier; - - #[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, tunnel_generation, name, ip, port, created_at, updated_at) -VALUES ('counter-node', 'test-generation-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(), - expected_tunnel_generation: Some("test-generation-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 - ); - } - - #[tokio::test] - async fn proxy_counter_enqueue_requires_the_expected_generation() { - 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 proxy_nodes (id, tunnel_generation, name, ip, port, created_at, updated_at) -VALUES ('generation-node', 'generation-a', 'generation node', '127.0.0.1', 8080, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("proxy node should seed"); - - assert!(!enqueue_proxy_node( - &pool, - ProxyNodeCounterDelta { - node_id: "generation-node".to_string(), - expected_tunnel_generation: None, - total_requests_delta: 1, - failed_requests_delta: 0, - dns_failures_delta: 0, - stream_errors_delta: 0, - }, - ) - .await - .expect("missing generation should be handled")); - assert!(!enqueue_proxy_node( - &pool, - ProxyNodeCounterDelta { - node_id: "generation-node".to_string(), - expected_tunnel_generation: Some("generation-b".to_string()), - total_requests_delta: 1, - failed_requests_delta: 0, - dns_failures_delta: 0, - stream_errors_delta: 0, - }, - ) - .await - .expect("mismatched generation should be handled")); - let pending: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_counter_deltas WHERE processed_at IS NULL", - ) - .fetch_one(&pool) - .await - .expect("pending rows should load"); - assert_eq!(pending, 0); - - assert!(enqueue_proxy_node( - &pool, - ProxyNodeCounterDelta { - node_id: "generation-node".to_string(), - expected_tunnel_generation: Some("generation-a".to_string()), - total_requests_delta: 1, - failed_requests_delta: 0, - dns_failures_delta: 0, - stream_errors_delta: 0, - }, - ) - .await - .expect("matching generation should enqueue")); - } - - #[tokio::test] - async fn proxy_counter_flush_does_not_cross_generation_reuse() { - 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 proxy_nodes (id, tunnel_generation, name, ip, port, created_at, updated_at) -VALUES ('reused-node', 'generation-a', 'reused node', '127.0.0.1', 8080, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("first proxy incarnation should seed"); - assert!(enqueue_proxy_node( - &pool, - ProxyNodeCounterDelta { - node_id: "reused-node".to_string(), - expected_tunnel_generation: Some("generation-a".to_string()), - total_requests_delta: 7, - failed_requests_delta: 3, - dns_failures_delta: 2, - stream_errors_delta: 1, - }, - ) - .await - .expect("first incarnation delta should enqueue")); - - // Simulate deletion followed by id reuse while the old outbox row is - // still pending. The generation predicate must retire the row without - // applying it to the replacement incarnation. - sqlx::query("DELETE FROM proxy_nodes WHERE id = 'reused-node'") - .execute(&pool) - .await - .expect("first incarnation should delete"); - sqlx::query( - r#" -INSERT INTO proxy_nodes (id, tunnel_generation, name, ip, port, created_at, updated_at) -VALUES ('reused-node', 'generation-b', 'reused node', '127.0.0.1', 8080, 2, 2) -"#, - ) - .execute(&pool) - .await - .expect("replacement incarnation should seed"); - - let summary = flush(&pool, 100) - .await - .expect("counter flush should succeed"); - assert_eq!(summary.rows_claimed, 1); - let counters: (i64, i64, i64, i64) = sqlx::query_as( - "SELECT total_requests, failed_requests, dns_failures, stream_errors FROM proxy_nodes WHERE id = 'reused-node'", - ) - .fetch_one(&pool) - .await - .expect("replacement counters should load"); - assert_eq!(counters, (0, 0, 0, 0)); - let pending: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_counter_deltas WHERE processed_at IS NULL", - ) - .fetch_one(&pool) - .await - .expect("pending rows should load"); - assert_eq!(pending, 0); - } - - #[tokio::test] - async fn proxy_counter_flush_and_delete_do_not_deadlock_or_cross_generation() { - let database_path = std::env::temp_dir().join(format!( - "aether-proxy-counter-delete-flush-{}.sqlite", - uuid::Uuid::new_v4().simple() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .busy_timeout(Duration::from_secs(10)); - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(8) - .connect_with(options) - .await - .expect("sqlite pool should connect"); - crate::run_migrations(&pool) - .await - .expect("sqlite migrations should run"); - - let node_id = "delete-flush-node"; - let generation_a = "delete-flush-generation-a"; - sqlx::query( - r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, is_manual, proxy_url, created_at, updated_at -) -VALUES (?, ?, 'delete/flush node', '127.0.0.91', 8091, 1, 'http://127.0.0.91:8091', 1, 1) -"#, - ) - .bind(node_id) - .bind(generation_a) - .execute(&pool) - .await - .expect("proxy node should seed"); - assert!(enqueue_proxy_node( - &pool, - ProxyNodeCounterDelta { - node_id: node_id.to_string(), - expected_tunnel_generation: Some(generation_a.to_string()), - total_requests_delta: 11, - failed_requests_delta: 3, - dns_failures_delta: 2, - stream_errors_delta: 1, - }, - ) - .await - .expect("counter delta should enqueue")); - - let repository = SqliteProxyNodeReadRepository::new(pool.clone()); - let barrier = Arc::new(Barrier::new(2)); - let flush_barrier = Arc::clone(&barrier); - let flush_pool = pool.clone(); - let flush_task = tokio::spawn(async move { - flush_barrier.wait().await; - tokio::time::timeout(Duration::from_secs(15), flush(&flush_pool, 100)) - .await - .expect("counter flush should not deadlock") - .expect("counter flush should succeed") - }); - let delete_barrier = Arc::clone(&barrier); - let delete_repository = repository.clone(); - let delete_task = tokio::spawn(async move { - delete_barrier.wait().await; - tokio::time::timeout( - Duration::from_secs(15), - delete_repository.delete_node(node_id), - ) - .await - .expect("proxy delete should not deadlock") - .expect("proxy delete should succeed") - }); - let (flush_summary, deleted) = tokio::join!(flush_task, delete_task); - let flush_summary = flush_summary.expect("flush task should join"); - let deleted = deleted.expect("delete task should join"); - assert!(deleted.is_some(), "the original node should be deleted"); - assert!(flush_summary.rows_claimed <= 1); - - // Reuse the id immediately. Any row that lost the race with the - // post-commit cleanup is still safe because flush checks generation. - let generation_b = "delete-flush-generation-b"; - sqlx::query( - r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, is_manual, proxy_url, created_at, updated_at -) -VALUES (?, ?, 'replacement node', '127.0.0.92', 8092, 1, 'http://127.0.0.92:8092', 2, 2) -"#, - ) - .bind(node_id) - .bind(generation_b) - .execute(&pool) - .await - .expect("replacement node should seed"); - flush(&pool, 100) - .await - .expect("stale counter flush should succeed"); - - let counters: (i64, i64, i64, i64) = sqlx::query_as( - "SELECT total_requests, failed_requests, dns_failures, stream_errors FROM proxy_nodes WHERE id = ?", - ) - .bind(node_id) - .fetch_one(&pool) - .await - .expect("replacement counters should load"); - assert_eq!(counters, (0, 0, 0, 0)); - let pending: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_counter_deltas WHERE processed_at IS NULL", - ) - .fetch_one(&pool) - .await - .expect("pending rows should load"); - assert_eq!(pending, 0); - - pool.close().await; - let _ = std::fs::remove_file(database_path); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs b/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs deleted file mode 100644 index 44b6a4580..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage/http_capture.rs +++ /dev/null @@ -1,745 +0,0 @@ -use aether_data_contracts::repository::usage::{ - canonical_usage_body_ref_for, parse_usage_body_ref, usage_body_ref, StoredRequestUsageAudit, - UpsertUsageRecord, UsageBodyCaptureState, UsageBodyField, -}; -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 { - payload_gzip: Option>, -} - -#[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, -} - -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(request_body_value)?; - let provider_request_body = prepare_body(provider_request_body_value)?; - let response_body = prepare_body(response_body_value)?; - let client_response_body = prepare_body(client_response_body_value)?; - - 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(value: Option<&Value>) -> Result { - if value.is_some() { - return Err(DataLayerError::InvalidInput( - "usage body persistence is disabled".to_string(), - )); - } - Ok(PreparedBody { payload_gzip: None }) -} - -fn resolved_write_ref( - explicit_ref: Option<&str>, - request_id: &str, - field: UsageBodyField, - has_blob: bool, -) -> Option { - explicit_ref - .and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, field)) - .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> { - let bodies = [ - &prepared.request_body, - &prepared.provider_request_body, - &prepared.response_body, - &prepared.client_response_body, - ]; - let contains_capture = prepared.request_headers.is_some() - || prepared.provider_request_headers.is_some() - || prepared.response_headers.is_some() - || prepared.client_response_headers.is_some() - || prepared.refs.any_present() - || bodies.iter().any(|body| body.payload_gzip.is_some()) - || prepared.capture_mode != "none"; - if contains_capture { - return Err(DataLayerError::InvalidInput( - "usage HTTP capture persistence is disabled".to_string(), - )); - } - - sqlx::query("DELETE FROM usage_http_audits WHERE request_id = ?") - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM usage_body_blobs WHERE request_id = ?") - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE "usage" -SET request_headers = NULL, - request_body = NULL, - provider_request_headers = NULL, - provider_request_body = NULL, - response_headers = NULL, - response_body = NULL, - client_response_headers = NULL, - client_response_body = NULL, - request_body_compressed = NULL, - provider_request_body_compressed = NULL, - response_body_compressed = NULL, - client_response_body_compressed = NULL -WHERE request_id = ? -"#, - ) - .bind(request_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - - 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() { - 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(()) -} - -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 - .and_then(|body_ref| canonical_usage_body_ref_for(&body_ref, request_id, field)) - .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) - .and_then(|body_ref| canonical_usage_body_ref_for(body_ref, request_id, 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 Some(body_ref) = canonical_usage_body_ref_for(body_ref, &usage.request_id, 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> { - let Some((request_id, field)) = parse_usage_body_ref(body_ref) else { - return Ok(None); - }; - let canonical_ref = usage_body_ref(&request_id, field); - if let Some(payload_gzip) = sqlx::query_scalar::<_, Vec>( - "SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? AND request_id = ? AND body_field = ? LIMIT 1", - ) - .bind(&canonical_ref) - .bind(&request_id) - .bind(field.as_storage_field()) - .fetch_optional(pool) - .await - .map_sql_err()? - { - return super::inflate_usage_json_value(&payload_gzip).map(Some); - } - 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 deleted file mode 100644 index 212e9e960..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage/snapshots.rs +++ /dev/null @@ -1,1047 +0,0 @@ -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 deleted file mode 100644 index 377d253bd..000000000 --- a/crates/aether-data/adapters/sqlite/src/usage/tests.rs +++ /dev/null @@ -1,2711 +0,0 @@ -use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository}; -use crate::run_migrations; -use aether_data_contracts::repository::usage::{ - ProviderApiKeyWindowUsageRequest, UpsertUsageRecord, UsageAuditAggregationGroupBy, - UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, - UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBreakdownGroupBy, - UsageBreakdownSummaryQuery, UsageCleanupExecutionMode, UsageCleanupTargets, UsageCleanupWindow, - UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, - UsageProviderPerformanceQuery, UsageReadRepository, UsageTimeSeriesGranularity, - UsageWriteRepository, -}; -use chrono::{DateTime, Utc}; -use flate2::{write::GzEncoder, Compression}; -use std::io::Write; - -fn gzip_json_for_test(value: &serde_json::Value) -> Vec { - let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); - encoder - .write_all(&serde_json::to_vec(value).expect("test JSON should serialize")) - .expect("test JSON should compress"); - encoder.finish().expect("test gzip should finish") -} - -#[test] -fn sqlite_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() { - for field in [ - "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", - "request_metadata", - "candidate_id", - "candidate_index", - "key_name", - "planner_kind", - "route_family", - "route_kind", - "execution_path", - "local_execution_runtime_miss_reason", - ] { - let assignment = format!("{field} = CASE WHEN ("); - assert!( - super::UPSERT_USAGE_SQL.contains(&assignment), - "missing lifecycle guard for {field}" - ); - assert!( - super::UPSERT_USAGE_SQL.contains(&format!("THEN \"usage\".{field}")), - "late lifecycle must preserve {field}" - ); - } - assert!(super::UPSERT_USAGE_SQL - .contains("OR (\"usage\".status = 'streaming' AND excluded.status = 'pending')")); -} - -#[test] -fn sqlite_first_byte_upsert_rejects_older_revision_in_sql() { - assert!(super::UPSERT_FIRST_BYTE_BATCH_UPDATE_SUFFIX_SQL.contains( - "excluded.updated_at_unix_secs >= COALESCE(\n NULLIF(\"usage\".updated_at_unix_secs, 0)," - )); -} - -#[tokio::test] -async fn sqlite_provider_performance_can_skip_timeline() { - 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; - - SqliteUsageWriteRepository::new(pool.clone()) - .upsert(sample_usage( - "provider-performance", - "completed", - "pending", - 1_000, - )) - .await - .expect("usage should upsert"); - let reader = SqliteUsageReadRepository::new(pool); - let mut query = UsageProviderPerformanceQuery { - created_from_unix_secs: 0, - created_until_unix_secs: 2_000, - granularity: UsageTimeSeriesGranularity::Hour, - tz_offset_minutes: 0, - limit: 1, - provider_id: None, - model: None, - api_format: None, - endpoint_kind: None, - is_stream: None, - has_format_conversion: None, - slow_threshold_ms: 10_000, - include_timeline: true, - }; - - let with_timeline = reader - .summarize_usage_provider_performance(&query) - .await - .expect("provider performance should load"); - assert_eq!(with_timeline.summary.request_count, 1); - assert_eq!(with_timeline.providers.len(), 1); - assert_eq!(with_timeline.timeline.len(), 1); - - query.include_timeline = false; - let without_timeline = reader - .summarize_usage_provider_performance(&query) - .await - .expect("provider performance without timeline should load"); - assert_eq!(without_timeline.summary, with_timeline.summary); - assert_eq!(without_timeline.providers, with_timeline.providers); - assert!(without_timeline.timeline.is_empty()); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_upserts_and_flushes_counter_deltas() { - 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()); - let record = repository - .upsert(sample_usage("request-1", "completed", "pending", 1_000)) - .await - .expect("usage should upsert"); - - assert_eq!(record.request_id, "request-1"); - assert_eq!(record.api_key_id.as_deref(), Some("api-key-1")); - assert_eq!(record.total_tokens, 5); - assert_eq!(record.cache_read_input_tokens, 2); - assert_eq!( - record.request_metadata.as_ref().unwrap()["trace_id"], - "trace-1" - ); - assert_eq!( - record.request_metadata.as_ref().unwrap()["upstream_is_stream"], - true - ); - let upstream_is_stream: Option = - sqlx::query_scalar("SELECT upstream_is_stream FROM \"usage\" WHERE request_id = ?") - .bind("request-1") - .fetch_one(&pool) - .await - .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 - .expect("usage should load") - .expect("usage should exist"); - assert_eq!( - loaded.provider_api_key_id.as_deref(), - Some("provider-key-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 = 'api-key-1'", - ) - .fetch_one(&pool) - .await - .expect("api key stats should load"); - 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'", - ) - .fetch_one(&pool) - .await - .expect("provider key stats should load"); - 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_is_not_persisted() { - 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!(stored.request_headers.is_none()); - assert!(stored.request_body.is_none()); - assert!(stored.provider_request_body.is_none()); - assert!(stored.response_body.is_none()); - assert!(stored.client_response_body.is_none()); - assert!(stored.request_body_state.is_none()); - assert!(stored.request_body_ref.is_none()); - 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_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_http_audits WHERE request_id = 'canonical-capture'", - ) - .fetch_one(&pool) - .await - .expect("canonical audits should count"); - assert_eq!(audit_count, 0); - 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, 0); - - let sparse = sample_usage("canonical-capture", "streaming", "pending", 1_001); - let sparse_stored = writer - .upsert(sparse) - .await - .expect("sparse lifecycle update should upsert"); - assert!(sparse_stored.request_headers.is_none()); - assert!(sparse_stored.request_body.is_none()); - assert!(sparse_stored.response_body.is_none()); - 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, 0); - - 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!(cleared.request_body_state.is_none()); - assert!(cleared.provider_request_body.is_none()); - 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, 0); - - 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!(resolved.is_none()); - let loaded = reader - .find_by_request_id("canonical-capture") - .await - .expect("canonical usage should load") - .expect("canonical usage should exist"); - assert!(loaded.provider_request_headers.is_none()); - assert!(loaded.provider_request_body.is_none()); -} - -#[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"); - 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":true}', - response_body = '{"inline":true}', - request_metadata = '{"request_body_ref":"usage://request/legacy-capture/request_body"}' -WHERE request_id = 'legacy-capture'; -"#, - ) - .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!(cleared.request_body_state.is_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_body_refs_enforce_request_and_field_ownership() { - 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()); - for (request_id, updated_at) in [("ref-target", 3_000), ("ref-owner", 3_001)] { - writer - .upsert(sample_usage(request_id, "completed", "settled", updated_at)) - .await - .expect("usage should seed"); - } - - let mismatched_payload = gzip_json_for_test(&serde_json::json!({ - "secret": "belongs to ref-owner response" - })); - sqlx::query( - "INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip) VALUES (?, ?, ?, ?)", - ) - .bind("usage://request/ref-target/request_body") - .bind("ref-owner") - .bind("response_body") - .bind(mismatched_payload) - .execute(&pool) - .await - .expect("mismatched legacy blob should seed"); - - let reader = SqliteUsageReadRepository::new(pool.clone()); - assert_eq!( - reader - .resolve_body_ref("usage://request/ref-target/request_body") - .await - .expect("mismatched blob lookup should remain safe"), - None - ); - - sqlx::query("DELETE FROM usage_body_blobs") - .execute(&pool) - .await - .expect("mismatched blob should clear"); - sqlx::query("UPDATE \"usage\" SET request_body = ? WHERE request_id = ?") - .bind(r#"{"secret":"belongs to ref-owner request"}"#) - .bind("ref-owner") - .execute(&pool) - .await - .expect("legacy owner body should seed"); - sqlx::query( - "INSERT INTO usage_http_audits (request_id, request_body_ref, body_capture_mode) VALUES (?, ?, ?)", - ) - .bind("ref-target") - .bind("usage://request/ref-owner/request_body") - .bind("ref_backed") - .execute(&pool) - .await - .expect("cross-request audit ref should seed"); - - let target = reader - .find_by_request_id("ref-target") - .await - .expect("target usage should load") - .expect("target usage should exist"); - assert!(target.request_body_ref.is_none()); - assert!(target.request_body.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_blobs: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'cleanup-detail'", - ) - .fetch_one(&pool) - .await - .expect("purged body blobs should count"); - assert_eq!(detail_blobs, 0); - 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_audits: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_http_audits WHERE request_id = 'cleanup-legacy'", - ) - .fetch_one(&pool) - .await - .expect("legacy audit refs should count"); - assert_eq!(legacy_audits, 0); - - 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()); - let audit_body_ref: Option = sqlx::query_scalar( - "SELECT request_body_ref FROM usage_http_audits WHERE request_id = 'cleanup-before-now'", - ) - .fetch_one(&pool) - .await - .expect("audit body ref should load"); - assert!(audit_body_ref.is_none()); - let body_blobs: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = 'cleanup-before-now'", - ) - .fetch_one(&pool) - .await - .expect("body blobs should count"); - assert_eq!(body_blobs, 0); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_does_not_regress_void_usage() { - 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("request-1", "failed", "void", 1_000)) - .await - .expect("void usage should upsert"); - let existing = repository - .upsert(sample_usage("request-1", "pending", "pending", 1_001)) - .await - .expect("stale usage should be ignored"); - - assert_eq!(existing.status, "failed"); - assert_eq!(existing.billing_status, "void"); - assert_eq!(existing.updated_at_unix_secs, 1_000); -} - -#[tokio::test] -async fn sqlite_stale_terminal_event_is_a_full_transaction_noop() { - 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()); - let mut newer = sample_usage("request-stale-terminal", "completed", "pending", 2_000); - newer.candidate_id = Some("candidate-new".to_string()); - newer.route_kind = Some("route-new".to_string()); - repository - .upsert(newer) - .await - .expect("newer terminal usage should upsert"); - - let counter_rows_before: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas") - .fetch_one(&pool) - .await - .expect("counter rows should count"); - let routing_before: (Option, Option) = sqlx::query_as( - "SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = ?", - ) - .bind("request-stale-terminal") - .fetch_one(&pool) - .await - .expect("routing snapshot should load"); - let settlement_before: (String, Option) = sqlx::query_as( - "SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = ?", - ) - .bind("request-stale-terminal") - .fetch_one(&pool) - .await - .expect("settlement snapshot should load"); - - let mut stale = sample_usage("request-stale-terminal", "failed", "void", 1_999); - stale.status_code = Some(503); - stale.total_cost_usd = Some(99.0); - stale.actual_total_cost_usd = Some(98.0); - stale.candidate_id = Some("candidate-stale".to_string()); - stale.route_kind = Some("route-stale".to_string()); - let stored = repository - .upsert(stale) - .await - .expect("stale terminal usage should be ignored"); - - assert_eq!(stored.status, "completed"); - assert_eq!(stored.billing_status, "pending"); - assert_eq!(stored.status_code, Some(200)); - assert_eq!(stored.total_cost_usd, 0.5); - assert_eq!(stored.routing_candidate_id(), Some("candidate-new")); - assert_eq!(stored.routing_route_kind(), Some("route-new")); - assert_eq!(stored.updated_at_unix_secs, 2_000); - - let counter_rows_after: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas") - .fetch_one(&pool) - .await - .expect("counter rows should count"); - let routing_after: (Option, Option) = sqlx::query_as( - "SELECT candidate_id, route_kind FROM usage_routing_snapshots WHERE request_id = ?", - ) - .bind("request-stale-terminal") - .fetch_one(&pool) - .await - .expect("routing snapshot should load"); - let settlement_after: (String, Option) = sqlx::query_as( - "SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = ?", - ) - .bind("request-stale-terminal") - .fetch_one(&pool) - .await - .expect("settlement snapshot should load"); - - assert_eq!(counter_rows_after, counter_rows_before); - assert_eq!(routing_after, routing_before); - assert_eq!(settlement_after, settlement_before); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_does_not_reopen_void_failure_from_late_streaming() { - 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); - for (request_id, status_code) in [ - ("request-late-active", None), - ("request-late-response-start", Some(200)), - ] { - let mut failed = sample_usage(request_id, "failed", "void", 1_000); - failed.status_code = Some(503); - repository - .upsert(failed) - .await - .expect("failed usage should upsert"); - - let mut late_streaming = sample_usage(request_id, "streaming", "pending", 1_001); - late_streaming.status_code = status_code; - late_streaming.finalized_at_unix_secs = None; - let current = repository - .upsert(late_streaming) - .await - .expect("late streaming usage should be ignored"); - - assert_eq!(current.status, "failed"); - assert_eq!(current.billing_status, "void"); - assert_eq!(current.status_code, Some(503)); - assert_eq!(current.finalized_at_unix_secs, Some(1_000)); - } -} - -#[tokio::test] -async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() { - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("sqlite pool should connect"); - run_migrations(&pool) - .await - .expect("sqlite migrations should run"); - seed_stats_targets(&pool).await; - - let repository = SqliteUsageWriteRepository::new(pool.clone()); - let mut terminal = sample_usage("request-1", "completed", "pending", 1_000); - terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal-trace"})); - repository - .upsert(terminal) - .await - .expect("terminal usage should upsert"); - - let mut late_streaming = sample_usage("request-1", "streaming", "pending", 1_001); - late_streaming.input_tokens = Some(0); - late_streaming.output_tokens = Some(0); - late_streaming.total_tokens = Some(0); - late_streaming.cache_read_input_tokens = Some(0); - late_streaming.cache_read_cost_usd = Some(0.0); - late_streaming.total_cost_usd = Some(0.0); - late_streaming.actual_total_cost_usd = Some(0.0); - late_streaming.response_time_ms = Some(9_999); - late_streaming.first_byte_time_ms = Some(9_999); - late_streaming.finalized_at_unix_secs = None; - late_streaming.provider_name = "Late Provider".to_string(); - late_streaming.model = "late-model".to_string(); - late_streaming.target_model = Some("late-target".to_string()); - late_streaming.request_type = Some("late-request".to_string()); - late_streaming.api_format = Some("late:api".to_string()); - late_streaming.api_family = Some("late-family".to_string()); - late_streaming.endpoint_kind = Some("late-endpoint".to_string()); - late_streaming.endpoint_api_format = Some("late:endpoint".to_string()); - late_streaming.provider_api_family = Some("late-provider-family".to_string()); - late_streaming.provider_endpoint_kind = Some("late-provider-endpoint".to_string()); - late_streaming.has_format_conversion = Some(false); - late_streaming.is_stream = Some(true); - late_streaming.candidate_id = Some("late-candidate".to_string()); - late_streaming.candidate_index = Some(99); - late_streaming.key_name = Some("late-key".to_string()); - late_streaming.planner_kind = Some("late-planner".to_string()); - late_streaming.route_family = Some("late-route-family".to_string()); - late_streaming.route_kind = Some("late-route-kind".to_string()); - late_streaming.execution_path = Some("late-path".to_string()); - late_streaming.local_execution_runtime_miss_reason = Some("late-miss".to_string()); - late_streaming.request_metadata = Some(serde_json::json!({ - "provider_service_tier": "priority", - "upstream_is_stream": true - })); - - let current = repository - .upsert(late_streaming) - .await - .expect("late streaming usage should not regress terminal usage"); - - assert_eq!(current.status, "completed"); - assert_eq!(current.billing_status, "pending"); - assert_eq!(current.total_tokens, 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); - assert_eq!(current.response_time_ms, Some(42)); - assert_eq!(current.first_byte_time_ms, Some(12)); - assert_eq!(current.finalized_at_unix_secs, Some(1_000)); - assert_eq!(current.updated_at_unix_secs, 1_000); - assert_eq!(current.provider_name, "Provider One"); - assert_eq!(current.model, "model-1"); - assert_eq!(current.target_model.as_deref(), Some("target-model")); - assert_eq!(current.request_type.as_deref(), Some("chat")); - assert_eq!(current.api_format.as_deref(), Some("openai")); - assert!(current.has_format_conversion); - assert!(!current.is_stream); - assert_eq!(current.candidate_id.as_deref(), Some("candidate-1")); - assert_eq!(current.candidate_index, Some(1)); - assert_eq!(current.key_name.as_deref(), Some("key-one")); - assert_eq!(current.planner_kind.as_deref(), Some("default")); - assert_eq!(current.route_family.as_deref(), Some("chat")); - assert_eq!(current.route_kind.as_deref(), Some("completion")); - assert_eq!(current.execution_path.as_deref(), Some("remote")); - assert_eq!(current.provider_service_tier(), None); - assert_eq!( - current - .request_metadata - .as_ref() - .and_then(|value| value.get("trace_id")) - .and_then(serde_json::Value::as_str), - Some("terminal-trace") - ); - - let listed = SqliteUsageReadRepository::new(pool) - .list_usage_audits(&UsageAuditListQuery { - limit: Some(10), - newest_first: true, - ..UsageAuditListQuery::default() - }) - .await - .expect("usage list should load") - .into_iter() - .find(|item| item.request_id == "request-1") - .expect("terminal usage should be listed"); - assert_eq!(listed.provider_name, "Provider One"); - assert_eq!(listed.model, "model-1"); - assert_eq!(listed.candidate_id.as_deref(), Some("candidate-1")); - assert_eq!(listed.provider_service_tier(), None); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_allows_authoritative_completed_recovery() { - 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); - repository - .upsert(sample_usage("request-recovery", "failed", "void", 1_000)) - .await - .expect("void failure should upsert"); - - let mut recovery = sample_usage("request-recovery", "completed", "pending", 1_001); - recovery.provider_name = "Recovered Provider".to_string(); - recovery.model = "recovered-model".to_string(); - recovery.target_model = Some("recovered-target".to_string()); - recovery.api_format = Some("recovered:api".to_string()); - recovery.candidate_id = Some("recovered-candidate".to_string()); - recovery.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"})); - let recovered = repository - .upsert(recovery) - .await - .expect("completed recovery should upsert"); - - assert_eq!(recovered.status, "completed"); - assert_eq!(recovered.billing_status, "pending"); - assert_eq!(recovered.provider_name, "Recovered Provider"); - assert_eq!(recovered.model, "recovered-model"); - assert_eq!(recovered.target_model.as_deref(), Some("recovered-target")); - assert_eq!(recovered.api_format.as_deref(), Some("recovered:api")); - assert_eq!( - recovered.candidate_id.as_deref(), - Some("recovered-candidate") - ); - assert_eq!( - recovered.provider_service_tier().as_deref(), - Some("priority") - ); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_preserves_streaming_response_start_from_late_active() { - 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); - repository - .upsert(sample_usage( - "request-late-active", - "streaming", - "pending", - 1_000, - )) - .await - .expect("response-start usage should upsert"); - - let mut late_active = sample_usage("request-late-active", "streaming", "pending", 1_001); - late_active.status_code = None; - late_active.response_time_ms = None; - late_active.first_byte_time_ms = None; - - let current = repository - .upsert(late_active) - .await - .expect("late active usage should not clear response-start fields"); - - assert_eq!(current.status, "streaming"); - assert_eq!(current.status_code, Some(200)); - assert_eq!(current.response_time_ms, Some(42)); - assert_eq!(current.first_byte_time_ms, Some(12)); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_keeps_streaming_capture_from_late_pending() { - 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); - let mut streaming = sample_usage("request-streaming-capture", "streaming", "pending", 1_000); - streaming.request_metadata = Some(serde_json::json!({"trace_id": "streaming-final"})); - repository - .upsert(streaming) - .await - .expect("streaming usage should upsert"); - - let mut late_pending = sample_usage("request-streaming-capture", "pending", "pending", 1_001); - late_pending.provider_name = "Late Provider".to_string(); - late_pending.model = "late-model".to_string(); - late_pending.candidate_id = Some("late-candidate".to_string()); - late_pending.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"})); - let current = repository - .upsert(late_pending) - .await - .expect("late pending usage should not regress streaming capture"); - - assert_eq!(current.status, "streaming"); - assert_eq!(current.provider_name, "Provider One"); - assert_eq!(current.model, "model-1"); - assert_eq!(current.candidate_id.as_deref(), Some("candidate-1")); - assert_eq!(current.provider_service_tier(), None); - assert_eq!( - current - .request_metadata - .as_ref() - .and_then(|value| value.get("trace_id")) - .and_then(serde_json::Value::as_str), - Some("streaming-final") - ); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_cleans_stale_pending_requests() { - 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("request-recovered", "streaming", "pending", 1)) - .await - .expect("streaming usage should upsert"); - repository - .upsert(sample_usage("request-failed", "pending", "pending", 1)) - .await - .expect("pending usage should upsert"); - - sqlx::query( - r#" -INSERT INTO request_candidates ( - id, - request_id, - candidate_index, - retry_index, - status, - is_cached, - created_at -) VALUES - ('candidate-recovered', 'request-recovered', 0, 0, 'streaming', 0, 1), - ('candidate-failed', 'request-failed', 0, 0, 'pending', 0, 1) -"#, - ) - .execute(&pool) - .await - .expect("request candidates should seed"); - - let summary = repository - .cleanup_stale_pending_requests(2, 10, 5, 1) - .await - .expect("cleanup should run"); - assert_eq!(summary.recovered, 1); - assert_eq!(summary.failed, 1); - - let recovered = repository - .find_by_request_id("request-recovered") - .await - .expect("recovered usage should load") - .expect("recovered usage should exist"); - assert_eq!(recovered.status, "completed"); - assert_eq!(recovered.status_code, Some(200)); - - let failed = repository - .find_by_request_id("request-failed") - .await - .expect("failed usage should load") - .expect("failed usage should exist"); - assert_eq!(failed.status, "failed"); - assert_eq!(failed.status_code, Some(504)); - assert_eq!(failed.billing_status, "void"); - assert_eq!(failed.total_cost_usd, 0.0); - assert_eq!(failed.finalized_at_unix_secs, Some(10)); - - let candidate_statuses = sqlx::query_as::<_, (String, String, Option)>( - r#" -SELECT request_id, status, finished_at -FROM request_candidates -ORDER BY request_id -"#, - ) - .fetch_all(&pool) - .await - .expect("candidate statuses should load"); - assert_eq!( - candidate_statuses, - vec![ - ( - "request-failed".to_string(), - "failed".to_string(), - Some(10_000) - ), - ( - "request-recovered".to_string(), - "success".to_string(), - Some(10_000) - ), - ] - ); - - let snapshot = sqlx::query_as::<_, (String, Option)>( - "SELECT billing_status, finalized_at FROM usage_settlement_snapshots WHERE request_id = 'request-failed'", - ) - .fetch_one(&pool) - .await - .expect("void settlement snapshot should load"); - assert_eq!(snapshot, ("void".to_string(), Some(10))); -} - -#[tokio::test] -async fn sqlite_usage_write_repository_cleanup_uses_failed_candidate_status_when_present() { - 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( - "request-upstream-reset", - "pending", - "pending", - 1, - )) - .await - .expect("pending usage should upsert"); - repository - .upsert(sample_usage("request-stuck", "pending", "pending", 1)) - .await - .expect("pending usage should upsert"); - - // request-upstream-reset has a failed candidate carrying a concrete 502 status. - // Cleanup keeps the status but must not copy the candidate's diagnostic text. - // request-stuck has only a still-pending candidate, so cleanup should fall back to 504. - sqlx::query( - r#" -INSERT INTO request_candidates ( - id, - request_id, - candidate_index, - retry_index, - status, - status_code, - error_message, - is_cached, - created_at, - started_at, - finished_at -) VALUES - ('candidate-reset', 'request-upstream-reset', 0, 0, 'failed', 502, 'upstream connection reset by peer', 0, 1, 2, 3), - ('candidate-stuck', 'request-stuck', 0, 0, 'pending', NULL, NULL, 0, 1, NULL, NULL) -"#, - ) - .execute(&pool) - .await - .expect("request candidates should seed"); - - let summary = repository - .cleanup_stale_pending_requests(2, 10, 5, 5) - .await - .expect("cleanup should run"); - assert_eq!(summary.recovered, 0); - assert_eq!(summary.failed, 2); - - let raw_messages = sqlx::query_as::<_, (String, Option, Option)>( - r#" -SELECT request_id, error_message, error_category -FROM "usage" -WHERE request_id IN ('request-upstream-reset', 'request-stuck') -ORDER BY request_id -"#, - ) - .fetch_all(&pool) - .await - .expect("raw stale usage diagnostics should load"); - assert_eq!( - raw_messages, - vec![ - ( - "request-stuck".to_string(), - None, - Some("server_error".to_string()) - ), - ( - "request-upstream-reset".to_string(), - None, - Some("server_error".to_string()) - ), - ] - ); - - let reset = repository - .find_by_request_id("request-upstream-reset") - .await - .expect("upstream-reset usage should load") - .expect("upstream-reset usage should exist"); - assert_eq!(reset.status, "failed"); - assert_eq!(reset.status_code, Some(502)); - assert_eq!(reset.error_message, None); - assert_eq!(reset.error_category.as_deref(), Some("server_error")); - - let stuck = repository - .find_by_request_id("request-stuck") - .await - .expect("stuck usage should load") - .expect("stuck usage should exist"); - assert_eq!(stuck.status, "failed"); - assert_eq!(stuck.status_code, Some(504)); - assert_eq!(stuck.error_message, None); - assert_eq!(stuck.error_category.as_deref(), Some("server_error")); -} - -#[tokio::test] -async fn sqlite_stale_cleanup_derives_category_from_candidate_status_code() { - 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( - "request-client-error-cleanup", - "pending", - "pending", - 1, - )) - .await - .expect("pending usage should upsert"); - sqlx::query( - r#" -INSERT INTO request_candidates ( - id, request_id, candidate_index, retry_index, status, status_code, - is_cached, created_at, started_at, finished_at -) VALUES ('candidate-client-error-cleanup', 'request-client-error-cleanup', 0, 0, - 'failed', 429, 0, 1, 2, 3) -"#, - ) - .execute(&pool) - .await - .expect("failed candidate should seed"); - - let summary = repository - .cleanup_stale_pending_requests(2, 10, 5, 10) - .await - .expect("cleanup should run"); - assert_eq!(summary.failed, 1); - - let stored = repository - .find_by_request_id("request-client-error-cleanup") - .await - .expect("usage should load") - .expect("usage should exist"); - assert_eq!(stored.status_code, Some(429)); - assert_eq!(stored.error_category.as_deref(), Some("client_error")); - assert_eq!(stored.error_message, None); -} - -#[tokio::test] -async fn sqlite_usage_read_repository_reads_usage_contract_views() { - 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()); - writer - .upsert(sample_usage("request-1", "completed", "settled", 1_000)) - .await - .expect("usage should upsert"); - writer - .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 - .find_by_request_id("request-1") - .await - .expect("usage should load") - .expect("usage should exist"); - 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 { - provider_name: Some("Provider One".to_string()), - newest_first: true, - ..UsageAuditListQuery::default() - }) - .await - .expect("usage list should load"); - assert_eq!(listed.len(), 2); - assert_eq!(listed[0].request_id, "request-2"); - - let summary = reader - .summarize_dashboard_usage(&UsageDashboardSummaryQuery { - created_from_unix_secs: 999, - created_until_unix_secs: 1_020, - user_id: Some("user-1".to_string()), - }) - .await - .expect("dashboard summary should load"); - assert_eq!(summary.total_requests, 2); - assert_eq!(summary.error_requests, 1); - assert_eq!(summary.total_tokens, 10); -} - -#[tokio::test] -async fn sqlite_usage_websocket_filter_applies_to_list_count_and_keyword_search() { - 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()); - writer - .upsert(sample_usage("request-http", "completed", "settled", 1_000)) - .await - .expect("HTTP usage should upsert"); - let mut websocket = sample_usage("request-ws", "completed", "void", 1_001); - websocket.request_metadata = Some(serde_json::json!({ - "websocket_mode": true, - "websocket_transport": "codex_live_direct", - "usage_available": false, - })); - websocket.input_tokens = None; - websocket.output_tokens = None; - websocket.total_tokens = None; - websocket.cache_creation_input_tokens = None; - websocket.cache_creation_ephemeral_5m_input_tokens = None; - websocket.cache_creation_ephemeral_1h_input_tokens = None; - websocket.cache_read_input_tokens = None; - websocket.cache_creation_cost_usd = None; - websocket.cache_read_cost_usd = None; - websocket.total_cost_usd = None; - websocket.actual_total_cost_usd = None; - writer - .upsert(websocket) - .await - .expect("WebSocket usage should upsert"); - - let reader = SqliteUsageReadRepository::new(pool); - let list_query = UsageAuditListQuery { - is_websocket: Some(true), - ..UsageAuditListQuery::default() - }; - let listed = reader - .list_usage_audits(&list_query) - .await - .expect("WebSocket list should load"); - assert_eq!(listed.len(), 1); - assert_eq!(listed[0].request_id, "request-ws"); - assert_eq!( - reader - .count_usage_audits(&list_query) - .await - .expect("WebSocket count should load"), - 1 - ); - - let keyword_query = UsageAuditKeywordSearchQuery { - is_websocket: Some(true), - keywords: vec!["model-1".to_string()], - ..UsageAuditKeywordSearchQuery::default() - }; - let keyword_matches = reader - .list_usage_audits_by_keyword_search(&keyword_query) - .await - .expect("WebSocket keyword list should load"); - assert_eq!(keyword_matches.len(), 1); - assert_eq!(keyword_matches[0].request_id, "request-ws"); - assert_eq!( - reader - .count_usage_audits_by_keyword_search(&keyword_query) - .await - .expect("WebSocket keyword count should load"), - 1 - ); - - let summary = reader - .summarize_usage_audits(&UsageAuditSummaryQuery { - created_from_unix_secs: 0, - created_until_unix_secs: 2_000, - ..UsageAuditSummaryQuery::default() - }) - .await - .expect("lifecycle summary should load"); - assert_eq!(summary.total_requests, 2); - assert_eq!(summary.recorded_total_tokens, 5); - - let provider_key_summaries = reader - .summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()]) - .await - .expect("provider key lifecycle summary should load"); - let provider_key_summary = provider_key_summaries - .get("provider-key-1") - .expect("provider key summary"); - assert_eq!(provider_key_summary.request_count, 2); - assert_eq!(provider_key_summary.total_tokens, 5); -} - -#[tokio::test] -async fn sqlite_usage_daily_heatmap_reads_imported_daily_aggregates() { - 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"); - - sqlx::query( - r#" -INSERT INTO stats_daily ( - id, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, is_complete, created_at, updated_at -) VALUES ( - 'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1 -); -INSERT INTO stats_user_daily ( - id, user_id, username, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) VALUES ( - 'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1 -); -"#, - ) - .execute(&pool) - .await - .expect("daily aggregates should seed"); - - let reader = SqliteUsageReadRepository::new(pool); - let admin = reader - .summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery { - created_from_unix_secs: 0, - user_id: None, - admin_mode: true, - }) - .await - .expect("admin heatmap should load"); - assert_eq!(admin.len(), 1); - assert_eq!(admin[0].date, "1970-01-02"); - assert_eq!(admin[0].requests, 9); - assert_eq!(admin[0].total_tokens, 37); - assert_eq!(admin[0].actual_total_cost_usd, 1.0); - - let user = reader - .summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery { - created_from_unix_secs: 0, - user_id: Some("user-1".to_string()), - admin_mode: false, - }) - .await - .expect("user heatmap should load"); - assert_eq!(user.len(), 1); - assert_eq!(user[0].date, "1970-01-02"); - assert_eq!(user[0].requests, 5); - assert_eq!(user[0].total_tokens, 18); - assert_eq!(user[0].actual_total_cost_usd, 0.75); -} - -#[tokio::test] -async fn sqlite_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() { - 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"); - - sqlx::query( - r#" -INSERT INTO stats_user_daily ( - id, user_id, username, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) VALUES ( - 'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1 -); -INSERT INTO "usage" ( - request_id, id, user_id, api_key_id, provider_name, model, total_tokens, - status, billing_status, created_at_unix_ms, updated_at_unix_secs -) VALUES - ('raw-before-cutoff', 'usage-1', 'user-1', 'api-key-1', 'Provider One', 'model-1', 99, - 'completed', 'settled', 90000, 90000), - ('raw-after-cutoff', 'usage-2', 'user-1', 'api-key-1', 'Provider One', 'model-1', 7, - 'completed', 'settled', 172800, 172800); -"#, - ) - .execute(&pool) - .await - .expect("usage totals fixtures should seed"); - - let reader = SqliteUsageReadRepository::new(pool); - let totals = reader - .summarize_usage_totals_by_user_ids(&["user-1".to_string()]) - .await - .expect("user totals should load"); - - assert_eq!(totals.len(), 1); - assert_eq!(totals[0].user_id, "user-1"); - assert_eq!(totals[0].request_count, 6); - assert_eq!(totals[0].total_tokens, 25); -} - -#[tokio::test] -async fn sqlite_dashboard_daily_stats_reads_imported_daily_aggregates() { - 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"); - - sqlx::query( - r#" -INSERT INTO stats_daily ( - id, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, is_complete, created_at, updated_at -) VALUES ( - 'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1 -); -"#, - ) - .execute(&pool) - .await - .expect("daily aggregates should seed"); - - let reader = SqliteUsageReadRepository::new(pool); - let summary = reader - .summarize_dashboard_usage(&UsageDashboardSummaryQuery { - created_from_unix_secs: 0, - created_until_unix_secs: 172800, - user_id: None, - }) - .await - .expect("dashboard summary should load"); - assert_eq!(summary.total_requests, 9); - assert_eq!(summary.total_tokens, 37); - - let rows = reader - .list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery { - created_from_unix_secs: 0, - created_until_unix_secs: 172800, - tz_offset_minutes: 480, - user_id: None, - }) - .await - .expect("dashboard daily breakdown should load"); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].date, "1970-01-02"); - assert_eq!(rows[0].model, "aggregate"); - assert_eq!(rows[0].requests, 9); - 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!(duplicate.request_body.is_none()); - assert!(duplicate.request_body_ref.is_none()); - - 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_first_byte_fast_path_rejects_stale_revision() { - 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()); - - let mut current = sample_usage("first-byte-stale", "streaming", "pending", 2_000); - current.finalized_at_unix_secs = None; - current.provider_name = "current-provider".to_string(); - current.model = "current-model".to_string(); - current.response_time_ms = Some(50); - repository - .upsert(current) - .await - .expect("current streaming usage should seed"); - - let mut stale = sample_usage("first-byte-stale", "streaming", "pending", 1_999); - stale.finalized_at_unix_secs = None; - stale.provider_name = "stale-provider".to_string(); - stale.model = "stale-model".to_string(); - stale.status_code = Some(503); - stale.response_time_ms = Some(999); - repository - .upsert_first_byte(stale) - .await - .expect("stale first-byte usage should be ignored"); - - let stored = repository - .find_by_request_id("first-byte-stale") - .await - .expect("streaming usage should load") - .expect("streaming usage should exist"); - assert_eq!(stored.updated_at_unix_secs, 2_000); - assert_eq!(stored.provider_name, "current-provider"); - assert_eq!(stored.model, "current-model"); - assert_eq!(stored.status_code, Some(200)); - assert_eq!(stored.response_time_ms, Some(50)); -} - -#[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_routing_snapshot - BEFORE INSERT ON usage_routing_snapshots -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_routing_snapshot") - .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, 0, 0, 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#" -INSERT INTO users (id, auth_source, created_at, updated_at) -VALUES ('user-1', 'local', 1, 1); -INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) -VALUES ('api-key-1', 'user-1', 'hash-1', 1, 1); -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) - .await - .expect("stats targets should seed"); -} - -fn sample_usage( - request_id: &str, - status: &str, - billing_status: &str, - updated_at: u64, -) -> UpsertUsageRecord { - UpsertUsageRecord { - request_id: request_id.to_string(), - user_id: Some("user-1".to_string()), - api_key_id: Some("api-key-1".to_string()), - username: Some("legacy-user".to_string()), - api_key_name: Some("legacy-key".to_string()), - provider_name: "Provider One".to_string(), - model: "model-1".to_string(), - target_model: Some("target-model".to_string()), - provider_id: Some("provider-1".to_string()), - provider_endpoint_id: Some("endpoint-1".to_string()), - provider_api_key_id: Some("provider-key-1".to_string()), - request_type: Some("chat".to_string()), - api_format: Some("openai".to_string()), - api_family: Some("chat".to_string()), - endpoint_kind: Some("chat".to_string()), - endpoint_api_format: Some("openai".to_string()), - provider_api_family: Some("chat".to_string()), - provider_endpoint_kind: Some("chat".to_string()), - has_format_conversion: Some(true), - is_stream: Some(false), - input_tokens: Some(2), - output_tokens: Some(3), - total_tokens: None, - cache_creation_input_tokens: None, - cache_creation_ephemeral_5m_input_tokens: Some(0), - cache_creation_ephemeral_1h_input_tokens: Some(0), - cache_read_input_tokens: Some(2), - cache_creation_cost_usd: Some(0.0), - cache_read_cost_usd: Some(0.1), - output_price_per_1m: Some(2.0), - total_cost_usd: Some(0.5), - actual_total_cost_usd: Some(0.4), - status_code: Some(200), - error_message: None, - error_category: None, - response_time_ms: Some(42), - first_byte_time_ms: Some(12), - status: status.to_string(), - billing_status: billing_status.to_string(), - request_headers: None, - request_body: None, - request_body_ref: None, - request_body_state: None, - provider_request_headers: None, - provider_request_body: None, - provider_request_body_ref: None, - provider_request_body_state: None, - response_headers: None, - response_body: None, - response_body_ref: None, - response_body_state: None, - client_response_headers: None, - client_response_body: None, - client_response_body_ref: None, - client_response_body_state: None, - candidate_id: Some("candidate-1".to_string()), - candidate_index: Some(1), - key_name: Some("key-one".to_string()), - planner_kind: Some("default".to_string()), - route_family: Some("chat".to_string()), - route_kind: Some("completion".to_string()), - execution_path: Some("remote".to_string()), - local_execution_runtime_miss_reason: None, - request_metadata: Some(serde_json::json!({ - "trace_id": "trace-1", - "upstream_is_stream": true, - })), - finalized_at_unix_secs: Some(updated_at), - created_at_unix_ms: Some(updated_at), - 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/users.rs b/crates/aether-data/adapters/sqlite/src/users.rs deleted file mode 100644 index f7c0eaf0b..000000000 --- a/crates/aether-data/adapters/sqlite/src/users.rs +++ /dev/null @@ -1,5336 +0,0 @@ -use async_trait::async_trait; -use chrono::{DateTime, TimeZone, Utc}; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::users::{ - is_valid_bcrypt_hash, last_oauth_unbind_denial, normalize_user_group_name, - BindUserOAuthLinkOutcome, BindUserOAuthLinkSessionExpectation, DeleteUserOAuthLinkOutcome, - LdapAuthUserProvisioningOutcome, ResolveOAuthLinkedUserOutcome, StoredUserAuthRecord, - StoredUserExportRow, StoredUserGroup, StoredUserGroupMember, StoredUserGroupMembership, - StoredUserOAuthLinkSummary, StoredUserPreferenceRecord, StoredUserSessionRecord, - StoredUserSummary, UpsertUserGroupRecord, UserExportListQuery, UserExportSortBy, - UserExportSummary, UserReadRepository, LAST_ACTIVE_ADMIN_DELETE_DENIED, - LAST_ACTIVE_ADMIN_UPDATE_DENIED, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const USER_SUMMARY_COLUMNS: &str = r#" -SELECT - id, - username, - email, - role, - is_active, - is_deleted -FROM users -"#; - -const SQLITE_ACTIVE_ADMIN_UPDATE_GUARD: &str = r#" - AND ( - ? = 0 - OR COALESCE(LOWER(role), '') != 'admin' - OR is_active = 0 - OR is_deleted != 0 - OR EXISTS ( - SELECT 1 - FROM users AS other_admin - WHERE other_admin.id != users.id - AND LOWER(other_admin.role) = 'admin' - AND other_admin.is_active = 1 - AND other_admin.is_deleted = 0 - ) - ) -"#; - -const SQLITE_DELETE_USER_SQL: &str = r#" -DELETE FROM users -WHERE id = ? - AND ( - COALESCE(LOWER(role), '') != 'admin' - OR is_active = 0 - OR is_deleted != 0 - OR EXISTS ( - SELECT 1 - FROM users AS other_admin - WHERE other_admin.id != users.id - AND LOWER(other_admin.role) = 'admin' - AND other_admin.is_active = 1 - AND other_admin.is_deleted = 0 - ) - ) -"#; - -const SQLITE_DELETE_USER_IF_WALLET_ABSENT_SQL: &str = r#" -DELETE FROM users -WHERE id = ? - AND NOT EXISTS ( - SELECT 1 - FROM wallets AS wallet - WHERE wallet.user_id = ? - OR EXISTS ( - SELECT 1 - FROM api_keys AS api_key - WHERE api_key.id = wallet.api_key_id - AND api_key.user_id = ? - ) - ) - AND ( - COALESCE(LOWER(role), '') != 'admin' - OR is_active = 0 - OR is_deleted != 0 - OR EXISTS ( - SELECT 1 - FROM users AS other_admin - WHERE other_admin.id != users.id - AND LOWER(other_admin.role) = 'admin' - AND other_admin.is_active = 1 - AND other_admin.is_deleted = 0 - ) - ) -"#; - -const SQLITE_DELETE_USER_API_KEYS_SQL: &str = "DELETE FROM api_keys WHERE user_id = ?"; - -const SQLITE_DELETE_USER_DEPENDENTS_SQL: &[&str] = &[ - "DELETE FROM usage_request_admissions WHERE subject_id = ?", - "DELETE FROM usage_cost_reservations WHERE subject_id = ?", - "DELETE FROM gemini_file_mappings WHERE user_id = ?", - "DELETE FROM api_key_provider_mappings WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - SQLITE_DELETE_USER_API_KEYS_SQL, - "DELETE FROM management_tokens WHERE user_id = ?", - "DELETE FROM user_sessions WHERE user_id = ?", - "DELETE FROM user_oauth_links WHERE user_id = ?", - "DELETE FROM user_group_members WHERE user_id = ?", - "DELETE FROM user_preferences WHERE user_id = ?", - "DELETE FROM user_invite_codes WHERE user_id = ?", - "DELETE FROM announcement_reads WHERE user_id = ?", -]; - -const SQLITE_PREPARE_USER_FACTS_FOR_DELETION_SQL: &[&str] = &[ - "UPDATE referral_rewards SET status = CASE WHEN status IN ('pending', 'failed', 'applying') THEN 'voided' ELSE status END, failure_reason = NULL, admin_note = NULL, updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE ? IN (inviter_user_id, invitee_user_id)", - "UPDATE referral_rewards SET failure_reason = NULL, admin_note = NULL, updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE admin_operator_id = ?", - "UPDATE user_referrals SET invite_code_snapshot = 'deleted-user', source_json = NULL, updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE ? IN (inviter_user_id, invitee_user_id)", - "UPDATE user_plan_entitlements SET status = CASE WHEN status = 'active' THEN 'revoked' ELSE status END, expires_at = MIN(expires_at, CAST(strftime('%s', 'now') AS INTEGER)), updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE user_id = ?", - "UPDATE wallets SET status = 'disabled', updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE user_id = ?", - "UPDATE wallets SET status = 'disabled', updated_at = CAST(strftime('%s', 'now') AS INTEGER) WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - "UPDATE audit_logs SET description = 'deleted user event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE user_id = ?", - "UPDATE audit_logs SET description = 'deleted API key event', ip_address = NULL, user_agent = NULL, event_metadata = NULL, error_message = NULL WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE user_id = ?)", - "UPDATE wallet_transactions SET description = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE wallet_transactions SET description = NULL WHERE operator_id = ?", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order WHERE history_order.user_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.user_id = ? AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_callbacks SET payload = NULL, error_message = NULL WHERE EXISTS (SELECT 1 FROM payment_orders AS history_order JOIN wallets AS history_wallet ON history_wallet.id = history_order.wallet_id WHERE history_wallet.api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?) AND (history_order.id = payment_callbacks.payment_order_id OR (payment_callbacks.order_no IS NOT NULL AND history_order.order_no = payment_callbacks.order_no)))", - "UPDATE payment_orders SET gateway_response = NULL WHERE user_id = ?", - "UPDATE payment_orders SET gateway_response = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE user_id = ?", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE ? IN (requested_by, approved_by, processed_by)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE user_id = ?)", - "UPDATE refund_requests SET reason = NULL, payout_reference = NULL, payout_proof = NULL, failure_reason = NULL WHERE wallet_id IN (SELECT id FROM wallets WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?))", - "UPDATE redeem_code_batches SET description = NULL WHERE created_by = ?", -]; - -const SQLITE_ANONYMIZE_USER_HISTORY_SQL: &[&str] = &[ - "UPDATE request_candidates SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE video_tasks SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE usage SET username = NULL, api_key_name = NULL WHERE user_id = ?", - "UPDATE stats_user_daily SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_summary SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_model SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_api_format SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_model_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_provider SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_model SET username = NULL WHERE user_id = ?", - "UPDATE stats_user_daily_cost_savings_model_provider SET username = NULL WHERE user_id = ?", -]; - -const SQLITE_ANONYMIZE_USER_API_KEY_HISTORY_SQL: &str = - "UPDATE stats_daily_api_key SET api_key_name = NULL WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id = ?)"; - -const USER_EXPORT_COLUMNS: &str = r#" -SELECT - id, - email, - email_verified, - username, - password_hash, - role, - auth_source, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - rate_limit, - rate_limit_mode, - model_capability_settings, - feature_settings, - is_active -FROM users -"#; - -const USER_AUTH_COLUMNS: &str = r#" -SELECT - id, - email, - email_verified, - username, - password_hash, - role, - auth_source, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - is_active, - is_deleted, - security_version, - created_at, - last_login_at -FROM users -"#; - -const USER_AUTH_COLUMNS_QUALIFIED: &str = r#" -SELECT - users.id AS id, - users.email AS email, - users.email_verified AS email_verified, - users.username AS username, - users.password_hash AS password_hash, - users.role AS role, - users.auth_source AS auth_source, - users.allowed_providers AS allowed_providers, - users.allowed_providers_mode AS allowed_providers_mode, - users.allowed_api_formats AS allowed_api_formats, - users.allowed_api_formats_mode AS allowed_api_formats_mode, - users.allowed_models AS allowed_models, - users.allowed_models_mode AS allowed_models_mode, - users.is_active AS is_active, - users.is_deleted AS is_deleted, - users.security_version AS security_version, - users.created_at AS created_at, - users.last_login_at AS last_login_at -FROM users -"#; - -const USER_OAUTH_LINK_SUMMARY_COLUMNS: &str = r#" -SELECT - user_oauth_links.provider_type, - oauth_providers.display_name, - user_oauth_links.provider_username, - user_oauth_links.provider_email, - user_oauth_links.linked_at, - user_oauth_links.last_login_at, - oauth_providers.is_enabled AS provider_enabled -FROM user_oauth_links -JOIN oauth_providers - ON oauth_providers.provider_type = user_oauth_links.provider_type -"#; - -const USER_PREFERENCES_COLUMNS: &str = r#" -SELECT - up.user_id, - up.avatar_url, - up.bio, - up.default_provider_id, - p.name AS default_provider_name, - up.theme, - up.language, - up.timezone, - up.email_notifications, - up.usage_alerts, - up.announcement_notifications -FROM user_preferences up -LEFT JOIN providers p - ON p.id = up.default_provider_id -"#; - -const USER_SESSION_COLUMNS: &str = r#" -SELECT - id, - user_id, - security_version, - client_device_id, - device_label, - refresh_token_hash, - prev_refresh_token_hash, - rotated_at, - last_seen_at, - expires_at, - revoked_at, - revoke_reason, - ip_address, - user_agent, - created_at, - updated_at -FROM user_sessions -"#; - -const USER_GROUP_COLUMNS: &str = r#" -SELECT - id, - name, - normalized_name, - description, - priority, - allowed_providers, - allowed_providers_mode, - allowed_api_formats, - allowed_api_formats_mode, - allowed_models, - allowed_models_mode, - rate_limit, - rate_limit_mode, - created_at, - updated_at -FROM user_groups -"#; - -const USER_GROUP_MEMBER_COLUMNS: &str = r#" -SELECT - user_group_members.group_id, - users.id AS user_id, - users.username, - users.email, - users.role, - users.is_active, - users.is_deleted, - user_group_members.created_at -FROM user_group_members -JOIN users ON users.id = user_group_members.user_id -"#; - -#[derive(Debug, Clone)] -pub struct SqliteUserReadRepository { - pool: SqlitePool, -} - -impl SqliteUserReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn fetch_summary_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_row).collect() - } - - async fn fetch_export_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_export_row).collect() - } - - async fn fetch_auth_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_auth_row).collect() - } - - async fn fetch_group_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_row).collect() - } - - async fn fetch_group_member_rows( - &self, - mut builder: QueryBuilder<'_, Sqlite>, - ) -> Result, DataLayerError> { - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_member_row).collect() - } - - async fn delete_local_auth_user_inner( - &self, - user_id: &str, - require_wallet_absent: bool, - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - if require_wallet_absent { - let wallet_exists: Option = sqlx::query_scalar( - r#" -SELECT 1 -FROM wallets AS wallet -WHERE wallet.user_id = ? - OR EXISTS ( - SELECT 1 - FROM api_keys AS api_key - WHERE api_key.id = wallet.api_key_id - AND api_key.user_id = ? - ) -LIMIT 1 - "#, - ) - .bind(user_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if wallet_exists.is_some() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - } - for sql in SQLITE_PREPARE_USER_FACTS_FOR_DELETION_SQL { - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - for sql in SQLITE_ANONYMIZE_USER_HISTORY_SQL { - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - sqlx::query(SQLITE_ANONYMIZE_USER_API_KEY_HISTORY_SQL) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - for sql in SQLITE_DELETE_USER_DEPENDENTS_SQL { - if require_wallet_absent && *sql == SQLITE_DELETE_USER_API_KEYS_SQL { - continue; - } - sqlx::query(sql) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let result = if require_wallet_absent { - sqlx::query(SQLITE_DELETE_USER_IF_WALLET_ABSENT_SQL) - .bind(user_id) - .bind(user_id) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query(SQLITE_DELETE_USER_SQL) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - }; - if result.rows_affected() > 0 { - if require_wallet_absent { - sqlx::query(SQLITE_DELETE_USER_API_KEYS_SQL) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - return Ok(true); - } - let blocked_active_admin: Option = sqlx::query_scalar( - "SELECT 1 FROM users WHERE id = ? AND LOWER(role) = 'admin' AND is_active = 1 AND is_deleted = 0 LIMIT 1", - ) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - tx.rollback().await.map_sql_err()?; - if blocked_active_admin.is_some() { - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_DELETE_DENIED.to_string(), - )); - } - Ok(false) - } -} - -#[async_trait] -impl UserReadRepository for SqliteUserReadRepository { - async fn list_users_by_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_SUMMARY_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push(") ORDER BY id ASC"); - self.fetch_summary_rows(builder).await - } - - async fn list_users_by_username_search( - &self, - username_search: &str, - ) -> Result, DataLayerError> { - let username_search = username_search.trim(); - if username_search.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_SUMMARY_COLUMNS); - builder - .push(" WHERE is_deleted = 0 AND LOWER(username) LIKE ") - .push_bind(format!("%{}%", username_search.to_ascii_lowercase())) - .push(" ORDER BY id ASC"); - self.fetch_summary_rows(builder).await - } - - async fn list_export_users(&self) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0 ORDER BY id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_export_users_page( - &self, - query: &UserExportListQuery, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0"); - if let Some(role) = query.role.as_deref() { - builder - .push(" AND LOWER(role) = ") - .push_bind(role.trim().to_ascii_lowercase()); - } - if let Some(is_active) = query.is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - if let Some(group_id) = query - .group_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = "); - builder.push_bind(group_id); - builder.push(")"); - } - 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(id) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(username) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(COALESCE(email, '')) LIKE ") - .push_bind(pattern) - .push(")"); - } - match query.sort_by { - UserExportSortBy::CreatedAt => { - builder - .push(" ORDER BY created_at ") - .push(if query.sort_order.is_desc() { - "DESC" - } else { - "ASC" - }) - .push(", id ASC"); - } - UserExportSortBy::Id => { - builder.push(" ORDER BY id ASC"); - } - } - - builder - .push(" LIMIT ") - .push_bind(i64::try_from(query.limit).map_err(|_| { - DataLayerError::InvalidInput(format!("invalid user export limit: {}", query.limit)) - })?) - .push(" OFFSET ") - .push_bind(i64::try_from(query.skip).map_err(|_| { - DataLayerError::InvalidInput(format!("invalid user export skip: {}", query.skip)) - })?); - self.fetch_export_rows(builder).await - } - - async fn count_export_users(&self, query: &UserExportListQuery) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(*) AS total FROM users"); - builder.push(" WHERE is_deleted = 0"); - if let Some(role) = query.role.as_deref() { - builder - .push(" AND LOWER(role) = ") - .push_bind(role.trim().to_ascii_lowercase()); - } - if let Some(is_active) = query.is_active { - builder.push(" AND is_active = ").push_bind(is_active); - } - if let Some(group_id) = query - .group_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - builder.push(" AND id IN (SELECT user_id FROM user_group_members WHERE group_id = "); - builder.push_bind(group_id); - builder.push(")"); - } - 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(id) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(username) LIKE ") - .push_bind(pattern.clone()) - .push(" OR LOWER(COALESCE(email, '')) LIKE ") - .push_bind(pattern) - .push(")"); - } - - let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?; - Ok(row.try_get::("total").map_sql_err()?.max(0) as u64) - } - - async fn summarize_export_users(&self) -> Result { - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total, - SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END) AS active -FROM users -WHERE is_deleted = 0 -"#, - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - Ok(UserExportSummary { - total: row.try_get::("total").map_sql_err()?.max(0) as u64, - active: row - .try_get::, _>("active") - .map_sql_err()? - .unwrap_or(0) - .max(0) as u64, - }) - } - - async fn find_export_user_by_id( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder - .push(" WHERE is_deleted = 0 AND id = ") - .push_bind(user_id) - .push(" LIMIT 1"); - Ok(self.fetch_export_rows(builder).await?.into_iter().next()) - } - - async fn list_non_admin_export_users( - &self, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_EXPORT_COLUMNS); - builder.push(" WHERE is_deleted = 0 AND LOWER(role) != 'admin' ORDER BY id ASC"); - self.fetch_export_rows(builder).await - } - - async fn list_user_groups(&self) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder.push(" ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn find_user_group_by_id( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder - .push(" WHERE id = ") - .push_bind(group_id) - .push(" LIMIT 1"); - Ok(self.fetch_group_rows(builder).await?.into_iter().next()) - } - - async fn list_user_groups_by_ids( - &self, - group_ids: &[String], - ) -> Result, DataLayerError> { - if group_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for group_id in group_ids { - separated.push_bind(group_id); - } - } - builder.push(") ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn create_user_group( - &self, - record: UpsertUserGroupRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let id = uuid::Uuid::new_v4().to_string(); - let name = normalize_user_group_name(&record.name); - let normalized_name = name.to_ascii_lowercase(); - let result = sqlx::query( - r#" -INSERT INTO user_groups ( - id, name, normalized_name, description, priority, - allowed_providers, allowed_providers_mode, - allowed_api_formats, allowed_api_formats_mode, - allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&id) - .bind(name) - .bind(normalized_name) - .bind(record.description) - .bind(record.priority) - .bind(json_string_from_option_vec( - record.allowed_providers.as_ref(), - )) - .bind(record.allowed_providers_mode) - .bind(json_string_from_option_vec( - record.allowed_api_formats.as_ref(), - )) - .bind(record.allowed_api_formats_mode) - .bind(json_string_from_option_vec(record.allowed_models.as_ref())) - .bind(record.allowed_models_mode) - .bind(record.rate_limit) - .bind(record.rate_limit_mode) - .bind(now) - .bind(now) - .execute(&self.pool) - .await; - match result { - Ok(_) => self.find_user_group_by_id(&id).await, - Err(sqlx::Error::Database(err)) if err.is_unique_violation() => Err( - DataLayerError::InvalidInput("duplicate user group name".to_string()), - ), - Err(err) => Err(err).map_sql_err(), - } - } - - async fn update_user_group( - &self, - group_id: &str, - record: UpsertUserGroupRecord, - ) -> Result, DataLayerError> { - let now = current_unix_secs(); - let name = normalize_user_group_name(&record.name); - let normalized_name = name.to_ascii_lowercase(); - let result = sqlx::query( - r#" -UPDATE user_groups -SET name = ?, - normalized_name = ?, - description = ?, - priority = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(name) - .bind(normalized_name) - .bind(record.description) - .bind(record.priority) - .bind(json_string_from_option_vec( - record.allowed_providers.as_ref(), - )) - .bind(record.allowed_providers_mode) - .bind(json_string_from_option_vec( - record.allowed_api_formats.as_ref(), - )) - .bind(record.allowed_api_formats_mode) - .bind(json_string_from_option_vec(record.allowed_models.as_ref())) - .bind(record.allowed_models_mode) - .bind(record.rate_limit) - .bind(record.rate_limit_mode) - .bind(now) - .bind(group_id) - .execute(&self.pool) - .await; - match result { - Ok(result) if result.rows_affected() == 0 => Ok(None), - Ok(_) => self.find_user_group_by_id(group_id).await, - Err(sqlx::Error::Database(err)) if err.is_unique_violation() => Err( - DataLayerError::InvalidInput("duplicate user group name".to_string()), - ), - Err(err) => Err(err).map_sql_err(), - } - } - - /// BEGIN IMMEDIATE serializes writers while the complete snapshot is - /// compared and restored, preventing a rollback from overwriting a newer - /// administrator update. - async fn restore_user_group_if_matches( - &self, - expected: &StoredUserGroup, - restored: &StoredUserGroup, - ) -> Result { - if expected.id != restored.id || expected.id.trim().is_empty() { - return Ok(false); - } - - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder.push(" WHERE id = ").push_bind(&expected.id); - let row = builder - .build() - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_user_group_row(&row)?; - if ¤t != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let result = sqlx::query( - r#" -UPDATE user_groups -SET name = ?, - normalized_name = ?, - description = ?, - priority = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - created_at = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&restored.name) - .bind(&restored.normalized_name) - .bind(&restored.description) - .bind(restored.priority) - .bind(json_string_from_option_vec( - restored.allowed_providers.as_ref(), - )) - .bind(&restored.allowed_providers_mode) - .bind(json_string_from_option_vec( - restored.allowed_api_formats.as_ref(), - )) - .bind(&restored.allowed_api_formats_mode) - .bind(json_string_from_option_vec( - restored.allowed_models.as_ref(), - )) - .bind(&restored.allowed_models_mode) - .bind(restored.rate_limit) - .bind(&restored.rate_limit_mode) - .bind(restored.created_at.map(|value| value.timestamp())) - .bind(restored.updated_at.map(|value| value.timestamp())) - .bind(&restored.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_user_group(&self, group_id: &str) -> Result { - let result = sqlx::query("DELETE FROM user_groups WHERE id = ?") - .bind(group_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn list_user_group_members( - &self, - group_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_MEMBER_COLUMNS); - builder - .push(" WHERE user_group_members.group_id = ") - .push_bind(group_id) - .push(" ORDER BY users.username ASC, users.id ASC"); - self.fetch_group_member_rows(builder).await - } - - async fn replace_user_group_members( - &self, - group_id: &str, - user_ids: &[String], - ) -> Result, DataLayerError> { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - sqlx::query("DELETE FROM user_group_members WHERE group_id = ?") - .bind(group_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for user_id in normalized_ids(user_ids) { - sqlx::query( - "INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.list_user_group_members(group_id).await - } - - async fn list_user_groups_for_user( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_GROUP_COLUMNS); - builder - .push(" WHERE id IN (SELECT group_id FROM user_group_members WHERE user_id = ") - .push_bind(user_id) - .push(") ORDER BY name ASC, id ASC"); - self.fetch_group_rows(builder).await - } - - async fn list_user_group_memberships_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - r#" -SELECT - user_group_members.user_id, - user_groups.id AS group_id, - user_groups.name AS group_name, - user_groups.priority AS group_priority, - user_group_members.created_at -FROM user_group_members -JOIN user_groups ON user_groups.id = user_group_members.group_id -WHERE user_group_members.user_id IN ( -"#, - ); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push( - ") ORDER BY user_group_members.user_id ASC, user_groups.name ASC, user_groups.id ASC", - ); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_group_membership_row).collect() - } - - async fn replace_user_groups_for_user( - &self, - user_id: &str, - group_ids: &[String], - ) -> Result, DataLayerError> { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(Vec::new()); - } - sqlx::query("DELETE FROM user_group_members WHERE user_id = ?") - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for group_id in normalized_ids(group_ids) { - sqlx::query( - "INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.list_user_groups_for_user(user_id).await - } - - async fn restore_user_groups_if_matches( - &self, - user_id: &str, - expected_group_ids: &[String], - restored_group_ids: &[String], - ) -> Result { - let expected = normalized_ids(expected_group_ids); - let restored = normalized_ids(restored_group_ids); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let current = sqlx::query_scalar::<_, String>( - "SELECT group_id FROM user_group_members WHERE user_id = ? ORDER BY group_id ASC", - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if current != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - if !restored.is_empty() { - let mut builder = QueryBuilder::::new( - "SELECT COUNT(*) AS count FROM user_groups WHERE id IN (", - ); - { - let mut separated = builder.separated(", "); - for group_id in &restored { - separated.push_bind(group_id); - } - } - builder.push(")"); - let count = builder - .build() - .fetch_one(&mut *tx) - .await - .map_sql_err()? - .try_get::("count") - .map_sql_err()?; - if count != restored.len() as i64 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - } - sqlx::query("DELETE FROM user_group_members WHERE user_id = ?") - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let now = current_unix_secs(); - for group_id in restored { - sqlx::query( - "INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn add_user_to_group( - &self, - group_id: &str, - user_id: &str, - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let result = sqlx::query( - "INSERT OR IGNORE INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, ?)", - ) - .bind(group_id) - .bind(user_id) - .bind(current_unix_secs()) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn find_user_auth_by_id( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE id = ") - .push_bind(user_id) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn list_user_auth_by_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder.push(" WHERE id IN ("); - { - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - } - builder.push(") ORDER BY id ASC"); - self.fetch_auth_rows(builder).await - } - - async fn find_user_auth_by_identifier( - &self, - identifier: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE email = ") - .push_bind(identifier) - .push(" OR username = ") - .push_bind(identifier) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_user_auth_by_email( - &self, - email: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE email = ") - .push_bind(email) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_active_user_auth_by_email_ci( - &self, - email: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE LOWER(email) = LOWER(") - .push_bind(email) - .push(") AND is_deleted = 0 LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn find_user_auth_by_username( - &self, - username: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS); - builder - .push(" WHERE username = ") - .push_bind(username) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn list_user_oauth_links( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_OAUTH_LINK_SUMMARY_COLUMNS); - builder - .push(" WHERE user_oauth_links.user_id = ") - .push_bind(user_id) - .push(" ORDER BY user_oauth_links.linked_at ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_oauth_link_summary_row).collect() - } - - async fn find_oauth_linked_user( - &self, - provider_type: &str, - provider_user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_AUTH_COLUMNS_QUALIFIED); - builder - .push(" JOIN user_oauth_links ON users.id = user_oauth_links.user_id") - .push(" WHERE user_oauth_links.provider_type = ") - .push_bind(provider_type) - .push(" AND user_oauth_links.provider_user_id = ") - .push_bind(provider_user_id) - .push(" LIMIT 1"); - Ok(self.fetch_auth_rows(builder).await?.into_iter().next()) - } - - async fn resolve_enabled_oauth_linked_user( - &self, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - verified_email: Option<&str>, - touched_at: DateTime, - _provider_enabled_snapshot: bool, - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let provider_enabled: Option = - sqlx::query_scalar("SELECT is_enabled FROM oauth_providers WHERE provider_type = ?") - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_enabled != Some(true) { - tx.rollback().await.map_sql_err()?; - return Ok(ResolveOAuthLinkedUserOutcome::ProviderUnavailable); - } - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS_QUALIFIED} JOIN user_oauth_links ON users.id = user_oauth_links.user_id WHERE user_oauth_links.provider_type = ? AND user_oauth_links.provider_user_id = ? LIMIT 1" - )) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(ResolveOAuthLinkedUserOutcome::NotLinked); - }; - let mut user = map_user_auth_row(&row)?; - sqlx::query( - "UPDATE user_oauth_links SET provider_username = COALESCE(?, provider_username), provider_email = COALESCE(?, provider_email), extra_data = COALESCE(?, extra_data), last_login_at = ? WHERE provider_type = ? AND provider_user_id = ?", - ) - .bind(provider_username) - .bind(provider_email) - .bind(optional_json_string(extra_data, "user_oauth_links.extra_data")?) - .bind(touched_at.timestamp()) - .bind(provider_type) - .bind(provider_user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if let Some(verified_email) = verified_email { - let result = sqlx::query( - "UPDATE users SET email_verified = 1, updated_at = ? WHERE id = ? AND email_verified = 0 AND LOWER(TRIM(email)) = LOWER(TRIM(?))", - ) - .bind(touched_at.timestamp()) - .bind(&user.id) - .bind(verified_email) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() == 1 { - user.email_verified = true; - } - } - tx.commit().await.map_sql_err()?; - Ok(ResolveOAuthLinkedUserOutcome::Linked(user)) - } - - async fn touch_oauth_link( - &self, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - touched_at: DateTime, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_oauth_links -SET provider_username = COALESCE(?, provider_username), - provider_email = COALESCE(?, provider_email), - extra_data = COALESCE(?, extra_data), - last_login_at = ? -WHERE provider_type = ? - AND provider_user_id = ? -"#, - ) - .bind(provider_username) - .bind(provider_email) - .bind(optional_json_string( - extra_data, - "user_oauth_links.extra_data", - )?) - .bind(touched_at.timestamp()) - .bind(provider_type) - .bind(provider_user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn create_oauth_auth_user( - &self, - email: Option, - email_verified: bool, - username: String, - created_at: DateTime, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, rate_limit_mode, - is_active, is_deleted, created_at, updated_at, last_login_at -) -VALUES (?, ?, ?, ?, NULL, 'user', 'oauth', 'inherit', 'inherit', 'inherit', 'inherit', 1, 0, ?, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(created_at.timestamp()) - .bind(created_at.timestamp()) - .bind(created_at.timestamp()) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn find_oauth_link_owner( - &self, - provider_type: &str, - provider_user_id: &str, - ) -> Result, DataLayerError> { - sqlx::query_scalar( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ? AND provider_user_id = ? LIMIT 1", - ) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err() - } - - async fn has_user_oauth_provider_link( - &self, - user_id: &str, - provider_type: &str, - ) -> Result { - let owner: Option = sqlx::query_scalar( - "SELECT user_id FROM user_oauth_links WHERE user_id = ? AND provider_type = ? LIMIT 1", - ) - .bind(user_id) - .bind(provider_type) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(owner.is_some()) - } - - async fn count_user_oauth_links(&self, user_id: &str) -> Result { - let total: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links WHERE user_id = ?") - .bind(user_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(total.max(0) as u64) - } - - async fn has_oauth_links_for_provider( - &self, - provider_type: &str, - ) -> Result { - let exists: Option = - sqlx::query_scalar("SELECT 1 FROM user_oauth_links WHERE provider_type = ? LIMIT 1") - .bind(provider_type) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - Ok(exists.is_some()) - } - - async fn bind_user_oauth_link_if_provider_enabled( - &self, - user_id: &str, - provider_type: &str, - provider_user_id: &str, - provider_username: Option<&str>, - provider_email: Option<&str>, - extra_data: Option, - linked_at: DateTime, - _provider_enabled_snapshot: bool, - session_expectation: Option<&BindUserOAuthLinkSessionExpectation>, - ) -> Result { - let extra_data = optional_json_string(extra_data, "user_oauth_links.extra_data")?; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let provider_enabled: Option = - sqlx::query_scalar("SELECT is_enabled FROM oauth_providers WHERE provider_type = ?") - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if provider_enabled.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::ProviderNotFound); - } - if provider_enabled != Some(true) { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::ProviderDisabled); - } - if let Some(expectation) = session_expectation { - let session_is_current: Option = sqlx::query_scalar( - r#" -SELECT 1 -FROM users -JOIN user_sessions - ON user_sessions.user_id = users.id -WHERE users.id = ? - AND users.is_active = 1 - AND users.is_deleted = 0 - AND users.security_version = ? - AND user_sessions.id = ? - AND user_sessions.security_version = ? - AND user_sessions.client_device_id = ? - AND user_sessions.revoked_at IS NULL - AND user_sessions.expires_at > MAX(?, CAST(strftime('%s', 'now') AS INTEGER)) -"#, - ) - .bind(user_id) - .bind(expectation.security_version) - .bind(&expectation.session_id) - .bind(expectation.security_version) - .bind(&expectation.client_device_id) - .bind(expectation.checked_at.timestamp()) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if session_is_current.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::SessionUnavailable); - } - } else { - let user_exists: Option = sqlx::query_scalar("SELECT 1 FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::UserNotFound); - } - } - if let Some(owner) = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ? AND provider_user_id = ? LIMIT 1", - ) - .bind(provider_type) - .bind(provider_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - { - tx.rollback().await.map_sql_err()?; - return Ok(if owner == user_id { - BindUserOAuthLinkOutcome::IdentityAlreadyBoundToUser - } else { - BindUserOAuthLinkOutcome::IdentityBoundToAnotherUser - }); - } - if sqlx::query_scalar::<_, i32>( - "SELECT 1 FROM user_oauth_links WHERE user_id = ? AND provider_type = ? LIMIT 1", - ) - .bind(user_id) - .bind(provider_type) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - .is_some() - { - tx.rollback().await.map_sql_err()?; - return Ok(BindUserOAuthLinkOutcome::UserAlreadyLinkedProvider); - } - sqlx::query( - r#" -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, provider_username, provider_email, - extra_data, linked_at, last_login_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(user_id) - .bind(provider_type) - .bind(provider_user_id) - .bind(provider_username) - .bind(provider_email) - .bind(extra_data.as_deref()) - .bind(linked_at.timestamp()) - .bind(linked_at.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(BindUserOAuthLinkOutcome::Bound) - } - - async fn upgrade_oauth_email_verification_if_matches( - &self, - user_id: &str, - verified_email: &str, - verified_at: DateTime, - ) -> Result { - let result = sqlx::query( - "UPDATE users SET email_verified = 1, updated_at = ? WHERE id = ? AND email_verified = 0 AND LOWER(TRIM(email)) = LOWER(TRIM(?))", - ) - .bind(verified_at.timestamp()) - .bind(user_id) - .bind(verified_email) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn delete_user_oauth_link( - &self, - user_id: &str, - provider_type: &str, - local_password_login_allowed: bool, - _enabled_provider_types_snapshot: &[String], - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let user = sqlx::query("SELECT auth_source, password_hash FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(user) = user else { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - }; - let auth_source = user.try_get::("auth_source").map_sql_err()?; - let password_hash = user - .try_get::, _>("password_hash") - .map_sql_err()?; - let provider_types = sqlx::query_scalar::<_, String>( - "SELECT provider_type FROM user_oauth_links WHERE user_id = ?", - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - if !provider_types.iter().any(|value| value == provider_type) { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - } - let enabled_provider_types = sqlx::query_scalar::<_, String>( - r#" -SELECT user_oauth_links.provider_type -FROM user_oauth_links -JOIN oauth_providers - ON oauth_providers.provider_type = user_oauth_links.provider_type -WHERE user_oauth_links.user_id = ? - AND oauth_providers.is_enabled = 1 -"#, - ) - .bind(user_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - let has_remaining_enabled_oauth_link = enabled_provider_types - .iter() - .any(|value| value != provider_type); - if !has_remaining_enabled_oauth_link { - if let Some(outcome) = last_oauth_unbind_denial( - &auth_source, - password_hash.as_deref(), - local_password_login_allowed, - ) { - tx.rollback().await.map_sql_err()?; - return Ok(outcome); - } - } - let result = - sqlx::query("DELETE FROM user_oauth_links WHERE user_id = ? AND provider_type = ?") - .bind(user_id) - .bind(provider_type) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(DeleteUserOAuthLinkOutcome::NotFound); - } - tx.commit().await.map_sql_err()?; - Ok(DeleteUserOAuthLinkOutcome::Deleted) - } - - async fn get_or_create_ldap_auth_user( - &self, - email: String, - username: String, - ldap_dn: Option, - ldap_username: Option, - logged_in_at: DateTime, - ) -> Result, DataLayerError> { - get_or_create_sqlite_ldap_auth_user( - &self.pool, - email, - username, - ldap_dn, - ldap_username, - logged_in_at, - ) - .await - } - - async fn touch_auth_user_last_login( - &self, - user_id: &str, - logged_in_at: DateTime, - ) -> Result { - let result = sqlx::query("UPDATE users SET last_login_at = ?, updated_at = ? WHERE id = ?") - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn update_local_auth_user_profile( - &self, - user_id: &str, - email_present: bool, - email: Option, - email_verified: Option, - username: Option, - ) -> Result, DataLayerError> { - let now = chrono::Utc::now().timestamp(); - let result = sqlx::query( - "UPDATE users SET email = CASE WHEN ? THEN ? ELSE email END, email_verified = COALESCE(?, email_verified), username = COALESCE(?, username), updated_at = ? WHERE id = ?", - ) - .bind(email_present) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(now) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn restore_local_auth_user_state_if_matches( - &self, - expected_auth: &StoredUserAuthRecord, - restored_auth: &StoredUserAuthRecord, - expected_export: &StoredUserExportRow, - restored_export: &StoredUserExportRow, - expected_model_capability_settings: Option<&serde_json::Value>, - restored_model_capability_settings: Option, - expected_feature_settings: Option<&serde_json::Value>, - restored_feature_settings: Option, - ) -> Result { - if expected_auth.id != restored_auth.id - || expected_export.id != expected_auth.id - || restored_export.id != restored_auth.id - { - return Ok(false); - } - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let auth_row = sqlx::query(&format!("{USER_AUTH_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(&expected_auth.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let export_row = sqlx::query(&format!("{USER_EXPORT_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(&expected_auth.id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (Some(auth_row), Some(export_row)) = (auth_row, export_row) else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current_auth = map_user_auth_row(&auth_row)?; - let current_export = map_user_export_row(&export_row)?; - if !current_auth.matches_restore_state(expected_auth) - || !current_export.matches_restore_state(expected_export) - || current_export.rate_limit != expected_export.rate_limit - || current_export.rate_limit_mode != expected_export.rate_limit_mode - || current_export.model_capability_settings.as_ref() - != expected_model_capability_settings - || current_export.feature_settings.as_ref() != expected_feature_settings - { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - - let removes_active_admin = current_auth.role.eq_ignore_ascii_case("admin") - && current_auth.is_active - && !current_auth.is_deleted - && (!restored_auth.role.eq_ignore_ascii_case("admin") || !restored_auth.is_active); - if removes_active_admin { - let active_admin_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM users WHERE LOWER(role) = 'admin' AND is_active = 1 AND is_deleted = 0", - ) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - if active_admin_count <= 1 { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_UPDATE_DENIED.to_string(), - )); - } - } - - let security_state_changed = expected_auth.role != restored_auth.role - || expected_auth.is_active != restored_auth.is_active; - let result = sqlx::query( - r#" -UPDATE users -SET email = ?, - email_verified = ?, - username = ?, - role = ?, - allowed_providers = ?, - allowed_providers_mode = ?, - allowed_api_formats = ?, - allowed_api_formats_mode = ?, - allowed_models = ?, - allowed_models_mode = ?, - rate_limit = ?, - rate_limit_mode = ?, - model_capability_settings = ?, - feature_settings = ?, - is_active = ?, - security_version = security_version + CASE WHEN ? THEN 1 ELSE 0 END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(restored_auth.email.as_deref()) - .bind(restored_auth.email_verified) - .bind(&restored_auth.username) - .bind(&restored_auth.role) - .bind(optional_string_list_json( - restored_auth.allowed_providers.clone(), - "users.allowed_providers", - )?) - .bind(&restored_auth.allowed_providers_mode) - .bind(optional_string_list_json( - restored_auth.allowed_api_formats.clone(), - "users.allowed_api_formats", - )?) - .bind(&restored_auth.allowed_api_formats_mode) - .bind(optional_string_list_json( - restored_auth.allowed_models.clone(), - "users.allowed_models", - )?) - .bind(&restored_auth.allowed_models_mode) - .bind(restored_export.rate_limit) - .bind(&restored_export.rate_limit_mode) - .bind(optional_json_string( - restored_model_capability_settings.clone(), - "users.model_capability_settings", - )?) - .bind(optional_json_string( - restored_feature_settings.clone(), - "users.feature_settings", - )?) - .bind(restored_auth.is_active) - .bind(security_state_changed) - .bind(current_unix_secs()) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - if security_state_changed { - let now = current_unix_secs(); - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'user_security_state_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(now) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE api_keys SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE management_tokens SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(now) - .bind(&expected_auth.id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn update_local_auth_user_password_hash( - &self, - user_id: &str, - password_hash: String, - updated_at: DateTime, - ) -> Result, DataLayerError> { - let result = sqlx::query( - "UPDATE users SET password_hash = ?, security_version = security_version + 1, updated_at = ? WHERE id = ?", - ) - .bind(password_hash) - .bind(updated_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn restore_local_auth_user_password_hash_if_matches( - &self, - user_id: &str, - expected_password_hash: Option<&str>, - password_hash: Option, - updated_at: DateTime, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE users -SET password_hash = ?, - security_version = security_version + 1, - updated_at = ? -WHERE id = ? - AND ((? IS NULL AND password_hash IS NULL) OR password_hash = ?) -"#, - ) - .bind(password_hash) - .bind(updated_at.timestamp()) - .bind(user_id) - .bind(expected_password_hash) - .bind(expected_password_hash) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() == 1) - } - - async fn reset_local_auth_user_password_and_revoke_sessions( - &self, - user_id: &str, - password_hash: String, - changed_at: DateTime, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let updated = sqlx::query( - "UPDATE users SET password_hash = ?, security_version = security_version + 1, updated_at = ? WHERE id = ? AND is_deleted = 0", - ) - .bind(password_hash) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'admin_password_reset', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(changed_at.timestamp()) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn change_local_auth_password_and_revoke_sessions( - &self, - user_id: &str, - current_session_id: &str, - expected_password_hash: Option<&str>, - next_password_hash: String, - changed_at: DateTime, - ) -> Result { - let mut tx = self.pool.begin().await.map_sql_err()?; - let updated = sqlx::query( - r#" -UPDATE users -SET password_hash = ?, security_version = security_version + 1, updated_at = ? -WHERE id = ? - AND is_active = 1 - AND is_deleted = 0 - AND ((? IS NULL AND password_hash IS NULL) OR password_hash = ?) - AND EXISTS ( - SELECT 1 FROM user_sessions - WHERE user_id = ? AND id = ? AND revoked_at IS NULL AND expires_at > ? - ) -"#, - ) - .bind(next_password_hash) - .bind(changed_at.timestamp()) - .bind(user_id) - .bind(expected_password_hash) - .bind(expected_password_hash) - .bind(user_id) - .bind(current_session_id) - .bind(changed_at.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() != 1 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let revoked = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'password_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(changed_at.timestamp()) - .bind(changed_at.timestamp()) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if revoked.rows_affected() == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn update_local_auth_user_admin_fields( - &self, - user_id: &str, - role: Option, - allowed_providers_present: bool, - allowed_providers: Option>, - allowed_api_formats_present: bool, - allowed_api_formats: Option>, - allowed_models_present: bool, - allowed_models: Option>, - rate_limit_present: bool, - rate_limit: Option, - is_active: Option, - ) -> Result, DataLayerError> { - let removes_active_admin = role - .as_deref() - .is_some_and(|value| !value.eq_ignore_ascii_case("admin")) - || is_active == Some(false); - let allowed_providers_mode = if allowed_providers - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_api_formats_mode = if allowed_api_formats - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_models_mode = if allowed_models - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let rate_limit_mode = if rate_limit.is_some() { - "custom" - } else { - "system" - }; - let update_sql = format!( - r#" -UPDATE users -SET role = CASE WHEN ? THEN COALESCE(?, role) ELSE role END, - allowed_providers = CASE WHEN ? THEN ? ELSE allowed_providers END, - allowed_providers_mode = CASE WHEN ? THEN ? ELSE allowed_providers_mode END, - allowed_api_formats = CASE WHEN ? THEN ? ELSE allowed_api_formats END, - allowed_api_formats_mode = CASE WHEN ? THEN ? ELSE allowed_api_formats_mode END, - allowed_models = CASE WHEN ? THEN ? ELSE allowed_models END, - allowed_models_mode = CASE WHEN ? THEN ? ELSE allowed_models_mode END, - rate_limit = CASE WHEN ? THEN ? ELSE rate_limit END, - rate_limit_mode = CASE WHEN ? THEN ? ELSE rate_limit_mode END, - is_active = CASE WHEN ? THEN ? ELSE is_active END, - security_version = security_version + CASE WHEN ? THEN 1 ELSE 0 END, - updated_at = ? -WHERE id = ? -{SQLITE_ACTIVE_ADMIN_UPDATE_GUARD} -"#, - ); - let mut tx = self.pool.begin().await.map_sql_err()?; - let current_security_state = sqlx::query("SELECT role, is_active FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let security_state_changed = current_security_state.as_ref().is_some_and(|row| { - role.as_deref().is_some_and(|next_role| { - row.try_get::("role") - .is_ok_and(|current_role| !current_role.eq_ignore_ascii_case(next_role)) - }) || is_active.is_some_and(|next_active| { - row.try_get::("is_active") - .is_ok_and(|current_active| current_active != next_active) - }) - }); - let result = sqlx::query(&update_sql) - .bind(role.is_some()) - .bind(role) - .bind(allowed_providers_present) - .bind(optional_string_list_json( - allowed_providers, - "users.allowed_providers", - )?) - .bind(allowed_providers_present) - .bind(allowed_providers_mode) - .bind(allowed_api_formats_present) - .bind(optional_string_list_json( - allowed_api_formats, - "users.allowed_api_formats", - )?) - .bind(allowed_api_formats_present) - .bind(allowed_api_formats_mode) - .bind(allowed_models_present) - .bind(optional_string_list_json( - allowed_models, - "users.allowed_models", - )?) - .bind(allowed_models_present) - .bind(allowed_models_mode) - .bind(rate_limit_present) - .bind(rate_limit) - .bind(rate_limit_present) - .bind(rate_limit_mode) - .bind(is_active.is_some()) - .bind(is_active) - .bind(security_state_changed) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .bind(removes_active_admin) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - let blocked_active_admin: Option = sqlx::query_scalar( - "SELECT 1 FROM users WHERE id = ? AND LOWER(role) = 'admin' AND is_active = 1 AND is_deleted = 0 LIMIT 1", - ) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - tx.rollback().await.map_sql_err()?; - if removes_active_admin && blocked_active_admin.is_some() { - return Err(DataLayerError::InvalidInput( - LAST_ACTIVE_ADMIN_UPDATE_DENIED.to_string(), - )); - } - return Ok(None); - } - if security_state_changed { - let revoked_at = chrono::Utc::now().timestamp(); - sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = 'user_security_state_changed', updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(revoked_at) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE api_keys SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - "UPDATE management_tokens SET is_active = 0, updated_at = ? WHERE user_id = ? AND is_active = 1", - ) - .bind(revoked_at) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - tx.commit().await.map_sql_err()?; - self.find_user_auth_by_id(user_id).await - } - - async fn update_local_auth_user_policy_modes( - &self, - user_id: &str, - allowed_providers_mode: Option, - allowed_api_formats_mode: Option, - allowed_models_mode: Option, - rate_limit_mode: Option, - ) -> Result, DataLayerError> { - let result = sqlx::query( - r#" -UPDATE users -SET allowed_providers_mode = CASE WHEN ? THEN ? ELSE allowed_providers_mode END, - allowed_api_formats_mode = CASE WHEN ? THEN ? ELSE allowed_api_formats_mode END, - allowed_models_mode = CASE WHEN ? THEN ? ELSE allowed_models_mode END, - rate_limit_mode = CASE WHEN ? THEN ? ELSE rate_limit_mode END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(allowed_providers_mode.is_some()) - .bind(allowed_providers_mode) - .bind(allowed_api_formats_mode.is_some()) - .bind(allowed_api_formats_mode) - .bind(allowed_models_mode.is_some()) - .bind(allowed_models_mode) - .bind(rate_limit_mode.is_some()) - .bind(rate_limit_mode) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find_user_auth_by_id(user_id).await - } - - async fn update_user_model_capability_settings( - &self, - user_id: &str, - settings: Option, - ) -> Result, DataLayerError> { - let normalized = normalize_optional_json_value(settings); - let result = sqlx::query( - "UPDATE users SET model_capability_settings = ?, updated_at = ? WHERE id = ?", - ) - .bind(optional_json_string( - normalized.clone(), - "users.model_capability_settings", - )?) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - Ok(normalized) - } - - async fn update_user_feature_settings( - &self, - user_id: &str, - settings: Option, - ) -> Result, DataLayerError> { - let normalized = normalize_optional_json_value(settings); - let result = - sqlx::query("UPDATE users SET feature_settings = ?, updated_at = ? WHERE id = ?") - .bind(optional_json_string( - normalized.clone(), - "users.feature_settings", - )?) - .bind(chrono::Utc::now().timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - Ok(normalized) - } - - async fn create_local_auth_user( - &self, - email: Option, - email_verified: bool, - username: String, - password_hash: String, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - let now = chrono::Utc::now().timestamp(); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers_mode, allowed_api_formats_mode, allowed_models_mode, rate_limit_mode, - is_active, is_deleted, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, 'user', 'local', 'inherit', 'inherit', 'inherit', 'inherit', 1, 0, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(password_hash) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn create_local_auth_user_with_settings( - &self, - email: Option, - email_verified: bool, - username: String, - password_hash: String, - role: String, - allowed_providers: Option>, - allowed_api_formats: Option>, - allowed_models: Option>, - rate_limit: Option, - ) -> Result, DataLayerError> { - let user_id = uuid::Uuid::new_v4().to_string(); - let now = chrono::Utc::now().timestamp(); - let allowed_providers_mode = if allowed_providers - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_api_formats_mode = if allowed_api_formats - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let allowed_models_mode = if allowed_models - .as_ref() - .is_some_and(|values| !values.is_empty()) - { - "specific" - } else { - "unrestricted" - }; - let rate_limit_mode = if rate_limit.is_some() { - "custom" - } else { - "system" - }; - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers, allowed_providers_mode, - allowed_api_formats, allowed_api_formats_mode, - allowed_models, allowed_models_mode, - rate_limit, rate_limit_mode, - is_active, is_deleted, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 'local', ?, ?, ?, ?, ?, ?, ?, ?, 1, 0, ?, ?) -"#, - ) - .bind(&user_id) - .bind(email) - .bind(email_verified) - .bind(username) - .bind(password_hash) - .bind(role) - .bind(optional_string_list_json( - allowed_providers, - "users.allowed_providers", - )?) - .bind(allowed_providers_mode) - .bind(optional_string_list_json( - allowed_api_formats, - "users.allowed_api_formats", - )?) - .bind(allowed_api_formats_mode) - .bind(optional_string_list_json( - allowed_models, - "users.allowed_models", - )?) - .bind(allowed_models_mode) - .bind(rate_limit) - .bind(rate_limit_mode) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.find_user_auth_by_id(&user_id).await - } - - async fn delete_local_auth_user(&self, user_id: &str) -> Result { - self.delete_local_auth_user_inner(user_id, false).await - } - - async fn delete_local_auth_user_if_wallet_absent( - &self, - user_id: &str, - ) -> Result { - self.delete_local_auth_user_inner(user_id, true).await - } - - async fn count_active_admin_users(&self) -> Result { - let total: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM users -WHERE LOWER(role) = 'admin' - AND is_deleted = 0 - AND is_active = 1 -"#, - ) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - Ok(total.max(0) as u64) - } - - async fn read_user_preferences( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_PREFERENCES_COLUMNS); - builder.push(" WHERE up.user_id = ").push_bind(user_id); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_preference_row).transpose() - } - - async fn write_user_preferences( - &self, - preferences: &StoredUserPreferenceRecord, - ) -> Result, DataLayerError> { - let now = Utc::now().timestamp(); - sqlx::query( - r#" -INSERT INTO user_preferences ( - id, user_id, avatar_url, bio, default_provider_id, theme, language, timezone, - email_notifications, usage_alerts, announcement_notifications, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(user_id) DO UPDATE SET - avatar_url = excluded.avatar_url, - bio = excluded.bio, - default_provider_id = excluded.default_provider_id, - theme = excluded.theme, - language = excluded.language, - timezone = excluded.timezone, - email_notifications = excluded.email_notifications, - usage_alerts = excluded.usage_alerts, - announcement_notifications = excluded.announcement_notifications, - updated_at = excluded.updated_at -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&preferences.user_id) - .bind(preferences.avatar_url.as_deref()) - .bind(preferences.bio.as_deref()) - .bind(preferences.default_provider_id.as_deref()) - .bind(&preferences.theme) - .bind(&preferences.language) - .bind(&preferences.timezone) - .bind(preferences.email_notifications) - .bind(preferences.usage_alerts) - .bind(preferences.announcement_notifications) - .bind(now) - .bind(now) - .execute(&self.pool) - .await - .map_sql_err()?; - self.read_user_preferences(&preferences.user_id).await - } - - async fn find_user_session( - &self, - user_id: &str, - session_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(user_id) - .push(" AND id = ") - .push_bind(session_id) - .push(" LIMIT 1"); - let row = builder - .build() - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_session_row).transpose() - } - - async fn list_user_sessions( - &self, - user_id: &str, - ) -> Result, DataLayerError> { - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(user_id) - .push(" AND revoked_at IS NULL AND expires_at > ") - .push_bind(Utc::now().timestamp()) - .push(" ORDER BY last_seen_at DESC, created_at DESC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_user_session_row).collect() - } - - async fn create_user_session( - &self, - session: &StoredUserSessionRecord, - ) -> Result, DataLayerError> { - let now = session - .created_at - .or(session.updated_at) - .or(session.last_seen_at) - .unwrap_or_else(Utc::now); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let user_is_eligible: Option = sqlx::query_scalar( - "SELECT security_version FROM users WHERE id = ? AND is_active = 1 AND is_deleted = 0 AND security_version = ?", - ) - .bind(&session.user_id) - .bind(session.security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_is_eligible.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query( - r#" -UPDATE user_sessions -SET revoked_at = ?, revoke_reason = 'replaced_by_new_login', updated_at = ? -WHERE user_id = ? AND client_device_id = ? AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(now.timestamp()) - .bind(now.timestamp()) - .bind(&session.user_id) - .bind(&session.client_device_id) - .bind(now.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -INSERT INTO user_sessions ( - id, user_id, security_version, client_device_id, device_label, device_type, ip_address, user_agent, - refresh_token_hash, last_seen_at, expires_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, 'unknown', ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&session.id) - .bind(&session.user_id) - .bind(session.security_version) - .bind(&session.client_device_id) - .bind(session.device_label.as_deref()) - .bind(session.ip_address.as_deref()) - .bind(session.user_agent.as_deref()) - .bind(&session.refresh_token_hash) - .bind(session.last_seen_at.unwrap_or(now).timestamp()) - .bind(session.expires_at.unwrap_or(now).timestamp()) - .bind(session.created_at.unwrap_or(now).timestamp()) - .bind(session.updated_at.unwrap_or(now).timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - let mut builder = QueryBuilder::::new(USER_SESSION_COLUMNS); - builder - .push(" WHERE user_id = ") - .push_bind(&session.user_id) - .push(" AND id = ") - .push_bind(&session.id) - .push(" LIMIT 1"); - let row = builder - .build() - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let created = row.as_ref().map(map_user_session_row).transpose()?; - tx.commit().await.map_sql_err()?; - Ok(created) - } - - async fn create_user_session_if_password_matches( - &self, - session: &StoredUserSessionRecord, - expected_password_hash: &str, - ) -> Result, DataLayerError> { - let now = session - .created_at - .or(session.updated_at) - .or(session.last_seen_at) - .unwrap_or_else(Utc::now); - let mut tx = self.pool.begin().await.map_sql_err()?; - let matched = sqlx::query_scalar::<_, String>( - r#" -SELECT password_hash FROM users -WHERE id = ? AND password_hash = ? AND LOWER(auth_source) = 'local' - AND is_active = 1 AND is_deleted = 0 AND security_version = ? -"#, - ) - .bind(&session.user_id) - .bind(expected_password_hash) - .bind(session.security_version) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if matched.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - sqlx::query("UPDATE users SET last_login_at = ? WHERE id = ?") - .bind(now.timestamp()) - .bind(&session.user_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE user_sessions -SET revoked_at = ?, revoke_reason = 'replaced_by_new_login', updated_at = ? -WHERE user_id = ? AND client_device_id = ? AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(now.timestamp()) - .bind(now.timestamp()) - .bind(&session.user_id) - .bind(&session.client_device_id) - .bind(now.timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -INSERT INTO user_sessions ( - id, user_id, security_version, client_device_id, device_label, device_type, ip_address, user_agent, - refresh_token_hash, last_seen_at, expires_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, 'unknown', ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(&session.id) - .bind(&session.user_id) - .bind(session.security_version) - .bind(&session.client_device_id) - .bind(session.device_label.as_deref()) - .bind(session.ip_address.as_deref()) - .bind(session.user_agent.as_deref()) - .bind(&session.refresh_token_hash) - .bind(session.last_seen_at.unwrap_or(now).timestamp()) - .bind(session.expires_at.unwrap_or(now).timestamp()) - .bind(session.created_at.unwrap_or(now).timestamp()) - .bind(session.updated_at.unwrap_or(now).timestamp()) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - self.find_user_session(&session.user_id, &session.id).await - } - - async fn touch_user_session( - &self, - user_id: &str, - session_id: &str, - touched_at: DateTime, - ip_address: Option<&str>, - user_agent: Option<&str>, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_sessions -SET last_seen_at = ?, ip_address = COALESCE(?, ip_address), - user_agent = COALESCE(?, user_agent), updated_at = ? -WHERE user_id = ? AND id = ? -"#, - ) - .bind(touched_at.timestamp()) - .bind(ip_address) - .bind(user_agent.map(|value| value.chars().take(1000).collect::())) - .bind(touched_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn update_user_session_device_label( - &self, - user_id: &str, - session_id: &str, - device_label: &str, - updated_at: DateTime, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_sessions -SET device_label = ?, updated_at = ? -WHERE user_id = ? AND id = ? -"#, - ) - .bind(device_label.chars().take(120).collect::()) - .bind(updated_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn rotate_user_session_refresh_token( - &self, - user_id: &str, - session_id: &str, - expected_refresh_token_hash: &str, - next_refresh_token_hash: &str, - rotated_at: DateTime, - expires_at: DateTime, - ip_address: Option<&str>, - user_agent: Option<&str>, - ) -> Result { - let result = sqlx::query( - r#" -UPDATE user_sessions -SET prev_refresh_token_hash = ?, rotated_at = ?, refresh_token_hash = ?, - expires_at = ?, last_seen_at = ?, ip_address = COALESCE(?, ip_address), - user_agent = COALESCE(?, user_agent), updated_at = ? -WHERE user_id = ? AND id = ? AND refresh_token_hash = ? - AND revoked_at IS NULL AND expires_at > ? -"#, - ) - .bind(expected_refresh_token_hash) - .bind(rotated_at.timestamp()) - .bind(next_refresh_token_hash) - .bind(expires_at.timestamp()) - .bind(rotated_at.timestamp()) - .bind(ip_address) - .bind(user_agent.map(|value| value.chars().take(1000).collect::())) - .bind(rotated_at.timestamp()) - .bind(user_id) - .bind(session_id) - .bind(expected_refresh_token_hash) - .bind(rotated_at.timestamp()) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn revoke_user_session( - &self, - user_id: &str, - session_id: &str, - revoked_at: DateTime, - reason: &str, - ) -> Result { - let result = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = ?, updated_at = ? WHERE user_id = ? AND id = ?", - ) - .bind(revoked_at.timestamp()) - .bind(reason.chars().take(100).collect::()) - .bind(revoked_at.timestamp()) - .bind(user_id) - .bind(session_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - async fn revoke_all_user_sessions( - &self, - user_id: &str, - revoked_at: DateTime, - reason: &str, - ) -> Result { - let result = sqlx::query( - "UPDATE user_sessions SET revoked_at = ?, revoke_reason = ?, updated_at = ? WHERE user_id = ? AND revoked_at IS NULL", - ) - .bind(revoked_at.timestamp()) - .bind(reason.chars().take(100).collect::()) - .bind(revoked_at.timestamp()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - Ok(result.rows_affected()) - } - - async fn count_active_local_admin_users_with_valid_password( - &self, - ) -> Result { - let hashes = sqlx::query_scalar::<_, String>( - r#" -SELECT password_hash -FROM users -WHERE LOWER(role) = 'admin' - AND LOWER(auth_source) = 'local' - AND is_deleted = 0 - AND is_active = 1 - AND password_hash IS NOT NULL -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - Ok(hashes - .iter() - .filter(|hash| is_valid_bcrypt_hash(hash)) - .count() as u64) - } -} - -fn optional_json_from_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn optional_string_list_json( - value: Option>, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} could not be serialized as JSON: {err}" - )) - }) - }) - .transpose() -} - -fn json_string_from_option_vec(value: Option<&Vec>) -> Option { - value.and_then(|items| serde_json::to_string(items).ok()) -} - -fn normalized_ids(values: &[String]) -> Vec { - values - .iter() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - .collect::>() - .into_iter() - .collect() -} - -fn current_unix_secs() -> i64 { - chrono::Utc::now().timestamp() -} - -fn optional_json_string( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::to_string(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} could not be serialized as JSON: {err}" - )) - }) - }) - .transpose() -} - -fn normalize_optional_json_value(value: Option) -> Option { - match value { - Some(serde_json::Value::Null) | None => None, - Some(value) => Some(value), - } -} - -async fn get_or_create_sqlite_ldap_auth_user( - pool: &SqlitePool, - email: String, - username: String, - ldap_dn: Option, - ldap_username: Option, - logged_in_at: DateTime, -) -> Result, DataLayerError> { - let existing = - find_sqlite_ldap_auth_user(pool, ldap_dn.as_deref(), ldap_username.as_deref(), &email) - .await?; - if let Some(existing) = existing { - if existing.is_deleted - || !existing.is_active - || !existing.auth_source.eq_ignore_ascii_case("ldap") - { - return Ok(None); - } - if existing.email.as_deref() != Some(email.as_str()) { - let taken: Option = - sqlx::query_scalar("SELECT 1 FROM users WHERE email = ? AND id <> ? LIMIT 1") - .bind(&email) - .bind(&existing.id) - .fetch_optional(pool) - .await - .map_sql_err()?; - if taken.is_some() { - return Ok(None); - } - } - sqlx::query("UPDATE users SET email = ?, email_verified = 1, ldap_dn = COALESCE(?, ldap_dn), ldap_username = COALESCE(?, ldap_username), last_login_at = ?, updated_at = ? WHERE id = ?") - .bind(&email) - .bind(ldap_dn.as_deref()) - .bind(ldap_username.as_deref()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(&existing.id) - .execute(pool) - .await - .map_sql_err()?; - let user = find_sqlite_auth_by_id(pool, &existing.id) - .await? - .ok_or_else(|| { - DataLayerError::UnexpectedValue("updated LDAP user disappeared".to_string()) - })?; - return Ok(Some(LdapAuthUserProvisioningOutcome { - user, - created: false, - })); - } - - let base_username = ldap_username - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or(username.as_str()) - .trim() - .to_string(); - let mut candidate_username = base_username.clone(); - for _attempt in 0..3 { - let taken: Option = - sqlx::query_scalar("SELECT 1 FROM users WHERE username = ? LIMIT 1") - .bind(&candidate_username) - .fetch_optional(pool) - .await - .map_sql_err()?; - if taken.is_some() { - let suffix = uuid::Uuid::new_v4().simple().to_string(); - candidate_username = format!( - "{}_ldap_{}{}", - base_username, - logged_in_at.timestamp(), - &suffix[..4] - ); - continue; - } - let user_id = uuid::Uuid::new_v4().to_string(); - sqlx::query("INSERT INTO users (id, email, email_verified, username, password_hash, role, auth_source, ldap_dn, ldap_username, is_active, is_deleted, created_at, updated_at, last_login_at) VALUES (?, ?, 1, ?, NULL, 'user', 'ldap', ?, ?, 1, 0, ?, ?, ?)") - .bind(&user_id) - .bind(&email) - .bind(&candidate_username) - .bind(ldap_dn.as_deref()) - .bind(ldap_username.as_deref()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .bind(logged_in_at.timestamp()) - .execute(pool) - .await - .map_sql_err()?; - let user = find_sqlite_auth_by_id(pool, &user_id) - .await? - .ok_or_else(|| { - DataLayerError::UnexpectedValue("created LDAP user disappeared".to_string()) - })?; - return Ok(Some(LdapAuthUserProvisioningOutcome { - user, - created: true, - })); - } - Ok(None) -} - -async fn find_sqlite_ldap_auth_user( - pool: &SqlitePool, - ldap_dn: Option<&str>, - ldap_username: Option<&str>, - email: &str, -) -> Result, DataLayerError> { - if let Some(ldap_dn) = ldap_dn.filter(|value| !value.trim().is_empty()) { - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS} WHERE auth_source = 'ldap' AND ldap_dn = ? LIMIT 1" - )) - .bind(ldap_dn) - .fetch_optional(pool) - .await - .map_sql_err()?; - if let Some(row) = row.as_ref() { - return map_user_auth_row(row).map(Some); - } - } - if let Some(ldap_username) = ldap_username.filter(|value| !value.trim().is_empty()) { - let row = sqlx::query(&format!( - "{USER_AUTH_COLUMNS} WHERE auth_source = 'ldap' AND ldap_username = ? LIMIT 1" - )) - .bind(ldap_username) - .fetch_optional(pool) - .await - .map_sql_err()?; - if let Some(row) = row.as_ref() { - return map_user_auth_row(row).map(Some); - } - } - let row = sqlx::query(&format!("{USER_AUTH_COLUMNS} WHERE email = ? LIMIT 1")) - .bind(email) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_auth_row).transpose() -} - -async fn find_sqlite_auth_by_id( - pool: &SqlitePool, - user_id: &str, -) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{USER_AUTH_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(user_id) - .fetch_optional(pool) - .await - .map_sql_err()?; - row.as_ref().map(map_user_auth_row).transpose() -} - -fn optional_datetime_from_unix_secs(value: Option) -> Option> { - value.and_then(|value| Utc.timestamp_opt(value, 0).single()) -} - -fn map_user_row(row: &SqliteRow) -> Result { - StoredUserSummary::new( - row.try_get("id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_deleted").map_sql_err()?, - ) -} - -fn map_user_export_row(row: &SqliteRow) -> Result { - let feature_settings = optional_json_from_string( - row.try_get("feature_settings").map_sql_err()?, - "users.feature_settings", - )?; - StoredUserExportRow::new( - row.try_get("id").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("email_verified").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("password_hash").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("auth_source").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("rate_limit").map_sql_err()?, - optional_json_from_string( - row.try_get("model_capability_settings").map_sql_err()?, - "users.model_capability_settings", - )?, - row.try_get("is_active").map_sql_err()?, - ) - .map(|record| record.with_feature_settings(feature_settings)) - .and_then(|record| { - record.with_policy_modes( - row.try_get("allowed_providers_mode").map_sql_err()?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - row.try_get("allowed_models_mode").map_sql_err()?, - row.try_get("rate_limit_mode").map_sql_err()?, - ) - }) -} - -fn map_user_auth_row(row: &SqliteRow) -> Result { - StoredUserAuthRecord::new( - row.try_get("id").map_sql_err()?, - row.try_get("email").map_sql_err()?, - row.try_get("email_verified").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("password_hash").map_sql_err()?, - row.try_get("role").map_sql_err()?, - row.try_get("auth_source").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "users.allowed_providers", - )?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "users.allowed_api_formats", - )?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "users.allowed_models", - )?, - row.try_get("is_active").map_sql_err()?, - row.try_get("is_deleted").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_login_at").map_sql_err()?), - ) - .and_then(|record| record.with_security_version(row.try_get("security_version").map_sql_err()?)) - .and_then(|record| { - record.with_policy_modes( - row.try_get("allowed_providers_mode").map_sql_err()?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - row.try_get("allowed_models_mode").map_sql_err()?, - ) - }) -} - -fn map_user_group_row(row: &SqliteRow) -> Result { - StoredUserGroup::new( - row.try_get("id").map_sql_err()?, - row.try_get("name").map_sql_err()?, - row.try_get("normalized_name").map_sql_err()?, - row.try_get("description").map_sql_err()?, - row.try_get("priority").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_providers").map_sql_err()?, - "user_groups.allowed_providers", - )?, - row.try_get("allowed_providers_mode").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_api_formats").map_sql_err()?, - "user_groups.allowed_api_formats", - )?, - row.try_get("allowed_api_formats_mode").map_sql_err()?, - optional_json_from_string( - row.try_get("allowed_models").map_sql_err()?, - "user_groups.allowed_models", - )?, - row.try_get("allowed_models_mode").map_sql_err()?, - row.try_get("rate_limit").map_sql_err()?, - row.try_get("rate_limit_mode").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("updated_at").map_sql_err()?), - ) -} - -fn map_user_group_member_row(row: &SqliteRow) -> Result { - Ok(StoredUserGroupMember { - group_id: row.try_get("group_id").map_sql_err()?, - user_id: row.try_get("user_id").map_sql_err()?, - username: row.try_get("username").map_sql_err()?, - email: row.try_get("email").map_sql_err()?, - role: row.try_get("role").map_sql_err()?, - is_active: row.try_get("is_active").map_sql_err()?, - is_deleted: row.try_get("is_deleted").map_sql_err()?, - created_at: optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - }) -} - -fn map_user_group_membership_row( - row: &SqliteRow, -) -> Result { - Ok(StoredUserGroupMembership { - user_id: row.try_get("user_id").map_sql_err()?, - group_id: row.try_get("group_id").map_sql_err()?, - group_name: row.try_get("group_name").map_sql_err()?, - group_priority: row.try_get("group_priority").map_sql_err()?, - created_at: optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - }) -} - -fn map_oauth_link_summary_row( - row: &SqliteRow, -) -> Result { - StoredUserOAuthLinkSummary::new( - row.try_get("provider_type").map_sql_err()?, - row.try_get("display_name").map_sql_err()?, - row.try_get("provider_username").map_sql_err()?, - row.try_get("provider_email").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("linked_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_login_at").map_sql_err()?), - row.try_get("provider_enabled").map_sql_err()?, - ) -} - -fn map_user_preference_row(row: &SqliteRow) -> Result { - let user_id: String = row.try_get("user_id").map_sql_err()?; - if user_id.trim().is_empty() { - return Err(DataLayerError::UnexpectedValue( - "user_preferences.user_id is empty".to_string(), - )); - } - - Ok(StoredUserPreferenceRecord { - user_id, - avatar_url: row.try_get("avatar_url").map_sql_err()?, - bio: row.try_get("bio").map_sql_err()?, - default_provider_id: row.try_get("default_provider_id").map_sql_err()?, - default_provider_name: row.try_get("default_provider_name").map_sql_err()?, - theme: row.try_get("theme").map_sql_err()?, - language: row.try_get("language").map_sql_err()?, - timezone: row.try_get("timezone").map_sql_err()?, - email_notifications: row.try_get("email_notifications").map_sql_err()?, - usage_alerts: row.try_get("usage_alerts").map_sql_err()?, - announcement_notifications: row.try_get("announcement_notifications").map_sql_err()?, - }) -} - -fn map_user_session_row(row: &SqliteRow) -> Result { - StoredUserSessionRecord::new( - row.try_get("id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("client_device_id").map_sql_err()?, - row.try_get("device_label").map_sql_err()?, - row.try_get("refresh_token_hash").map_sql_err()?, - row.try_get("prev_refresh_token_hash").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("rotated_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("last_seen_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("expires_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("revoked_at").map_sql_err()?), - row.try_get("revoke_reason").map_sql_err()?, - row.try_get("ip_address").map_sql_err()?, - row.try_get("user_agent").map_sql_err()?, - optional_datetime_from_unix_secs(row.try_get("created_at").map_sql_err()?), - optional_datetime_from_unix_secs(row.try_get("updated_at").map_sql_err()?), - ) - .and_then(|record| record.with_security_version(row.try_get("security_version").map_sql_err()?)) -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::time::Duration; - - use super::{ - SqliteUserReadRepository, SQLITE_ANONYMIZE_USER_API_KEY_HISTORY_SQL, - SQLITE_ANONYMIZE_USER_HISTORY_SQL, - }; - use crate::run_migrations; - use aether_data_contracts::repository::users::{ - is_last_active_admin_delete_denied, is_last_active_admin_update_denied, - BindUserOAuthLinkOutcome, DeleteUserOAuthLinkOutcome, StoredUserPreferenceRecord, - StoredUserSessionRecord, UserExportListQuery, UserReadRepository, - }; - use sqlx::Row; - - const USER_HISTORY_TABLES: &[&str] = &[ - "request_candidates", - "video_tasks", - "usage", - "stats_user_daily", - "stats_user_summary", - "stats_user_daily_model", - "stats_user_daily_provider", - "stats_user_daily_api_format", - "stats_user_daily_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", - ]; - const STABLE_USER_ID_ONLY_TABLES: &[&str] = &[ - "stats_hourly_user", - "stats_hourly_user_model", - "user_model_usage_counts", - ]; - - async fn seed_sqlite_admin(pool: &crate::SqlitePool, id: &str, username: &str) { - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES (?, ?, 1, ?, NULL, 'admin', 'local', 1, 0, 1, 1) -"#, - ) - .bind(id) - .bind(format!("{username}@example.com")) - .bind(username) - .execute(pool) - .await - .expect("admin should insert"); - } - - fn test_user_session( - id: &str, - user_id: &str, - client_device_id: &str, - refresh_token: &str, - now: chrono::DateTime, - ) -> StoredUserSessionRecord { - StoredUserSessionRecord::new( - id.to_string(), - user_id.to_string(), - client_device_id.to_string(), - None, - StoredUserSessionRecord::hash_refresh_token(refresh_token), - None, - None, - Some(now), - Some(now + chrono::Duration::hours(1)), - None, - None, - None, - None, - Some(now), - Some(now), - ) - .expect("session should build") - } - - #[test] - fn hard_delete_anonymizes_every_sqlite_history_snapshot() { - assert_eq!( - SQLITE_ANONYMIZE_USER_HISTORY_SQL.len(), - USER_HISTORY_TABLES.len() - ); - for table in USER_HISTORY_TABLES { - let statement = SQLITE_ANONYMIZE_USER_HISTORY_SQL - .iter() - .find(|sql| sql.starts_with(&format!("UPDATE {table} "))) - .unwrap_or_else(|| panic!("missing history anonymization for {table}")); - assert!(statement.contains("username = NULL")); - assert!(statement.ends_with("WHERE user_id = ?")); - } - for table in ["request_candidates", "video_tasks", "usage"] { - let statement = SQLITE_ANONYMIZE_USER_HISTORY_SQL - .iter() - .find(|sql| sql.starts_with(&format!("UPDATE {table} "))) - .expect("identity snapshot table should be covered"); - assert!(statement.contains("api_key_name = NULL")); - } - assert!(SQLITE_ANONYMIZE_USER_API_KEY_HISTORY_SQL - .starts_with("UPDATE stats_daily_api_key SET api_key_name = NULL")); - assert!(SQLITE_ANONYMIZE_USER_API_KEY_HISTORY_SQL - .contains("SELECT id FROM api_keys WHERE user_id = ?")); - } - - #[tokio::test] - async fn sqlite_hard_delete_preserves_history_ids_and_anonymizes_snapshots() { - 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"); - sqlx::raw_sql( - r#" -INSERT INTO users ( - id, username, password_hash, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES ('history-user', 'history-name', 'history-hash', 'user', 'local', 1, 0, 1, 1); - -INSERT INTO request_candidates ( - id, request_id, user_id, api_key_id, username, api_key_name, - candidate_index, status, created_at -) VALUES ( - 'history-row', 'history-request-candidate', 'history-user', 'history-key', - 'history-name', 'history-key-name', 0, 'success', 1 -); - -INSERT INTO video_tasks ( - id, request_id, user_id, api_key_id, username, api_key_name, created_at, updated_at -) VALUES ( - 'history-row', 'history-video-request', 'history-user', 'history-key', - 'history-name', 'history-key-name', 1, 1 -); - -INSERT INTO usage ( - request_id, id, user_id, api_key_id, username, api_key_name -) VALUES ( - 'history-usage-request', 'history-row', 'history-user', 'history-key', - 'history-name', 'history-key-name' -); - -INSERT INTO stats_user_daily (id, user_id, date, username, created_at, updated_at) -VALUES ('history-row', 'history-user', 1, 'history-name', 1, 1); -INSERT INTO stats_user_summary (id, user_id, username, cutoff_date, created_at, updated_at) -VALUES ('history-row', 'history-user', 'history-name', 1, 1, 1); -INSERT INTO stats_user_daily_model (id, user_id, username, date, model, created_at, updated_at) -VALUES ('history-row', 'history-user', 'history-name', 1, 'history-model', 1, 1); -INSERT INTO stats_user_daily_provider ( - id, user_id, username, date, provider_name, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-name', 1, 'history-provider', 1, 1); -INSERT INTO stats_user_daily_api_format ( - id, user_id, username, date, api_format, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-name', 1, 'history-format', 1, 1); -INSERT INTO stats_user_daily_model_provider ( - id, user_id, username, date, model, provider_name, created_at, updated_at -) VALUES ( - 'history-row', 'history-user', 'history-name', 1, - 'history-model', 'history-provider', 1, 1 -); -INSERT INTO stats_user_daily_cost_savings ( - id, user_id, username, date, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-name', 1, 1, 1); -INSERT INTO stats_user_daily_cost_savings_provider ( - id, user_id, username, date, provider_name, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-name', 1, 'history-provider', 1, 1); -INSERT INTO stats_user_daily_cost_savings_model ( - id, user_id, username, date, model, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-name', 1, 'history-model', 1, 1); -INSERT INTO stats_user_daily_cost_savings_model_provider ( - id, user_id, username, date, model, provider_name, created_at, updated_at -) VALUES ( - 'history-row', 'history-user', 'history-name', 1, - 'history-model', 'history-provider', 1, 1 -); -INSERT INTO stats_hourly_user ( - id, hour_utc, user_id, created_at, updated_at -) VALUES ('history-row', 1, 'history-user', 1, 1); -INSERT INTO stats_hourly_user_model ( - id, hour_utc, user_id, model, created_at, updated_at -) VALUES ('history-row', 1, 'history-user', 'history-model', 1, 1); -INSERT INTO user_model_usage_counts ( - id, user_id, model, created_at, updated_at -) VALUES ('history-row', 'history-user', 'history-model', 1, 1); -INSERT INTO stats_daily_api_key ( - id, api_key_id, date, api_key_name, created_at, updated_at -) VALUES ('history-row', 'history-key', 1, 'history-key-name', 1, 1); -INSERT INTO api_keys ( - id, user_id, name, key_hash, created_at, updated_at -) VALUES ('history-key', 'history-user', 'history-key-name', 'history-key-hash', 1, 1); - -INSERT INTO users ( - id, username, password_hash, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES ('history-inviter', 'history-inviter', 'history-hash', 'user', 'local', 1, 0, 1, 1); - -INSERT INTO wallets ( - id, user_id, balance, gift_balance, status, created_at, updated_at -) VALUES ('history-wallet', 'history-user', 12, 3, 'active', 1, 1); - -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, - operator_id, description, created_at -) VALUES ( - 'history-wallet-tx', 'history-wallet', 'adjust', 'manual', 1, - 14, 15, 11, 12, 3, 3, - 'history-user', 'private wallet note', 1 -); - -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, payment_method, - gateway_response, status, created_at -) VALUES ( - 'history-order', 'history-order-no', 'history-wallet', 'history-user', 12, - 'test', '{"customer_email":"history@example.com"}', 'credited', 1 -); - -INSERT INTO payment_callbacks ( - id, payment_order_id, payment_method, callback_key, order_no, - payload_hash, signature_valid, status, payload, error_message, created_at -) VALUES ( - 'history-callback', 'history-order', 'test', 'history-callback-key', - 'history-order-no', 'history-payload-hash', 1, 'processed', - '{"customer_email":"history@example.com"}', 'private callback error', 1 -); - -INSERT INTO billing_plans ( - id, title, price_amount, duration_unit, duration_value, - entitlements_json, created_at, updated_at -) VALUES ('history-plan', 'history plan', 12, 'day', 30, '[]', 1, 1); - -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) VALUES ( - 'history-entitlement', 'history-user', 'history-plan', 'history-order', - 'active', 1, 4102444800, '[]', 1, 1 -); - -INSERT INTO entitlement_usage_ledgers ( - id, user_entitlement_id, user_id, request_id, amount_usd, - balance_before, balance_after, usage_date, created_at -) VALUES ( - 'history-entitlement-ledger', 'history-entitlement', 'history-user', - 'history-entitlement-request', 1, 12, 11, '2026-08-27', 1 -); - -INSERT INTO user_referrals ( - id, inviter_user_id, invitee_user_id, invite_code_snapshot, source_json, - first_paid_order_id, first_paid_at, created_at, updated_at -) VALUES ( - 'history-referral', 'history-inviter', 'history-user', 'PRIVATE-CODE', - '{"ip":"192.0.2.1"}', 'history-order', 1, 1, 1 -); - -INSERT INTO referral_rewards ( - id, referral_id, inviter_user_id, invitee_user_id, reward_type, - trigger_point, source_order_id, idempotency_key, amount_usd, status, - failure_reason, admin_operator_id, admin_note, created_at, updated_at -) VALUES ( - 'history-reward', 'history-referral', 'history-inviter', 'history-user', - 'percent', 'paid_order', 'history-order', 'history-reward-key', 1, 'failed', - 'private failure', 'history-user', 'private admin note', 1, 1 -); - -INSERT INTO audit_logs ( - id, event_type, user_id, description, ip_address, user_agent, - event_metadata, error_message, created_at -) VALUES ( - 'history-audit', 'history_event', 'history-user', 'private description', - '192.0.2.2', 'private agent', '{"private":true}', 'private error', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("history fixtures should insert"); - - let repository = SqliteUserReadRepository::new(pool.clone()); - assert!(repository - .delete_local_auth_user("history-user") - .await - .expect("hard delete should succeed")); - - for table in USER_HISTORY_TABLES { - let row = sqlx::query(&format!( - "SELECT user_id, username FROM {table} WHERE id = 'history-row'" - )) - .fetch_one(&pool) - .await - .unwrap_or_else(|error| panic!("{table} history should remain: {error}")); - assert_eq!( - row.try_get::, _>("user_id") - .expect("user_id should decode") - .as_deref(), - Some("history-user"), - "{table} user_id must remain stable" - ); - assert_eq!( - row.try_get::, _>("username") - .expect("username should decode"), - None, - "{table} username snapshot must be removed" - ); - } - for table in ["request_candidates", "video_tasks", "usage"] { - let row = sqlx::query(&format!( - "SELECT api_key_id, api_key_name FROM {table} WHERE id = 'history-row'" - )) - .fetch_one(&pool) - .await - .unwrap_or_else(|error| panic!("{table} identity snapshot should remain: {error}")); - assert_eq!( - row.try_get::, _>("api_key_id") - .expect("api_key_id should decode") - .as_deref(), - Some("history-key"), - "{table} api_key_id must remain stable" - ); - assert_eq!( - row.try_get::, _>("api_key_name") - .expect("api_key_name should decode"), - None, - "{table} API key name snapshot must be removed" - ); - } - for table in STABLE_USER_ID_ONLY_TABLES { - let user_id: String = sqlx::query_scalar(&format!( - "SELECT user_id FROM {table} WHERE id = 'history-row'" - )) - .fetch_one(&pool) - .await - .unwrap_or_else(|error| panic!("{table} fact should remain: {error}")); - assert_eq!( - user_id, "history-user", - "{table} user_id must remain stable" - ); - } - let api_key_fact = sqlx::query( - "SELECT api_key_id, api_key_name FROM stats_daily_api_key WHERE id = 'history-row'", - ) - .fetch_one(&pool) - .await - .expect("API key aggregate should remain"); - assert_eq!( - api_key_fact - .try_get::("api_key_id") - .expect("aggregate api_key_id should decode"), - "history-key" - ); - assert_eq!( - api_key_fact - .try_get::, _>("api_key_name") - .expect("aggregate api_key_name should decode"), - None - ); - let wallet_fact = - sqlx::query("SELECT user_id, status FROM wallets WHERE id = 'history-wallet'") - .fetch_one(&pool) - .await - .expect("wallet fact should remain"); - assert_eq!( - wallet_fact - .try_get::, _>("user_id") - .expect("wallet user_id should decode") - .as_deref(), - Some("history-user") - ); - assert_eq!( - wallet_fact - .try_get::("status") - .expect("wallet status should decode"), - "disabled" - ); - let order_fact = sqlx::query( - "SELECT user_id, gateway_response FROM payment_orders WHERE id = 'history-order'", - ) - .fetch_one(&pool) - .await - .expect("payment order fact should remain"); - assert_eq!( - order_fact - .try_get::, _>("user_id") - .expect("order user_id should decode") - .as_deref(), - Some("history-user") - ); - assert_eq!( - order_fact - .try_get::, _>("gateway_response") - .expect("gateway response should decode"), - None - ); - let callback_fact = sqlx::query( - "SELECT payment_order_id, order_no, payload, error_message FROM payment_callbacks WHERE id = 'history-callback'", - ) - .fetch_one(&pool) - .await - .expect("payment callback fact should remain"); - assert_eq!( - callback_fact - .try_get::, _>("payment_order_id") - .expect("callback order id should decode") - .as_deref(), - Some("history-order") - ); - assert_eq!( - callback_fact - .try_get::, _>("order_no") - .expect("callback order number should decode") - .as_deref(), - Some("history-order-no") - ); - assert_eq!( - callback_fact - .try_get::, _>("payload") - .expect("callback payload should decode"), - None - ); - assert_eq!( - callback_fact - .try_get::, _>("error_message") - .expect("callback error should decode"), - None - ); - let entitlement_fact = sqlx::query( - "SELECT user_id, status FROM user_plan_entitlements WHERE id = 'history-entitlement'", - ) - .fetch_one(&pool) - .await - .expect("entitlement fact should remain"); - assert_eq!( - entitlement_fact - .try_get::("user_id") - .expect("entitlement user_id should decode"), - "history-user" - ); - assert_eq!( - entitlement_fact - .try_get::("status") - .expect("entitlement status should decode"), - "revoked" - ); - let ledger_user_id: String = sqlx::query_scalar( - "SELECT user_id FROM entitlement_usage_ledgers WHERE id = 'history-entitlement-ledger'", - ) - .fetch_one(&pool) - .await - .expect("entitlement ledger should remain"); - assert_eq!(ledger_user_id, "history-user"); - let referral_fact = sqlx::query( - "SELECT invitee_user_id, invite_code_snapshot, source_json FROM user_referrals WHERE id = 'history-referral'", - ) - .fetch_one(&pool) - .await - .expect("referral fact should remain"); - assert_eq!( - referral_fact - .try_get::("invitee_user_id") - .expect("invitee user_id should decode"), - "history-user" - ); - assert_eq!( - referral_fact - .try_get::("invite_code_snapshot") - .expect("invite code snapshot should decode"), - "deleted-user" - ); - assert_eq!( - referral_fact - .try_get::, _>("source_json") - .expect("referral source should decode"), - None - ); - let reward_fact = sqlx::query( - "SELECT invitee_user_id, status, failure_reason, admin_note FROM referral_rewards WHERE id = 'history-reward'", - ) - .fetch_one(&pool) - .await - .expect("referral reward fact should remain"); - assert_eq!( - reward_fact - .try_get::("invitee_user_id") - .expect("reward invitee should decode"), - "history-user" - ); - assert_eq!( - reward_fact - .try_get::("status") - .expect("reward status should decode"), - "voided" - ); - assert_eq!( - reward_fact - .try_get::, _>("failure_reason") - .expect("reward failure reason should decode"), - None - ); - assert_eq!( - reward_fact - .try_get::, _>("admin_note") - .expect("reward admin note should decode"), - None - ); - let audit_fact = sqlx::query( - "SELECT user_id, ip_address, user_agent, event_metadata, error_message FROM audit_logs WHERE id = 'history-audit'", - ) - .fetch_one(&pool) - .await - .expect("audit fact should remain"); - assert_eq!( - audit_fact - .try_get::, _>("user_id") - .expect("audit user_id should decode") - .as_deref(), - Some("history-user") - ); - for column in [ - "ip_address", - "user_agent", - "event_metadata", - "error_message", - ] { - assert_eq!( - audit_fact - .try_get::, _>(column) - .unwrap_or_else(|error| panic!("audit {column} should decode: {error}")), - None, - "audit {column} must be removed" - ); - } - let transaction_description: Option = sqlx::query_scalar( - "SELECT description FROM wallet_transactions WHERE id = 'history-wallet-tx'", - ) - .fetch_one(&pool) - .await - .expect("wallet transaction should remain"); - assert_eq!(transaction_description, None); - let user_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM users WHERE id = 'history-user'") - .fetch_one(&pool) - .await - .expect("deleted user count should load"); - assert_eq!(user_count, 0); - } - - #[tokio::test] - async fn sqlite_atomic_user_delete_requires_wallet_absence() { - 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"); - sqlx::query( - "INSERT INTO users (id, username, email, role, auth_source, is_active, is_deleted, created_at, updated_at) VALUES ('atomic-user', 'atomic-user', 'atomic@example.com', 'user', 'local', 1, 0, 1, 1)", - ) - .execute(&pool) - .await - .expect("user should seed"); - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, created_at, updated_at) VALUES ('atomic-wallet', 'atomic-user', 0, 0, 'active', 1, 1)", - ) - .execute(&pool) - .await - .expect("wallet should seed"); - - let repository = SqliteUserReadRepository::new(pool.clone()); - assert!(!repository - .delete_local_auth_user_if_wallet_absent("atomic-user") - .await - .expect("wallet guard should resolve")); - assert!(repository - .find_user_auth_by_id("atomic-user") - .await - .expect("user lookup should succeed") - .is_some()); - - sqlx::query("DELETE FROM wallets WHERE id = 'atomic-wallet'") - .execute(&pool) - .await - .expect("wallet should remove"); - assert!(repository - .delete_local_auth_user_if_wallet_absent("atomic-user") - .await - .expect("wallet-free user should delete")); - assert!(repository - .find_user_auth_by_id("atomic-user") - .await - .expect("user lookup should succeed") - .is_none()); - } - - #[tokio::test] - async fn sqlite_atomic_user_delete_detects_api_key_wallet() { - 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"); - sqlx::query( - "INSERT INTO users (id, username, email, role, auth_source, is_active, is_deleted, created_at, updated_at) VALUES ('api-wallet-user', 'api-wallet-user', 'api-wallet@example.com', 'user', 'local', 1, 0, 1, 1)", - ) - .execute(&pool) - .await - .expect("user should seed"); - sqlx::query( - "INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) VALUES ('api-wallet-key', 'api-wallet-user', 'api-wallet-key-hash', 1, 1)", - ) - .execute(&pool) - .await - .expect("api key should seed"); - sqlx::query( - "INSERT INTO wallets (id, api_key_id, balance, gift_balance, status, created_at, updated_at) VALUES ('api-wallet', 'api-wallet-key', 25, 0, 'active', 1, 1)", - ) - .execute(&pool) - .await - .expect("api key wallet should seed"); - - let repository = SqliteUserReadRepository::new(pool.clone()); - assert!(!repository - .delete_local_auth_user_if_wallet_absent("api-wallet-user") - .await - .expect("api key wallet guard should resolve")); - - assert!(repository - .find_user_auth_by_id("api-wallet-user") - .await - .expect("user lookup should succeed") - .is_some()); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM api_keys WHERE id = 'api-wallet-key'", - ) - .fetch_one(&pool) - .await - .expect("api key count should query"), - 1 - ); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallets WHERE id = 'api-wallet' AND api_key_id = 'api-wallet-key'", - ) - .fetch_one(&pool) - .await - .expect("wallet count should query"), - 1 - ); - - sqlx::query("DELETE FROM wallets WHERE id = 'api-wallet'") - .execute(&pool) - .await - .expect("api key wallet should remove"); - assert!(repository - .delete_local_auth_user_if_wallet_absent("api-wallet-user") - .await - .expect("wallet-free api key owner should delete")); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM api_keys WHERE id = 'api-wallet-key'", - ) - .fetch_one(&pool) - .await - .expect("api key count should query after delete"), - 0 - ); - } - - #[tokio::test] - async fn sqlite_atomically_preserves_last_active_admin_and_revokes_on_security_change() { - 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_sqlite_admin(&pool, "admin-1", "admin_one").await; - sqlx::query( - "INSERT INTO management_tokens (id, user_id, name, token_hash, created_at, updated_at) VALUES ('token-admin-1', 'admin-1', 'admin token', 'token-hash-admin-1', 1, 1)", - ) - .execute(&pool) - .await - .expect("management token should insert"); - let repository = SqliteUserReadRepository::new(pool.clone()); - - let update_error = repository - .update_local_auth_user_admin_fields( - "admin-1", - Some("audit_admin".to_string()), - false, - None, - false, - None, - false, - None, - false, - None, - None, - ) - .await - .expect_err("last active admin demotion must be rejected"); - assert!(is_last_active_admin_update_denied(&update_error)); - assert_eq!( - repository - .find_user_auth_by_id("admin-1") - .await - .expect("admin lookup should succeed") - .expect("admin should remain") - .role, - "admin" - ); - - let delete_error = repository - .delete_local_auth_user("admin-1") - .await - .expect_err("last active admin delete must be rejected"); - assert!(is_last_active_admin_delete_denied(&delete_error)); - let token_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM management_tokens WHERE user_id = 'admin-1'") - .fetch_one(&pool) - .await - .expect("management token count should load"); - assert_eq!( - token_count, 1, - "rejected delete must roll back credential cleanup" - ); - - seed_sqlite_admin(&pool, "admin-2", "admin_two").await; - let now = chrono::Utc::now(); - let session = test_user_session("session-admin-1", "admin-1", "device-1", "refresh", now); - repository - .create_user_session(&session) - .await - .expect("admin session should create") - .expect("admin session should exist"); - sqlx::query( - "INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) VALUES ('key-admin-1', 'admin-1', 'key-hash-admin-1', 1, 1)", - ) - .execute(&pool) - .await - .expect("api key should insert"); - sqlx::query( - "INSERT INTO api_key_provider_mappings (id, api_key_id, provider_id, created_at, updated_at) VALUES ('mapping-admin-1', 'key-admin-1', 'provider-1', 1, 1)", - ) - .execute(&pool) - .await - .expect("api key mapping should insert"); - sqlx::query( - "INSERT INTO user_oauth_links (id, user_id, provider_type, provider_user_id, linked_at) VALUES ('oauth-admin-1', 'admin-1', 'test', 'subject-admin-1', 1)", - ) - .execute(&pool) - .await - .expect("oauth link should insert"); - sqlx::query( - "INSERT INTO user_group_members (group_id, user_id, created_at) VALUES ('00000000-0000-0000-0000-000000000001', 'admin-1', 1)", - ) - .execute(&pool) - .await - .expect("group membership should insert"); - sqlx::query( - "INSERT INTO user_preferences (id, user_id, created_at, updated_at) VALUES ('preferences-admin-1', 'admin-1', 1, 1)", - ) - .execute(&pool) - .await - .expect("preferences should insert"); - sqlx::query( - "INSERT INTO announcements (id, title, content, created_at, updated_at) VALUES ('announcement-1', 'notice', 'content', 1, 1)", - ) - .execute(&pool) - .await - .expect("announcement should insert"); - sqlx::query( - "INSERT INTO announcement_reads (id, user_id, announcement_id, read_at) VALUES ('read-admin-1', 'admin-1', 'announcement-1', 1)", - ) - .execute(&pool) - .await - .expect("announcement read should insert"); - let updated = repository - .update_local_auth_user_admin_fields( - "admin-1", - Some("audit_admin".to_string()), - false, - None, - false, - None, - false, - None, - false, - None, - None, - ) - .await - .expect("demotion with another active admin should succeed") - .expect("admin should exist"); - assert_eq!(updated.role, "audit_admin"); - let revoked = repository - .find_user_session("admin-1", "session-admin-1") - .await - .expect("session lookup should succeed") - .expect("session should remain as audit record"); - assert!(revoked.revoked_at.is_some()); - assert_eq!( - revoked.revoke_reason.as_deref(), - Some("user_security_state_changed") - ); - assert!(repository - .delete_local_auth_user("admin-1") - .await - .expect("non-full-admin delete should succeed")); - for (table, predicate) in [ - ("api_key_provider_mappings", "api_key_id = 'key-admin-1'"), - ("api_keys", "user_id = 'admin-1'"), - ("management_tokens", "user_id = 'admin-1'"), - ("user_sessions", "user_id = 'admin-1'"), - ("user_oauth_links", "user_id = 'admin-1'"), - ("user_group_members", "user_id = 'admin-1'"), - ("user_preferences", "user_id = 'admin-1'"), - ("announcement_reads", "user_id = 'admin-1'"), - ] { - let count: i64 = - sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table} WHERE {predicate}")) - .fetch_one(&pool) - .await - .expect("dependent row count should load"); - assert_eq!(count, 0, "{table} credentials must be removed"); - } - } - - async fn seed_active_session_test_user(pool: &crate::SqlitePool, user_id: &str) { - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES (?, ?, 1, ?, NULL, 'user', 'oauth', 1, 0, ?, ?) -"#, - ) - .bind(user_id) - .bind(format!("{user_id}@example.com")) - .bind(user_id) - .bind(chrono::Utc::now().timestamp()) - .bind(chrono::Utc::now().timestamp()) - .execute(pool) - .await - .expect("session test user should insert"); - } - - #[tokio::test] - async fn sqlite_repository_reads_user_contract_views() { - 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"); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers, allowed_api_formats, allowed_models, model_capability_settings, - rate_limit, is_active, is_deleted, created_at, updated_at, last_login_at -) VALUES - ( - 'admin-1', 'admin@example.com', 1, 'admin', NULL, 'admin', 'local', - NULL, NULL, NULL, NULL, 100, 1, 0, 1, 1, NULL - ), - ( - 'user-1', 'user-1@example.com', 1, 'alice', 'hash', 'user', 'local', - '["openai"]', '["openai:chat"]', '["gpt-4.1"]', '{"gpt-4.1":{"cache_1h":true}}', - 60, 1, 0, 2, 2, 3 - ), - ( - 'user-2', NULL, 0, 'deleted', NULL, 'user', 'local', - NULL, NULL, NULL, NULL, NULL, 0, 1, 4, 4, NULL - ) -"#, - ) - .execute(&pool) - .await - .expect("seed users should insert"); - let valid_hash = "$2b$12$4qL4tdcsFwVaDTw5Ck3xzu8GpNdre56DiNR6Dnw7t6gCXaEnqAe7G".to_string(); - sqlx::query( - r#" -INSERT INTO users ( - id, email, email_verified, username, password_hash, role, auth_source, - allowed_providers, allowed_api_formats, allowed_models, model_capability_settings, - rate_limit, is_active, is_deleted, created_at, updated_at, last_login_at -) VALUES ( - 'admin-2', 'admin-2@example.com', 1, 'admin2', ?, 'admin', 'local', - NULL, NULL, NULL, NULL, 100, 1, 0, 5, 5, NULL -) -"#, - ) - .bind(valid_hash) - .execute(&pool) - .await - .expect("valid local admin should insert"); - - let repository = SqliteUserReadRepository::new(pool.clone()); - let summaries = repository - .list_users_by_ids(&["user-1".to_string(), "admin-1".to_string()]) - .await - .expect("summaries should load"); - assert_eq!(summaries.len(), 2); - assert_eq!(summaries[0].id, "admin-1"); - - let searched = repository - .list_users_by_username_search("ali") - .await - .expect("username search should load"); - assert_eq!(searched.len(), 1); - assert_eq!(searched[0].id, "user-1"); - - let exports = repository - .list_non_admin_export_users() - .await - .expect("non-admin exports should load"); - assert_eq!(exports.len(), 1); - assert_eq!(exports[0].allowed_models, Some(vec!["gpt-4.1".to_string()])); - - let page = repository - .list_export_users_page(&UserExportListQuery { - skip: 0, - limit: 10, - role: Some("user".to_string()), - is_active: Some(true), - search: None, - group_id: None, - ..Default::default() - }) - .await - .expect("export page should load"); - assert_eq!(page.len(), 1); - assert_eq!(page[0].id, "user-1"); - - let summary = repository - .summarize_export_users() - .await - .expect("export summary should load"); - assert_eq!(summary.total, 3); - assert_eq!(summary.active, 3); - - let auth = repository - .find_user_auth_by_identifier("user-1@example.com") - .await - .expect("auth lookup should load") - .expect("auth user should exist"); - assert_eq!(auth.id, "user-1"); - assert_eq!(auth.last_login_at.expect("last login").timestamp(), 3); - let logged_in_at = chrono::DateTime::from_timestamp(123, 0).expect("valid time"); - assert!(repository - .touch_auth_user_last_login("user-1", logged_in_at) - .await - .expect("last login touch should update")); - assert!(!repository - .touch_auth_user_last_login("missing-user", logged_in_at) - .await - .expect("missing last login touch should be harmless")); - let touched_auth = repository - .find_user_auth_by_id("user-1") - .await - .expect("auth lookup should load") - .expect("auth user should exist"); - assert_eq!( - touched_auth.last_login_at.expect("last login").timestamp(), - 123 - ); - let profile_updated = repository - .update_local_auth_user_profile( - "user-1", - true, - Some("user-1b@example.com".to_string()), - Some(true), - Some("alice-b".to_string()), - ) - .await - .expect("profile update should succeed") - .expect("profile update should return user"); - assert_eq!( - profile_updated.email.as_deref(), - Some("user-1b@example.com") - ); - assert!(profile_updated.email_verified); - assert_eq!(profile_updated.username, "alice-b"); - let password_updated = repository - .update_local_auth_user_password_hash( - "user-1", - "new-password-hash".to_string(), - logged_in_at, - ) - .await - .expect("password update should succeed") - .expect("password update should return user"); - assert_eq!( - password_updated.password_hash.as_deref(), - Some("new-password-hash") - ); - let created = repository - .create_local_auth_user_with_settings( - Some("created@example.com".to_string()), - true, - "created-user".to_string(), - "created-hash".to_string(), - "admin".to_string(), - Some(vec!["openai".to_string()]), - Some(vec!["chat".to_string()]), - Some(vec!["gpt-4.1".to_string()]), - Some(25), - ) - .await - .expect("local user create should succeed") - .expect("local user create should return user"); - assert_eq!(created.email.as_deref(), Some("created@example.com")); - assert_eq!(created.username, "created-user"); - assert_eq!(created.role, "admin"); - assert_eq!(created.allowed_providers, Some(vec!["openai".to_string()])); - assert_eq!(created.allowed_api_formats, Some(vec!["chat".to_string()])); - assert_eq!(created.allowed_models, Some(vec!["gpt-4.1".to_string()])); - let admin_updated = repository - .update_local_auth_user_admin_fields( - &created.id, - Some("user".to_string()), - true, - None, - true, - Some(vec!["responses".to_string()]), - true, - Some(vec!["gpt-4.1-mini".to_string()]), - true, - Some(5), - Some(false), - ) - .await - .expect("admin fields update should succeed") - .expect("admin fields update should return user"); - assert_eq!(admin_updated.role, "user"); - assert_eq!(admin_updated.allowed_providers, None); - assert_eq!( - admin_updated.allowed_api_formats, - Some(vec!["responses".to_string()]) - ); - assert_eq!( - admin_updated.allowed_models, - Some(vec!["gpt-4.1-mini".to_string()]) - ); - assert!(!admin_updated.is_active); - assert_eq!( - repository - .update_user_model_capability_settings( - &created.id, - Some(serde_json::json!({"gpt-4.1-mini": {"enabled": true}})), - ) - .await - .expect("model settings update should succeed"), - Some(serde_json::json!({"gpt-4.1-mini": {"enabled": true}})) - ); - assert_eq!( - repository - .update_user_model_capability_settings(&created.id, Some(serde_json::Value::Null)) - .await - .expect("model settings clear should succeed"), - None - ); - - let by_email = repository - .find_user_auth_by_email("user-1b@example.com") - .await - .expect("email lookup should load") - .expect("email lookup should find user"); - assert_eq!(by_email.id, "user-1"); - let by_username = repository - .find_user_auth_by_username("alice-b") - .await - .expect("username lookup should load") - .expect("username lookup should find user"); - assert_eq!(by_username.id, "user-1"); - assert!(repository - .find_user_auth_by_email("alice") - .await - .expect("email lookup should load") - .is_none()); - let cleared_profile = repository - .update_local_auth_user_profile("user-1", true, None, Some(false), None) - .await - .expect("nullable email update should succeed") - .expect("profile should remain"); - assert!(cleared_profile.email.is_none()); - assert!(!cleared_profile.email_verified); - assert_eq!( - repository - .count_active_admin_users() - .await - .expect("active admin count should load"), - 2 - ); - assert_eq!( - repository - .count_active_local_admin_users_with_valid_password() - .await - .expect("valid local admin count should load"), - 1 - ); - let preferences = StoredUserPreferenceRecord { - user_id: "user-1".to_string(), - avatar_url: Some("https://example.test/avatar.png".to_string()), - bio: Some("hello".to_string()), - default_provider_id: None, - default_provider_name: None, - theme: "dark".to_string(), - language: "en-US".to_string(), - timezone: "UTC".to_string(), - email_notifications: false, - usage_alerts: true, - announcement_notifications: false, - }; - assert_eq!( - repository - .write_user_preferences(&preferences) - .await - .expect("preferences should write"), - Some(preferences.clone()) - ); - assert_eq!( - repository - .read_user_preferences("user-1") - .await - .expect("preferences should read"), - Some(preferences) - ); - let now = chrono::Utc::now(); - let session = StoredUserSessionRecord::new( - "session-1".to_string(), - "user-1".to_string(), - "device-1".to_string(), - Some("Laptop".to_string()), - StoredUserSessionRecord::hash_refresh_token("refresh-1"), - None, - None, - Some(now), - Some(now + chrono::Duration::hours(1)), - None, - None, - Some("127.0.0.1".to_string()), - Some("agent".to_string()), - Some(now), - Some(now), - ) - .expect("session should build") - .with_security_version(1) - .expect("session security version should be valid"); - assert_eq!( - repository - .create_user_session(&session) - .await - .expect("session should create") - .map(|session| session.id), - Some("session-1".to_string()) - ); - assert_eq!( - repository - .list_user_sessions("user-1") - .await - .expect("sessions should list") - .len(), - 1 - ); - assert!(repository - .revoke_user_session("user-1", "session-1", now, "logout") - .await - .expect("session should revoke")); - assert!(repository - .list_user_sessions("user-1") - .await - .expect("sessions should list") - .is_empty()); - - let by_ids = repository - .list_user_auth_by_ids(&["user-1".to_string()]) - .await - .expect("auth list should load"); - assert_eq!(by_ids.len(), 1); - assert_eq!(by_ids[0].username, "alice-b"); - assert!(repository - .delete_local_auth_user("user-1") - .await - .expect("delete should succeed")); - assert!(!repository - .delete_local_auth_user("user-1") - .await - .expect("second delete should succeed")); - assert!(repository - .find_user_auth_by_id("user-1") - .await - .expect("deleted auth lookup should load") - .is_none()); - - assert!(repository - .find_export_user_by_id("user-2") - .await - .expect("deleted user lookup should run") - .is_none()); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn non_password_session_replacement_is_atomic_under_concurrency() { - let database_path = std::env::temp_dir().join(format!( - "aether-sqlite-session-replacement-race-{}.db", - uuid::Uuid::new_v4() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .foreign_keys(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .busy_timeout(Duration::from_secs(30)); - 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_active_session_test_user(&pool, "session-race-user").await; - - let repository = SqliteUserReadRepository::new(pool.clone()); - let now = chrono::Utc::now(); - let first = test_user_session( - "session-race-first", - "session-race-user", - "shared-device", - "refresh-first", - now, - ); - let second = test_user_session( - "session-race-second", - "session-race-user", - "shared-device", - "refresh-second", - now, - ); - let barrier = Arc::new(tokio::sync::Barrier::new(2)); - let first_repository = repository.clone(); - let first_barrier = Arc::clone(&barrier); - let first_create = tokio::spawn(async move { - first_barrier.wait().await; - first_repository.create_user_session(&first).await - }); - let second_repository = repository.clone(); - let second_barrier = Arc::clone(&barrier); - let second_create = tokio::spawn(async move { - second_barrier.wait().await; - second_repository.create_user_session(&second).await - }); - - assert!(first_create - .await - .expect("first login task should join") - .expect("first login should succeed") - .is_some()); - assert!(second_create - .await - .expect("second login task should join") - .expect("second login should succeed") - .is_some()); - let active = repository - .list_user_sessions("session-race-user") - .await - .expect("active sessions should list"); - assert_eq!(active.len(), 1); - assert_eq!(active[0].client_device_id, "shared-device"); - - pool.close().await; - let _ = std::fs::remove_file(database_path); - } - - #[tokio::test] - async fn failed_non_password_session_insert_rolls_back_device_revocation() { - 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_active_session_test_user(&pool, "session-rollback-user").await; - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let current = test_user_session( - "session-duplicate-id", - "session-rollback-user", - "shared-device", - "refresh-current", - now, - ); - repository - .create_user_session(¤t) - .await - .expect("initial session should create") - .expect("initial session should exist"); - - let duplicate = test_user_session( - "session-duplicate-id", - "session-rollback-user", - "shared-device", - "refresh-duplicate", - now + chrono::Duration::seconds(1), - ); - assert!(repository.create_user_session(&duplicate).await.is_err()); - - let active = repository - .list_user_sessions("session-rollback-user") - .await - .expect("active sessions should list"); - assert_eq!(active.len(), 1); - assert_eq!(active[0].refresh_token_hash, current.refresh_token_hash); - assert!(!active[0].is_revoked()); - } - - #[tokio::test] - async fn security_state_changes_revoke_sessions_without_reactivation() { - 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_active_session_test_user(&pool, "sqlite-security-state-user").await; - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let session = test_user_session( - "sqlite-security-state-session", - "sqlite-security-state-user", - "sqlite-security-state-device", - "sqlite-security-state-refresh", - now, - ); - repository - .create_user_session(&session) - .await - .expect("initial session should create") - .expect("initial session should exist"); - - repository - .update_local_auth_user_admin_fields( - "sqlite-security-state-user", - None, - false, - None, - false, - None, - false, - None, - false, - None, - Some(false), - ) - .await - .expect("disable should succeed") - .expect("user should exist"); - let revoked = repository - .find_user_session( - "sqlite-security-state-user", - "sqlite-security-state-session", - ) - .await - .expect("revoked session should load") - .expect("revoked session should remain stored"); - assert!(revoked.is_revoked()); - assert_eq!( - revoked.revoke_reason.as_deref(), - Some("user_security_state_changed") - ); - - repository - .update_local_auth_user_admin_fields( - "sqlite-security-state-user", - None, - false, - None, - false, - None, - false, - None, - false, - None, - Some(true), - ) - .await - .expect("reactivation should succeed") - .expect("user should exist"); - assert!(repository - .list_user_sessions("sqlite-security-state-user") - .await - .expect("sessions should list") - .is_empty()); - - let replacement = test_user_session( - "sqlite-security-state-replacement", - "sqlite-security-state-user", - "sqlite-security-state-device", - "sqlite-security-state-replacement-refresh", - now + chrono::Duration::seconds(1), - ); - let replacement = replacement - .with_security_version( - repository - .find_user_auth_by_id("sqlite-security-state-user") - .await - .expect("user lookup should succeed") - .expect("user should exist") - .security_version, - ) - .expect("security version should be valid"); - let created = repository - .create_user_session(&replacement) - .await - .expect("replacement session should create") - .expect("replacement session should exist"); - assert_eq!(created.security_version, replacement.security_version); - let persisted = repository - .find_user_session( - "sqlite-security-state-user", - "sqlite-security-state-replacement", - ) - .await - .expect("replacement session should load") - .expect("replacement session should remain stored"); - assert_eq!(persisted.security_version, replacement.security_version); - let active = repository - .list_user_sessions("sqlite-security-state-user") - .await - .expect("replacement session should list"); - assert_eq!(active.len(), 1); - assert_eq!(active[0].security_version, replacement.security_version); - repository - .update_local_auth_user_admin_fields( - "sqlite-security-state-user", - Some("audit_admin".to_string()), - false, - None, - false, - None, - false, - None, - false, - None, - None, - ) - .await - .expect("role update should succeed") - .expect("user should exist"); - assert!(repository - .list_user_sessions("sqlite-security-state-user") - .await - .expect("sessions should list") - .is_empty()); - } - - #[tokio::test] - async fn unchanged_security_state_preserves_sqlite_sessions() { - 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_active_session_test_user(&pool, "sqlite-unchanged-security-user").await; - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let session = test_user_session( - "sqlite-unchanged-security-session", - "sqlite-unchanged-security-user", - "sqlite-unchanged-security-device", - "sqlite-unchanged-security-refresh", - now, - ); - repository - .create_user_session(&session) - .await - .expect("initial session should create") - .expect("initial session should exist"); - - repository - .update_local_auth_user_admin_fields( - "sqlite-unchanged-security-user", - Some("USER".to_string()), - false, - None, - false, - None, - false, - None, - false, - None, - Some(true), - ) - .await - .expect("idempotent security update should succeed") - .expect("user should exist"); - assert_eq!( - repository - .list_user_sessions("sqlite-unchanged-security-user") - .await - .expect("sessions should list") - .len(), - 1 - ); - } - - #[tokio::test] - async fn sqlite_repository_manages_oauth_users_and_links() { - 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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES - ('linuxdo', 'Linux.do', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1), - ('github', 'GitHub', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("provider should insert"); - - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let user = repository - .create_oauth_auth_user( - Some("OAuth@Example.com".to_string()), - false, - "oauth_user".to_string(), - now, - ) - .await - .expect("oauth user should create") - .expect("oauth user should exist"); - assert_eq!(user.auth_source, "oauth"); - assert!(!user.email_verified); - assert_eq!( - repository - .find_active_user_auth_by_email_ci("oauth@example.com") - .await - .expect("ci lookup should work") - .map(|user| user.id), - Some(user.id.clone()) - ); - - assert!(!repository - .upgrade_oauth_email_verification_if_matches(&user.id, "different@example.com", now,) - .await - .expect("mismatched verification should resolve")); - assert!(repository - .upgrade_oauth_email_verification_if_matches(&user.id, "oauth@example.com", now) - .await - .expect("matching verification should resolve")); - assert!( - repository - .find_user_auth_by_id(&user.id) - .await - .expect("user should load") - .expect("user should exist") - .email_verified - ); - - assert_eq!( - repository - .bind_user_oauth_link( - &user.id, - "linuxdo", - "subject-1", - Some("alice"), - Some("alice@example.com"), - Some(serde_json::json!({"sub": "subject-1"})), - now, - ) - .await - .expect("oauth link should bind"), - BindUserOAuthLinkOutcome::Bound - ); - assert_eq!( - repository - .find_oauth_link_owner("linuxdo", "subject-1") - .await - .expect("owner lookup should work"), - Some(user.id.clone()) - ); - assert!(repository - .find_oauth_linked_user("linuxdo", "subject-1") - .await - .expect("linked user should load") - .is_some()); - assert_eq!( - repository - .list_user_oauth_links(&user.id) - .await - .expect("links should list") - .len(), - 1 - ); - - assert!(repository - .touch_oauth_link( - "linuxdo", - "subject-1", - Some("alice2"), - None, - Some(serde_json::json!({"sub": "subject-1", "fresh": true})), - now + chrono::Duration::seconds(10), - ) - .await - .expect("link should touch")); - assert_eq!( - repository - .count_user_oauth_links(&user.id) - .await - .expect("link count should load"), - 1 - ); - assert_eq!( - repository - .delete_user_oauth_link(&user.id, "linuxdo", false, &[]) - .await - .expect("last link deletion should resolve"), - DeleteUserOAuthLinkOutcome::LastOAuthBinding - ); - repository - .bind_user_oauth_link( - &user.id, - "github", - "subject-2", - Some("alice"), - Some("alice@example.com"), - None, - now, - ) - .await - .expect("second link should upsert"); - assert_eq!( - repository - .delete_user_oauth_link(&user.id, "linuxdo", false, &[]) - .await - .expect("link should delete"), - DeleteUserOAuthLinkOutcome::Deleted - ); - assert_eq!( - repository - .count_user_oauth_links(&user.id) - .await - .expect("link count should load"), - 1 - ); - } - - #[tokio::test] - async fn concurrent_sqlite_oauth_unbinds_preserve_one_login_method() { - let database_path = std::env::temp_dir().join(format!( - "aether-oauth-unbind-{}.sqlite", - uuid::Uuid::new_v4() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .foreign_keys(true) - .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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES - ('linuxdo', 'Linux.do', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1), - ('github', 'GitHub', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("providers should insert"); - - let repository = Arc::new(SqliteUserReadRepository::new(pool.clone())); - let now = chrono::Utc::now(); - let user = repository - .create_oauth_auth_user( - Some("concurrent-oauth@example.com".to_string()), - true, - "concurrent-oauth".to_string(), - now, - ) - .await - .expect("oauth user should create") - .expect("oauth user should exist"); - for (provider_type, subject) in [("linuxdo", "subject-1"), ("github", "subject-2")] { - repository - .bind_user_oauth_link(&user.id, provider_type, subject, None, None, None, now) - .await - .expect("oauth link should upsert"); - } - - let barrier = Arc::new(tokio::sync::Barrier::new(3)); - let first_repository = Arc::clone(&repository); - let first_barrier = Arc::clone(&barrier); - let first_user_id = user.id.clone(); - let first = tokio::spawn(async move { - first_barrier.wait().await; - first_repository - .delete_user_oauth_link(&first_user_id, "linuxdo", false, &[]) - .await - .expect("first unlink should resolve") - }); - let second_repository = Arc::clone(&repository); - let second_barrier = Arc::clone(&barrier); - let second_user_id = user.id.clone(); - let second = tokio::spawn(async move { - second_barrier.wait().await; - second_repository - .delete_user_oauth_link(&second_user_id, "github", false, &[]) - .await - .expect("second unlink should resolve") - }); - barrier.wait().await; - let outcomes = [ - first.await.expect("first unlink task should join"), - second.await.expect("second unlink task should join"), - ]; - - assert_eq!( - outcomes - .iter() - .filter(|outcome| **outcome == DeleteUserOAuthLinkOutcome::Deleted) - .count(), - 1 - ); - assert_eq!( - outcomes - .iter() - .filter(|outcome| **outcome == DeleteUserOAuthLinkOutcome::LastOAuthBinding) - .count(), - 1 - ); - assert_eq!( - repository - .count_user_oauth_links(&user.id) - .await - .expect("remaining links should count"), - 1 - ); - - drop(repository); - pool.close().await; - let _ = std::fs::remove_file(database_path); - } - - #[tokio::test] - async fn concurrent_sqlite_oauth_binds_preserve_single_identity_owner() { - let database_path = - std::env::temp_dir().join(format!("aether-oauth-bind-{}.sqlite", uuid::Uuid::new_v4())); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .foreign_keys(true) - .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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://example.test/callback', - 'https://example.test/app', 1, 1, 1 -) -"#, - ) - .execute(&pool) - .await - .expect("provider should insert"); - - let repository = Arc::new(SqliteUserReadRepository::new(pool.clone())); - let now = chrono::Utc::now(); - let first_user = repository - .create_oauth_auth_user(None, false, "bind-first".to_string(), now) - .await - .expect("first user should create") - .expect("first user should exist"); - let second_user = repository - .create_oauth_auth_user(None, false, "bind-second".to_string(), now) - .await - .expect("second user should create") - .expect("second user should exist"); - let barrier = Arc::new(tokio::sync::Barrier::new(3)); - - let first_repository = Arc::clone(&repository); - let first_barrier = Arc::clone(&barrier); - let first_id = first_user.id.clone(); - let first = tokio::spawn(async move { - first_barrier.wait().await; - first_repository - .bind_user_oauth_link( - &first_id, - "linuxdo", - "shared-subject", - None, - None, - None, - now, - ) - .await - .expect("first bind should resolve") - }); - let second_repository = Arc::clone(&repository); - let second_barrier = Arc::clone(&barrier); - let second_id = second_user.id.clone(); - let second = tokio::spawn(async move { - second_barrier.wait().await; - second_repository - .bind_user_oauth_link( - &second_id, - "linuxdo", - "shared-subject", - None, - None, - None, - now, - ) - .await - .expect("second bind should resolve") - }); - barrier.wait().await; - let outcomes = [ - first.await.expect("first bind task should join"), - second.await.expect("second bind task should join"), - ]; - - assert_eq!( - outcomes - .iter() - .filter(|outcome| **outcome == BindUserOAuthLinkOutcome::Bound) - .count(), - 1 - ); - assert_eq!( - outcomes - .iter() - .filter(|outcome| { - **outcome == BindUserOAuthLinkOutcome::IdentityBoundToAnotherUser - }) - .count(), - 1 - ); - let link_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_oauth_links WHERE provider_type = 'linuxdo' AND provider_user_id = 'shared-subject'", - ) - .fetch_one(&pool) - .await - .expect("link count should load"); - assert_eq!(link_count, 1); - - drop(repository); - pool.close().await; - let _ = std::fs::remove_file(database_path); - } - - #[tokio::test] - async fn sqlite_oauth_bind_rejects_disabled_provider() { - 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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://example.test/callback', - 'https://example.test/app', 0, 1, 1 -) -"#, - ) - .execute(&pool) - .await - .expect("disabled provider should insert"); - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let user = repository - .create_oauth_auth_user(None, false, "disabled-provider-user".to_string(), now) - .await - .expect("user should create") - .expect("user should exist"); - - assert_eq!( - repository - .bind_user_oauth_link(&user.id, "linuxdo", "subject", None, None, None, now) - .await - .expect("bind should resolve"), - BindUserOAuthLinkOutcome::ProviderDisabled - ); - assert!(!repository - .has_user_oauth_provider_link(&user.id, "linuxdo") - .await - .expect("link lookup should succeed")); - } - - #[tokio::test] - async fn sqlite_oauth_unbind_only_counts_enabled_provider_links() { - 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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES - ('linuxdo', 'Linux.do', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1), - ('github', 'GitHub', 'client', 'https://example.test/callback', 'https://example.test/app', 1, 1, 1) -"#, - ) - .execute(&pool) - .await - .expect("providers should insert"); - let repository = SqliteUserReadRepository::new(pool.clone()); - let now = chrono::Utc::now(); - let user = repository - .create_oauth_auth_user( - Some("enabled-link@example.com".to_string()), - true, - "enabled-link".to_string(), - now, - ) - .await - .expect("oauth user should create") - .expect("oauth user should exist"); - for (provider_type, subject) in [("linuxdo", "subject-1"), ("github", "subject-2")] { - assert_eq!( - repository - .bind_user_oauth_link(&user.id, provider_type, subject, None, None, None, now) - .await - .expect("oauth link should bind"), - BindUserOAuthLinkOutcome::Bound - ); - } - sqlx::query("UPDATE oauth_providers SET is_enabled = 0 WHERE provider_type = 'github'") - .execute(&pool) - .await - .expect("provider should disable after its link is created"); - - assert_eq!( - repository - .delete_user_oauth_link(&user.id, "linuxdo", false, &[]) - .await - .expect("enabled link deletion should resolve"), - DeleteUserOAuthLinkOutcome::LastOAuthBinding - ); - assert_eq!( - repository - .delete_user_oauth_link(&user.id, "github", false, &[]) - .await - .expect("disabled link deletion should resolve"), - DeleteUserOAuthLinkOutcome::Deleted - ); - assert!(repository - .has_user_oauth_provider_link(&user.id, "linuxdo") - .await - .expect("enabled provider link lookup should work")); - } - - #[tokio::test] - async fn sqlite_oauth_unbind_respects_ldap_exclusive_local_login_policy() { - 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"); - sqlx::query( - r#" -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://example.test/callback', - 'https://example.test/app', 1, 1, 1 -) -"#, - ) - .execute(&pool) - .await - .expect("provider should insert"); - let repository = SqliteUserReadRepository::new(pool); - let now = chrono::Utc::now(); - let valid_hash = "$2b$12$4qL4tdcsFwVaDTw5Ck3xzu8GpNdre56DiNR6Dnw7t6gCXaEnqAe7G".to_string(); - let user = repository - .create_local_auth_user( - Some("ldap-exclusive-local@example.com".to_string()), - true, - "ldap-exclusive-local".to_string(), - valid_hash, - ) - .await - .expect("local user should create") - .expect("local user should exist"); - repository - .bind_user_oauth_link(&user.id, "linuxdo", "subject-1", None, None, None, now) - .await - .expect("oauth link should upsert"); - - assert_eq!( - repository - .delete_user_oauth_link(&user.id, "linuxdo", false, &[]) - .await - .expect("unlink should resolve"), - DeleteUserOAuthLinkOutcome::LastLoginMethod - ); - assert!(repository - .has_user_oauth_provider_link(&user.id, "linuxdo") - .await - .expect("oauth link lookup should work")); - } -} diff --git a/crates/aether-data/adapters/sqlite/src/video_tasks.rs b/crates/aether-data/adapters/sqlite/src/video_tasks.rs deleted file mode 100644 index a00c0f85d..000000000 --- a/crates/aether-data/adapters/sqlite/src/video_tasks.rs +++ /dev/null @@ -1,1168 +0,0 @@ -use async_trait::async_trait; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -use aether_data_contracts::repository::video_tasks::{ - StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount, - VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskStatusCount, - VideoTaskWriteRepository, -}; -use aether_data_contracts::DataLayerError; - -use crate::error::SqlResultExt; -use crate::SqlitePool; - -const VIDEO_TASK_COLUMNS: &str = r#" -SELECT - id, - short_id, - request_id, - user_id, - api_key_id, - username, - api_key_name, - external_task_id, - provider_id, - endpoint_id, - key_id, - client_api_format, - provider_api_format, - format_converted, - model, - prompt, - original_request_body, - duration_seconds, - resolution, - aspect_ratio, - size, - status, - progress_percent, - progress_message, - retry_count, - poll_interval_seconds, - next_poll_at AS next_poll_at_unix_secs, - poll_count, - max_poll_count, - created_at AS created_at_unix_ms, - submitted_at AS submitted_at_unix_secs, - completed_at AS completed_at_unix_secs, - updated_at AS updated_at_unix_secs, - error_code, - error_message, - video_url, - request_metadata -FROM video_tasks -"#; - -#[derive(Debug, Clone)] -pub struct SqliteVideoTaskRepository { - pool: SqlitePool, -} - -impl SqliteVideoTaskRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - async fn find_by_id(&self, id: &str) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{VIDEO_TASK_COLUMNS} WHERE id = ? LIMIT 1")) - .bind(id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_id_for_user( - &self, - id: &str, - user_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE id = ? AND user_id = ? LIMIT 1" - )) - .bind(id) - .bind(user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_short_id( - &self, - short_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!("{VIDEO_TASK_COLUMNS} WHERE short_id = ? LIMIT 1")) - .bind(short_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_short_id_for_user( - &self, - short_id: &str, - user_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE short_id = ? AND user_id = ? LIMIT 1" - )) - .bind(short_id) - .bind(user_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } - - async fn find_by_user_external( - &self, - user_id: &str, - external_task_id: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE user_id = ? AND external_task_id = ? LIMIT 1" - )) - .bind(user_id) - .bind(external_task_id) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_video_task_row).transpose() - } -} - -#[async_trait] -impl VideoTaskReadRepository for SqliteVideoTaskRepository { - async fn find( - &self, - key: VideoTaskLookupKey<'_>, - ) -> Result, DataLayerError> { - match key { - VideoTaskLookupKey::Id(id) => self.find_by_id(id).await, - VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await, - VideoTaskLookupKey::UserExternal { - user_id, - external_task_id, - } => self.find_by_user_external(user_id, external_task_id).await, - } - } - - async fn find_for_user( - &self, - key: VideoTaskLookupKey<'_>, - user_id: &str, - ) -> Result, DataLayerError> { - match key { - VideoTaskLookupKey::Id(id) => self.find_by_id_for_user(id, user_id).await, - VideoTaskLookupKey::ShortId(short_id) => { - self.find_by_short_id_for_user(short_id, user_id).await - } - VideoTaskLookupKey::UserExternal { - user_id: lookup_user_id, - external_task_id, - } if lookup_user_id == user_id => { - self.find_by_user_external(user_id, external_task_id).await - } - VideoTaskLookupKey::UserExternal { .. } => Ok(None), - } - } - - async fn list_active(&self, limit: usize) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} WHERE status IN ('pending', 'submitted', 'queued', 'processing') ORDER BY updated_at DESC LIMIT ?" - )) - .bind(limit_i64(limit, "active video task limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_due( - &self, - now_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let rows = sqlx::query(&format!( - "{VIDEO_TASK_COLUMNS} 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 ?" - )) - .bind(u64_to_i64(now_unix_secs, "video task now")?) - .bind(limit_i64(limit, "due video task limit")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_page( - &self, - filter: &VideoTaskQueryFilter, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new(VIDEO_TASK_COLUMNS); - push_filter(&mut builder, filter, None); - builder - .push(" ORDER BY created_at DESC, updated_at DESC LIMIT ") - .push_bind(limit_i64(limit, "video task page limit")?) - .push(" OFFSET ") - .push_bind(limit_i64(offset, "video task page offset")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter().map(map_video_task_row).collect() - } - - async fn list_page_summary( - &self, - filter: &VideoTaskQueryFilter, - offset: usize, - limit: usize, - ) -> Result, DataLayerError> { - self.list_page(filter, offset, limit).await - } - - async fn count(&self, filter: &VideoTaskQueryFilter) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - count_query(builder, &self.pool).await - } - - async fn count_by_status( - &self, - filter: &VideoTaskQueryFilter, - ) -> Result, DataLayerError> { - let mut builder = - QueryBuilder::::new("SELECT status, COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - builder.push(" GROUP BY status ORDER BY status ASC"); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(VideoTaskStatusCount { - status: VideoTaskStatus::from_database( - row.try_get::("status").map_sql_err()?.as_str(), - )?, - count: count_value(row.try_get("total").map_sql_err()?)?, - }) - }) - .collect() - } - - async fn count_distinct_users( - &self, - filter: &VideoTaskQueryFilter, - ) -> Result { - let mut builder = - QueryBuilder::::new("SELECT COUNT(DISTINCT user_id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - push_clause(&mut builder, "user_id IS NOT NULL"); - push_clause(&mut builder, "user_id <> ''"); - count_query(builder, &self.pool).await - } - - async fn top_models( - &self, - filter: &VideoTaskQueryFilter, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - let mut builder = - QueryBuilder::::new("SELECT model, COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, None); - push_clause(&mut builder, "model IS NOT NULL"); - push_clause(&mut builder, "model <> ''"); - builder - .push(" GROUP BY model ORDER BY total DESC, model ASC LIMIT ") - .push_bind(limit_i64(limit, "video task top models limit")?); - let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?; - rows.iter() - .map(|row| { - Ok(VideoTaskModelCount { - model: row.try_get("model").map_sql_err()?, - count: count_value(row.try_get("total").map_sql_err()?)?, - }) - }) - .collect() - } - - async fn count_created_since( - &self, - filter: &VideoTaskQueryFilter, - created_since_unix_secs: u64, - ) -> Result { - let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); - push_filter(&mut builder, filter, Some(created_since_unix_secs)); - count_query(builder, &self.pool).await - } -} - -#[async_trait] -impl VideoTaskWriteRepository for SqliteVideoTaskRepository { - async fn upsert(&self, mut task: UpsertVideoTask) -> Result { - task.sanitize_for_persistence(); - let id = task.id.clone(); - let expected_identity = task.clone(); - bind_task(sqlx::query(UPSERT_SQL), task, true, false)? - .execute(&self.pool) - .await - .map_sql_err()?; - let stored = self.find_by_id(&id).await?.ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "video task {id} conflicts with persisted immutable identity" - )) - })?; - stored.ensure_immutable_identity_matches(&expected_identity)?; - Ok(stored) - } - - async fn update_if_active( - &self, - task: UpsertVideoTask, - ) -> Result, DataLayerError> { - let id = task.id.clone(); - let rows_affected = bind_task(sqlx::query(UPDATE_IF_ACTIVE_SQL), task, false, true)? - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); - if rows_affected == 0 { - return Ok(None); - } - self.find_by_id(&id).await - } - - async fn claim_due( - &self, - now_unix_secs: u64, - claim_until_unix_secs: u64, - limit: usize, - ) -> Result, DataLayerError> { - if limit == 0 { - return Ok(Vec::new()); - } - 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()); - } - - 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) - } -} - -const UPSERT_SQL: &str = r#" -INSERT INTO video_tasks ( - id, short_id, request_id, user_id, api_key_id, username, api_key_name, - external_task_id, provider_id, endpoint_id, key_id, client_api_format, - provider_api_format, format_converted, model, prompt, original_request_body, - duration_seconds, resolution, aspect_ratio, size, status, progress_percent, - progress_message, retry_count, poll_interval_seconds, next_poll_at, poll_count, - max_poll_count, video_url, error_code, error_message, request_metadata, - created_at, submitted_at, completed_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT(id) DO UPDATE SET - short_id = excluded.short_id, - request_id = excluded.request_id, - user_id = excluded.user_id, - api_key_id = excluded.api_key_id, - username = excluded.username, - api_key_name = excluded.api_key_name, - external_task_id = excluded.external_task_id, - provider_id = excluded.provider_id, - endpoint_id = excluded.endpoint_id, - key_id = excluded.key_id, - client_api_format = excluded.client_api_format, - provider_api_format = excluded.provider_api_format, - format_converted = excluded.format_converted, - model = excluded.model, - prompt = excluded.prompt, - original_request_body = excluded.original_request_body, - duration_seconds = excluded.duration_seconds, - resolution = excluded.resolution, - aspect_ratio = excluded.aspect_ratio, - size = excluded.size, - status = excluded.status, - progress_percent = excluded.progress_percent, - progress_message = excluded.progress_message, - retry_count = excluded.retry_count, - poll_interval_seconds = excluded.poll_interval_seconds, - next_poll_at = excluded.next_poll_at, - poll_count = excluded.poll_count, - max_poll_count = excluded.max_poll_count, - video_url = excluded.video_url, - error_code = excluded.error_code, - error_message = excluded.error_message, - request_metadata = excluded.request_metadata, - created_at = COALESCE(video_tasks.created_at, excluded.created_at), - submitted_at = excluded.submitted_at, - completed_at = excluded.completed_at, - updated_at = excluded.updated_at -WHERE video_tasks.short_id IS excluded.short_id - AND video_tasks.request_id IS excluded.request_id - AND video_tasks.user_id IS excluded.user_id - AND video_tasks.api_key_id IS excluded.api_key_id - AND video_tasks.external_task_id IS excluded.external_task_id - AND video_tasks.provider_id IS excluded.provider_id - AND video_tasks.endpoint_id IS excluded.endpoint_id - AND video_tasks.key_id IS excluded.key_id - AND video_tasks.client_api_format IS excluded.client_api_format - AND video_tasks.provider_api_format IS excluded.provider_api_format - AND video_tasks.format_converted IS excluded.format_converted - AND video_tasks.model IS excluded.model - AND video_tasks.duration_seconds IS excluded.duration_seconds - AND video_tasks.resolution IS excluded.resolution - AND video_tasks.aspect_ratio IS excluded.aspect_ratio - AND video_tasks.size IS excluded.size -"#; - -const UPDATE_IF_ACTIVE_SQL: &str = r#" -UPDATE video_tasks SET - short_id = ?, - request_id = ?, - user_id = ?, - api_key_id = ?, - username = ?, - api_key_name = ?, - external_task_id = ?, - provider_id = ?, - endpoint_id = ?, - key_id = ?, - client_api_format = ?, - provider_api_format = ?, - format_converted = ?, - model = ?, - prompt = ?, - original_request_body = ?, - duration_seconds = ?, - resolution = ?, - aspect_ratio = ?, - size = ?, - status = ?, - progress_percent = ?, - progress_message = ?, - retry_count = ?, - poll_interval_seconds = ?, - next_poll_at = ?, - poll_count = ?, - max_poll_count = ?, - video_url = ?, - error_code = ?, - error_message = ?, - request_metadata = ?, - created_at = COALESCE(created_at, ?), - submitted_at = ?, - completed_at = ?, - updated_at = ? -WHERE id = ? - AND status IN ('pending', 'submitted', 'queued', 'processing') - AND short_id IS ? - AND request_id IS ? - AND user_id IS ? - AND api_key_id IS ? - AND external_task_id IS ? - AND provider_id IS ? - AND endpoint_id IS ? - AND key_id IS ? - AND client_api_format IS ? - AND provider_api_format IS ? - AND format_converted IS ? - AND model IS ? - AND duration_seconds IS ? - AND resolution IS ? - AND aspect_ratio IS ? - AND size IS ? -"#; - -fn bind_task<'q>( - query: sqlx::query::Query<'q, Sqlite, sqlx::sqlite::SqliteArguments<'q>>, - mut task: UpsertVideoTask, - include_insert_id: bool, - include_update_id: bool, -) -> Result>, DataLayerError> { - task.sanitize_for_persistence(); - let identity = task.clone(); - let original_request_body = json_to_string(&task.original_request_body)?; - let request_metadata = json_to_string(&task.request_metadata)?; - let query = if include_insert_id { - query.bind(task.id.clone()) - } else { - query - }; - let bound = query - .bind(task.short_id) - .bind(task.request_id) - .bind(task.user_id) - .bind(task.api_key_id) - .bind(task.username) - .bind(task.api_key_name) - .bind(task.external_task_id) - .bind(task.provider_id) - .bind(task.endpoint_id) - .bind(task.key_id) - .bind(task.client_api_format) - .bind(task.provider_api_format) - .bind(task.format_converted) - .bind(task.model) - .bind(task.prompt) - .bind(original_request_body) - .bind(optional_u32_to_i32( - task.duration_seconds, - "video task duration_seconds", - )?) - .bind(task.resolution) - .bind(task.aspect_ratio) - .bind(task.size) - .bind(status_to_database(task.status)) - .bind(i32::from(task.progress_percent)) - .bind(task.progress_message) - .bind(u32_to_i32(task.retry_count, "video task retry_count")?) - .bind(u32_to_i32( - task.poll_interval_seconds, - "video task poll_interval_seconds", - )?) - .bind(optional_u64_to_i64( - task.next_poll_at_unix_secs, - "video task next_poll_at", - )?) - .bind(u32_to_i32(task.poll_count, "video task poll_count")?) - .bind(u32_to_i32( - task.max_poll_count, - "video task max_poll_count", - )?) - .bind(task.video_url) - .bind(task.error_code) - .bind(task.error_message) - .bind(request_metadata) - .bind(u64_to_i64( - task.created_at_unix_ms, - "video task created_at", - )?) - .bind(optional_u64_to_i64( - task.submitted_at_unix_secs, - "video task submitted_at", - )?) - .bind(optional_u64_to_i64( - task.completed_at_unix_secs, - "video task completed_at", - )?) - .bind(u64_to_i64( - task.updated_at_unix_secs, - "video task updated_at", - )?); - if include_update_id { - bind_identity_guard(bound.bind(task.id), identity) - } else { - Ok(bound) - } -} - -fn bind_identity_guard<'q>( - query: sqlx::query::Query<'q, Sqlite, sqlx::sqlite::SqliteArguments<'q>>, - identity: UpsertVideoTask, -) -> Result>, DataLayerError> { - Ok(query - .bind(identity.short_id) - .bind(identity.request_id) - .bind(identity.user_id) - .bind(identity.api_key_id) - .bind(identity.external_task_id) - .bind(identity.provider_id) - .bind(identity.endpoint_id) - .bind(identity.key_id) - .bind(identity.client_api_format) - .bind(identity.provider_api_format) - .bind(identity.format_converted) - .bind(identity.model) - .bind(optional_u32_to_i32( - identity.duration_seconds, - "video task duration_seconds", - )?) - .bind(identity.resolution) - .bind(identity.aspect_ratio) - .bind(identity.size)) -} - -fn push_filter<'args>( - builder: &mut QueryBuilder<'args, Sqlite>, - filter: &'args VideoTaskQueryFilter, - created_since_unix_secs: Option, -) { - if let Some(user_id) = filter.user_id.as_deref() { - push_clause(builder, "user_id = "); - builder.push_bind(user_id); - } - if let Some(status) = filter.status { - push_clause(builder, "status = "); - builder.push_bind(status_to_database(status)); - } - if let Some(model_substring) = filter.model_substring.as_deref() { - push_clause(builder, "LOWER(model) LIKE "); - builder.push_bind(format!( - "%{}%", - escape_like_pattern(&model_substring.trim().to_ascii_lowercase()) - )); - builder.push(" ESCAPE '\\'"); - } - if let Some(client_api_format) = filter.client_api_format.as_deref() { - push_clause(builder, "client_api_format = "); - builder.push_bind(client_api_format); - } - if let Some(created_since_unix_secs) = created_since_unix_secs { - push_clause(builder, "created_at >= "); - builder.push_bind(created_since_unix_secs as i64); - } -} - -fn push_clause<'args>(builder: &mut QueryBuilder<'args, Sqlite>, clause: &str) { - let sql = builder.sql(); - if sql.contains(" WHERE ") || sql.contains("\nWHERE ") { - builder.push(" AND "); - } else { - builder.push(" WHERE "); - } - builder.push(clause); -} - -async fn count_query( - mut builder: QueryBuilder<'_, Sqlite>, - pool: &SqlitePool, -) -> Result { - let row = builder.build().fetch_one(pool).await.map_sql_err()?; - count_value(row.try_get("total").map_sql_err()?) -} - -fn count_value(value: i64) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("invalid video task count result: {value}")) - }) -} - -fn map_video_task_row(row: &SqliteRow) -> Result { - StoredVideoTask::new( - row.try_get("id").map_sql_err()?, - row.try_get("short_id").map_sql_err()?, - row.try_get("request_id").map_sql_err()?, - row.try_get("user_id").map_sql_err()?, - row.try_get("api_key_id").map_sql_err()?, - row.try_get("username").map_sql_err()?, - row.try_get("api_key_name").map_sql_err()?, - row.try_get("external_task_id").map_sql_err()?, - row.try_get("provider_id").map_sql_err()?, - row.try_get("endpoint_id").map_sql_err()?, - row.try_get("key_id").map_sql_err()?, - row.try_get("client_api_format").map_sql_err()?, - row.try_get("provider_api_format").map_sql_err()?, - row.try_get("format_converted").map_sql_err()?, - row.try_get("model").map_sql_err()?, - row.try_get("prompt").map_sql_err()?, - parse_json(row.try_get("original_request_body").ok().flatten())?, - row.try_get("duration_seconds").map_sql_err()?, - row.try_get("resolution").map_sql_err()?, - row.try_get("aspect_ratio").map_sql_err()?, - row.try_get("size").map_sql_err()?, - VideoTaskStatus::from_database(row.try_get::("status").map_sql_err()?.as_str())?, - row.try_get("progress_percent").map_sql_err()?, - row.try_get("progress_message").map_sql_err()?, - row.try_get("retry_count").map_sql_err()?, - row.try_get("poll_interval_seconds").map_sql_err()?, - row.try_get("next_poll_at_unix_secs").map_sql_err()?, - row.try_get("poll_count").map_sql_err()?, - row.try_get("max_poll_count").map_sql_err()?, - row.try_get("created_at_unix_ms").map_sql_err()?, - row.try_get("submitted_at_unix_secs").map_sql_err()?, - row.try_get("completed_at_unix_secs").map_sql_err()?, - row.try_get("updated_at_unix_secs").map_sql_err()?, - row.try_get("error_code").map_sql_err()?, - row.try_get("error_message").map_sql_err()?, - row.try_get("video_url").map_sql_err()?, - parse_json(row.try_get("request_metadata").ok().flatten())?, - ) -} - -fn parse_json(value: Option) -> Result, DataLayerError> { - value - .filter(|value| !value.trim().is_empty()) - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("video task JSON field is invalid: {err}")) - }) - }) - .transpose() -} - -fn json_to_string(value: &Option) -> Result, DataLayerError> { - value - .as_ref() - .map(|value| { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "video task JSON field is unserializable: {err}" - )) - }) - }) - .transpose() -} - -fn status_to_database(status: VideoTaskStatus) -> &'static str { - match status { - VideoTaskStatus::Pending => "pending", - VideoTaskStatus::Submitted => "submitted", - VideoTaskStatus::Queued => "queued", - VideoTaskStatus::Processing => "processing", - VideoTaskStatus::Completed => "completed", - VideoTaskStatus::Failed => "failed", - VideoTaskStatus::Cancelled => "cancelled", - VideoTaskStatus::Expired => "expired", - VideoTaskStatus::Deleted => "deleted", - } -} - -fn escape_like_pattern(value: &str) -> String { - value - .replace('\\', "\\\\") - .replace('%', "\\%") - .replace('_', "\\_") -} - -fn limit_i64(value: usize, name: &str) -> Result { - i64::try_from(value) - .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {name}: {value}"))) -} - -fn u64_to_i64(value: u64, name: &str) -> Result { - i64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u64_to_i64(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u64_to_i64(value, name)).transpose() -} - -fn u32_to_i32(value: u32, name: &str) -> Result { - i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow"))) -} - -fn optional_u32_to_i32(value: Option, name: &str) -> Result, DataLayerError> { - value.map(|value| u32_to_i32(value, name)).transpose() -} - -#[cfg(test)] -mod tests { - use super::{SqliteVideoTaskRepository, UPDATE_IF_ACTIVE_SQL, UPSERT_SQL}; - use crate::run_migrations; - use aether_data_contracts::repository::video_tasks::{ - UpsertVideoTask, VideoTaskLookupKey, VideoTaskQueryFilter, VideoTaskReadRepository, - VideoTaskStatus, VideoTaskWriteRepository, - }; - use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; - use std::{sync::Arc, time::Duration}; - - #[test] - fn sqlite_write_sql_guards_every_immutable_identity_field() { - for column in [ - "short_id", - "request_id", - "user_id", - "api_key_id", - "external_task_id", - "provider_id", - "endpoint_id", - "key_id", - "client_api_format", - "provider_api_format", - "format_converted", - "model", - "duration_seconds", - "resolution", - "aspect_ratio", - "size", - ] { - assert!( - UPSERT_SQL.contains(&format!("video_tasks.{column} IS excluded.{column}")), - "upsert should guard {column}" - ); - assert!( - UPDATE_IF_ACTIVE_SQL.contains(&format!("{column} IS ?")), - "active update should guard {column}" - ); - } - assert!(UPSERT_SQL - .contains("created_at = COALESCE(video_tasks.created_at, excluded.created_at)")); - assert!(UPDATE_IF_ACTIVE_SQL.contains("created_at = COALESCE(created_at, ?)")); - } - - #[tokio::test] - async fn sqlite_repository_writes_and_reads_video_tasks() { - 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 = SqliteVideoTaskRepository::new(pool); - repository - .upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100)) - .await - .expect("task should insert"); - repository - .upsert(UpsertVideoTask { - user_id: Some("user-2".to_string()), - model: Some("veo-3-fast".to_string()), - client_api_format: Some("gemini:video".to_string()), - created_at_unix_ms: 260, - updated_at_unix_secs: 260, - status: VideoTaskStatus::Completed, - ..sample_task("task-2", VideoTaskStatus::Completed, 260) - }) - .await - .expect("task should insert"); - - assert!(repository - .find(VideoTaskLookupKey::ShortId("short-task-1")) - .await - .expect("short lookup should load") - .is_some()); - assert!(repository - .find(VideoTaskLookupKey::UserExternal { - user_id: "user-1", - external_task_id: "ext-task-1", - }) - .await - .expect("user/external lookup should load") - .is_some()); - assert!(repository - .find_for_user(VideoTaskLookupKey::Id("task-1"), "user-1") - .await - .expect("owner id lookup should load") - .is_some()); - assert!(repository - .find_for_user(VideoTaskLookupKey::Id("task-1"), "user-2") - .await - .expect("foreign id lookup should run") - .is_none()); - assert!(repository - .find_for_user(VideoTaskLookupKey::ShortId("short-task-1"), "user-1") - .await - .expect("owner short id lookup should load") - .is_some()); - assert!(repository - .find_for_user(VideoTaskLookupKey::ShortId("short-task-1"), "user-2") - .await - .expect("foreign short id lookup should run") - .is_none()); - - let due = repository - .list_due(100, 10) - .await - .expect("due tasks should load"); - assert_eq!(due.len(), 1); - - let claimed = repository - .claim_due(100, 130, 10) - .await - .expect("due tasks should claim"); - assert_eq!(claimed.len(), 1); - assert_eq!(claimed[0].next_poll_at_unix_secs, Some(130)); - - let filter = VideoTaskQueryFilter { - user_id: Some("user-2".to_string()), - status: Some(VideoTaskStatus::Completed), - model_substring: Some("veo".to_string()), - client_api_format: Some("gemini:video".to_string()), - }; - assert_eq!( - repository.count(&filter).await.expect("count should load"), - 1 - ); - assert_eq!( - repository - .count_by_status(&filter) - .await - .expect("status counts should load")[0] - .count, - 1 - ); - assert_eq!( - repository - .top_models(&filter, 10) - .await - .expect("top models should load")[0] - .model, - "veo-3-fast" - ); - - let updated = repository - .update_if_active(UpsertVideoTask { - status: VideoTaskStatus::Processing, - progress_percent: 50, - created_at_unix_ms: 90, - submitted_at_unix_secs: Some(90), - ..sample_task("task-1", VideoTaskStatus::Processing, 150) - }) - .await - .expect("active task should update") - .expect("active task should exist"); - assert_eq!(updated.progress_percent, 50); - } - - #[tokio::test] - async fn sqlite_rejects_identity_conflicts_without_modifying_task_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"); - - let repository = SqliteVideoTaskRepository::new(pool); - let original = sample_task("task-owned", VideoTaskStatus::Submitted, 100); - repository - .upsert(original.clone()) - .await - .expect("original task should insert"); - - let conflict = repository - .upsert(UpsertVideoTask { - user_id: Some("attacker".to_string()), - key_id: Some("attacker-key".to_string()), - status: VideoTaskStatus::Completed, - progress_percent: 100, - completed_at_unix_secs: Some(200), - updated_at_unix_secs: 200, - ..original.clone() - }) - .await - .expect_err("conflicting owner should be rejected"); - assert!(conflict.to_string().contains("immutable field user_id")); - - let after_upsert = repository - .find(VideoTaskLookupKey::Id("task-owned")) - .await - .expect("task lookup should succeed") - .expect("original task should remain"); - assert_eq!(after_upsert.user_id.as_deref(), Some("user-1")); - assert_eq!(after_upsert.key_id.as_deref(), Some("provider-key-1")); - assert_eq!(after_upsert.status, VideoTaskStatus::Submitted); - assert_eq!(after_upsert.progress_percent, 0); - assert_eq!(after_upsert.completed_at_unix_secs, None); - assert_eq!(after_upsert.updated_at_unix_secs, 100); - - let active_conflict = repository - .update_if_active(UpsertVideoTask { - request_id: "attacker-request".to_string(), - status: VideoTaskStatus::Failed, - updated_at_unix_secs: 300, - ..original.clone() - }) - .await - .expect("guarded active update should execute"); - assert!(active_conflict.is_none()); - let after_active_conflict = repository - .find(VideoTaskLookupKey::Id("task-owned")) - .await - .expect("task lookup should succeed") - .expect("original task should remain"); - assert_eq!(after_active_conflict.request_id, "request-task-owned"); - assert_eq!(after_active_conflict.status, VideoTaskStatus::Submitted); - assert_eq!(after_active_conflict.updated_at_unix_secs, 100); - - let updated = repository - .upsert(UpsertVideoTask { - status: VideoTaskStatus::Processing, - progress_percent: 50, - poll_count: 2, - created_at_unix_ms: 999, - updated_at_unix_secs: 200, - ..original - }) - .await - .expect("same owner state update should succeed"); - assert_eq!(updated.status, VideoTaskStatus::Processing); - assert_eq!(updated.progress_percent, 50); - assert_eq!(updated.poll_count, 2); - assert_eq!(updated.created_at_unix_ms, 90); - } - - #[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, - updated_at_unix_secs: u64, - ) -> UpsertVideoTask { - UpsertVideoTask { - id: id.to_string(), - short_id: Some(format!("short-{id}")), - request_id: format!("request-{id}"), - user_id: Some("user-1".to_string()), - api_key_id: Some("api-key-1".to_string()), - username: Some("user".to_string()), - api_key_name: Some("primary".to_string()), - external_task_id: Some(format!("ext-{id}")), - provider_id: Some("provider-1".to_string()), - endpoint_id: Some("endpoint-1".to_string()), - key_id: Some("provider-key-1".to_string()), - 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("hello".to_string()), - original_request_body: Some(serde_json::json!({"prompt": "hello"})), - duration_seconds: Some(4), - resolution: Some("720p".to_string()), - aspect_ratio: Some("16:9".to_string()), - size: Some("1280x720".to_string()), - status, - progress_percent: 0, - progress_message: None, - retry_count: 0, - poll_interval_seconds: 10, - next_poll_at_unix_secs: Some(updated_at_unix_secs), - poll_count: 0, - max_poll_count: 360, - created_at_unix_ms: updated_at_unix_secs.saturating_sub(10), - submitted_at_unix_secs: Some(updated_at_unix_secs.saturating_sub(10)), - completed_at_unix_secs: None, - updated_at_unix_secs, - error_code: None, - error_message: None, - video_url: None, - request_metadata: Some(serde_json::json!({"request": id})), - } - } -} diff --git a/crates/aether-data/adapters/sqlite/src/wallet.rs b/crates/aether-data/adapters/sqlite/src/wallet.rs deleted file mode 100644 index e1c8c34fc..000000000 --- a/crates/aether-data/adapters/sqlite/src/wallet.rs +++ /dev/null @@ -1,6698 +0,0 @@ -use crate::error::SqlResultExt; -use crate::{sqlite_optional_real, sqlite_real, SqlitePool}; -use aether_data_contracts::repository::billing::{ - checked_plan_duration_days_from_snapshot, entitlements_have_replacement_selector, - entitlements_should_replace_existing, -}; -use aether_data_contracts::repository::wallet::{ - canonicalize_payment_method, canonicalize_wallet_refund_fields, - payment_callback_amount_matches_order, payment_callback_method_matches_order, - payment_callback_provider_matches_order, payment_order_is_failed_wallet_checkout_placeholder, - payment_order_is_uncertain_wallet_checkout_placeholder, - payment_order_refund_amounts_are_consistent, - payment_order_stripe_client_secret_cas_replacement, project_wallet_gateway_response, - project_wallet_recharge_gateway_response, redeem_code_payment_method, - redeem_code_refundable_amount, validate_admin_redeem_code_batch_input, - validate_manual_wallet_recharge, validate_payment_order_credit_amounts, - validate_plan_purchase_order_input, validate_plan_wallet_credit_entitlements, - validate_redeem_wallet_credit, validate_wallet_recharge_order_input, - wallet_recharge_replay_matches, -}; -use aether_data_contracts::repository::wallet::{ - wallet_recharge_checkout_claim_response, wallet_recharge_checkout_claim_token, - wallet_recharge_checkout_failed_response, wallet_recharge_checkout_uncertain_response, - wallet_recharge_order_is_checkout_placeholder, - wallet_recharge_order_is_reclaimable_placeholder, - wallet_recharge_response_is_checkout_placeholder, wallet_refund_proof_is_success, - AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, - AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery, - AdminWalletRefundRequestListQuery, CompareAndSwapPaymentOrderStripeClientSecretInput, - CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput, - CreateAdminRedeemCodeBatchResult, CreateManualWalletRechargeInput, - CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput, - CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput, - CreateWalletRefundRequestOutcome, CreatedAdminRedeemCodePlaintext, - CreditAdminPaymentOrderInput, DeleteAdminRedeemCodeBatchInput, - DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput, - FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome, ProcessAdminWalletRefundInput, - ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, ReclaimWalletRechargeCheckoutInput, - RedeemWalletCodeInput, RedeemWalletCodeOutcome, StoredAdminPaymentCallback, - StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, - StoredAdminRedeemCode, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, - StoredAdminRedeemCodePage, StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage, - StoredAdminWalletListItem, StoredAdminWalletListPage, StoredAdminWalletRefund, - StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem, - StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction, - StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger, - StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput, - UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome, - WalletReadRepository, WalletWriteRepository, -}; -use aether_data_contracts::DataLayerError; -use async_trait::async_trait; -use chrono::Utc; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; - -#[derive(Debug, Clone)] -pub struct SqliteWalletReadRepository { - pool: SqlitePool, -} - -impl SqliteWalletReadRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl WalletReadRepository for SqliteWalletReadRepository { - async fn find( - &self, - key: WalletLookupKey<'_>, - ) -> Result, DataLayerError> { - 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( - &self, - user_id: &str, - limit_mode: &str, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE wallets SET limit_mode = ?, updated_at = ? WHERE user_id = ?") - .bind(limit_mode) - .bind(current_unix_secs_i64()) - .bind(user_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find(WalletLookupKey::UserId(user_id)).await - } - - async fn update_auth_api_key_wallet_limit_mode( - &self, - api_key_id: &str, - limit_mode: &str, - ) -> Result, DataLayerError> { - let result = - sqlx::query("UPDATE wallets SET limit_mode = ?, updated_at = ? WHERE api_key_id = ?") - .bind(limit_mode) - .bind(current_unix_secs_i64()) - .bind(api_key_id) - .execute(&self.pool) - .await - .map_sql_err()?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.find(WalletLookupKey::ApiKeyId(api_key_id)).await - } - - async fn initialize_auth_user_wallet( - &self, - user_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_sqlite_auth_wallet(&self.pool, Some(user_id), None, initial_gift_usd, unlimited) - .await - .map(|result| result.map(|(wallet, _created)| wallet)) - } - - async fn initialize_auth_user_wallet_with_outcome( - &self, - user_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_sqlite_auth_wallet(&self.pool, Some(user_id), None, initial_gift_usd, unlimited) - .await - .map(|result| { - result.map(|(wallet, created)| InitializeAuthWalletOutcome { wallet, created }) - }) - } - - async fn initialize_auth_api_key_wallet( - &self, - api_key_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_sqlite_auth_wallet( - &self.pool, - None, - Some(api_key_id), - initial_gift_usd, - unlimited, - ) - .await - .map(|result| result.map(|(wallet, _created)| wallet)) - } - - async fn initialize_auth_api_key_wallet_with_outcome( - &self, - api_key_id: &str, - initial_gift_usd: f64, - unlimited: bool, - ) -> Result, DataLayerError> { - initialize_sqlite_auth_wallet( - &self.pool, - None, - Some(api_key_id), - initial_gift_usd, - unlimited, - ) - .await - .map(|result| { - result.map(|(wallet, created)| InitializeAuthWalletOutcome { wallet, created }) - }) - } - - async fn update_auth_user_wallet_snapshot( - &self, - user_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, - ) -> Result, DataLayerError> { - update_sqlite_wallet_snapshot( - &self.pool, - "user_id", - user_id, - balance, - gift_balance, - limit_mode, - currency, - status, - total_recharged, - total_consumed, - total_refunded, - total_adjusted, - updated_at_unix_secs, - ) - .await?; - self.find(WalletLookupKey::UserId(user_id)).await - } - - async fn update_auth_api_key_wallet_snapshot( - &self, - api_key_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, - ) -> Result, DataLayerError> { - update_sqlite_wallet_snapshot( - &self.pool, - "api_key_id", - api_key_id, - balance, - gift_balance, - limit_mode, - currency, - status, - total_recharged, - total_consumed, - total_refunded, - total_adjusted, - updated_at_unix_secs, - ) - .await?; - self.find(WalletLookupKey::ApiKeyId(api_key_id)).await - } - - async fn list_wallets_by_user_ids( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - if user_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - 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 -WHERE user_id IN ( -"#, - ); - let mut separated = builder.separated(", "); - for user_id in user_ids { - separated.push_bind(user_id); - } - separated.push_unseparated(")"); - 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> { - if api_key_ids.is_empty() { - return Ok(Vec::new()); - } - let mut builder = QueryBuilder::::new( - 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 -WHERE api_key_id IN ( -"#, - ); - let mut separated = builder.separated(", "); - for api_key_id in api_key_ids { - separated.push_bind(api_key_id); - } - separated.push_unseparated(")"); - 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 { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM wallets -WHERE (? IS NULL OR status = ?) - AND ( - ? IS NULL - OR (? = 'user' AND user_id IS NOT NULL) - OR (? = 'api_key' AND api_key_id IS NOT NULL) - ) -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let rows = sqlx::query( - r#" -SELECT - 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 (? IS NULL OR w.status = ?) - AND ( - ? IS NULL - OR (? = 'user' AND w.user_id IS NOT NULL) - OR (? = 'api_key' AND w.api_key_id IS NOT NULL) - ) -ORDER BY w.updated_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(i64_from_usize(query.limit, "wallet limit")?) - .bind(i64_from_usize(query.offset, "wallet offset")?) - .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 { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM wallet_transactions tx -JOIN wallets w ON w.id = tx.wallet_id -WHERE (? IS NULL OR tx.category = ?) - AND (? IS NULL OR tx.reason_code = ?) - AND ( - ? IS NULL - OR (? = 'user' AND w.user_id IS NOT NULL) - OR (? = 'api_key' AND w.api_key_id IS NOT NULL) - ) -"#, - ) - .bind(query.category.as_deref()) - .bind(query.category.as_deref()) - .bind(query.reason_code.as_deref()) - .bind(query.reason_code.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .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, - 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 (? IS NULL OR tx.category = ?) - AND (? IS NULL OR tx.reason_code = ?) - AND ( - ? IS NULL - OR (? = 'user' AND w.user_id IS NOT NULL) - OR (? = 'api_key' AND w.api_key_id IS NOT NULL) - ) -ORDER BY tx.created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.category.as_deref()) - .bind(query.category.as_deref()) - .bind(query.reason_code.as_deref()) - .bind(query.reason_code.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(query.owner_type.as_deref()) - .bind(i64_from_usize(query.limit, "wallet ledger limit")?) - .bind(i64_from_usize(query.offset, "wallet ledger offset")?) - .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 { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM refund_requests rr -JOIN wallets w ON w.id = rr.wallet_id -WHERE (? IS NULL OR rr.status = ?) - AND w.user_id IS NOT NULL -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let rows = sqlx::query( - 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 (? IS NULL OR rr.status = ?) - AND w.user_id IS NOT NULL -ORDER BY rr.created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .bind(i64_from_usize(query.limit, "wallet refund request limit")?) - .bind(i64_from_usize( - query.offset, - "wallet refund request offset", - )?) - .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( - &self, - wallet_id: &str, - limit: usize, - offset: usize, - ) -> Result { - 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 -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( - &self, - wallet_id: &str, - billing_timezone: &str, - ) -> Result, DataLayerError> { - let billing_date = current_billing_date(billing_timezone)?; - let sql = daily_usage_select_sql("AND billing_date = ? LIMIT 1"); - let row = sqlx::query(&sql) - .bind(wallet_id) - .bind(billing_timezone) - .bind(billing_date) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_daily_usage_row).transpose() - } - - async fn list_wallet_daily_usage_history( - &self, - wallet_id: &str, - billing_timezone: &str, - limit: usize, - ) -> Result { - let billing_date = current_billing_date(billing_timezone)?; - let total: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_timezone = ? - AND billing_date < ? -"#, - ) - .bind(wallet_id) - .bind(billing_timezone) - .bind(&billing_date) - .fetch_one(&self.pool) - .await - .map_sql_err()?; - - let sql = daily_usage_select_sql("AND billing_date < ? ORDER BY billing_date DESC LIMIT ?"); - let rows = sqlx::query(&sql) - .bind(wallet_id) - .bind(billing_timezone) - .bind(billing_date) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_daily_usage_row) - .collect::, _>>()?; - Ok(StoredWalletDailyUsageLedgerPage { - items, - total: total.max(0) as u64, - }) - } - - async fn list_admin_wallet_refunds( - &self, - wallet_id: &str, - limit: usize, - offset: usize, - ) -> Result { - 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 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 { - let now = current_unix_secs_i64(); - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM payment_orders -WHERE (? IS NULL OR payment_method = ?) - AND ( - ? IS NULL - OR ( - CASE - WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at <= ? THEN 'expired' - ELSE status - END - ) = ? - ) -"#, - ) - .bind(query.payment_method.as_deref()) - .bind(query.payment_method.as_deref()) - .bind(query.status.as_deref()) - .bind(now) - .bind(query.status.as_deref()) - .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, 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 (? IS NULL OR payment_method = ?) - AND ( - ? IS NULL - OR ( - CASE - WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at <= ? THEN 'expired' - ELSE status - END - ) = ? - ) -ORDER BY created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.payment_method.as_deref()) - .bind(query.payment_method.as_deref()) - .bind(query.status.as_deref()) - .bind(now) - .bind(query.status.as_deref()) - .bind(i64_from_usize(query.limit, "payment order limit")?) - .bind(i64_from_usize(query.offset, "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 find_admin_payment_order( - &self, - order_id: &str, - ) -> Result, DataLayerError> { - 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( - &self, - user_id: &str, - limit: usize, - offset: usize, - ) -> Result { - let total = read_count_row( - sqlx::query("SELECT COUNT(*) AS total FROM payment_orders WHERE user_id = ? AND order_kind = 'wallet_recharge'") - .bind(user_id) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let now = current_unix_secs_i64(); - 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 = ? - AND order_kind = 'wallet_recharge' -ORDER BY created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(now) - .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 { - 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 { - 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( - &self, - user_id: &str, - order_id: &str, - ) -> Result, DataLayerError> { - 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 = ? - AND order_kind = 'wallet_recharge' -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_wallet_recharge_order_by_order_no( - &self, - user_id: &str, - order_no: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql( - "WHERE user_id = ? AND order_no = ? AND order_kind = 'wallet_recharge' LIMIT 1", - ); - let row = sqlx::query(&sql) - .bind(user_id) - .bind(order_no) - .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( - &self, - user_id: &str, - product_id: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql( - r#" -WHERE user_id = ? - AND product_id = ? - AND order_kind = 'plan_purchase' - AND status = 'pending' - AND expires_at > ? -ORDER BY created_at DESC -LIMIT 1 -"#, - ); - let row = sqlx::query(&sql) - .bind(user_id) - .bind(product_id) - .bind(current_unix_secs_i64()) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_payment_order_row).transpose() - } - - async fn find_payment_order_by_order_no( - &self, - order_no: &str, - ) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE order_no = ? LIMIT 1"); - let row = sqlx::query(&sql) - .bind(order_no) - .fetch_optional(&self.pool) - .await - .map_sql_err()?; - row.as_ref().map(map_payment_order_row).transpose() - } - - async fn find_wallet_refund( - &self, - wallet_id: &str, - refund_id: &str, - ) -> Result, DataLayerError> { - 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( - &self, - payment_method: Option<&str>, - limit: usize, - offset: usize, - ) -> Result { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM payment_callbacks -WHERE (? IS NULL OR payment_method = ?) -"#, - ) - .bind(payment_method) - .bind(payment_method) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let rows = sqlx::query( - 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 -WHERE (? IS NULL OR payment_method = ?) -ORDER BY created_at DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(payment_method) - .bind(payment_method) - .bind(i64_from_usize(limit, "payment callback limit")?) - .bind(i64_from_usize(offset, "payment callback offset")?) - .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 { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM redeem_code_batches -WHERE (? IS NULL OR status = ?) -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let rows = sqlx::query( - r#" -SELECT - batches.id, batches.name, batches.amount_usd, batches.currency, - batches.balance_bucket, batches.total_count, - COALESCE(stats.redeemed_count, 0) AS redeemed_count, - COALESCE(stats.active_count, 0) 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 ( - 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 (? IS NULL OR batches.status = ?) -ORDER BY batches.created_at DESC, batches.id DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .bind(i64_from_usize(query.limit, "redeem code batch limit")?) - .bind(i64_from_usize(query.offset, "redeem code batch offset")?) - .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> { - let row = sqlx::query( - r#" -SELECT - batches.id, batches.name, batches.amount_usd, batches.currency, - batches.balance_bucket, batches.total_count, - COALESCE(stats.redeemed_count, 0) AS redeemed_count, - COALESCE(stats.active_count, 0) 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 ( - 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 batches.id = ? -LIMIT 1 -"#, - ) - .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 { - let total = read_count_row( - sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM redeem_codes -WHERE batch_id = ? - AND (? IS NULL OR status = ?) -"#, - ) - .bind(&query.batch_id) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .fetch_one(&self.pool) - .await - .map_sql_err()?, - )?; - let sql = redeem_code_select_sql( - "WHERE codes.batch_id = ? AND (? IS NULL OR codes.status = ?) ORDER BY codes.created_at DESC, codes.id DESC LIMIT ? OFFSET ?", - ); - let rows = sqlx::query(&sql) - .bind(&query.batch_id) - .bind(query.status.as_deref()) - .bind(query.status.as_deref()) - .bind(i64_from_usize(query.limit, "redeem code limit")?) - .bind(i64_from_usize(query.offset, "redeem code offset")?) - .fetch_all(&self.pool) - .await - .map_sql_err()?; - let items = rows - .iter() - .map(map_redeem_code_row) - .collect::, _>>()?; - Ok(StoredAdminRedeemCodePage { items, total }) - } -} - -#[async_trait] -impl WalletWriteRepository for SqliteWalletReadRepository { - async fn delete_wallet_if_unreferenced( - &self, - wallet_id: &str, - owner: WalletLookupKey<'_>, - ) -> Result { - if wallet_id.trim().is_empty() { - return Ok(false); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet compensation requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - // The reference predicates and the delete must run while holding SQLite's single - // writer slot. A deferred transaction could observe an unreferenced wallet, then let a - // concurrent writer attach a financial row before the delete is issued. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let select_sql = format!( - r#" -SELECT id -FROM wallets -WHERE id = ? - AND {owner_clause} - AND balance = 0 - AND gift_balance = 0 - AND total_recharged = 0 - AND total_consumed = 0 - AND total_refunded = 0 - AND total_adjusted = 0 - AND limit_mode IN ('finite', 'unlimited') - AND currency = 'USD' - AND status = 'active' - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = wallets.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM "usage" u WHERE u.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = wallets.id - ) -LIMIT 1 - "# - ); - let found = sqlx::query_scalar::<_, String>(&select_sql) - .bind(wallet_id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(found_id) = found else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let removed = sqlx::query("DELETE FROM wallets WHERE id = ?") - .bind(&found_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn delete_wallet_if_snapshot_matches_and_unreferenced( - &self, - expected: &StoredWalletSnapshot, - owner: WalletLookupKey<'_>, - ) -> Result { - if expected.id.trim().is_empty() { - return Ok(false); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet compensation requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - // Serialize the snapshot check, reference check, and delete behind SQLite's writer lock. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let select_sql = wallet_select_sql(&format!( - r#"WHERE id = ? - AND {owner_clause} - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = wallets.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM "usage" u WHERE u.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = wallets.id - ) - AND NOT EXISTS (SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = wallets.id) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = wallets.id - ) -LIMIT 1"# - )); - let row = sqlx::query(&select_sql) - .bind(&expected.id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_wallet_row(&row)?; - if ¤t != expected { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let removed = sqlx::query("DELETE FROM wallets WHERE id = ?") - .bind(&expected.id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn restore_wallet_if_snapshot_matches( - &self, - before: &StoredWalletSnapshot, - after: &StoredWalletSnapshot, - owner: WalletLookupKey<'_>, - ) -> Result { - if before.id.trim().is_empty() || after.id.trim().is_empty() { - return Ok(false); - } - if before.id != after.id { - return Err(DataLayerError::InvalidInput( - "wallet restore snapshots must reference the same wallet".to_string(), - )); - } - let (owner_clause, owner_id) = match owner { - WalletLookupKey::UserId(user_id) if !user_id.trim().is_empty() => { - ("user_id = ? AND api_key_id IS NULL", user_id) - } - WalletLookupKey::ApiKeyId(api_key_id) if !api_key_id.trim().is_empty() => { - ("api_key_id = ? AND user_id IS NULL", api_key_id) - } - WalletLookupKey::WalletId(_) => { - return Err(DataLayerError::InvalidInput( - "wallet restore requires an explicit user or API-key owner".to_string(), - )) - } - _ => return Ok(false), - }; - let owner_matches = match owner { - WalletLookupKey::UserId(user_id) => { - before.user_id.as_deref() == Some(user_id) - && after.user_id.as_deref() == Some(user_id) - && before.api_key_id.is_none() - && after.api_key_id.is_none() - } - WalletLookupKey::ApiKeyId(api_key_id) => { - before.api_key_id.as_deref() == Some(api_key_id) - && after.api_key_id.as_deref() == Some(api_key_id) - && before.user_id.is_none() - && after.user_id.is_none() - } - WalletLookupKey::WalletId(_) => false, - }; - if !owner_matches { - return Ok(false); - } - let before_updated_at = i64::try_from(before.updated_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput( - "wallet restore timestamp is outside the supported range".to_string(), - ) - })?; - - // BEGIN IMMEDIATE serializes the snapshot check and replacement with all SQLite wallet - // writers. No caller can change the row after it is read but before it is restored. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let select_sql = wallet_select_sql(&format!("WHERE id = ? AND {owner_clause} LIMIT 1")); - let row = sqlx::query(&select_sql) - .bind(&after.id) - .bind(owner_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(row) = row else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - let current = map_wallet_row(&row)?; - if current != *after { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - let updated = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - limit_mode = ?, - currency = ?, - status = ?, - total_recharged = ?, - total_consumed = ?, - total_refunded = ?, - total_adjusted = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(before.balance) - .bind(before.gift_balance) - .bind(&before.limit_mode) - .bind(&before.currency) - .bind(&before.status) - .bind(before.total_recharged) - .bind(before.total_consumed) - .bind(before.total_refunded) - .bind(before.total_adjusted) - .bind(before_updated_at) - .bind(&before.id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - if updated == 0 { - tx.rollback().await.map_sql_err()?; - return Ok(false); - } - tx.commit().await.map_sql_err()?; - Ok(true) - } - - async fn delete_provisional_auth_user_wallet( - &self, - wallet_id: &str, - user_id: &str, - ) -> Result { - if wallet_id.trim().is_empty() || user_id.trim().is_empty() { - return Ok(false); - } - // Keep the eligibility read and the compensating delete atomic with respect to any new - // wallet/order/ledger writes (SQLite transactions are deferred by default). - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let found_wallet_id = sqlx::query_scalar::<_, String>( - r#" -SELECT w.id -FROM wallets AS w -WHERE w.id = ? - AND w.user_id = ? - AND w.api_key_id IS NULL - AND w.balance = 0 - AND w.gift_balance >= 0 - AND w.total_recharged = 0 - AND w.total_consumed = 0 - AND w.total_refunded = 0 - AND w.total_adjusted = w.gift_balance - AND w.limit_mode IN ('finite', 'unlimited') - AND w.currency = 'USD' - AND w.status = 'active' - AND NOT EXISTS (SELECT 1 FROM payment_orders p WHERE p.wallet_id = w.id) - AND NOT EXISTS (SELECT 1 FROM refund_requests r WHERE r.wallet_id = w.id) - AND NOT EXISTS ( - SELECT 1 FROM wallet_daily_usage_ledgers d WHERE d.wallet_id = w.id - ) - AND NOT EXISTS (SELECT 1 FROM "usage" u WHERE u.wallet_id = w.id) - AND NOT EXISTS ( - SELECT 1 FROM usage_settlement_snapshots s WHERE s.wallet_id = w.id - ) - AND NOT EXISTS ( - SELECT 1 FROM redeem_codes c WHERE c.redeemed_wallet_id = w.id - ) - AND ( - (w.gift_balance = 0 AND NOT EXISTS ( - SELECT 1 FROM wallet_transactions t WHERE t.wallet_id = w.id - )) - OR - (w.gift_balance > 0 - AND (SELECT COUNT(*) FROM wallet_transactions t WHERE t.wallet_id = w.id) = 1 - AND EXISTS ( - SELECT 1 FROM wallet_transactions t - WHERE t.wallet_id = w.id - AND t.category = 'gift' - AND t.reason_code = 'gift_initial' - AND t.amount = w.gift_balance - AND t.balance_before = 0 - AND t.balance_after = w.gift_balance - AND t.recharge_balance_before = 0 - AND t.recharge_balance_after = 0 - AND t.gift_balance_before = 0 - AND t.gift_balance_after = w.gift_balance - AND t.link_type = 'system_task' - AND t.link_id = w.user_id - AND t.operator_id IS NULL - )) - ) -LIMIT 1 - "#, - ) - .bind(wallet_id) - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(found_wallet_id) = found_wallet_id else { - tx.rollback().await.map_sql_err()?; - return Ok(false); - }; - sqlx::query("DELETE FROM wallet_transactions WHERE wallet_id = ?") - .bind(&found_wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let removed = sqlx::query("DELETE FROM wallets WHERE id = ? AND user_id = ?") - .bind(&found_wallet_id) - .bind(user_id) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected() - > 0; - tx.commit().await.map_sql_err()?; - Ok(removed) - } - - async fn create_wallet_recharge_order( - &self, - mut input: CreateWalletRechargeOrderInput, - ) -> Result { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - validate_wallet_recharge_order_input(&input).map_err(DataLayerError::InvalidInput)?; - if !input.amount_usd.is_finite() - || input.amount_usd <= 0.0 - || input - .pay_amount - .is_some_and(|value| !value.is_finite() || value <= 0.0) - || input - .exchange_rate - .is_some_and(|value| !value.is_finite() || value <= 0.0) - { - return Err(DataLayerError::InvalidInput( - "invalid wallet recharge numeric fields".to_string(), - )); - } - let projected_gateway_response = - project_wallet_recharge_gateway_response(&input.gateway_response) - .map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("wallet recharge expires_at overflow".to_string()) - })?; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - // Serialize wallet creation and order idempotency checks with the insert. A deferred - // transaction can let concurrent callers race through the read-before-write section. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - - // The wallet schema intentionally keeps `user_id` nullable for - // deleted-user history, so it cannot protect this creation path with - // a mandatory foreign key. Validate the owner while holding the - // writer transaction before creating either the wallet or order. - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput("user not found".to_string())); - } - - let wallet_row = sqlx::query( - r#" -SELECT id, status -FROM wallets -WHERE user_id = ? -LIMIT 1 -"#, - ) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (wallet_id, wallet_status, created_wallet) = if let Some(row) = wallet_row { - ( - get::(&row, "id")?, - get::(&row, "status")?, - false, - ) - } else { - let requested_wallet_id = input - .preferred_wallet_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?) -ON CONFLICT DO NOTHING -"#, - ) - .bind(&requested_wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if insert_result.rows_affected() == 0 { - let Some(row) = sqlite_wallet_by_user_id(&mut tx, &input.user_id).await? else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - }; - ( - get::(&row, "id")?, - get::(&row, "status")?, - false, - ) - } else { - (requested_wallet_id, "active".to_string(), true) - } - }; - if wallet_status != "active" { - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::WalletInactive); - } - - if let Some(existing_row) = - sqlite_payment_order_by_order_no(&mut tx, &input.order_no).await? - { - let existing_user_id: Option = existing_row.try_get("user_id").map_sql_err()?; - let existing_kind: String = existing_row.try_get("order_kind").map_sql_err()?; - if existing_user_id.as_deref() == Some(input.user_id.as_str()) - && existing_kind == "wallet_recharge" - { - if !sqlite_wallet_recharge_replay_matches(&existing_row, &wallet_id, &input)? { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge replay changes immutable order fields".to_string(), - )); - } - let existing = map_payment_order_row(&existing_row)?; - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::Existing(existing)); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment order number already belongs to another order".to_string(), - )); - } - - let order_id = uuid::Uuid::new_v4().to_string(); - let insert_result = sqlx::query( - r#" -INSERT INTO payment_orders ( - 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, fulfillment_status, - gateway_order_id, gateway_response, status, created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?, 'wallet_recharge', 'pending', ?, ?, 'pending', ?, ?) -ON CONFLICT DO NOTHING -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(input.amount_usd) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(&input.payment_method) - .bind(input.payment_provider.as_deref()) - .bind(input.payment_channel.as_deref()) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(now) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - if insert_result.rows_affected() == 0 { - if let Some(existing_row) = - sqlite_payment_order_by_order_no(&mut tx, &input.order_no).await? - { - let existing_user_id: Option = - existing_row.try_get("user_id").map_sql_err()?; - let existing_kind: String = existing_row.try_get("order_kind").map_sql_err()?; - let existing = map_payment_order_row(&existing_row)?; - if existing_user_id.as_deref() == Some(input.user_id.as_str()) - && existing_kind == "wallet_recharge" - { - if !sqlite_wallet_recharge_replay_matches(&existing_row, &wallet_id, &input)? { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge replay changes immutable order fields".to_string(), - )); - } - if created_wallet { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - return Ok(CreateWalletRechargeOrderOutcome::Existing(existing)); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment order number already belongs to another order".to_string(), - )); - } - if sqlite_payment_order_by_gateway_order_id( - &mut tx, - &input.payment_method, - &input.gateway_order_id, - ) - .await? - .is_some() - { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "payment gateway order already belongs to another order".to_string(), - )); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet recharge order could not be created".to_string(), - )); - } - - let row = sqlite_payment_order_by_id(&mut tx, &order_id).await?; - tx.commit().await.map_sql_err()?; - Ok(CreateWalletRechargeOrderOutcome::Created( - map_payment_order_row(&row)?, - )) - } - - async fn update_wallet_recharge_checkout( - &self, - input: UpdateWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - if input.order_id.trim().is_empty() || input.gateway_order_id.trim().is_empty() { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout identifiers are required".to_string(), - )); - } - let projected_gateway_response = - match project_wallet_recharge_gateway_response(&input.gateway_response) { - Ok(value) => value, - Err(error) => return Ok(WalletMutationOutcome::Invalid(error)), - }; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(current_row) = - sqlite_payment_order_by_id_optional(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order_kind: Option = get(¤t_row, "order_kind")?; - if order_kind.as_deref() != Some("wallet_recharge") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not a wallet recharge".to_string(), - )); - } - let current = map_payment_order_row(¤t_row)?; - let current_is_checkout_placeholder = - wallet_recharge_order_is_checkout_placeholder(¤t); - let current_token = current - .gateway_response - .as_ref() - .and_then(wallet_recharge_checkout_claim_token); - let requested_token = wallet_recharge_checkout_claim_token(&projected_gateway_response); - if current_token.is_some() && current_token != requested_token { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - if current.status != "pending" { - if current.gateway_order_id.as_deref() == Some(input.gateway_order_id.as_str()) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(current)); - } - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is no longer pending".to_string(), - )); - } - let now = current_unix_secs_i64(); - if current - .expires_at_unix_secs - .is_none_or(|expires_at| expires_at <= now.max(0) as u64) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is expired".to_string(), - )); - } - // A newly-created row uses order_no as a temporary gateway id. Once - // the provider checkout is stored, do not let a concurrent request - // replace that checkout evidence. - if current.gateway_order_id.as_deref().is_some_and(|existing| { - existing != input.gateway_order_id.as_str() - && existing != current.order_no.as_str() - && !current_is_checkout_placeholder - }) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout is already bound".to_string(), - )); - } - let conflict = sqlite_payment_order_by_gateway_order_id( - &mut tx, - ¤t.payment_method, - &input.gateway_order_id, - ) - .await? - .is_some_and(|row| { - row.try_get::("id").ok().as_deref() != Some(input.order_id.as_str()) - }); - if conflict { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment gateway order already belongs to another order".to_string(), - )); - } - let updated = sqlx::query( - "UPDATE payment_orders SET gateway_order_id = ?, gateway_response = ? WHERE id = ? AND status = 'pending' AND expires_at > ?", - ) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(&input.order_id) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge order is expired or no longer pending".to_string(), - )); - } - let updated = - map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn compare_and_swap_payment_order_stripe_client_secret( - &self, - input: CompareAndSwapPaymentOrderStripeClientSecretInput, - ) -> Result { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(row) = sqlite_payment_order_by_id_optional(&mut tx, &input.order_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(false); - }; - let current = map_payment_order_row(&row)?; - let Some(replacement) = - payment_order_stripe_client_secret_cas_replacement(¤t, &input) - .map_err(DataLayerError::InvalidInput)? - else { - tx.commit().await.map_sql_err()?; - return Ok(false); - }; - let replacement = json_string(&replacement, "payment_orders.gateway_response")?; - let updated = sqlx::query("UPDATE payment_orders SET gateway_response = ? WHERE id = ?") - .bind(replacement) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(updated.rows_affected() == 1) - } - - async fn fail_wallet_recharge_checkout( - &self, - input: FailWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - if input.order_id.trim().is_empty() - || input.claim_token.trim().is_empty() - || input.claim_token.len() > 128 - { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout failure identifiers are required".to_string(), - )); - } - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(row) = sqlite_payment_order_by_id_optional(&mut tx, &input.order_id).await? else { - tx.rollback().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if !wallet_recharge_order_is_checkout_placeholder(&order) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not a checkout placeholder".to_string(), - )); - } - let current_token = order - .gateway_response - .as_ref() - .and_then(wallet_recharge_checkout_claim_token); - if current_token != Some(input.claim_token.trim()) { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - if order.status != "pending" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(order)); - } - let failed = if input.provider_request_may_have_succeeded { - wallet_recharge_checkout_uncertain_response( - order.gateway_response.as_ref(), - &input.reason, - current_unix_secs_i64().max(0) as u64, - ) - } else { - wallet_recharge_checkout_failed_response( - order.gateway_response.as_ref(), - &input.reason, - current_unix_secs_i64().max(0) as u64, - ) - }; - let failed = json_string(&failed, "payment_orders.gateway_response")?; - let updated = sqlx::query( - "UPDATE payment_orders SET status = 'failed', gateway_response = ? WHERE id = ? AND status = 'pending'", - ) - .bind(failed) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout claim is no longer current".to_string(), - )); - } - let updated = - map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn reclaim_wallet_recharge_checkout( - &self, - input: ReclaimWalletRechargeCheckoutInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64().max(0) as u64; - if input.order_id.trim().is_empty() - || input.claim_token.trim().is_empty() - || input.claim_token.len() > 128 - || input.expires_at_unix_secs <= now - || input.expires_at_unix_secs > i64::MAX as u64 - { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout reclaim identifiers are invalid".to_string(), - )); - } - if !wallet_recharge_response_is_checkout_placeholder(&input.gateway_response) { - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge reclaim response must be a placeholder".to_string(), - )); - } - let response = wallet_recharge_checkout_claim_response( - &input.gateway_response, - &input.claim_token, - now, - ) - .map_err(DataLayerError::InvalidInput)?; - let response = json_string(&response, "payment_orders.gateway_response")?; - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("wallet recharge expires_at overflow".to_string()) - })?; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(row) = sqlite_payment_order_by_id_optional(&mut tx, &input.order_id).await? else { - tx.rollback().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order_kind: Option = get(&row, "order_kind")?; - let order = map_payment_order_row(&row)?; - if order_kind.as_deref() != Some("wallet_recharge") - || !wallet_recharge_order_is_reclaimable_placeholder(&order, now) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout is still in progress or already completed".to_string(), - )); - } - let updated = sqlx::query( - "UPDATE payment_orders SET gateway_order_id = ?, gateway_response = ?, status = 'pending', expires_at = ? WHERE id = ?", - ) - .bind(&order.order_no) - .bind(response) - .bind(expires_at) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if updated.rows_affected() == 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet recharge checkout reclaim lost the order race".to_string(), - )); - } - let updated = - map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn create_plan_purchase_order( - &self, - mut input: CreatePlanPurchaseOrderInput, - ) -> Result { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - validate_plan_purchase_order_input(&input).map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let expires_at = i64::try_from(input.expires_at_unix_secs).map_err(|_| { - DataLayerError::InvalidInput("plan purchase expires_at overflow".to_string()) - })?; - let projected_gateway_response = project_wallet_gateway_response(&input.gateway_response) - .map_err(DataLayerError::InvalidInput)?; - let gateway_response = json_string( - &projected_gateway_response, - "payment_orders.gateway_response", - )?; - let product_snapshot = - json_string(&input.product_snapshot, "payment_orders.product_snapshot")?; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - - // The baseline SQLite schema does not enforce the wallet -> user - // relationship. Validate the order owner before creating an automatic - // wallet, so an invalid checkout cannot leave a financial row behind. - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput("user not found".to_string())); - } - - let wallet_row = sqlx::query( - r#" -SELECT id, status -FROM wallets -WHERE user_id = ? -LIMIT 1 -"#, - ) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let (wallet_id, wallet_status) = if let Some(row) = wallet_row { - (get::(&row, "id")?, get::(&row, "status")?) - } else { - let requested_wallet_id = input - .preferred_wallet_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - // `user_id` is unique in the portable schema, but the explicit - // read above can still race with another initializer on a - // connection that started before it committed. Treat either a - // user or wallet-id conflict as a no-op, then resolve the winner - // by owner instead of leaking a database constraint error. - let insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, ?, ?) -ON CONFLICT DO NOTHING -"#, - ) - .bind(&requested_wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if insert_result.rows_affected() == 0 { - let Some(row) = sqlite_wallet_by_user_id(&mut tx, &input.user_id).await? else { - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - }; - (get::(&row, "id")?, get::(&row, "status")?) - } else { - (requested_wallet_id, "active".to_string()) - } - }; - if wallet_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(CreatePlanPurchaseOrderOutcome::WalletInactive); - } - - let purchase_limit_scope = plan_purchase_limit_scope(&input.product_snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&input.product_snapshot); - let mut active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' -"#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? -"#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - active_count += sqlx::query_scalar::<_, i64>( - r#" - SELECT COUNT(*) - FROM payment_orders - WHERE user_id = ? - AND product_id = ? - AND order_kind = 'plan_purchase' - AND status = 'pending' - AND expires_at > ? - "#, - ) - .bind(&input.user_id) - .bind(&input.product_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - if active_count >= max_active_per_user { - tx.commit().await.map_sql_err()?; - return Ok(CreatePlanPurchaseOrderOutcome::ActivePlanLimitReached); - } - } - - let order_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO payment_orders ( - 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, - fulfillment_status, gateway_order_id, gateway_response, status, created_at, expires_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?, 'plan_purchase', ?, ?, 'pending', ?, ?, 'pending', ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(input.amount_usd) - .bind(input.pay_amount) - .bind(&input.pay_currency) - .bind(input.exchange_rate) - .bind(&input.payment_method) - .bind(input.payment_provider.as_deref()) - .bind(input.payment_channel.as_deref()) - .bind(&input.product_id) - .bind(product_snapshot) - .bind(&input.gateway_order_id) - .bind(gateway_response) - .bind(now) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let row = sqlite_payment_order_by_id(&mut tx, &order_id).await?; - tx.commit().await.map_sql_err()?; - Ok(CreatePlanPurchaseOrderOutcome::Created( - map_payment_order_row(&row)?, - )) - } - - async fn create_wallet_refund_request( - &self, - input: CreateWalletRefundRequestInput, - ) -> Result { - if !input.amount_usd.is_finite() || input.amount_usd <= 0.0 { - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "refund amount must be finite and greater than zero".to_string(), - )); - } - let now = current_unix_secs_i64(); - // Hold SQLite's single writer slot while checking existing reservations and inserting - // the new request. A deferred transaction would allow two callers to both observe the - // same available balance before either reservation is written. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - - let Some(wallet_row) = sqlx::query( - r#" -SELECT id, balance -FROM wallets -WHERE id = ? - AND user_id = ? -LIMIT 1 -"#, - ) - .bind(&input.wallet_id) - .bind(&input.user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::WalletMissing); - }; - - if let Some(idempotency_key) = input.idempotency_key.as_deref() { - let existing = - sqlite_refund_by_idempotency(&mut tx, &input.user_id, idempotency_key).await?; - if let Some(row) = existing { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::Duplicate(map_refund_row( - &row, - )?)); - } - } - - let wallet_recharge_balance = sqlite_real(&wallet_row, "balance")?; - if !wallet_recharge_balance.is_finite() { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "wallet recharge balance is invalid".to_string(), - )); - } - let wallet_reserved_amount = sqlx::query_scalar::<_, Option>( - r#" -SELECT amount_usd -FROM refund_requests -WHERE wallet_id = ? - AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(&input.wallet_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()? - .into_iter() - .try_fold(0.0_f64, |total, amount| { - let amount = amount?; - if !amount.is_finite() || amount <= 0.0 { - return None; - } - let next = total + amount; - next.is_finite().then_some(next) - }); - let Some(wallet_reserved_amount) = wallet_reserved_amount else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "wallet refund reservation is invalid".to_string(), - )); - }; - if input.amount_usd > (wallet_recharge_balance - wallet_reserved_amount) { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::RefundAmountExceedsAvailableBalance); - } - - let mut payment_order_id = None; - let mut resolved_payment_method = None; - if let Some(order_id) = input.payment_order_id.as_deref() { - let Some(order_row) = sqlx::query( - r#" -SELECT id, status, payment_method, amount_usd, refunded_amount_usd, refundable_amount_usd -FROM payment_orders -WHERE id = ? - AND wallet_id = ? -LIMIT 1 -"#, - ) - .bind(order_id) - .bind(&input.wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::PaymentOrderNotFound); - }; - let status: String = get(&order_row, "status")?; - if status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::PaymentOrderNotRefundable); - } - let order_reserved_amount = sqlx::query_scalar::<_, Option>( - r#" -SELECT amount_usd -FROM refund_requests -WHERE payment_order_id = ? - AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(order_id) - .fetch_all(&mut *tx) - .await - .map_sql_err()? - .into_iter() - .try_fold(0.0_f64, |total, amount| { - let amount = amount?; - if !amount.is_finite() || amount <= 0.0 { - return None; - } - let next = total + amount; - next.is_finite().then_some(next) - }); - let Some(order_reserved_amount) = order_reserved_amount else { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "payment order refund reservation is invalid".to_string(), - )); - }; - let order_amount = sqlite_real(&order_row, "amount_usd")?; - let refunded_amount = sqlite_real(&order_row, "refunded_amount_usd")?; - let refundable_amount = sqlite_real(&order_row, "refundable_amount_usd")?; - if !payment_order_refund_amounts_are_consistent( - order_amount, - refunded_amount, - refundable_amount, - ) { - tx.commit().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::InvalidInput( - "payment order refund amounts are invalid".to_string(), - )); - } - if input.amount_usd > (refundable_amount - order_reserved_amount) { - tx.commit().await.map_sql_err()?; - return Ok( - CreateWalletRefundRequestOutcome::RefundAmountExceedsAvailableOrderAmount, - ); - } - payment_order_id = Some(order_id.to_string()); - resolved_payment_method = Some(get::(&order_row, "payment_method")?); - } - - let canonical = canonicalize_wallet_refund_fields( - payment_order_id.as_deref(), - input.source_type.as_deref(), - input.source_id.as_deref(), - input.refund_mode.as_deref(), - resolved_payment_method.as_deref(), - ) - .map_err(DataLayerError::InvalidInput)?; - let source_type = canonical.source_type; - let source_id = canonical.source_id; - let refund_mode = canonical.refund_mode; - - let refund_id = uuid::Uuid::new_v4().to_string(); - let insert = sqlx::query( - r#" -INSERT INTO refund_requests ( - id, refund_no, wallet_id, user_id, payment_order_id, source_type, source_id, - refund_mode, amount_usd, status, reason, requested_by, idempotency_key, - created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending_approval', ?, ?, ?, ?, ?) -"#, - ) - .bind(&refund_id) - .bind(&input.refund_no) - .bind(&input.wallet_id) - .bind(&input.user_id) - .bind(payment_order_id.as_deref()) - .bind(&source_type) - .bind(source_id.as_deref()) - .bind(&refund_mode) - .bind(input.amount_usd) - .bind(input.reason.as_deref()) - .bind(&input.user_id) - .bind(input.idempotency_key.as_deref()) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await; - - if let Err(err) = insert { - if input.idempotency_key.is_some() - && err - .as_database_error() - .is_some_and(|database_error| database_error.is_unique_violation()) - { - tx.rollback().await.map_sql_err()?; - return Ok(CreateWalletRefundRequestOutcome::DuplicateRejected); - } - return Err(DataLayerError::sql(err)); - } - - let row = sqlite_refund_by_id(&mut tx, &refund_id).await?; - tx.commit().await.map_sql_err()?; - Ok(CreateWalletRefundRequestOutcome::Created(map_refund_row( - &row, - )?)) - } - - async fn process_payment_callback( - &self, - mut input: ProcessPaymentCallbackInput, - ) -> Result { - input - .canonicalize_and_validate() - .map_err(DataLayerError::InvalidInput)?; - if input.callback_key.trim().is_empty() - || input.callback_key.chars().count() > 128 - || input.payload_hash.trim().is_empty() - || !input.amount_usd.is_finite() - || input.amount_usd <= 0.0 - || input - .pay_amount - .is_some_and(|value| !value.is_finite() || value <= 0.0) - || input - .exchange_rate - .is_some_and(|value| !value.is_finite() || value <= 0.0) - { - return Err(DataLayerError::InvalidInput( - "invalid payment callback numeric or identity fields".to_string(), - )); - } - let now = current_unix_secs_i64(); - let payload = json_string(&input.payload, "payment_callbacks.payload")?; - // Callback processing reads and then mutates the callback, order, and - // wallet rows. Acquire SQLite's writer lock before the first read so - // a deferred transaction cannot observe `pending` and later fail while - // upgrading after another callback or admin credit has committed. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - - // Register the callback atomically. A preceding SELECT allows two - // concurrent first deliveries to both observe a missing key and then - // race on the unique constraint. Insert first and re-read by the - // key so every caller uses the row that actually won the race. - let candidate_callback_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -INSERT INTO payment_callbacks ( - id, payment_order_id, payment_method, callback_key, order_no, gateway_order_id, - payload_hash, signature_valid, status, payload, error_message, created_at, processed_at -) -VALUES (?, NULL, ?, ?, ?, ?, ?, ?, 'received', NULL, NULL, ?, NULL) -ON CONFLICT(callback_key) DO NOTHING -"#, - ) - .bind(&candidate_callback_id) - .bind(&input.payment_method) - .bind(&input.callback_key) - .bind(input.order_no.as_deref()) - .bind(input.gateway_order_id.as_deref()) - .bind(&input.payload_hash) - .bind(sqlite_bool(input.signature_valid)) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let callback_row = sqlx::query( - r#" -SELECT id, payment_order_id, payment_method, payload_hash, status, order_no, gateway_order_id -FROM payment_callbacks -WHERE callback_key = ? -LIMIT 1 -"#, - ) - .bind(&input.callback_key) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - let callback_id: String = get(&callback_row, "id")?; - let duplicate = callback_id != candidate_callback_id; - let callback_order_no: Option = get(&callback_row, "order_no")?; - let callback_gateway_order_id: Option = get(&callback_row, "gateway_order_id")?; - - let stored_method: String = get(&callback_row, "payment_method")?; - let stored_hash: Option = get(&callback_row, "payload_hash")?; - if !stored_method.eq_ignore_ascii_case(&input.payment_method) - || stored_hash.as_deref() != Some(input.payload_hash.as_str()) - { - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate: true, - error: "callback key reused with different payment payload".to_string(), - }); - } - let status: String = get(&callback_row, "status")?; - if status == "processed" { - let order_id: Option = get(&callback_row, "payment_order_id")?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::DuplicateProcessed { order_id }); - } - - if !input.signature_valid { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "invalid callback signature", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "invalid callback signature".to_string(), - }); - } - - let lookup_order_no = input.order_no.clone().or_else(|| callback_order_no.clone()); - let lookup_gateway_order_id = input - .gateway_order_id - .clone() - .or_else(|| callback_gateway_order_id.clone()); - let order_row = if let Some(order_no) = lookup_order_no.as_deref() { - sqlite_payment_order_by_order_no(&mut tx, order_no).await? - } else if let Some(gateway_order_id) = lookup_gateway_order_id.as_deref() { - sqlite_payment_order_by_gateway_order_id( - &mut tx, - &input.payment_method, - gateway_order_id, - ) - .await? - } else { - None - }; - let Some(order_row) = order_row else { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order not found".to_string(), - }); - }; - - let order_id: String = get(&order_row, "id")?; - let order_no: String = get(&order_row, "order_no")?; - let order_wallet_id: String = get(&order_row, "wallet_id")?; - let order_payment_method: String = get(&order_row, "payment_method")?; - let order_payment_provider: Option = get(&order_row, "payment_provider")?; - let order_payment_channel: Option = get(&order_row, "payment_channel")?; - let order_pay_currency: Option = get(&order_row, "pay_currency")?; - let order_gateway_order_id: Option = get(&order_row, "gateway_order_id")?; - let order_kind: String = get(&order_row, "order_kind")?; - let order_amount_usd = sqlite_real(&order_row, "amount_usd")?; - let order_pay_amount = sqlite_optional_real(&order_row, "pay_amount")?; - let order_exchange_rate = sqlite_optional_real(&order_row, "exchange_rate")?; - let order_status: String = get(&order_row, "status")?; - let expires_at_unix_secs: Option = get(&order_row, "expires_at_unix_secs")?; - let order_gateway_response = if order_status.eq_ignore_ascii_case("failed") { - optional_json( - get(&order_row, "gateway_response")?, - "payment_orders.gateway_response", - )? - } else { - None - }; - let failed_checkout_recoverable = payment_order_is_failed_wallet_checkout_placeholder( - &order_status, - &order_kind, - order_gateway_response.as_ref(), - ); - let uncertain_checkout = payment_order_is_uncertain_wallet_checkout_placeholder( - &order_status, - &order_kind, - order_gateway_response.as_ref(), - ); - if !order_amount_usd.is_finite() - || order_amount_usd <= 0.0 - || order_pay_amount.is_some_and(|value| !value.is_finite() || value <= 0.0) - { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order amount is invalid", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order amount is invalid".to_string(), - }); - } - - // A payment order must credit the wallet that belongs to the same - // live user. Do this before binding a gateway id or changing any - // entitlement, wallet, or order state. Legacy rows may violate the - // wallet owner XOR invariant, so reject every ambiguous shape here. - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(order_user_id) = order_user_id - .as_deref() - .filter(|value| !value.trim().is_empty()) - else { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order user missing", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order user missing".to_string(), - }); - }; - let Some(wallet_owner_row) = sqlx::query( - r#" -SELECT - w.user_id AS wallet_user_id, - w.api_key_id AS wallet_api_key_id, - api_keys.user_id AS api_key_user_id -FROM wallets AS w -LEFT JOIN api_keys ON api_keys.id = w.api_key_id -WHERE w.id = ? -LIMIT 1 - "#, - ) - .bind(&order_wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet not found".to_string(), - }); - }; - let wallet_user_id: Option = get(&wallet_owner_row, "wallet_user_id")?; - let wallet_api_key_id: Option = get(&wallet_owner_row, "wallet_api_key_id")?; - let api_key_user_id: Option = get(&wallet_owner_row, "api_key_user_id")?; - let wallet_owner_matches = match ( - wallet_user_id.as_deref(), - wallet_api_key_id.as_deref(), - api_key_user_id.as_deref(), - ) { - (Some(wallet_user_id), None, _) if !wallet_user_id.trim().is_empty() => { - wallet_user_id == order_user_id - } - (None, Some(wallet_api_key_id), Some(api_key_user_id)) - if !wallet_api_key_id.trim().is_empty() && !api_key_user_id.trim().is_empty() => - { - api_key_user_id == order_user_id - } - _ => false, - }; - if !wallet_owner_matches { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order wallet owner mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order wallet owner mismatch".to_string(), - }); - } - - // The lookup identifier is not proof that the callback belongs to - // this order: order_no takes precedence over gateway_order_id. Check - // every identifier supplied by this delivery (and any persisted - // fallback from the callback row) before changing the order or - // wallet. Orders created before the gateway returns a provider - // transaction id store order_no as a placeholder; that value may be - // replaced by a verified callback, but a real id must never be - // rebound to another order. - if input - .order_no - .as_deref() - .is_some_and(|value| value != order_no) - || callback_order_no - .as_deref() - .is_some_and(|value| value != order_no) - { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order number mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order number mismatch".to_string(), - }); - } - let input_gateway_order_id = input - .gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - let callback_gateway_order_id = callback_gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - let stored_real_gateway_order_id = order_gateway_order_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty() && *value != order_no); - let input_real_gateway_order_id = input_gateway_order_id.filter(|value| *value != order_no); - let callback_real_gateway_order_id = - callback_gateway_order_id.filter(|value| *value != order_no); - let effective_gateway_order_id = input_real_gateway_order_id - .or(callback_real_gateway_order_id) - .or(stored_real_gateway_order_id); - if let Some(expected_gateway_order_id) = stored_real_gateway_order_id { - if input_gateway_order_id.is_some_and(|value| value != expected_gateway_order_id) - || callback_gateway_order_id.is_some_and(|value| value != expected_gateway_order_id) - { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order mismatch".to_string(), - }); - } - } else if let (Some(input_gateway), Some(callback_gateway)) = - (input_real_gateway_order_id, callback_real_gateway_order_id) - { - if input_gateway != callback_gateway { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order identifier mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order identifier mismatch".to_string(), - }); - } - } - if stored_real_gateway_order_id.is_none() { - if let Some(gateway_order_id) = effective_gateway_order_id { - let conflicting_order_id: Option = sqlx::query_scalar( - "SELECT id FROM payment_orders WHERE payment_method = ? AND gateway_order_id = ? AND id <> ? LIMIT 1", - ) - .bind(&order_payment_method) - .bind(gateway_order_id) - .bind(&order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if conflicting_order_id.is_some() { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order belongs to another payment order", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order belongs to another payment order".to_string(), - }); - } - } - } - - let amount_matches = payment_callback_amount_matches_order( - order_amount_usd, - order_pay_amount, - order_pay_currency.as_deref(), - order_exchange_rate, - input.amount_usd, - input.pay_amount, - ); - if !amount_matches { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "callback amount mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "callback amount mismatch".to_string(), - }); - } - if !payment_callback_method_matches_order( - &order_payment_method, - order_payment_provider.as_deref(), - &input.payment_method, - input.payment_provider.as_deref(), - ) { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment method mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment method mismatch".to_string(), - }); - } - let payment_provider_matches = payment_callback_provider_matches_order( - &order_payment_method, - order_payment_provider.as_deref(), - &input.payment_method, - input.payment_provider.as_deref(), - ); - if !payment_provider_matches { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment provider mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment provider mismatch".to_string(), - }); - } - let currency_matches = match (input.pay_currency.as_deref(), order_pay_currency.as_deref()) - { - (Some(callback), Some(order)) => order.eq_ignore_ascii_case(callback), - (None, None) => input.pay_amount.is_none() && order_pay_amount.is_none(), - _ => false, - }; - if !currency_matches { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment currency mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment currency mismatch".to_string(), - }); - } - if let Some(expected_channel) = input.payment_channel.as_deref() { - let stored_channel = order_payment_channel.as_deref().or_else(|| { - (order_payment_provider.is_none() - && ["alipay", "wxpay"] - .iter() - .any(|method| method.eq_ignore_ascii_case(&order_payment_method))) - .then_some(order_payment_method.as_str()) - }); - if !stored_channel.is_some_and(|value| value.eq_ignore_ascii_case(expected_channel)) { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment channel mismatch", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment channel mismatch".to_string(), - }); - } - } - if order_status == "credited" { - mark_sqlite_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::AlreadyCredited { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - }); - } - if !matches!(order_status.as_str(), "pending" | "paid") && !failed_checkout_recoverable { - let error = format!("payment order is not creditable: {order_status}"); - update_sqlite_payment_callback_failure(&mut tx, &callback_id, &input, &payload, &error) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { duplicate, error }); - } - if (order_status == "pending" || (failed_checkout_recoverable && !uncertain_checkout)) - && expires_at_unix_secs.is_some_and(|value| value <= now) - { - sqlx::query("UPDATE payment_orders SET status = 'expired' WHERE id = ?") - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order expired", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order expired".to_string(), - }); - } - - if stored_real_gateway_order_id.is_none() { - if let Some(gateway_order_id) = effective_gateway_order_id { - if !sqlite_bind_payment_gateway_order_id(&mut tx, &order_id, gateway_order_id) - .await? - { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment gateway order belongs to another payment order", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment gateway order belongs to another payment order".to_string(), - }); - } - } - } - - if order_kind == "plan_purchase" { - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(user_id) = order_user_id else { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "payment order user missing", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "payment order user missing".to_string(), - }); - }; - let product_id: Option = get(&order_row, "product_id")?; - let snapshot = optional_json( - get::>(&order_row, "product_snapshot")?, - "payment_orders.product_snapshot", - )? - .unwrap_or_else(|| serde_json::json!({})); - let plan_id = product_id.unwrap_or_else(|| { - snapshot - .get("id") - .and_then(|value| value.as_str()) - .unwrap_or("unknown") - .to_string() - }); - let entitlements = plan_entitlements_snapshot(&snapshot); - let existing_entitlement_id = sqlx::query_scalar::<_, String>( - "SELECT id FROM user_plan_entitlements WHERE payment_order_id = ? LIMIT 1", - ) - .bind(&order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if existing_entitlement_id.is_none() { - let purchase_limit_scope = plan_purchase_limit_scope(&snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&snapshot); - let active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - if active_count >= max_active_per_user { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "plan purchase limit reached", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "plan purchase limit reached".to_string(), - }); - } - } - replace_matching_plan_entitlements_sqlite(&mut tx, &user_id, &snapshot, now) - .await?; - sqlx::query( - r#" -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) -VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&user_id) - .bind(&plan_id) - .bind(&order_id) - .bind(now) - .bind(plan_expires_at_unix(&snapshot, now)?) - .bind(json_string( - &entitlements, - "user_plan_entitlements.entitlements_snapshot", - )?) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - apply_plan_wallet_credit_sqlite( - &mut tx, - &order_wallet_id, - &order_id, - &input.payment_method, - &entitlements, - now, - ) - .await?; - } - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = COALESCE(?, gateway_order_id), - gateway_response = ?, - pay_amount = COALESCE(pay_amount, ?), - pay_currency = COALESCE(pay_currency, ?), - exchange_rate = COALESCE(exchange_rate, ?), - status = 'credited', - fulfillment_status = 'fulfilled', - fulfillment_error = NULL, - paid_at = COALESCE(paid_at, ?), - credited_at = ?, - refundable_amount_usd = 0 -WHERE id = ? - "#, - ) - .bind(effective_gateway_order_id) - .bind(json_string( - &input.gateway_response_projection(&order_no, effective_gateway_order_id), - "payment_orders.gateway_response", - )?) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(now) - .bind(now) - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated_order_row = sqlite_payment_order_by_id(&mut tx, &order_id).await?; - mark_sqlite_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Applied { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - order: map_payment_order_row(&updated_order_row)?, - }); - } - - let Some(wallet_row) = sqlx::query( - r#" -SELECT id, status, balance, gift_balance, total_recharged -FROM wallets -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(&order_wallet_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet not found", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet not found".to_string(), - }); - }; - let wallet_status: String = get(&wallet_row, "status")?; - if wallet_status != "active" { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet is not active", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet is not active".to_string(), - }); - } - - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let total_recharged = sqlite_real(&wallet_row, "total_recharged")?; - // Finite recharge balances may be negative: usage settlement permits a - // finite wallet to overdraft, and a later recharge must be able to - // restore that balance. Reject only malformed values and arithmetic - // overflow here. - if !before_recharge.is_finite() - || !before_gift.is_finite() - || before_gift < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - || !(total_recharged + order_amount_usd).is_finite() - || !(before_recharge + before_gift + order_amount_usd).is_finite() - { - update_sqlite_payment_callback_failure( - &mut tx, - &callback_id, - &input, - &payload, - "wallet balance is invalid", - ) - .await?; - tx.commit().await.map_sql_err()?; - return Ok(ProcessPaymentCallbackOutcome::Failed { - duplicate, - error: "wallet balance is invalid".to_string(), - }); - } - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + order_amount_usd; - let after_total = after_recharge + before_gift; - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = total_recharged + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(order_amount_usd) - .bind(now) - .bind(&order_wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_gateway', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&order_wallet_id) - .bind(order_amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&order_id) - .bind(format!("充值到账({})", input.payment_method)) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = COALESCE(?, gateway_order_id), - gateway_response = ?, - pay_amount = COALESCE(pay_amount, ?), - pay_currency = COALESCE(pay_currency, ?), - exchange_rate = COALESCE(exchange_rate, ?), - status = 'credited', - paid_at = COALESCE(paid_at, ?), - credited_at = ?, - refundable_amount_usd = amount_usd -WHERE id = ? -"#, - ) - .bind(effective_gateway_order_id) - .bind(json_string( - &input.gateway_response_projection(&order_no, effective_gateway_order_id), - "payment_orders.gateway_response", - )?) - .bind(input.pay_amount) - .bind(input.pay_currency.as_deref()) - .bind(input.exchange_rate) - .bind(now) - .bind(now) - .bind(&order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated_order_row = sqlite_payment_order_by_id(&mut tx, &order_id).await?; - mark_sqlite_payment_callback_processed( - &mut tx, - &callback_id, - &input, - &payload, - &order_id, - &order_no, - ) - .await?; - tx.commit().await.map_sql_err()?; - Ok(ProcessPaymentCallbackOutcome::Applied { - duplicate, - order_id, - order_no, - wallet_id: order_wallet_id, - order: map_payment_order_row(&updated_order_row)?, - }) - } - - async fn adjust_wallet_balance( - &self, - input: AdjustWalletBalanceInput, - ) -> Result, DataLayerError> { - if !input.amount_usd.is_finite() || input.amount_usd == 0.0 { - return Err(DataLayerError::InvalidInput( - "adjustment amount must be finite and non-zero".to_string(), - )); - } - let now = current_unix_secs_i64(); - // Admin credit follows the same read/validate/write sequence as a - // provider callback. Serialize it at transaction start to prevent two - // writers from both crediting a pending order. - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(row) = sqlite_wallet_by_id_optional(&mut tx, &input.wallet_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let before_recharge = sqlite_real(&row, "balance")?; - let before_gift = sqlite_real(&row, "gift_balance")?; - let before_total = before_recharge + before_gift; - let before_total_adjusted = sqlite_real(&row, "total_adjusted")?; - if !before_recharge.is_finite() - || !before_gift.is_finite() - || !before_total.is_finite() - || !before_total_adjusted.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance is invalid".to_string(), - )); - } - let mut after_recharge = before_recharge; - let mut after_gift = before_gift; - apply_admin_balance_adjustment( - input.amount_usd, - &input.balance_type, - &mut after_recharge, - &mut after_gift, - ); - let after_total = after_recharge + after_gift; - let after_total_adjusted = before_total_adjusted + input.amount_usd; - if !after_recharge.is_finite() - || !after_gift.is_finite() - || !after_total.is_finite() - || !after_total_adjusted.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance overflow during admin adjustment".to_string(), - )); - } - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_adjusted = total_adjusted + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(input.amount_usd) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - - let transaction_id = uuid::Uuid::new_v4().to_string(); - let description = input - .description - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or("管理员调账") - .to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'adjust', 'adjust_admin', ?, ?, ?, ?, ?, ?, ?, 'admin_action', ?, ?, ?, ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(input.amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(after_gift) - .bind(&input.wallet_id) - .bind(input.operator_id.as_deref()) - .bind(&description) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - - Ok(Some(( - wallet, - StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "adjust".to_string(), - reason_code: "adjust_admin".to_string(), - amount: input.amount_usd, - balance_before: before_total, - balance_after: after_total, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: after_gift, - link_type: Some("admin_action".to_string()), - link_id: Some(input.wallet_id), - operator_id: input.operator_id, - operator_name: None, - operator_email: None, - description: Some(description), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }, - ))) - } - - async fn create_manual_wallet_recharge( - &self, - mut input: CreateManualWalletRechargeInput, - ) -> Result, DataLayerError> { - input.payment_method = canonicalize_payment_method(&input.payment_method) - .map_err(DataLayerError::InvalidInput)?; - if !input.amount_usd.is_finite() || input.amount_usd <= 0.0 { - return Err(DataLayerError::InvalidInput( - "manual recharge amount must be finite and positive".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(wallet_row) = sqlite_wallet_by_id_optional(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(None); - }; - - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let before_total_recharged = sqlite_real(&wallet_row, "total_recharged")?; - let (after_recharge, after_total_recharged) = validate_manual_wallet_recharge( - input.amount_usd, - before_recharge, - before_gift, - before_total_recharged, - ) - .map_err(DataLayerError::InvalidInput)?; - let user_id: Option = get(&wallet_row, "user_id")?; - let order_id = uuid::Uuid::new_v4().to_string(); - let gateway_response = json_string( - &serde_json::json!({ - "source": "manual", - "operator_id": input.operator_id, - "description": input.description, - }), - "payment_orders.gateway_response", - )?; - - sqlx::query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, status, gateway_response, - created_at, paid_at, credited_at -) -VALUES (?, ?, ?, ?, ?, 0, ?, ?, 'credited', ?, ?, ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&input.wallet_id) - .bind(user_id.as_deref()) - .bind(input.amount_usd) - .bind(input.amount_usd) - .bind(&input.payment_method) - .bind(&gateway_response) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_recharged) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let reason_code = if matches!( - input.payment_method.as_str(), - "card_code" | "gift_code" | "card_recharge" - ) { - "topup_card_code" - } else { - "topup_admin_manual" - }; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', ?, ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, ?, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&input.wallet_id) - .bind(reason_code) - .bind(input.amount_usd) - .bind(before_recharge + before_gift) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&order_id) - .bind(input.operator_id.as_deref()) - .bind( - input - .description - .as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or("管理员手动充值"), - ) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let order = map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(Some((wallet, order))) - } - - async fn process_admin_wallet_refund( - &self, - input: ProcessAdminWalletRefundInput, - ) -> Result< - WalletMutationOutcome<( - StoredWalletSnapshot, - StoredAdminWalletRefund, - StoredAdminWalletTransaction, - )>, - DataLayerError, - > { - let now = current_unix_secs_i64(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(refund_row) = - sqlite_refund_by_id_and_wallet(&mut tx, &input.refund_id, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(&refund_row)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if !matches!(refund.status.as_str(), "approved" | "pending_approval") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status is not approvable".to_string(), - )); - } - - let Some(wallet_row) = sqlite_wallet_by_id_optional(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let before_total_refunded = sqlite_real(&wallet_row, "total_refunded")?; - let amount_usd = refund.amount_usd; - let after_recharge = before_recharge - amount_usd; - let before_total = before_recharge + before_gift; - let after_total = after_recharge + before_gift; - let after_total_refunded = before_total_refunded + amount_usd; - if !before_recharge.is_finite() - || before_recharge < 0.0 - || !before_gift.is_finite() - || before_gift < 0.0 - || !before_total_refunded.is_finite() - || before_total_refunded < 0.0 - || !before_total.is_finite() - || !after_recharge.is_finite() - || !after_total.is_finite() - || !after_total_refunded.is_finite() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid".to_string(), - )); - } - if after_recharge < 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount exceeds refundable recharge balance".to_string(), - )); - } - - if let Some(payment_order_id) = refund.payment_order_id.as_deref() { - let Some(order_row) = - sqlite_payment_order_by_id_optional(&mut tx, payment_order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order not found".to_string(), - )); - }; - let order_wallet_id: String = get(&order_row, "wallet_id")?; - let order_status: String = get(&order_row, "status")?; - if order_wallet_id != input.wallet_id || order_status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not refundable for this wallet".to_string(), - )); - } - let order_amount = sqlite_real(&order_row, "amount_usd")?; - let refunded_before = sqlite_real(&order_row, "refunded_amount_usd")?; - let refundable_before = sqlite_real(&order_row, "refundable_amount_usd")?; - let refunded_after = refunded_before + amount_usd; - let refundable_after = refundable_before - amount_usd; - if !payment_order_refund_amounts_are_consistent( - order_amount, - refunded_before, - refundable_before, - ) || amount_usd > refundable_before - || !refunded_after.is_finite() - || refunded_after < 0.0 - || refunded_after > order_amount - || !refundable_after.is_finite() - || refundable_after < 0.0 - || refundable_after > order_amount - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order refund amounts are invalid".to_string(), - )); - } - let result = sqlx::query( - r#" -UPDATE payment_orders -SET refunded_amount_usd = ?, - refundable_amount_usd = ? -WHERE id = ? -"#, - ) - .bind(refunded_after) - .bind(refundable_after) - .bind(payment_order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "payment order disappeared during refund processing".to_string(), - )); - } - } - - let result = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_refunded = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_refunded) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "wallet disappeared during refund processing".to_string(), - )); - } - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - - let transaction_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'refund', 'refund_out', ?, ?, ?, ?, ?, ?, ?, 'refund_request', ?, ?, '退款占款', ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(-amount_usd) - .bind(before_total) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.refund_id) - .bind(input.operator_id.as_deref()) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let refund_update = sqlx::query( - r#" - UPDATE refund_requests -SET status = 'processing', - approved_by = ?, - processed_by = ?, - processed_at = ?, - updated_at = ? -WHERE id = ? - AND wallet_id = ? - AND status IN ('approved', 'pending_approval') -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if refund_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during refund processing".to_string(), - )); - } - let refund = map_refund_row(&sqlite_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(( - wallet, - refund, - StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "refund".to_string(), - reason_code: "refund_out".to_string(), - amount: -amount_usd, - balance_before: before_total, - balance_after: after_recharge + before_gift, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: before_gift, - link_type: Some("refund_request".to_string()), - link_id: Some(input.refund_id.clone()), - operator_id: input.operator_id.clone(), - operator_name: None, - operator_email: None, - description: Some("退款占款".to_string()), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }, - ))) - } - - async fn complete_admin_wallet_refund( - &self, - input: CompleteAdminWalletRefundInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(current_refund) = - sqlite_refund_by_id_and_wallet(&mut tx, &input.refund_id, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(¤t_refund)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if let (Some(existing_id), Some(incoming_id)) = ( - refund.gateway_refund_id.as_deref(), - input.gateway_refund_id.as_deref(), - ) { - if existing_id != incoming_id { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier conflicts with existing evidence".to_string(), - )); - } - } - if refund.status == "succeeded" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(refund)); - } - if refund.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status must be processing before completion".to_string(), - )); - } - // Preserve a processing proof for ordinary replays, but allow an - // explicit successful gateway proof to upgrade it at completion. - let selected_payout_proof = input - .payout_proof - .as_ref() - .filter(|proof| refund.payout_proof.is_none() || wallet_refund_proof_is_success(proof)) - .cloned() - .or_else(|| refund.payout_proof.clone()); - let payout_proof = selected_payout_proof - .as_ref() - .map(|value| json_string(value, "refund_requests.payout_proof")) - .transpose()?; - - let refund_update = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'succeeded', - gateway_refund_id = COALESCE(gateway_refund_id, ?), - payout_reference = COALESCE(payout_reference, ?), - payout_proof = ?, - completed_at = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(input.gateway_refund_id.as_deref()) - .bind(input.payout_reference.as_deref()) - .bind(payout_proof.as_deref()) - .bind(now) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if refund_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during refund completion".to_string(), - )); - } - let refund = map_refund_row(&sqlite_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(refund)) - } - - async fn update_admin_wallet_refund_gateway( - &self, - input: UpdateAdminWalletRefundGatewayInput, - ) -> Result, DataLayerError> { - if input.gateway_refund_id.trim().is_empty() || input.gateway_refund_id.len() > 128 { - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier is invalid".to_string(), - )); - } - if input - .payout_proof - .as_ref() - .is_some_and(|proof| !proof.is_object()) - { - return Ok(WalletMutationOutcome::Invalid( - "gateway refund proof must be an object".to_string(), - )); - } - let now = current_unix_secs_i64(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(current_row) = - sqlite_refund_by_id_and_wallet(&mut tx, &input.refund_id, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let current = map_refund_row(¤t_row)?; - if !current.amount_usd.is_finite() || current.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - if let Some(existing_id) = current.gateway_refund_id.as_deref() { - if existing_id != input.gateway_refund_id { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "gateway refund identifier conflicts with existing evidence".to_string(), - )); - } - } - if current.status == "succeeded" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied(current)); - } - if current.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund status must be processing before gateway update".to_string(), - )); - } - // Do not overwrite durable processing evidence with an arbitrary - // replay; only a terminal success proof may replace it. - let selected_payout_proof = input - .payout_proof - .as_ref() - .filter(|proof| current.payout_proof.is_none() || wallet_refund_proof_is_success(proof)) - .cloned() - .or_else(|| current.payout_proof.clone()); - let proof = selected_payout_proof - .as_ref() - .map(|value| json_string(value, "refund_requests.payout_proof")) - .transpose()?; - let gateway_update = sqlx::query( - r#" -UPDATE refund_requests -SET gateway_refund_id = COALESCE(gateway_refund_id, ?), - payout_proof = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(&input.gateway_refund_id) - .bind(proof.as_deref()) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if gateway_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during gateway evidence update".to_string(), - )); - } - let updated = map_refund_row(&sqlite_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn fail_admin_wallet_refund( - &self, - input: FailAdminWalletRefundInput, - ) -> Result< - WalletMutationOutcome<( - StoredWalletSnapshot, - StoredAdminWalletRefund, - Option, - )>, - DataLayerError, - > { - let now = current_unix_secs_i64(); - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_sql_err()?; - let Some(refund_row) = - sqlite_refund_by_id_and_wallet(&mut tx, &input.refund_id, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let refund = map_refund_row(&refund_row)?; - if !refund.amount_usd.is_finite() || refund.amount_usd <= 0.0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "refund amount must be finite and greater than zero".to_string(), - )); - } - - if matches!(refund.status.as_str(), "pending_approval" | "approved") { - let refund_update = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'failed', - failure_reason = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status IN ('pending_approval', 'approved') -"#, - ) - .bind(&input.reason) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if refund_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during refund failure".to_string(), - )); - } - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let refund = map_refund_row(&sqlite_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((wallet, refund, None))); - } - - if refund.status != "processing" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "cannot fail refund in status: {}", - refund.status - ))); - } - - // Only an explicitly offline payout can be released without external - // settlement evidence. An original-channel refund may still be in - // flight between the provider request and the evidence update. - if refund.gateway_refund_id.is_some() - || refund.payout_proof.is_some() - || !refund - .refund_mode - .trim() - .eq_ignore_ascii_case("offline_payout") - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "cannot fail refund while gateway settlement is processing".to_string(), - )); - } - - let Some(wallet_row) = sqlite_wallet_by_id_optional(&mut tx, &input.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let amount_usd = refund.amount_usd; - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let before_total_refunded = sqlite_real(&wallet_row, "total_refunded")?; - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + amount_usd; - let after_total = after_recharge + before_gift; - let after_total_refunded = before_total_refunded - amount_usd; - if !before_recharge.is_finite() - || before_recharge < 0.0 - || !before_gift.is_finite() - || before_gift < 0.0 - || !before_total_refunded.is_finite() - || before_total_refunded < 0.0 - || before_total_refunded < amount_usd - || !before_total.is_finite() - || !after_recharge.is_finite() - || !after_total.is_finite() - || !after_total_refunded.is_finite() - || after_total_refunded < 0.0 - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid for refund recovery".to_string(), - )); - } - - let mut order_amounts = None; - if let Some(payment_order_id) = refund.payment_order_id.as_deref() { - let Some(order_row) = - sqlite_payment_order_by_id_optional(&mut tx, payment_order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order not found".to_string(), - )); - }; - let order = map_payment_order_row(&order_row)?; - if order.wallet_id != input.wallet_id || order.status != "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order is not refundable for this wallet".to_string(), - )); - } - let refunded_before = order.refunded_amount_usd; - let refundable_before = order.refundable_amount_usd; - let refunded_after = refunded_before - amount_usd; - let refundable_after = refundable_before + amount_usd; - if !payment_order_refund_amounts_are_consistent( - order.amount_usd, - refunded_before, - refundable_before, - ) || refunded_before < amount_usd - || !refunded_after.is_finite() - || refunded_after < 0.0 - || !refundable_after.is_finite() - || refundable_after < 0.0 - || refundable_after > order.amount_usd - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order refund amounts are invalid".to_string(), - )); - } - order_amounts = Some(( - payment_order_id.to_string(), - refunded_after, - refundable_after, - )); - } - - let wallet_update = sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_refunded = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_total_refunded) - .bind(now) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if wallet_update.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "wallet disappeared during refund recovery".to_string(), - )); - } - - let transaction_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'refund', 'refund_revert', ?, ?, ?, ?, ?, ?, ?, 'refund_request', ?, ?, '退款失败回补', ?) -"#, - ) - .bind(&transaction_id) - .bind(&input.wallet_id) - .bind(amount_usd) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.refund_id) - .bind(input.operator_id.as_deref()) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - if let Some((payment_order_id, refunded_after, refundable_after)) = order_amounts { - let result = sqlx::query( - r#" -UPDATE payment_orders -SET refunded_amount_usd = ?, - refundable_amount_usd = ? -WHERE id = ? -"#, - ) - .bind(refunded_after) - .bind(refundable_after) - .bind(payment_order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "payment order disappeared during refund recovery".to_string(), - )); - } - } - - let result = sqlx::query( - r#" -UPDATE refund_requests -SET status = 'failed', - failure_reason = ?, - updated_at = ? -WHERE id = ? AND wallet_id = ? AND status = 'processing' -"#, - ) - .bind(&input.reason) - .bind(now) - .bind(&input.refund_id) - .bind(&input.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - if result.rows_affected() != 1 { - return Err(DataLayerError::UnexpectedValue( - "refund status changed during recovery".to_string(), - )); - } - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &input.wallet_id).await?)?; - let refund = map_refund_row(&sqlite_refund_by_id(&mut tx, &input.refund_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(( - wallet, - refund, - Some(StoredAdminWalletTransaction { - id: transaction_id, - wallet_id: input.wallet_id.clone(), - category: "refund".to_string(), - reason_code: "refund_revert".to_string(), - amount: amount_usd, - balance_before: before_total, - balance_after: after_total, - recharge_balance_before: before_recharge, - recharge_balance_after: after_recharge, - gift_balance_before: before_gift, - gift_balance_after: before_gift, - link_type: Some("refund_request".to_string()), - link_id: Some(input.refund_id.clone()), - operator_id: input.operator_id.clone(), - operator_name: None, - operator_email: None, - description: Some("退款失败回补".to_string()), - created_at_unix_ms: Some(timestamp(now, "wallet_transactions.created_at")?), - }), - ))) - } - - async fn expire_admin_payment_order( - &self, - order_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = sqlite_payment_order_by_id_optional(&mut tx, order_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "credited order cannot be expired".to_string(), - )); - } - if order.status == "expired" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, false))); - } - if order.status != "pending" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "only pending order can be expired: {}", - order.status - ))); - } - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - gateway_response.insert( - "expire_reason".to_string(), - serde_json::Value::String("admin_mark_expired".to_string()), - ); - gateway_response.insert( - "expired_at".to_string(), - serde_json::Value::String(Utc::now().to_rfc3339()), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - sqlx::query( - "UPDATE payment_orders SET status = 'expired', gateway_response = ? WHERE id = ?", - ) - .bind(gateway_response) - .bind(order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied((updated, true))) - } - - async fn fail_admin_payment_order( - &self, - order_id: &str, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(row) = sqlite_payment_order_by_id_optional(&mut tx, order_id).await? else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "credited order cannot be failed".to_string(), - )); - } - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - gateway_response.insert( - "failure_reason".to_string(), - serde_json::Value::String("admin_mark_failed".to_string()), - ); - gateway_response.insert( - "failed_at".to_string(), - serde_json::Value::String(Utc::now().to_rfc3339()), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - sqlx::query( - "UPDATE payment_orders SET status = 'failed', gateway_response = ? WHERE id = ?", - ) - .bind(gateway_response) - .bind(order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let updated = map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(updated)) - } - - async fn credit_admin_payment_order( - &self, - input: CreditAdminPaymentOrderInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(order_row) = sqlite_payment_order_by_id_optional(&mut tx, &input.order_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let order = map_payment_order_row(&order_row)?; - if order.status == "credited" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, false))); - } - if matches!(order.status.as_str(), "failed" | "expired" | "refunded") { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid(format!( - "payment order is not creditable: {}", - order.status - ))); - } - if order - .expires_at_unix_secs - .is_some_and(|value| value <= now as u64) - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order expired".to_string(), - )); - } - let order_kind: String = get(&order_row, "order_kind")?; - let order_payment_provider: Option = get(&order_row, "payment_provider")?; - let order_payment_channel: Option = get(&order_row, "payment_channel")?; - if validate_payment_order_credit_amounts( - &order_kind, - &order.payment_method, - order_payment_provider.as_deref(), - order_payment_channel.as_deref(), - order.amount_usd, - order.pay_amount, - ) - .is_err() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order amount is invalid".to_string(), - )); - } - if order_kind == "plan_purchase" { - let order_user_id: Option = get(&order_row, "user_id")?; - let Some(user_id) = order_user_id else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "payment order user missing".to_string(), - )); - }; - let product_id: Option = get(&order_row, "product_id")?; - let snapshot = optional_json( - get::>(&order_row, "product_snapshot")?, - "payment_orders.product_snapshot", - )? - .unwrap_or_else(|| serde_json::json!({})); - let plan_id = product_id.unwrap_or_else(|| { - snapshot - .get("id") - .and_then(|value| value.as_str()) - .unwrap_or("unknown") - .to_string() - }); - let entitlements = plan_entitlements_snapshot(&snapshot); - let existing_entitlement_id = sqlx::query_scalar::<_, String>( - "SELECT id FROM user_plan_entitlements WHERE payment_order_id = ? LIMIT 1", - ) - .bind(&input.order_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if existing_entitlement_id.is_none() { - let purchase_limit_scope = plan_purchase_limit_scope(&snapshot); - if purchase_limit_scope != "unlimited" { - let max_active_per_user = plan_max_active_per_user(&snapshot); - let active_count = if purchase_limit_scope == "lifetime" { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - } else { - sqlx::query_scalar::<_, i64>( - r#" -SELECT COUNT(*) -FROM user_plan_entitlements -WHERE user_id = ? - AND plan_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(&user_id) - .bind(&plan_id) - .bind(now) - .fetch_one(&mut *tx) - .await - .map_sql_err()? - }; - if active_count >= max_active_per_user { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "plan purchase limit reached".to_string(), - )); - } - } - replace_matching_plan_entitlements_sqlite(&mut tx, &user_id, &snapshot, now) - .await?; - sqlx::query( - r#" -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) -VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&user_id) - .bind(&plan_id) - .bind(&input.order_id) - .bind(now) - .bind(plan_expires_at_unix(&snapshot, now)?) - .bind(json_string( - &entitlements, - "user_plan_entitlements.entitlements_snapshot", - )?) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - apply_plan_wallet_credit_sqlite( - &mut tx, - &order.wallet_id, - &input.order_id, - &order.payment_method, - &entitlements, - now, - ) - .await?; - } - - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - if let Some(serde_json::Value::Object(map)) = input.gateway_response_patch.clone() { - gateway_response.extend(map); - } - gateway_response.insert("manual_credit".to_string(), serde_json::Value::Bool(true)); - gateway_response.insert( - "credited_by".to_string(), - input - .operator_id - .clone() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - let next_gateway_order_id = input.gateway_order_id.clone().or(order.gateway_order_id); - let next_pay_amount = input.pay_amount.or(order.pay_amount); - let next_pay_currency = input.pay_currency.clone().or(order.pay_currency); - let next_exchange_rate = input.exchange_rate.or(order.exchange_rate); - let next_paid_at = order.paid_at_unix_secs.unwrap_or(now as u64) as i64; - - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = ?, - gateway_response = ?, - pay_amount = ?, - pay_currency = ?, - exchange_rate = ?, - status = 'credited', - fulfillment_status = 'fulfilled', - fulfillment_error = NULL, - paid_at = ?, - credited_at = ?, - refundable_amount_usd = 0 -WHERE id = ? -"#, - ) - .bind(next_gateway_order_id.as_deref()) - .bind(&gateway_response) - .bind(next_pay_amount) - .bind(next_pay_currency.as_deref()) - .bind(next_exchange_rate) - .bind(next_paid_at) - .bind(now) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let order = map_payment_order_row( - &sqlite_payment_order_by_id(&mut tx, &input.order_id).await?, - )?; - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Applied((order, true))); - } - - let Some(wallet_row) = sqlite_wallet_by_id_optional(&mut tx, &order.wallet_id).await? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet not found".to_string(), - )); - }; - let wallet_status: String = get(&wallet_row, "status")?; - if wallet_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet is not active".to_string(), - )); - } - - let before_recharge = sqlite_real(&wallet_row, "balance")?; - let before_gift = sqlite_real(&wallet_row, "gift_balance")?; - let total_recharged = sqlite_real(&wallet_row, "total_recharged")?; - if !before_recharge.is_finite() - || !before_gift.is_finite() - || before_gift < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - || !(total_recharged + order.amount_usd).is_finite() - || !(before_recharge + before_gift + order.amount_usd).is_finite() - { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "wallet balance is invalid".to_string(), - )); - } - let before_total = before_recharge + before_gift; - let after_recharge = before_recharge + order.amount_usd; - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - total_recharged = total_recharged + ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(order.amount_usd) - .bind(now) - .bind(&order.wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_gateway', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&order.wallet_id) - .bind(order.amount_usd) - .bind(before_total) - .bind(after_recharge + before_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(before_gift) - .bind(&input.order_id) - .bind(format!("充值到账({})", order.payment_method)) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let mut gateway_response = payment_gateway_response_map(order.gateway_response.clone()); - if let Some(serde_json::Value::Object(map)) = input.gateway_response_patch { - gateway_response.extend(map); - } - gateway_response.insert("manual_credit".to_string(), serde_json::Value::Bool(true)); - gateway_response.insert( - "credited_by".to_string(), - input - .operator_id - .clone() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - ); - let gateway_response = json_string( - &serde_json::Value::Object(gateway_response), - "payment_orders.gateway_response", - )?; - let next_gateway_order_id = input.gateway_order_id.or(order.gateway_order_id); - let next_pay_amount = input.pay_amount.or(order.pay_amount); - let next_pay_currency = input.pay_currency.or(order.pay_currency); - let next_exchange_rate = input.exchange_rate.or(order.exchange_rate); - let next_paid_at = order.paid_at_unix_secs.unwrap_or(now as u64) as i64; - - sqlx::query( - r#" -UPDATE payment_orders -SET gateway_order_id = ?, - gateway_response = ?, - pay_amount = ?, - pay_currency = ?, - exchange_rate = ?, - status = 'credited', - paid_at = ?, - credited_at = ?, - refundable_amount_usd = amount_usd -WHERE id = ? -"#, - ) - .bind(next_gateway_order_id.as_deref()) - .bind(&gateway_response) - .bind(next_pay_amount) - .bind(next_pay_currency.as_deref()) - .bind(next_exchange_rate) - .bind(next_paid_at) - .bind(now) - .bind(&input.order_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let order = - map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &input.order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied((order, true))) - } - - async fn create_admin_redeem_code_batch( - &self, - input: CreateAdminRedeemCodeBatchInput, - ) -> Result { - validate_admin_redeem_code_batch_input(&input).map_err(DataLayerError::InvalidInput)?; - let now = current_unix_secs_i64(); - let batch_id = uuid::Uuid::new_v4().to_string(); - let expires_at = input - .expires_at_unix_secs - .map(|value| { - i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput( - "redeem code batch expires_at overflow".to_string(), - ) - }) - }) - .transpose()?; - let mut tx = self.pool.begin().await.map_sql_err()?; - - sqlx::query( - r#" -INSERT INTO redeem_code_batches ( - id, name, amount_usd, currency, balance_bucket, total_count, status, - description, created_by, expires_at, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) -"#, - ) - .bind(&batch_id) - .bind(&input.name) - .bind(input.amount_usd) - .bind(&input.currency) - .bind(&input.balance_bucket) - .bind(i64::try_from(input.total_count).map_err(|_| { - DataLayerError::InvalidInput(format!( - "invalid redeem code count: {}", - input.total_count - )) - })?) - .bind(input.description.as_deref()) - .bind(input.created_by.as_deref()) - .bind(expires_at) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let mut codes = Vec::with_capacity(input.total_count); - for _ in 0..input.total_count { - let (code_id, code, masked_code, code_hash, prefix, suffix) = - generate_redeem_code_candidate(); - sqlx::query( - r#" -INSERT INTO redeem_codes ( - id, batch_id, code_hash, code_prefix, code_suffix, status, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, 'active', ?, ?) -"#, - ) - .bind(&code_id) - .bind(&batch_id) - .bind(&code_hash) - .bind(&prefix) - .bind(&suffix) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - codes.push(CreatedAdminRedeemCodePlaintext { - code_id, - code, - masked_code, - }); - } - - let batch = StoredAdminRedeemCodeBatch { - id: batch_id, - name: input.name, - amount_usd: input.amount_usd, - currency: input.currency, - balance_bucket: input.balance_bucket, - total_count: input.total_count as u64, - redeemed_count: 0, - active_count: input.total_count as u64, - status: "active".to_string(), - description: input.description, - created_by: input.created_by, - expires_at_unix_secs: input.expires_at_unix_secs, - created_at_unix_ms: timestamp(now, "redeem_code_batches.created_at")?, - updated_at_unix_secs: timestamp(now, "redeem_code_batches.updated_at")?, - }; - tx.commit().await.map_sql_err()?; - Ok(CreateAdminRedeemCodeBatchResult { batch, codes }) - } - - async fn disable_admin_redeem_code_batch( - &self, - input: DisableAdminRedeemCodeBatchInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_batch) = sqlx::query( - r#" -SELECT status -FROM redeem_code_batches -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(&input.batch_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let status: String = get(¤t_batch, "status")?; - if status != "disabled" { - sqlx::query( - r#" -UPDATE redeem_code_batches -SET status = 'disabled', - updated_at = ? -WHERE id = ? -"#, - ) - .bind(now) - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'disabled', - disabled_by = COALESCE(?, disabled_by), - updated_at = ? -WHERE batch_id = ? - AND status = 'active' -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - let batch = - map_redeem_batch_row(&sqlite_redeem_batch_by_id(&mut tx, &input.batch_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(batch)) - } - - async fn delete_admin_redeem_code_batch( - &self, - input: DeleteAdminRedeemCodeBatchInput, - ) -> Result, DataLayerError> { - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_batch) = sqlx::query( - r#" -SELECT status -FROM redeem_code_batches -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(&input.batch_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let status: String = get(¤t_batch, "status")?; - if status != "disabled" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "only disabled redeem code batch can be deleted".to_string(), - )); - } - - let redeemed_count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM redeem_codes -WHERE batch_id = ? - AND status = 'redeemed' -"#, - ) - .bind(&input.batch_id) - .fetch_one(&mut *tx) - .await - .map_sql_err()?; - if redeemed_count > 0 { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "redeemed batch cannot be deleted".to_string(), - )); - } - - let batch = - map_redeem_batch_row(&sqlite_redeem_batch_by_id(&mut tx, &input.batch_id).await?)?; - let _ = input.operator_id; - sqlx::query("DELETE FROM redeem_codes WHERE batch_id = ?") - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("DELETE FROM redeem_code_batches WHERE id = ?") - .bind(&input.batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(batch)) - } - - async fn disable_admin_redeem_code( - &self, - input: DisableAdminRedeemCodeInput, - ) -> Result, DataLayerError> { - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - let Some(current_code) = sqlx::query( - r#" -SELECT batch_id, status -FROM redeem_codes -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(&input.code_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::NotFound); - }; - let batch_id: String = get(¤t_code, "batch_id")?; - let status: String = get(¤t_code, "status")?; - if status == "redeemed" { - tx.commit().await.map_sql_err()?; - return Ok(WalletMutationOutcome::Invalid( - "redeemed code cannot be disabled".to_string(), - )); - } - if status != "disabled" { - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'disabled', - disabled_by = COALESCE(?, disabled_by), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(input.operator_id.as_deref()) - .bind(now) - .bind(&input.code_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - - sqlx::query("UPDATE redeem_code_batches SET updated_at = ? WHERE id = ?") - .bind(now) - .bind(&batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - let code = map_redeem_code_row(&sqlite_redeem_code_by_id(&mut tx, &input.code_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(WalletMutationOutcome::Applied(code)) - } - - async fn redeem_wallet_code( - &self, - input: RedeemWalletCodeInput, - ) -> Result { - let Some(normalized) = normalize_redeem_code(&input.code) else { - return Ok(RedeemWalletCodeOutcome::InvalidCode); - }; - let code_hash = hash_redeem_code(&normalized); - let now = current_unix_secs_i64(); - let mut tx = self.pool.begin().await.map_sql_err()?; - - let Some(code_row) = sqlx::query( - r#" -SELECT - codes.id AS code_id, - codes.status AS code_status, - codes.batch_id, - batches.name AS batch_name, - batches.status AS batch_status, - batches.balance_bucket, - batches.amount_usd, - batches.expires_at AS batch_expires_at -FROM redeem_codes AS codes -JOIN redeem_code_batches AS batches ON batches.id = codes.batch_id -WHERE codes.code_hash = ? -LIMIT 1 -"#, - ) - .bind(&code_hash) - .fetch_optional(&mut *tx) - .await - .map_sql_err()? - else { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeNotFound); - }; - - let code_status: String = get(&code_row, "code_status")?; - match code_status.as_str() { - "disabled" => { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeDisabled); - } - "redeemed" => { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeRedeemed); - } - _ => {} - } - let batch_status: String = get(&code_row, "batch_status")?; - if batch_status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::BatchDisabled); - } - let batch_expires_at: Option = get(&code_row, "batch_expires_at")?; - if batch_expires_at.is_some_and(|value| value <= now) { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::CodeExpired); - } - - let code_id: String = get(&code_row, "code_id")?; - let batch_id: String = get(&code_row, "batch_id")?; - let batch_name: String = get(&code_row, "batch_name")?; - let balance_bucket: String = get(&code_row, "balance_bucket")?; - let amount_usd = sqlite_real(&code_row, "amount_usd")?; - - let wallet_row = sqlite_wallet_by_user_id(&mut tx, &input.user_id).await?; - let wallet_id = if let Some(row) = wallet_row.as_ref() { - let status: String = get(row, "status")?; - if status != "active" { - tx.commit().await.map_sql_err()?; - return Ok(RedeemWalletCodeOutcome::WalletInactive); - } - get(row, "id")? - } else { - uuid::Uuid::new_v4().to_string() - }; - - let (before_recharge, before_gift, before_total_recharged) = - if let Some(row) = wallet_row.as_ref() { - ( - sqlite_real(row, "balance")?, - sqlite_real(row, "gift_balance")?, - sqlite_real(row, "total_recharged")?, - ) - } else { - sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, 0.0, 0.0, 'finite', 'USD', 'active', 0.0, 0.0, 0.0, 0.0, ?, ?) -"#, - ) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - (0.0, 0.0, 0.0) - }; - let (after_recharge, after_gift, after_total_recharged) = validate_redeem_wallet_credit( - &balance_bucket, - amount_usd, - before_recharge, - before_gift, - before_total_recharged, - ) - .map_err(DataLayerError::UnexpectedValue)?; - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(after_recharge) - .bind(after_gift) - .bind(after_total_recharged) - .bind(now) - .bind(&wallet_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let payment_method = redeem_code_payment_method(&balance_bucket); - let order_id = uuid::Uuid::new_v4().to_string(); - let gateway_order_id = format!("card_{}", uuid::Uuid::new_v4().simple()); - let gateway_response = json_string( - &serde_json::json!({ - "source": "redeem_code", - "batch_id": batch_id, - "batch_name": batch_name, - "balance_bucket": balance_bucket, - }), - "payment_orders.gateway_response", - )?; - sqlx::query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, - exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, - gateway_order_id, gateway_response, status, created_at, paid_at, credited_at -) -VALUES (?, ?, ?, ?, ?, NULL, NULL, NULL, 0, ?, ?, ?, ?, 'credited', ?, ?, ?) -"#, - ) - .bind(&order_id) - .bind(&input.order_no) - .bind(&wallet_id) - .bind(&input.user_id) - .bind(amount_usd) - .bind(redeem_code_refundable_amount(&balance_bucket, amount_usd)) - .bind(payment_method) - .bind(&gateway_order_id) - .bind(&gateway_response) - .bind(now) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'topup_card_code', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&wallet_id) - .bind(amount_usd) - .bind(before_recharge + before_gift) - .bind(after_recharge + after_gift) - .bind(before_recharge) - .bind(after_recharge) - .bind(before_gift) - .bind(after_gift) - .bind(&order_id) - .bind("兑换码充值") - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -UPDATE redeem_codes -SET status = 'redeemed', - redeemed_by_user_id = ?, - redeemed_wallet_id = ?, - redeemed_payment_order_id = ?, - redeemed_at = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&input.user_id) - .bind(&wallet_id) - .bind(&order_id) - .bind(now) - .bind(now) - .bind(&code_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - sqlx::query("UPDATE redeem_code_batches SET updated_at = ? WHERE id = ?") - .bind(now) - .bind(&batch_id) - .execute(&mut *tx) - .await - .map_sql_err()?; - - let wallet = map_wallet_row(&sqlite_wallet_by_id(&mut tx, &wallet_id).await?)?; - let order = map_payment_order_row(&sqlite_payment_order_by_id(&mut tx, &order_id).await?)?; - tx.commit().await.map_sql_err()?; - Ok(RedeemWalletCodeOutcome::Redeemed { - wallet, - order, - amount_usd, - batch_name, - }) - } -} - -fn daily_usage_select_sql(suffix: &'static str) -> String { - format!( - r#" -SELECT - id, billing_date, billing_timezone, total_cost_usd, total_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - first_finalized_at AS first_finalized_at_unix_secs, - last_finalized_at AS last_finalized_at_unix_secs, - aggregated_at AS aggregated_at_unix_secs -FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_timezone = ? - {suffix} -"# - ) -} - -fn current_billing_date(billing_timezone: &str) -> Result { - let timezone = billing_timezone.parse::().map_err(|err| { - DataLayerError::InvalidInput(format!("invalid wallet billing timezone: {err}")) - })?; - Ok(Utc::now().with_timezone(&timezone).date_naive().to_string()) -} - -fn map_wallet_row(row: &SqliteRow) -> Result { - StoredWalletSnapshot::new( - get(row, "id")?, - get(row, "user_id")?, - get(row, "api_key_id")?, - sqlite_real(row, "balance")?, - sqlite_real(row, "gift_balance")?, - get(row, "limit_mode")?, - get(row, "currency")?, - get(row, "status")?, - sqlite_real(row, "total_recharged")?, - sqlite_real(row, "total_consumed")?, - sqlite_real(row, "total_refunded")?, - sqlite_real(row, "total_adjusted")?, - get(row, "updated_at_unix_secs")?, - ) -} - -fn map_admin_wallet_list_item_row( - row: &SqliteRow, -) -> Result { - Ok(StoredAdminWalletListItem { - id: get(row, "id")?, - user_id: get(row, "user_id")?, - api_key_id: get(row, "api_key_id")?, - balance: sqlite_real(row, "balance")?, - gift_balance: sqlite_real(row, "gift_balance")?, - limit_mode: get(row, "limit_mode")?, - currency: get(row, "currency")?, - status: get(row, "status")?, - total_recharged: sqlite_real(row, "total_recharged")?, - total_consumed: sqlite_real(row, "total_consumed")?, - total_refunded: sqlite_real(row, "total_refunded")?, - total_adjusted: sqlite_real(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: &SqliteRow, -) -> Result { - Ok(StoredAdminWalletLedgerItem { - id: get(row, "id")?, - wallet_id: get(row, "wallet_id")?, - category: get(row, "category")?, - reason_code: get(row, "reason_code")?, - amount: sqlite_real(row, "amount")?, - balance_before: sqlite_real(row, "balance_before")?, - balance_after: sqlite_real(row, "balance_after")?, - recharge_balance_before: sqlite_real(row, "recharge_balance_before")?, - recharge_balance_after: sqlite_real(row, "recharge_balance_after")?, - gift_balance_before: sqlite_real(row, "gift_balance_before")?, - gift_balance_after: sqlite_real(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: &SqliteRow, -) -> 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: sqlite_real(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 sqlite_bool(value: bool) -> i64 { - if value { - 1 - } else { - 0 - } -} - -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}")) - }) -} - -fn plan_entitlements_snapshot(snapshot: &serde_json::Value) -> serde_json::Value { - snapshot - .get("entitlements") - .or_else(|| snapshot.get("entitlements_json")) - .cloned() - .unwrap_or_else(|| serde_json::json!([])) -} - -fn plan_max_active_per_user(snapshot: &serde_json::Value) -> i64 { - snapshot - .get("max_active_per_user") - .and_then(|value| value.as_i64()) - .unwrap_or(1) - .max(1) -} - -fn plan_purchase_limit_scope(snapshot: &serde_json::Value) -> &str { - match snapshot - .get("purchase_limit_scope") - .and_then(|value| value.as_str()) - { - Some("lifetime") => "lifetime", - Some("unlimited") => "unlimited", - _ => "active_period", - } -} - -async fn replace_matching_plan_entitlements_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - user_id: &str, - snapshot: &serde_json::Value, - now: i64, -) -> Result<(), DataLayerError> { - let incoming_entitlements = plan_entitlements_snapshot(snapshot); - if !entitlements_have_replacement_selector(&incoming_entitlements) { - return Ok(()); - } - - let rows = sqlx::query( - r#" -SELECT id, entitlements_snapshot -FROM user_plan_entitlements -WHERE user_id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(user_id) - .bind(now) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - - for row in rows { - let entitlements = optional_json( - get::>(&row, "entitlements_snapshot")?, - "user_plan_entitlements.entitlements_snapshot", - )? - .unwrap_or_else(|| serde_json::json!([])); - let should_replace = - entitlements_should_replace_existing(&incoming_entitlements, &entitlements); - if !should_replace { - continue; - } - let entitlement_id: String = get(&row, "id")?; - sqlx::query( - r#" -UPDATE user_plan_entitlements -SET status = 'replaced', - expires_at = CASE WHEN expires_at > ? THEN ? ELSE expires_at END, - updated_at = ? -WHERE id = ? - AND status = 'active' - AND expires_at > ? - "#, - ) - .bind(now) - .bind(now) - .bind(now) - .bind(entitlement_id) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - Ok(()) -} - -fn plan_expires_at_unix( - snapshot: &serde_json::Value, - starts_at_unix_secs: i64, -) -> Result { - let days = - checked_plan_duration_days_from_snapshot(snapshot).map_err(DataLayerError::InvalidInput)?; - let seconds = days.checked_mul(86_400).ok_or_else(|| { - DataLayerError::InvalidInput("plan duration exceeds the supported range".to_string()) - })?; - starts_at_unix_secs.checked_add(seconds).ok_or_else(|| { - DataLayerError::InvalidInput("plan expiration exceeds the supported range".to_string()) - }) -} - -async fn apply_plan_wallet_credit_sqlite( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - wallet_id: &str, - order_id: &str, - payment_method: &str, - entitlements: &serde_json::Value, - now: i64, -) -> Result<(), DataLayerError> { - validate_plan_wallet_credit_entitlements(entitlements).map_err(DataLayerError::InvalidInput)?; - let credits = entitlements - .as_array() - .into_iter() - .flatten() - .filter(|item| { - item.get("type") - .and_then(|value| value.as_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("wallet_credit")) - }) - .filter_map(|item| { - let amount = item.get("amount_usd").and_then(|value| value.as_f64())?; - if amount <= 0.0 || !amount.is_finite() { - return None; - } - let bucket = item - .get("balance_bucket") - .and_then(|value| value.as_str()) - .unwrap_or("gift") - .trim() - .to_ascii_lowercase(); - Some((amount, bucket)) - }) - .collect::>(); - if credits.is_empty() { - return Ok(()); - } - let Some(wallet_row) = - sqlx::query("SELECT id, status, balance, gift_balance, total_recharged FROM wallets WHERE id = ? LIMIT 1") - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - else { - return Err(DataLayerError::UnexpectedValue( - "wallet not found for plan wallet_credit".to_string(), - )); - }; - let status: String = get(&wallet_row, "status")?; - if status != "active" { - return Err(DataLayerError::UnexpectedValue( - "wallet is not active for plan wallet_credit".to_string(), - )); - } - let mut recharge_balance = sqlite_real(&wallet_row, "balance")?; - let mut gift_balance = sqlite_real(&wallet_row, "gift_balance")?; - let mut total_recharged = sqlite_real(&wallet_row, "total_recharged")?; - if !recharge_balance.is_finite() - || !gift_balance.is_finite() - || gift_balance < 0.0 - || !total_recharged.is_finite() - || total_recharged < 0.0 - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance is invalid for plan wallet_credit".to_string(), - )); - } - for (amount, bucket) in credits { - let before_recharge = recharge_balance; - let before_gift = gift_balance; - let before_total = before_recharge + before_gift; - let credits_recharge = bucket == "recharge"; - if credits_recharge { - recharge_balance += amount; - total_recharged += amount; - } else { - gift_balance += amount; - } - let after_total = recharge_balance + gift_balance; - if !before_total.is_finite() - || !recharge_balance.is_finite() - || !gift_balance.is_finite() - || !total_recharged.is_finite() - || !after_total.is_finite() - { - return Err(DataLayerError::UnexpectedValue( - "wallet balance overflow for plan wallet_credit".to_string(), - )); - } - sqlx::query( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - total_recharged = ?, - updated_at = ? -WHERE id = ? - "#, - ) - .bind(recharge_balance) - .bind(gift_balance) - .bind(total_recharged) - .bind(now) - .bind(wallet_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'recharge', 'plan_wallet_credit', ?, ?, ?, ?, ?, ?, ?, 'payment_order', ?, NULL, ?, ?) - "#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(wallet_id) - .bind(amount) - .bind(before_total) - .bind(after_total) - .bind(before_recharge) - .bind(recharge_balance) - .bind(before_gift) - .bind(gift_balance) - .bind(order_id) - .bind(format!("套餐附赠余额({payment_method})")) - .bind(now) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - Ok(()) -} - -fn payment_gateway_response_map( - value: Option, -) -> serde_json::Map { - match value { - Some(serde_json::Value::Object(map)) => map, - _ => serde_json::Map::new(), - } -} - -fn normalize_redeem_code(value: &str) -> Option { - let normalized = value - .chars() - .filter(|ch| ch.is_ascii_alphanumeric()) - .map(|ch| ch.to_ascii_uppercase()) - .collect::(); - if normalized.len() < 16 { - None - } else { - Some(normalized) - } -} - -fn hash_redeem_code(normalized: &str) -> String { - use sha2::Digest; - - format!("{:x}", sha2::Sha256::digest(normalized.as_bytes())) -} - -fn format_redeem_code(normalized: &str) -> String { - normalized - .as_bytes() - .chunks(8) - .map(|chunk| std::str::from_utf8(chunk).unwrap_or_default()) - .collect::>() - .join("-") -} - -fn generate_redeem_code_candidate() -> (String, String, String, String, String, String) { - let normalized = uuid::Uuid::new_v4() - .simple() - .to_string() - .to_ascii_uppercase(); - let code = format_redeem_code(&normalized); - let code_id = uuid::Uuid::new_v4().to_string(); - let prefix = normalized.chars().take(4).collect::(); - let suffix = normalized - .chars() - .rev() - .take(4) - .collect::>() - .into_iter() - .rev() - .collect::(); - let masked_code = mask_redeem_code(&prefix, &suffix); - let code_hash = hash_redeem_code(&normalized); - (code_id, code, masked_code, code_hash, prefix, suffix) -} - -fn wallet_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -fn apply_admin_balance_adjustment( - amount_usd: f64, - balance_type: &str, - recharge_balance: &mut f64, - gift_balance: &mut f64, -) { - if amount_usd > 0.0 { - if balance_type.eq_ignore_ascii_case("gift") { - *gift_balance += amount_usd; - } else { - *recharge_balance += amount_usd; - } - return; - } - - let mut remaining = -amount_usd; - let consume_positive_bucket = |balance: &mut f64, to_consume: &mut f64| { - if *to_consume <= 0.0 { - return; - } - let available = (*balance).max(0.0); - let consumed = available.min(*to_consume); - *balance -= consumed; - *to_consume -= consumed; - }; - if balance_type.eq_ignore_ascii_case("gift") { - consume_positive_bucket(gift_balance, &mut remaining); - consume_positive_bucket(recharge_balance, &mut remaining); - } else { - consume_positive_bucket(recharge_balance, &mut remaining); - consume_positive_bucket(gift_balance, &mut remaining); - } - if remaining > 0.0 { - *recharge_balance -= remaining; - } -} - -async fn sqlite_wallet_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - wallet_id: &str, -) -> Result { - let sql = wallet_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(wallet_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_wallet_by_id_optional( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - wallet_id: &str, -) -> Result, DataLayerError> { - let sql = wallet_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_wallet_by_user_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - user_id: &str, -) -> Result, DataLayerError> { - let sql = wallet_select_sql("WHERE user_id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -fn payment_order_select_sql(where_clause: &str) -> String { - format!( - 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, 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_clause} -"# - ) -} - -async fn sqlite_payment_order_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - order_id: &str, -) -> Result { - let sql = payment_order_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(order_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_payment_order_by_id_optional( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - order_id: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(order_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_payment_order_by_order_no( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - order_no: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE order_no = ? LIMIT 1"); - sqlx::query(&sql) - .bind(order_no) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_payment_order_by_gateway_order_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - payment_method: &str, - gateway_order_id: &str, -) -> Result, DataLayerError> { - let sql = payment_order_select_sql("WHERE payment_method = ? AND gateway_order_id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(payment_method) - .bind(gateway_order_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_bind_payment_gateway_order_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - order_id: &str, - gateway_order_id: &str, -) -> Result { - sqlx::query("SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - let bind_result = sqlx::query("UPDATE payment_orders SET gateway_order_id = ? WHERE id = ?") - .bind(gateway_order_id) - .bind(order_id) - .execute(&mut **tx) - .await; - match bind_result { - Ok(_) => { - sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(true) - } - Err(error) - if error - .as_database_error() - .is_some_and(|database_error| database_error.is_unique_violation()) => - { - sqlx::query("ROLLBACK TO SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(false) - } - Err(error) => { - let _ = sqlx::query("ROLLBACK TO SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await; - let _ = sqlx::query("RELEASE SAVEPOINT payment_gateway_order_binding") - .execute(&mut **tx) - .await; - Err(DataLayerError::sql(error)) - } - } -} - -fn refund_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -async fn sqlite_refund_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - refund_id: &str, -) -> Result { - let sql = refund_select_sql("WHERE id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(refund_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_refund_by_id_and_wallet( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - refund_id: &str, - wallet_id: &str, -) -> Result, DataLayerError> { - let sql = refund_select_sql("WHERE id = ? AND wallet_id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(refund_id) - .bind(wallet_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -async fn sqlite_refund_by_idempotency( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - user_id: &str, - idempotency_key: &str, -) -> Result, DataLayerError> { - let sql = refund_select_sql("WHERE user_id = ? AND idempotency_key = ? LIMIT 1"); - sqlx::query(&sql) - .bind(user_id) - .bind(idempotency_key) - .fetch_optional(&mut **tx) - .await - .map_sql_err() -} - -fn redeem_batch_select_sql(where_clause: &str) -> String { - format!( - r#" -SELECT - batches.id, batches.name, batches.amount_usd, batches.currency, - batches.balance_bucket, batches.total_count, - COALESCE(SUM(CASE WHEN codes.status = 'redeemed' THEN 1 ELSE 0 END), 0) AS redeemed_count, - COALESCE(SUM(CASE WHEN codes.status = 'active' THEN 1 ELSE 0 END), 0) 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 -{where_clause} -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 -"# - ) -} - -async fn sqlite_redeem_batch_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - batch_id: &str, -) -> Result { - let sql = redeem_batch_select_sql("WHERE batches.id = ?"); - sqlx::query(&sql) - .bind(batch_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -fn redeem_code_select_sql(where_clause: &str) -> String { - format!( - 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 -{where_clause} -"# - ) -} - -async fn sqlite_redeem_code_by_id( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - code_id: &str, -) -> Result { - let sql = redeem_code_select_sql("WHERE codes.id = ? LIMIT 1"); - sqlx::query(&sql) - .bind(code_id) - .fetch_one(&mut **tx) - .await - .map_sql_err() -} - -async fn update_sqlite_payment_callback_failure( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - callback_id: &str, - input: &ProcessPaymentCallbackInput, - _payload: &str, - error: &str, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -UPDATE payment_callbacks -SET signature_valid = ?, - status = 'failed', - error_message = ?, - payload_hash = ?, - payload = NULL, - processed_at = ? -WHERE id = ? -"#, - ) - .bind(sqlite_bool(input.signature_valid)) - .bind(error) - .bind(&input.payload_hash) - .bind(current_unix_secs_i64()) - .bind(callback_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -async fn mark_sqlite_payment_callback_processed( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - callback_id: &str, - input: &ProcessPaymentCallbackInput, - _payload: &str, - order_id: &str, - order_no: &str, -) -> Result<(), DataLayerError> { - sqlx::query( - r#" -UPDATE payment_callbacks -SET payment_order_id = ?, - signature_valid = 1, - status = 'processed', - error_message = NULL, - payload_hash = ?, - payload = NULL, - processed_at = ?, - order_no = ?, - gateway_order_id = COALESCE(?, gateway_order_id) -WHERE id = ? -"#, - ) - .bind(order_id) - .bind(&input.payload_hash) - .bind(current_unix_secs_i64()) - .bind(order_no) - .bind(input.gateway_order_id.as_deref()) - .bind(callback_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -async fn update_sqlite_wallet_snapshot( - pool: &SqlitePool, - owner_column: &str, - owner_id: &str, - balance: f64, - gift_balance: f64, - limit_mode: &str, - currency: &str, - status: &str, - total_recharged: f64, - total_consumed: f64, - total_refunded: f64, - total_adjusted: f64, - updated_at_unix_secs: Option, -) -> Result<(), DataLayerError> { - let owner_predicate = match owner_column { - "user_id" => "user_id = ?", - "api_key_id" => "api_key_id = ?", - _ => { - return Err(DataLayerError::UnexpectedValue(format!( - "unsupported wallet owner column: {owner_column}" - ))); - } - }; - let sql = format!( - r#" -UPDATE wallets -SET balance = ?, - gift_balance = ?, - limit_mode = ?, - currency = ?, - status = ?, - total_recharged = ?, - total_consumed = ?, - total_refunded = ?, - total_adjusted = ?, - updated_at = ? -WHERE {owner_predicate} -"# - ); - sqlx::query(&sql) - .bind(balance) - .bind(gift_balance) - .bind(limit_mode) - .bind(currency) - .bind(status) - .bind(total_recharged) - .bind(total_consumed) - .bind(total_refunded) - .bind(total_adjusted) - .bind( - updated_at_unix_secs - .map(|value| value as i64) - .unwrap_or_else(current_unix_secs_i64), - ) - .bind(owner_id) - .execute(pool) - .await - .map_sql_err()?; - Ok(()) -} - -async fn initialize_sqlite_auth_wallet( - pool: &SqlitePool, - user_id: Option<&str>, - api_key_id: Option<&str>, - initial_gift_usd: f64, - unlimited: bool, -) -> Result, DataLayerError> { - let owner = user_id - .or(api_key_id) - .filter(|value| !value.trim().is_empty()); - if owner.is_none() || (user_id.is_some() && api_key_id.is_some()) { - return Err(DataLayerError::InvalidInput( - "wallet owner must be exactly one non-empty user or API-key id".to_string(), - )); - } - if !initial_gift_usd.is_finite() { - return Err(DataLayerError::InvalidInput( - "initial gift amount must be finite".to_string(), - )); - } - let gift_amount = if unlimited { - 0.0 - } else { - initial_gift_usd.max(0.0) - }; - let now = current_unix_secs_i64(); - let wallet = StoredWalletSnapshot::new( - uuid::Uuid::new_v4().to_string(), - user_id.map(str::to_string), - api_key_id.map(str::to_string), - 0.0, - gift_amount, - if unlimited { "unlimited" } else { "finite" }.to_string(), - "USD".to_string(), - "active".to_string(), - 0.0, - 0.0, - 0.0, - gift_amount, - now, - )?; - // Keep the owner lookup and insert in one writer transaction so concurrent - // initialization retries cannot mint duplicate wallets or gift entries. - let mut tx = pool.begin_with("BEGIN IMMEDIATE").await.map_sql_err()?; - let owner_column = if user_id.is_some() { - "user_id" - } else { - "api_key_id" - }; - let owner_value = owner.expect("validated wallet owner"); - - // SQLite's baseline schema does not declare owner foreign keys on wallets. - // Validate the owner while holding the writer transaction so a wallet (and - // its initial gift journal entry) can never be created for a missing auth - // record. - if let Some(user_id) = user_id { - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - } else { - let api_key_user_id: Option = - sqlx::query_scalar("SELECT user_id FROM api_keys WHERE id = ?") - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - let Some(api_key_user_id) = api_key_user_id else { - tx.rollback().await.map_sql_err()?; - return Ok(None); - }; - let user_exists: Option = sqlx::query_scalar("SELECT id FROM users WHERE id = ?") - .bind(&api_key_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if user_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - let api_key_exists: Option = - sqlx::query_scalar("SELECT id FROM api_keys WHERE id = ? AND user_id = ?") - .bind(owner_value) - .bind(&api_key_user_id) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if api_key_exists.is_none() { - tx.rollback().await.map_sql_err()?; - return Ok(None); - } - } - - let existing_row = sqlx::query(&wallet_select_sql(&format!( - "WHERE {owner_column} = ? LIMIT 1" - ))) - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if let Some(row) = existing_row { - let existing = map_wallet_row(&row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some((existing, false))); - } - - let insert_result = sqlx::query( - r#" -INSERT INTO wallets ( - id, user_id, api_key_id, balance, gift_balance, limit_mode, currency, - status, total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) -VALUES (?, ?, ?, 0, ?, ?, 'USD', 'active', 0, 0, 0, ?, ?, ?) -ON CONFLICT DO NOTHING -"#, - ) - .bind(&wallet.id) - .bind(user_id) - .bind(api_key_id) - .bind(gift_amount) - .bind(&wallet.limit_mode) - .bind(gift_amount) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - if insert_result.rows_affected() == 0 { - // Another initializer may have won the owner race. If no owner row is - // visible, the generated wallet id collided with a different owner. - let row = sqlx::query(&wallet_select_sql(&format!( - "WHERE {owner_column} = ? LIMIT 1" - ))) - .bind(owner_value) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if let Some(row) = row { - let existing = map_wallet_row(&row)?; - tx.commit().await.map_sql_err()?; - return Ok(Some((existing, false))); - } - tx.rollback().await.map_sql_err()?; - return Err(DataLayerError::InvalidInput( - "wallet identifier already belongs to another owner".to_string(), - )); - } - if gift_amount > 0.0 { - let link_id = user_id.or(api_key_id).unwrap_or_default(); - let description = if api_key_id.is_some() { - "独立余额 Key 初始赠款" - } else { - "用户初始赠款" - }; - sqlx::query( - r#" -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, link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'gift', 'gift_initial', ?, 0, ?, 0, 0, 0, ?, 'system_task', ?, NULL, ?, ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&wallet.id) - .bind(gift_amount) - .bind(gift_amount) - .bind(gift_amount) - .bind(link_id) - .bind(description) - .bind(now) - .execute(&mut *tx) - .await - .map_sql_err()?; - } - let row = sqlite_wallet_by_id(&mut tx, &wallet.id).await?; - let wallet = map_wallet_row(&row)?; - tx.commit().await.map_sql_err()?; - Ok(Some((wallet, true))) -} - -fn sqlite_wallet_recharge_replay_matches( - row: &SqliteRow, - wallet_id: &str, - input: &CreateWalletRechargeOrderInput, -) -> Result { - let existing_wallet_id: String = get(row, "wallet_id")?; - let pay_currency: Option = get(row, "pay_currency")?; - let payment_method: String = get(row, "payment_method")?; - let payment_provider: Option = get(row, "payment_provider")?; - let payment_channel: Option = get(row, "payment_channel")?; - Ok(wallet_recharge_replay_matches( - &existing_wallet_id, - sqlite_real(row, "amount_usd")?, - sqlite_optional_real(row, "pay_amount")?, - pay_currency.as_deref(), - sqlite_optional_real(row, "exchange_rate")?, - &payment_method, - payment_provider.as_deref(), - payment_channel.as_deref(), - wallet_id, - input, - )) -} - -fn map_payment_order_row(row: &SqliteRow) -> Result { - Ok(StoredAdminPaymentOrder { - id: get(row, "id")?, - order_no: get(row, "order_no")?, - wallet_id: get(row, "wallet_id")?, - user_id: get(row, "user_id")?, - amount_usd: sqlite_real(row, "amount_usd")?, - pay_amount: sqlite_optional_real(row, "pay_amount")?, - pay_currency: get(row, "pay_currency")?, - exchange_rate: sqlite_optional_real(row, "exchange_rate")?, - refunded_amount_usd: sqlite_real(row, "refunded_amount_usd")?, - refundable_amount_usd: sqlite_real(row, "refundable_amount_usd")?, - payment_method: get(row, "payment_method")?, - payment_provider: get(row, "payment_provider")?, - order_kind: get(row, "order_kind")?, - gateway_order_id: get(row, "gateway_order_id")?, - gateway_response: optional_json( - get(row, "gateway_response")?, - "payment_orders.gateway_response", - )?, - status: get(row, "status")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "payment_orders.created_at", - )?, - paid_at_unix_secs: optional_timestamp( - get(row, "paid_at_unix_secs")?, - "payment_orders.paid_at", - )?, - credited_at_unix_secs: optional_timestamp( - get(row, "credited_at_unix_secs")?, - "payment_orders.credited_at", - )?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "payment_orders.expires_at", - )?, - }) -} - -fn map_payment_callback_row(row: &SqliteRow) -> Result { - Ok(StoredAdminPaymentCallback { - id: get(row, "id")?, - payment_order_id: get(row, "payment_order_id")?, - payment_method: get(row, "payment_method")?, - callback_key: get(row, "callback_key")?, - order_no: get(row, "order_no")?, - gateway_order_id: get(row, "gateway_order_id")?, - payload_hash: get(row, "payload_hash")?, - signature_valid: get::(row, "signature_valid")? != 0, - status: get(row, "status")?, - payload: optional_json(get(row, "payload")?, "payment_callbacks.payload")?, - error_message: get(row, "error_message")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "payment_callbacks.created_at", - )?, - processed_at_unix_secs: optional_timestamp( - get(row, "processed_at_unix_secs")?, - "payment_callbacks.processed_at", - )?, - }) -} - -fn map_wallet_transaction_row( - row: &SqliteRow, -) -> Result { - Ok(StoredAdminWalletTransaction { - id: get(row, "id")?, - wallet_id: get(row, "wallet_id")?, - category: get(row, "category")?, - reason_code: get(row, "reason_code")?, - amount: sqlite_real(row, "amount")?, - balance_before: sqlite_real(row, "balance_before")?, - balance_after: sqlite_real(row, "balance_after")?, - recharge_balance_before: sqlite_real(row, "recharge_balance_before")?, - recharge_balance_after: sqlite_real(row, "recharge_balance_after")?, - gift_balance_before: sqlite_real(row, "gift_balance_before")?, - gift_balance_after: sqlite_real(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")?, - created_at_unix_ms: optional_timestamp( - get(row, "created_at_unix_ms")?, - "wallet_transactions.created_at", - )?, - }) -} - -fn map_refund_row(row: &SqliteRow) -> Result { - Ok(StoredAdminWalletRefund { - 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: sqlite_real(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")?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "refund_requests.created_at", - )?, - updated_at_unix_secs: 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 map_redeem_batch_row(row: &SqliteRow) -> Result { - Ok(StoredAdminRedeemCodeBatch { - id: get(row, "id")?, - name: get(row, "name")?, - amount_usd: sqlite_real(row, "amount_usd")?, - currency: get(row, "currency")?, - balance_bucket: get(row, "balance_bucket")?, - total_count: nonnegative_u64(get(row, "total_count")?, "redeem_code_batches.total_count")?, - redeemed_count: nonnegative_u64( - get(row, "redeemed_count")?, - "redeem_codes.redeemed_count", - )?, - active_count: nonnegative_u64(get(row, "active_count")?, "redeem_codes.active_count")?, - status: get(row, "status")?, - description: get(row, "description")?, - created_by: get(row, "created_by")?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "redeem_code_batches.expires_at", - )?, - created_at_unix_ms: timestamp( - get(row, "created_at_unix_ms")?, - "redeem_code_batches.created_at", - )?, - updated_at_unix_secs: timestamp( - get(row, "updated_at_unix_secs")?, - "redeem_code_batches.updated_at", - )?, - }) -} - -fn map_redeem_code_row(row: &SqliteRow) -> Result { - let code_prefix: String = get(row, "code_prefix")?; - let code_suffix: String = get(row, "code_suffix")?; - Ok(StoredAdminRedeemCode { - id: get(row, "id")?, - batch_id: get(row, "batch_id")?, - batch_name: get(row, "batch_name")?, - masked_code: mask_redeem_code(&code_prefix, &code_suffix), - code_prefix, - code_suffix, - status: get(row, "status")?, - redeemed_by_user_id: get(row, "redeemed_by_user_id")?, - redeemed_by_user_name: get(row, "redeemed_by_user_name")?, - redeemed_wallet_id: get(row, "redeemed_wallet_id")?, - redeemed_payment_order_id: get(row, "redeemed_payment_order_id")?, - redeemed_order_no: get(row, "redeemed_order_no")?, - redeemed_at_unix_secs: optional_timestamp( - get(row, "redeemed_at_unix_secs")?, - "redeem_codes.redeemed_at", - )?, - disabled_by: get(row, "disabled_by")?, - expires_at_unix_secs: optional_timestamp( - get(row, "expires_at_unix_secs")?, - "redeem_code_batches.expires_at", - )?, - created_at_unix_ms: timestamp(get(row, "created_at_unix_ms")?, "redeem_codes.created_at")?, - updated_at_unix_secs: timestamp( - get(row, "updated_at_unix_secs")?, - "redeem_codes.updated_at", - )?, - }) -} - -fn map_daily_usage_row(row: &SqliteRow) -> Result { - Ok(StoredWalletDailyUsageLedger { - id: get(row, "id")?, - billing_date: get(row, "billing_date")?, - billing_timezone: get(row, "billing_timezone")?, - total_cost_usd: sqlite_real(row, "total_cost_usd")?, - total_requests: nonnegative_u64( - get(row, "total_requests")?, - "wallet_daily_usage_ledgers.total_requests", - )?, - input_tokens: nonnegative_u64( - get(row, "input_tokens")?, - "wallet_daily_usage_ledgers.input_tokens", - )?, - output_tokens: nonnegative_u64( - get(row, "output_tokens")?, - "wallet_daily_usage_ledgers.output_tokens", - )?, - cache_creation_tokens: nonnegative_u64( - get(row, "cache_creation_tokens")?, - "wallet_daily_usage_ledgers.cache_creation_tokens", - )?, - cache_read_tokens: nonnegative_u64( - get(row, "cache_read_tokens")?, - "wallet_daily_usage_ledgers.cache_read_tokens", - )?, - first_finalized_at_unix_secs: optional_timestamp( - get(row, "first_finalized_at_unix_secs")?, - "wallet_daily_usage_ledgers.first_finalized_at", - )?, - last_finalized_at_unix_secs: optional_timestamp( - get(row, "last_finalized_at_unix_secs")?, - "wallet_daily_usage_ledgers.last_finalized_at", - )?, - aggregated_at_unix_secs: optional_timestamp( - get(row, "aggregated_at_unix_secs")?, - "wallet_daily_usage_ledgers.aggregated_at", - )?, - }) -} - -fn get(row: &SqliteRow, field: &str) -> Result -where - for<'r> T: sqlx::Decode<'r, sqlx::Sqlite> + sqlx::Type, -{ - row.try_get(field).map_sql_err() -} - -fn read_count_row(row: SqliteRow) -> Result { - nonnegative_u64(get(&row, "total")?, "count total") -} - -fn optional_json( - value: Option, - field_name: &str, -) -> Result, DataLayerError> { - value - .map(|value| { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!( - "{field_name} contains invalid JSON: {err}" - )) - }) - }) - .transpose() -} - -fn timestamp(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("{field_name} contains a negative timestamp")) - }) -} - -fn optional_timestamp(value: Option, field_name: &str) -> Result, DataLayerError> { - value.map(|value| timestamp(value, field_name)).transpose() -} - -fn nonnegative_u64(value: i64, field_name: &str) -> Result { - u64::try_from(value).map_err(|_| { - DataLayerError::UnexpectedValue(format!("{field_name} contains a negative value")) - }) -} - -fn mask_redeem_code(prefix: &str, suffix: &str) -> String { - format!("{prefix}****{suffix}") -} - -#[cfg(test)] -mod tests; diff --git a/crates/aether-data/adapters/sqlite/src/wallet/tests.rs b/crates/aether-data/adapters/sqlite/src/wallet/tests.rs deleted file mode 100644 index 7ac05fc50..000000000 --- a/crates/aether-data/adapters/sqlite/src/wallet/tests.rs +++ /dev/null @@ -1,5478 +0,0 @@ -use super::{replace_matching_plan_entitlements_sqlite, SqliteWalletReadRepository}; -use crate::run_migrations; -use aether_data_contracts::repository::wallet::{ - AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeListQuery, - AdminWalletListQuery, CompareAndSwapPaymentOrderStripeClientSecretInput, - CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput, - CreateManualWalletRechargeInput, CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, - CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome, - CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, - DeleteAdminRedeemCodeBatchInput, DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, - FailAdminWalletRefundInput, FailWalletRechargeCheckoutInput, ProcessAdminWalletRefundInput, - ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, RedeemWalletCodeInput, - RedeemWalletCodeOutcome, UpdateAdminWalletRefundGatewayInput, - UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome, - WalletReadRepository, WalletWriteRepository, -}; -use aether_data_contracts::DataLayerError; -use serde_json::json; -use std::{sync::Arc, time::Duration}; - -async fn ensure_test_user(pool: &sqlx::SqlitePool, user_id: &str) { - let username = format!("wallet-test-{user_id}"); - let email = format!("{user_id}@wallet-test.example"); - sqlx::query( - "INSERT OR IGNORE INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, 'local', 1, 1)", - ) - .bind(user_id) - .bind(username) - .bind(email) - .execute(pool) - .await - .expect("test user should seed"); -} - -async fn ensure_test_users(pool: &sqlx::SqlitePool, user_ids: &[&str]) { - for user_id in user_ids { - ensure_test_user(pool, user_id).await; - } -} - -#[tokio::test] -async fn sqlite_stripe_secret_cas_requires_the_exact_locked_row() { - 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"); - ensure_test_user(&pool, "stripe-cas-user").await; - let repository = SqliteWalletReadRepository::new(pool.clone()); - let wallet = repository - .initialize_auth_user_wallet("stripe-cas-user", 0.0, false) - .await - .expect("wallet initialization should run") - .expect("wallet should exist"); - let legacy = "gAAAAABsqlite-legacy"; - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some(wallet.id), - user_id: "stripe-cas-user".to_string(), - amount_usd: 10.0, - pay_amount: Some(10.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "pi-sqlite-cas".to_string(), - gateway_response: json!({ - "gateway": "stripe", - "publishable_key": "pk_test_public", - "_stripe_client_secret_encrypted": legacy, - }), - order_no: "po-sqlite-cas".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("order creation should run") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - other => panic!("unexpected order creation outcome: {other:?}"), - }; - let observed = order - .gateway_response - .clone() - .expect("created order should contain a response"); - let replacement = concat!( - "aether-payment-order-stripe-client-secret-v2:", - "aether-runtime-secret-v1:gAAAAABsqlite-replacement" - ); - let input = CompareAndSwapPaymentOrderStripeClientSecretInput { - order_id: order.id.clone(), - order_no: order.order_no.clone(), - wallet_id: order.wallet_id.clone(), - user_id: order.user_id.clone(), - payment_method: order.payment_method.clone(), - payment_provider: order.payment_provider.clone(), - order_kind: order.order_kind.clone(), - gateway_order_id: order.gateway_order_id.clone(), - expected_status: order.status.clone(), - expected_expires_at_unix_secs: order.expires_at_unix_secs, - expected_gateway_response: observed, - expected_client_secret_encrypted: legacy.to_string(), - replacement_client_secret_encrypted: replacement.to_string(), - }; - - let mut foreign = input.clone(); - foreign.user_id = Some("stripe-cas-foreign-user".to_string()); - assert!(!repository - .compare_and_swap_payment_order_stripe_client_secret(foreign) - .await - .expect("identity mismatch should be a normal CAS miss")); - assert!(repository - .compare_and_swap_payment_order_stripe_client_secret(input.clone()) - .await - .expect("exact CAS should succeed")); - assert!(!repository - .compare_and_swap_payment_order_stripe_client_secret(input) - .await - .expect("stale CAS must not replace the new value")); - - let stored: String = - sqlx::query_scalar("SELECT gateway_response FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("stored gateway response should query"); - let stored: serde_json::Value = - serde_json::from_str(&stored).expect("stored response should be valid JSON"); - assert_eq!( - stored["_stripe_client_secret_encrypted"].as_str(), - Some(replacement) - ); - assert_eq!(stored["publishable_key"], "pk_test_public"); -} - -#[tokio::test] -async fn sqlite_admin_balance_adjustment_rejects_invalid_numbers_without_writes() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "invalid-adjustment-user").await; - let wallet = repository - .initialize_auth_user_wallet("invalid-adjustment-user", 0.0, false) - .await - .expect("wallet initialization should run") - .expect("wallet should exist"); - - for amount_usd in [0.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { - let error = repository - .adjust_wallet_balance(AdjustWalletBalanceInput { - wallet_id: wallet.id.clone(), - amount_usd, - balance_type: "recharge".to_string(), - operator_id: Some("admin-1".to_string()), - description: None, - }) - .await - .expect_err("invalid adjustment should fail before writing"); - assert!(matches!(error, DataLayerError::InvalidInput(_))); - } - - sqlx::query("UPDATE wallets SET balance = ? WHERE id = ?") - .bind(f64::MAX) - .bind(&wallet.id) - .execute(&pool) - .await - .expect("overflow fixture should update"); - let error = repository - .adjust_wallet_balance(AdjustWalletBalanceInput { - wallet_id: wallet.id.clone(), - amount_usd: f64::MAX, - balance_type: "recharge".to_string(), - operator_id: Some("admin-1".to_string()), - description: None, - }) - .await - .expect_err("overflowing adjustment should fail before writing"); - assert!(matches!(error, DataLayerError::UnexpectedValue(_))); - - let stored_balance = sqlx::query_scalar::<_, f64>("SELECT balance FROM wallets WHERE id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("wallet balance should query"); - assert_eq!(stored_balance, f64::MAX); - let adjustment_count = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ? AND category = 'adjust'", - ) - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("adjustment ledger count should query"); - assert_eq!(adjustment_count, 0); -} - -#[tokio::test] -async fn sqlite_manual_wallet_recharge_rejects_invalid_numbers_without_writes() { - 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"); - ensure_test_user(&pool, "invalid-manual-recharge-user").await; - let repository = SqliteWalletReadRepository::new(pool.clone()); - let wallet = repository - .initialize_auth_user_wallet("invalid-manual-recharge-user", 0.0, false) - .await - .expect("wallet initialization should run") - .expect("wallet should exist"); - let initial_order_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM payment_orders WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("payment order count should query"); - let initial_transaction_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("wallet transaction count should query"); - - for (index, amount_usd) in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] - .into_iter() - .enumerate() - { - let error = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id.clone(), - amount_usd, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-invalid-recharge".to_string()), - description: None, - order_no: format!("invalid-manual-recharge-{index}"), - }) - .await - .expect_err("invalid manual recharge should fail before writing"); - assert!(matches!(error, DataLayerError::InvalidInput(_))); - } - - let unchanged = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet should query") - .expect("wallet should remain present"); - assert_eq!(unchanged, wallet); - let order_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM payment_orders WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("payment order count should query"); - let transaction_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("wallet transaction count should query"); - assert_eq!(order_count, initial_order_count); - assert_eq!(transaction_count, initial_transaction_count); - - sqlx::query("UPDATE wallets SET balance = ?, total_recharged = ? WHERE id = ?") - .bind(f64::MAX) - .bind(f64::MAX) - .bind(&wallet.id) - .execute(&pool) - .await - .expect("overflow fixture should update"); - let overflow_fixture = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("overflow fixture should query") - .expect("wallet should remain present"); - let error = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id.clone(), - amount_usd: f64::MAX, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-overflow-recharge".to_string()), - description: None, - order_no: "overflow-manual-recharge".to_string(), - }) - .await - .expect_err("overflowing manual recharge should fail before writing"); - assert!(matches!(error, DataLayerError::InvalidInput(_))); - - let unchanged = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet should query") - .expect("wallet should remain present"); - assert_eq!(unchanged, overflow_fixture); - let order_count_after_overflow: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM payment_orders WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("payment order count should query"); - let transaction_count_after_overflow: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?") - .bind(&wallet.id) - .fetch_one(&pool) - .await - .expect("wallet transaction count should query"); - assert_eq!(order_count_after_overflow, initial_order_count); - assert_eq!(transaction_count_after_overflow, initial_transaction_count); -} - -#[tokio::test] -async fn sqlite_wallet_initialization_rejects_missing_owners_without_writes() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - assert!(repository - .initialize_auth_user_wallet("missing-wallet-user", 5.0, false) - .await - .expect("missing user initialization should resolve") - .is_none()); - assert!(repository - .initialize_auth_api_key_wallet("missing-wallet-api-key", 5.0, false) - .await - .expect("missing api key initialization should resolve") - .is_none()); - - let wallet_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM wallets") - .fetch_one(&pool) - .await - .expect("wallet count should query"); - assert_eq!(wallet_count, 0); - let transaction_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions") - .fetch_one(&pool) - .await - .expect("transaction count should query"); - assert_eq!(transaction_count, 0); - - ensure_test_user(&pool, "wallet-owner-user").await; - sqlx::query( - "INSERT INTO api_keys (id, user_id, key_hash, created_at, updated_at) VALUES (?, ?, ?, ?, ?)", - ) - .bind("wallet-owner-api-key") - .bind("wallet-owner-user") - .bind("wallet-owner-api-key-hash") - .bind(1_i64) - .bind(1_i64) - .execute(&pool) - .await - .expect("api key should seed"); - let initialized = repository - .initialize_auth_api_key_wallet("wallet-owner-api-key", 5.0, false) - .await - .expect("valid api key initialization should resolve") - .expect("valid api key wallet should be created"); - assert_eq!( - initialized.api_key_id.as_deref(), - Some("wallet-owner-api-key") - ); - assert_eq!(initialized.gift_balance, 5.0); -} - -#[tokio::test] -async fn sqlite_wallet_read_repository_reads_wallet_contract_views() { - 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_rows(&pool).await; - - let repository = SqliteWalletReadRepository::new(pool); - let wallet = repository - .find(WalletLookupKey::UserId("user-1")) - .await - .expect("wallet find should query") - .expect("wallet should exist"); - assert_eq!(wallet.total_adjusted, 3.0); - - let page = repository - .list_admin_wallets(&AdminWalletListQuery { - status: Some("active".to_string()), - owner_type: Some("user".to_string()), - limit: 10, - offset: 0, - }) - .await - .expect("admin wallets should list"); - assert_eq!(page.total, 1); - assert_eq!(page.items[0].total_adjusted, 3.0); - - let orders = repository - .list_admin_payment_orders(&AdminPaymentOrderListQuery { - status: Some("credited".to_string()), - payment_method: Some("redeem_code".to_string()), - limit: 10, - offset: 0, - }) - .await - .expect("payment orders should list"); - assert_eq!(orders.total, 1); - assert_eq!( - orders.items[0].gateway_response.as_ref().unwrap()["ok"], - true - ); - - let refunds = repository - .list_admin_wallet_refunds("wallet-1", 10, 0) - .await - .expect("refunds should list"); - assert_eq!(refunds.total, 1); - assert_eq!( - refunds.items[0].payout_proof.as_ref().unwrap()["proof"], - "ok" - ); - - let callbacks = repository - .list_admin_payment_callbacks(Some("redeem_code"), 10, 0) - .await - .expect("callbacks should list"); - assert_eq!(callbacks.total, 1); - assert!(callbacks.items[0].signature_valid); - - let codes = repository - .list_admin_redeem_codes(&AdminRedeemCodeListQuery { - batch_id: "batch-1".to_string(), - status: Some("redeemed".to_string()), - limit: 10, - offset: 0, - }) - .await - .expect("redeem codes should list"); - assert_eq!(codes.total, 1); - assert_eq!(codes.items[0].masked_code, "ABCD****WXYZ"); - - 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'") - .bind(today) - .execute(repository.pool()) - .await - .expect("daily row should update"); - let daily = repository - .find_wallet_today_usage("wallet-1", "UTC") - .await - .expect("daily usage should query") - .expect("daily usage should exist"); - assert_eq!(daily.total_requests, 2); -} - -#[tokio::test] -async fn sqlite_provisional_wallet_cleanup_is_activity_guarded() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - ensure_test_user(&pool, "provisional-user").await; - let provisional_wallet = repository - .initialize_auth_user_wallet("provisional-user", 10.0, false) - .await - .expect("wallet initialization should succeed") - .expect("provisional wallet should exist"); - assert!(repository - .delete_provisional_auth_user_wallet(&provisional_wallet.id, "provisional-user") - .await - .expect("provisional cleanup should succeed")); - assert!(repository - .find(WalletLookupKey::UserId("provisional-user")) - .await - .expect("wallet lookup should succeed") - .is_none()); - - ensure_test_user(&pool, "active-user").await; - repository - .initialize_auth_user_wallet("active-user", 10.0, false) - .await - .expect("wallet initialization should succeed"); - let active_wallet = repository - .find(WalletLookupKey::UserId("active-user")) - .await - .expect("wallet lookup should succeed") - .expect("active wallet should exist"); - sqlx::query( - r#"INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at - ) VALUES (?, ?, '2000-01-01', 'UTC', 1.0, 1, 1, 1, 0, 0, 1, 1, 1)"#, - ) - .bind("active-daily") - .bind(&active_wallet.id) - .execute(&pool) - .await - .expect("activity row should insert"); - assert!(!repository - .delete_provisional_auth_user_wallet(&active_wallet.id, "active-user") - .await - .expect("provisional cleanup should succeed")); - assert!(repository - .find(WalletLookupKey::UserId("active-user")) - .await - .expect("wallet lookup should succeed") - .is_some()); -} - -#[tokio::test] -async fn sqlite_wallet_compensation_delete_is_owner_and_reference_guarded() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - ensure_test_user(&pool, "funded-compensation-user").await; - let funded_wallet = repository - .initialize_auth_user_wallet("funded-compensation-user", 0.0, false) - .await - .expect("wallet initialization should succeed") - .expect("wallet should exist"); - sqlx::query( - "UPDATE wallets SET balance = ?, total_recharged = ?, total_adjusted = ? WHERE id = ?", - ) - .bind(1.0) - .bind(1.0) - .bind(1.0) - .bind(&funded_wallet.id) - .execute(&pool) - .await - .expect("funded wallet update should succeed"); - assert!(!repository - .delete_wallet_if_unreferenced( - &funded_wallet.id, - WalletLookupKey::UserId("funded-compensation-user"), - ) - .await - .expect("funded wallet must not be deleted")); - assert!(repository - .find(WalletLookupKey::UserId("funded-compensation-user")) - .await - .expect("wallet lookup should succeed") - .is_some()); - - ensure_test_user(&pool, "compensation-user").await; - let wallet = repository - .initialize_auth_user_wallet("compensation-user", 0.0, false) - .await - .expect("wallet initialization should succeed") - .expect("wallet should exist"); - let wallet_id = wallet.id.clone(); - - // A caller with a different owner must not be able to reclaim this wallet, even when the - // wallet id is known. - assert!(!repository - .delete_wallet_if_unreferenced(&wallet_id, WalletLookupKey::ApiKeyId("different-api-key"),) - .await - .expect("owner mismatch should be handled cleanly")); - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES (?, ?, '2000-01-01', 'UTC', 0, 1, 0, 0, 0, 0, 1, 1, 1) -"#, - ) - .bind("compensation-daily") - .bind(&wallet_id) - .execute(&pool) - .await - .expect("usage reference should insert"); - - assert!(!repository - .delete_wallet_if_unreferenced(&wallet_id, WalletLookupKey::UserId("compensation-user"),) - .await - .expect("referenced wallet should be retained")); - assert!(repository - .find(WalletLookupKey::UserId("compensation-user")) - .await - .expect("wallet lookup should succeed") - .is_some()); - - sqlx::query("DELETE FROM wallet_daily_usage_ledgers WHERE id = ?") - .bind("compensation-daily") - .execute(&pool) - .await - .expect("usage reference should delete"); - assert!(repository - .delete_wallet_if_unreferenced(&wallet_id, WalletLookupKey::UserId("compensation-user"),) - .await - .expect("unreferenced wallet should be deleted")); - assert!(repository - .find(WalletLookupKey::UserId("compensation-user")) - .await - .expect("wallet lookup should succeed") - .is_none()); -} - -#[tokio::test] -async fn sqlite_wallet_snapshot_compensation_deletes_funded_match_only() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - ensure_test_user(&pool, "snapshot-compensation-user").await; - let wallet = repository - .initialize_auth_user_wallet("snapshot-compensation-user", 0.0, false) - .await - .expect("wallet initialization should succeed") - .expect("wallet should exist"); - sqlx::query( - "UPDATE wallets SET balance = ?, total_recharged = ?, total_adjusted = ? WHERE id = ?", - ) - .bind(12.5) - .bind(12.5) - .bind(0.0) - .bind(&wallet.id) - .execute(&pool) - .await - .expect("wallet funding fixture should update"); - let expected = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet lookup should succeed") - .expect("funded wallet should exist"); - assert!(repository - .delete_wallet_if_snapshot_matches_and_unreferenced( - &expected, - WalletLookupKey::UserId("snapshot-compensation-user"), - ) - .await - .expect("matching snapshot delete should succeed")); - - ensure_test_user(&pool, "snapshot-compensation-changed").await; - let wallet = repository - .initialize_auth_user_wallet("snapshot-compensation-changed", 0.0, false) - .await - .expect("wallet initialization should succeed") - .expect("wallet should exist"); - let expected = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet lookup should succeed") - .expect("wallet should exist"); - sqlx::query("UPDATE wallets SET balance = 1.0 WHERE id = ?") - .bind(&wallet.id) - .execute(&pool) - .await - .expect("concurrent wallet change fixture should update"); - assert!(!repository - .delete_wallet_if_snapshot_matches_and_unreferenced( - &expected, - WalletLookupKey::UserId("snapshot-compensation-changed"), - ) - .await - .expect("changed snapshot should be retained")); - assert!(repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet lookup should succeed") - .is_some()); - - ensure_test_user(&pool, "snapshot-compensation-referenced").await; - let wallet = repository - .initialize_auth_user_wallet("snapshot-compensation-referenced", 0.0, false) - .await - .expect("wallet initialization should succeed") - .expect("wallet should exist"); - let expected = repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet lookup should succeed") - .expect("wallet should exist"); - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES (?, ?, '2000-01-01', 'UTC', 0, 1, 0, 0, 0, 0, 1, 1, 1) -"#, - ) - .bind("snapshot-compensation-reference") - .bind(&wallet.id) - .execute(&pool) - .await - .expect("usage reference should insert"); - assert!(!repository - .delete_wallet_if_snapshot_matches_and_unreferenced( - &expected, - WalletLookupKey::UserId("snapshot-compensation-referenced"), - ) - .await - .expect("referenced snapshot should be retained")); - assert!(repository - .find(WalletLookupKey::WalletId(&wallet.id)) - .await - .expect("wallet lookup should succeed") - .is_some()); -} - -#[tokio::test] -async fn sqlite_wallet_refund_rejects_invalid_amounts_and_foreign_wallets() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - ensure_test_user(&pool, "refund-security-user").await; - let wallet = repository - .initialize_auth_user_wallet("refund-security-user", 0.0, false) - .await - .expect("setup wallet initialization should run") - .expect("setup wallet should exist"); - let (wallet, _) = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id, - amount_usd: 10.0, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-1".to_string()), - description: Some("refund security setup".to_string()), - order_no: "refund-security-recharge".to_string(), - }) - .await - .expect("manual recharge should run") - .expect("setup wallet should exist"); - - for (index, amount_usd) in [0.0, -1.0, f64::NAN, f64::INFINITY].into_iter().enumerate() { - let outcome = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id.clone(), - user_id: wallet - .user_id - .clone() - .expect("setup wallet should have an owner"), - amount_usd, - payment_order_id: None, - source_type: None, - source_id: None, - refund_mode: None, - reason: None, - idempotency_key: Some(format!("refund-security-invalid-{index}")), - refund_no: format!("refund-security-invalid-{index}"), - }) - .await - .expect("invalid refund should be rejected cleanly"); - assert!(matches!( - outcome, - CreateWalletRefundRequestOutcome::InvalidInput(_) - )); - } - - let foreign_outcome = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id, - user_id: "different-user".to_string(), - amount_usd: 1.0, - payment_order_id: None, - source_type: None, - source_id: None, - refund_mode: None, - reason: None, - idempotency_key: Some("refund-security-foreign-wallet".to_string()), - refund_no: "refund-security-foreign-wallet".to_string(), - }) - .await - .expect("foreign wallet refund should be rejected cleanly"); - assert!(matches!( - foreign_outcome, - CreateWalletRefundRequestOutcome::WalletMissing - )); - - let refund_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM refund_requests") - .fetch_one(&pool) - .await - .expect("refund count should query"); - assert_eq!(refund_count, 0); -} - -#[tokio::test] -async fn sqlite_wallet_refund_rejects_invalid_reserved_amounts() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - ensure_test_user(&pool, "refund-reservation-user").await; - let wallet = repository - .initialize_auth_user_wallet("refund-reservation-user", 0.0, false) - .await - .expect("reservation wallet initialization should run") - .expect("reservation wallet should exist"); - let (wallet, order) = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id, - amount_usd: 10.0, - payment_method: "admin_manual".to_string(), - operator_id: Some("reservation-admin".to_string()), - description: Some("reservation setup".to_string()), - order_no: "reservation-order".to_string(), - }) - .await - .expect("reservation recharge should run") - .expect("reservation wallet should still exist"); - - for (id, status, amount_usd) in [ - ("reservation-valid", "pending_approval", 2.0), - ("reservation-negative", "approved", -100.0), - ("reservation-infinite", "pending_approval", f64::INFINITY), - ] { - sqlx::query( - r#" -INSERT INTO refund_requests ( - id, refund_no, wallet_id, user_id, payment_order_id, source_type, - refund_mode, amount_usd, status, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, 'payment_order', 'offline_payout', ?, ?, 1, 1) -"#, - ) - .bind(id) - .bind(format!("{id}-no")) - .bind(&wallet.id) - .bind("refund-reservation-user") - .bind(&order.id) - .bind(amount_usd) - .bind(status) - .execute(&pool) - .await - .expect("corrupt reservation row should insert"); - } - - let outcome = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id, - user_id: "refund-reservation-user".to_string(), - amount_usd: 8.0, - payment_order_id: Some(order.id), - source_type: None, - source_id: None, - refund_mode: None, - reason: Some("reservation regression".to_string()), - idempotency_key: Some("reservation-regression-idempotency".to_string()), - refund_no: "reservation-regression-refund".to_string(), - }) - .await - .expect("reservation request should run"); - assert!(matches!( - outcome, - CreateWalletRefundRequestOutcome::InvalidInput(_) - )); - let refund_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM refund_requests WHERE idempotency_key = ?") - .bind("reservation-regression-idempotency") - .fetch_one(&pool) - .await - .expect("refund count should query"); - assert_eq!(refund_count, 0); -} - -#[derive(Debug, PartialEq)] -struct RefundMutationSnapshot { - wallet_balance: f64, - wallet_total_refunded: f64, - order_refunded_amount: f64, - order_refundable_amount: f64, - refund_status: String, - refund_failure_reason: Option, - refund_gateway_id: Option, - refund_payout_reference: Option, - refund_payout_proof: Option, - wallet_transaction_count: i64, -} - -struct RefundMutationFixture { - wallet_id: String, - payment_order_id: String, - refund_id: String, -} - -async fn create_refund_mutation_fixture( - repository: &SqliteWalletReadRepository, - case_name: &str, -) -> RefundMutationFixture { - let user_id = format!("refund-corruption-{case_name}"); - ensure_test_user(repository.pool(), &user_id).await; - let wallet = repository - .initialize_auth_user_wallet(&user_id, 0.0, false) - .await - .expect("refund corruption wallet initialization should run") - .expect("refund corruption wallet should exist"); - let (wallet, order) = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id, - amount_usd: 10.0, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-refund-corruption".to_string()), - description: Some("refund corruption setup".to_string()), - order_no: format!("refund-corruption-order-{case_name}"), - }) - .await - .expect("refund corruption recharge should run") - .expect("refund corruption wallet should still exist"); - let refund = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id.clone(), - user_id, - amount_usd: 2.0, - payment_order_id: Some(order.id.clone()), - source_type: None, - source_id: None, - refund_mode: None, - reason: Some("refund corruption setup".to_string()), - idempotency_key: Some(format!("refund-corruption-idempotency-{case_name}")), - refund_no: format!("refund-corruption-refund-{case_name}"), - }) - .await - .expect("refund corruption request should run"); - let CreateWalletRefundRequestOutcome::Created(refund) = refund else { - panic!("refund corruption request should be created"); - }; - - RefundMutationFixture { - wallet_id: wallet.id, - payment_order_id: order.id, - refund_id: refund.id, - } -} - -async fn process_refund_mutation_fixture( - repository: &SqliteWalletReadRepository, - fixture: &RefundMutationFixture, -) { - let outcome = repository - .process_admin_wallet_refund(ProcessAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - operator_id: Some("admin-refund-corruption".to_string()), - }) - .await - .expect("refund corruption setup should process"); - assert!(matches!(outcome, WalletMutationOutcome::Applied(_))); -} - -async fn corrupt_refund_amount(pool: &sqlx::SqlitePool, refund_id: &str, amount_usd: f64) { - let result = sqlx::query("UPDATE refund_requests SET amount_usd = ? WHERE id = ?") - .bind(amount_usd) - .bind(refund_id) - .execute(pool) - .await - .expect("persisted refund amount should be corruptible for the regression test"); - assert_eq!(result.rows_affected(), 1); - - let stored_amount: f64 = - sqlx::query_scalar("SELECT amount_usd FROM refund_requests WHERE id = ?") - .bind(refund_id) - .fetch_one(pool) - .await - .expect("corrupted refund amount should query"); - if amount_usd.is_infinite() { - assert!(stored_amount.is_infinite()); - assert_eq!( - stored_amount.is_sign_positive(), - amount_usd.is_sign_positive() - ); - } else { - assert_eq!(stored_amount, amount_usd); - } -} - -async fn refund_mutation_snapshot( - pool: &sqlx::SqlitePool, - fixture: &RefundMutationFixture, -) -> RefundMutationSnapshot { - let (wallet_balance, wallet_total_refunded): (f64, f64) = - sqlx::query_as("SELECT balance, total_refunded FROM wallets WHERE id = ?") - .bind(&fixture.wallet_id) - .fetch_one(pool) - .await - .expect("refund corruption wallet should query"); - let (order_refunded_amount, order_refundable_amount): (f64, f64) = sqlx::query_as( - "SELECT refunded_amount_usd, refundable_amount_usd FROM payment_orders WHERE id = ?", - ) - .bind(&fixture.payment_order_id) - .fetch_one(pool) - .await - .expect("refund corruption payment order should query"); - let ( - refund_status, - refund_failure_reason, - refund_gateway_id, - refund_payout_reference, - refund_payout_proof, - ): ( - String, - Option, - Option, - Option, - Option, - ) = sqlx::query_as( - r#" -SELECT status, failure_reason, gateway_refund_id, payout_reference, payout_proof -FROM refund_requests -WHERE id = ? -"#, - ) - .bind(&fixture.refund_id) - .fetch_one(pool) - .await - .expect("refund corruption request should query"); - let wallet_transaction_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?") - .bind(&fixture.wallet_id) - .fetch_one(pool) - .await - .expect("refund corruption transactions should count"); - - RefundMutationSnapshot { - wallet_balance, - wallet_total_refunded, - order_refunded_amount, - order_refundable_amount, - refund_status, - refund_failure_reason, - refund_gateway_id, - refund_payout_reference, - refund_payout_proof, - wallet_transaction_count, - } -} - -#[tokio::test] -async fn sqlite_process_refund_rejects_corrupt_persisted_amount_without_side_effects() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - for (index, amount_usd) in [0.0, -1.0, f64::INFINITY].into_iter().enumerate() { - let fixture = - create_refund_mutation_fixture(&repository, &format!("process-{index}")).await; - corrupt_refund_amount(&pool, &fixture.refund_id, amount_usd).await; - let before = refund_mutation_snapshot(&pool, &fixture).await; - - let outcome = repository - .process_admin_wallet_refund(ProcessAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - operator_id: Some("admin-refund-corruption".to_string()), - }) - .await - .expect("corrupted refund process should be rejected cleanly"); - - assert!(matches!(outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, &fixture).await, before); - } -} - -#[tokio::test] -async fn sqlite_complete_refund_rejects_corrupt_persisted_amount_without_side_effects() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - for (index, amount_usd) in [0.0, -1.0, f64::INFINITY].into_iter().enumerate() { - let fixture = - create_refund_mutation_fixture(&repository, &format!("complete-{index}")).await; - process_refund_mutation_fixture(&repository, &fixture).await; - corrupt_refund_amount(&pool, &fixture.refund_id, amount_usd).await; - let before = refund_mutation_snapshot(&pool, &fixture).await; - - let outcome = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: Some("must-not-be-stored".to_string()), - payout_reference: Some("must-not-be-stored".to_string()), - payout_proof: Some(json!({ "proof": "must-not-be-stored" })), - }) - .await - .expect("corrupted refund completion should be rejected cleanly"); - - assert!(matches!(outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, &fixture).await, before); - } -} - -#[tokio::test] -async fn sqlite_fail_refund_rejects_corrupt_persisted_amount_without_side_effects() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - for (index, amount_usd) in [0.0, -1.0, f64::INFINITY].into_iter().enumerate() { - let fixture = create_refund_mutation_fixture(&repository, &format!("fail-{index}")).await; - process_refund_mutation_fixture(&repository, &fixture).await; - corrupt_refund_amount(&pool, &fixture.refund_id, amount_usd).await; - let before = refund_mutation_snapshot(&pool, &fixture).await; - - let outcome = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - reason: "must not be stored".to_string(), - operator_id: Some("admin-refund-corruption".to_string()), - }) - .await - .expect("corrupted refund failure should be rejected cleanly"); - - assert!(matches!(outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, &fixture).await, before); - } -} - -#[tokio::test] -async fn sqlite_fail_refund_rejects_negative_recharge_balance_without_side_effects() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let fixture = create_refund_mutation_fixture(&repository, "fail-negative-balance").await; - process_refund_mutation_fixture(&repository, &fixture).await; - - sqlx::query("UPDATE wallets SET balance = ? WHERE id = ?") - .bind(-1.0_f64) - .bind(&fixture.wallet_id) - .execute(&pool) - .await - .expect("negative wallet balance should be seedable for the regression test"); - let before = refund_mutation_snapshot(&pool, &fixture).await; - - let outcome = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - reason: "must not recover a corrupt wallet".to_string(), - operator_id: Some("admin-refund-corruption".to_string()), - }) - .await - .expect("negative wallet balance failure should resolve"); - - assert!(matches!(outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, &fixture).await, before); -} - -#[tokio::test] -async fn sqlite_pending_gateway_refund_evidence_is_durable_and_cannot_be_reverted() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let fixture = create_refund_mutation_fixture(&repository, "pending-gateway").await; - process_refund_mutation_fixture(&repository, &fixture).await; - - let proof = json!({ - "gateway": "wxpay", - "status": "processing", - "refund_no": "provider-refund-1" - }); - let recorded = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-1".to_string(), - payout_proof: Some(proof.clone()), - }) - .await - .expect("gateway evidence update should run"); - let WalletMutationOutcome::Applied(recorded_refund) = recorded else { - panic!("gateway evidence should be recorded"); - }; - assert_eq!(recorded_refund.status, "processing"); - assert_eq!( - recorded_refund.gateway_refund_id.as_deref(), - Some("provider-refund-1") - ); - assert_eq!(recorded_refund.payout_proof, Some(proof.clone())); - - let replay = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-1".to_string(), - payout_proof: Some(json!({ "status": "different" })), - }) - .await - .expect("same gateway evidence replay should run"); - let WalletMutationOutcome::Applied(replayed_refund) = replay else { - panic!("same gateway evidence replay should be accepted"); - }; - assert_eq!(replayed_refund.payout_proof, Some(proof)); - - let conflict = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-attacker".to_string(), - payout_proof: None, - }) - .await - .expect("conflicting gateway evidence should resolve"); - assert!(matches!(conflict, WalletMutationOutcome::Invalid(_))); - - let before_fail = refund_mutation_snapshot(&pool, &fixture).await; - let fail = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - reason: "provider still processing".to_string(), - operator_id: Some("admin-refund-pending".to_string()), - }) - .await - .expect("processing refund failure should resolve"); - assert!(matches!(fail, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, &fixture).await, before_fail); - - let completed = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: None, - payout_reference: None, - payout_proof: None, - }) - .await - .expect("completion should preserve provider evidence"); - let WalletMutationOutcome::Applied(completed_refund) = completed else { - panic!("processing refund should complete"); - }; - assert_eq!(completed_refund.status, "succeeded"); - assert_eq!( - completed_refund.gateway_refund_id.as_deref(), - Some("provider-refund-1") - ); - assert_eq!( - completed_refund.payout_proof, - Some(json!({ - "gateway": "wxpay", - "status": "processing", - "refund_no": "provider-refund-1" - })) - ); - - let terminal_update_conflict = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-attacker".to_string(), - payout_proof: None, - }) - .await - .expect("terminal gateway evidence conflict should resolve"); - assert!(matches!( - terminal_update_conflict, - WalletMutationOutcome::Invalid(_) - )); - - let terminal_complete_conflict = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: fixture.wallet_id, - refund_id: fixture.refund_id, - gateway_refund_id: Some("provider-refund-attacker".to_string()), - payout_reference: None, - payout_proof: None, - }) - .await - .expect("terminal completion conflict should resolve"); - assert!(matches!( - terminal_complete_conflict, - WalletMutationOutcome::Invalid(_) - )); -} - -#[tokio::test] -async fn sqlite_success_gateway_refund_proof_upgrades_processing_evidence() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let fixture = create_refund_mutation_fixture(&repository, "success-proof-upgrade").await; - process_refund_mutation_fixture(&repository, &fixture).await; - - let processing_proof = json!({ - "gateway": "wxpay", - "id": "provider-refund-upgrade", - "status": "processing" - }); - let recorded = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-upgrade".to_string(), - payout_proof: Some(processing_proof.clone()), - }) - .await - .expect("processing evidence should persist"); - assert!(matches!(recorded, WalletMutationOutcome::Applied(_))); - - let replay = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-upgrade".to_string(), - payout_proof: Some(json!({ - "gateway": "wxpay", - "id": "provider-refund-upgrade", - "status": "processing", - "attempt": 2 - })), - }) - .await - .expect("processing replay should resolve"); - let WalletMutationOutcome::Applied(replayed) = replay else { - panic!("processing replay should be accepted"); - }; - assert_eq!(replayed.payout_proof, Some(processing_proof)); - - let success_proof = json!({ - "gateway": "wxpay", - "id": "provider-refund-upgrade", - "status": "success", - "processed_at": "2026-08-29T00:00:00Z" - }); - let upgraded = repository - .update_admin_wallet_refund_gateway(UpdateAdminWalletRefundGatewayInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - gateway_refund_id: "provider-refund-upgrade".to_string(), - payout_proof: Some(success_proof.clone()), - }) - .await - .expect("success evidence should upgrade processing proof"); - let WalletMutationOutcome::Applied(upgraded) = upgraded else { - panic!("success evidence should be accepted"); - }; - assert_eq!(upgraded.payout_proof, Some(success_proof.clone())); - - let completed = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id, - gateway_refund_id: Some("provider-refund-upgrade".to_string()), - payout_reference: None, - payout_proof: None, - }) - .await - .expect("refund should complete"); - let WalletMutationOutcome::Applied(completed) = completed else { - panic!("refund should complete"); - }; - assert_eq!(completed.status, "succeeded"); - assert_eq!(completed.payout_proof, Some(success_proof)); -} - -#[tokio::test] -async fn sqlite_offline_processing_refund_failure_releases_reservation() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let fixture = create_refund_mutation_fixture(&repository, "offline-failure").await; - process_refund_mutation_fixture(&repository, &fixture).await; - let before = refund_mutation_snapshot(&pool, &fixture).await; - - let outcome = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - reason: "offline payout was not sent".to_string(), - operator_id: Some("admin-offline-failure".to_string()), - }) - .await - .expect("offline processing refund failure should resolve"); - let WalletMutationOutcome::Applied((wallet, refund, transaction)) = outcome else { - panic!("offline processing refund should be released"); - }; - let transaction = transaction.expect("refund recovery transaction should be recorded"); - assert_eq!(refund.status, "failed"); - assert_eq!( - refund.failure_reason.as_deref(), - Some("offline payout was not sent") - ); - assert_eq!(transaction.reason_code, "refund_revert"); - assert_eq!(wallet.balance, 10.0); - assert_eq!(wallet.total_refunded, 0.0); - - let after = refund_mutation_snapshot(&pool, &fixture).await; - assert_eq!(after.wallet_balance, before.wallet_balance + 2.0); - assert_eq!(after.wallet_total_refunded, 0.0); - assert_eq!(after.order_refunded_amount, 0.0); - assert_eq!(after.order_refundable_amount, 10.0); - assert_eq!(after.refund_status, "failed"); - assert_eq!( - after.wallet_transaction_count, - before.wallet_transaction_count + 1 - ); -} - -#[tokio::test] -async fn sqlite_processing_refund_requires_offline_mode_without_gateway_evidence() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - let original_channel = create_refund_mutation_fixture(&repository, "original-channel").await; - process_refund_mutation_fixture(&repository, &original_channel).await; - sqlx::query("UPDATE refund_requests SET refund_mode = 'original_channel' WHERE id = ?") - .bind(&original_channel.refund_id) - .execute(&pool) - .await - .expect("refund mode should update for the regression test"); - - let proof_only = create_refund_mutation_fixture(&repository, "proof-only").await; - process_refund_mutation_fixture(&repository, &proof_only).await; - sqlx::query("UPDATE refund_requests SET payout_proof = ? WHERE id = ?") - .bind(r#"{"gateway":"manual-settlement"}"#) - .bind(&proof_only.refund_id) - .execute(&pool) - .await - .expect("proof-only evidence should update for the regression test"); - - for fixture in [&original_channel, &proof_only] { - let before = refund_mutation_snapshot(&pool, fixture).await; - let outcome = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - reason: "must remain reserved".to_string(), - operator_id: Some("admin-preserve-reservation".to_string()), - }) - .await - .expect("protected processing refund failure should resolve"); - assert!(matches!(outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!(refund_mutation_snapshot(&pool, fixture).await, before); - } -} - -#[tokio::test] -async fn sqlite_refund_rejects_foreign_or_uncredited_payment_order_without_side_effects() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let fixture = create_refund_mutation_fixture(&repository, "order-integrity").await; - - ensure_test_user(&pool, "refund-order-integrity-other").await; - let other_wallet = repository - .initialize_auth_user_wallet("refund-order-integrity-other", 0.0, false) - .await - .expect("other wallet initialization should run") - .expect("other wallet should exist"); - let (_, other_order) = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: other_wallet.id, - amount_usd: 10.0, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-order-integrity".to_string()), - description: Some("order integrity setup".to_string()), - order_no: "refund-order-integrity-other-order".to_string(), - }) - .await - .expect("other recharge should run") - .expect("other recharge should create an order"); - - sqlx::query("UPDATE refund_requests SET payment_order_id = ? WHERE id = ?") - .bind(&other_order.id) - .bind(&fixture.refund_id) - .execute(&pool) - .await - .expect("foreign payment order should be assignable for regression setup"); - let before_foreign = refund_mutation_snapshot(&pool, &fixture).await; - let foreign_outcome = repository - .process_admin_wallet_refund(ProcessAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - operator_id: Some("admin-order-integrity".to_string()), - }) - .await - .expect("foreign payment order refund should resolve"); - assert!(matches!(foreign_outcome, WalletMutationOutcome::Invalid(_))); - assert_eq!( - refund_mutation_snapshot(&pool, &fixture).await, - before_foreign - ); - - sqlx::query("UPDATE refund_requests SET payment_order_id = ? WHERE id = ?") - .bind(&fixture.payment_order_id) - .bind(&fixture.refund_id) - .execute(&pool) - .await - .expect("original payment order should be restored for regression setup"); - sqlx::query("UPDATE payment_orders SET status = 'pending' WHERE id = ?") - .bind(&fixture.payment_order_id) - .execute(&pool) - .await - .expect("payment order status should be corruptible for regression setup"); - let before_uncredited = refund_mutation_snapshot(&pool, &fixture).await; - let uncredited_outcome = repository - .process_admin_wallet_refund(ProcessAdminWalletRefundInput { - wallet_id: fixture.wallet_id.clone(), - refund_id: fixture.refund_id.clone(), - operator_id: Some("admin-order-integrity".to_string()), - }) - .await - .expect("uncredited payment order refund should resolve"); - assert!(matches!( - uncredited_outcome, - WalletMutationOutcome::Invalid(_) - )); - assert_eq!( - refund_mutation_snapshot(&pool, &fixture).await, - before_uncredited - ); -} - -#[tokio::test] -async fn sqlite_wallet_write_repository_handles_public_recharge_callback_and_refund() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users( - &pool, - &[ - "user-write-1", - "user-credit-1", - "user-expire-1", - "user-fail-1", - ], - ) - .await; - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-write-1".to_string()), - user_id: "user-write-1".to_string(), - amount_usd: 12.5, - pay_amount: Some(12.5), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-order-write-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-no-write-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new wallet order should not already exist") - } - }; - assert_eq!(order.status, "pending"); - - let callback = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-write-1".to_string(), - order_no: Some("order-no-write-1".to_string()), - gateway_order_id: Some("gateway-order-write-1".to_string()), - amount_usd: 12.5, - pay_amount: Some(12.5), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-hash-write-1".to_string(), - payload: json!({ - "status": "paid", - "client_secret": "pi_1_secret_replayable", - "customer": {"email": "payer@example.com"}, - "authorization": "Bearer upstream-secret", - }), - signature_valid: true, - }) - .await - .expect("payment callback should process"); - let ProcessPaymentCallbackOutcome::Applied { - wallet_id, order, .. - } = callback - else { - panic!("callback should credit the order"); - }; - assert_eq!(wallet_id, "wallet-write-1"); - assert_eq!(order.status, "credited"); - let stored_callback_payload: Option = - sqlx::query_scalar("SELECT payload FROM payment_callbacks WHERE callback_key = ?") - .bind("callback-write-1") - .fetch_one(&pool) - .await - .expect("callback payload should query"); - assert_eq!(stored_callback_payload, None); - let stored_gateway_response: String = - sqlx::query_scalar("SELECT gateway_response FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("gateway response should query"); - let stored_gateway_response: serde_json::Value = - serde_json::from_str(&stored_gateway_response).expect("gateway response should be JSON"); - assert_eq!(stored_gateway_response["gateway"], "alipay"); - assert_eq!(stored_gateway_response["payment_provider"], "alipay"); - assert_eq!(stored_gateway_response["payment_channel"], "alipay"); - assert_eq!(stored_gateway_response["order_no"], "order-no-write-1"); - assert_eq!(stored_gateway_response["amount_usd"], 12.5); - assert!(stored_gateway_response.get("status").is_none()); - let stored_settlement_binding: (String, Option, Option) = sqlx::query_as( - "SELECT payment_method, payment_provider, payment_channel FROM payment_orders WHERE id = ?", - ) - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("settlement binding should query"); - assert_eq!( - stored_settlement_binding, - ( - "alipay".to_string(), - Some("alipay".to_string()), - Some("alipay".to_string()), - ) - ); - let encoded_gateway_response = stored_gateway_response.to_string(); - for forbidden in [ - "client_secret", - "replayable", - "payer@example.com", - "authorization", - "upstream-secret", - ] { - assert!( - !encoded_gateway_response.contains(forbidden), - "persisted {forbidden}" - ); - } - - let wallet = repository - .find(WalletLookupKey::UserId("user-write-1")) - .await - .expect("wallet should query") - .expect("wallet should exist"); - assert_eq!(wallet.balance, 12.5); - assert_eq!(wallet.total_recharged, 12.5); - - let refund = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id.clone(), - user_id: "user-write-1".to_string(), - amount_usd: 4.0, - payment_order_id: Some(order.id.clone()), - source_type: None, - source_id: None, - refund_mode: None, - reason: Some("requested".to_string()), - idempotency_key: Some("idem-refund-write-1".to_string()), - refund_no: "refund-no-write-1".to_string(), - }) - .await - .expect("refund request should create"); - let CreateWalletRefundRequestOutcome::Created(refund) = refund else { - panic!("refund request should be created"); - }; - - let duplicate = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id.clone(), - user_id: "user-write-1".to_string(), - amount_usd: 4.0, - payment_order_id: Some(order.id.clone()), - source_type: None, - source_id: None, - refund_mode: None, - reason: Some("requested".to_string()), - idempotency_key: Some("idem-refund-write-1".to_string()), - refund_no: "refund-no-write-duplicate".to_string(), - }) - .await - .expect("duplicate refund request should resolve"); - assert!(matches!( - duplicate, - CreateWalletRefundRequestOutcome::Duplicate(_) - )); - - let processed = repository - .process_admin_wallet_refund(ProcessAdminWalletRefundInput { - wallet_id: wallet.id.clone(), - refund_id: refund.id.clone(), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("refund should process"); - let WalletMutationOutcome::Applied((wallet, refund, transaction)) = processed else { - panic!("refund should be processed"); - }; - assert_eq!(wallet.balance, 8.5); - assert_eq!(refund.status, "processing"); - assert_eq!(transaction.reason_code, "refund_out"); - - let completed = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: wallet.id.clone(), - refund_id: refund.id.clone(), - gateway_refund_id: Some("gateway-refund-write-1".to_string()), - payout_reference: Some("payout-ref-write-1".to_string()), - payout_proof: Some(json!({ "proof": "ok" })), - }) - .await - .expect("refund should complete"); - let WalletMutationOutcome::Applied(completed_refund) = completed else { - panic!("refund should be completed"); - }; - assert_eq!(completed_refund.status, "succeeded"); - assert_eq!( - completed_refund.payout_proof.as_ref().unwrap()["proof"], - "ok" - ); - let repeated_completion = repository - .complete_admin_wallet_refund(CompleteAdminWalletRefundInput { - wallet_id: wallet.id.clone(), - refund_id: refund.id.clone(), - gateway_refund_id: Some("gateway-refund-attacker".to_string()), - payout_reference: Some("payout-ref-attacker".to_string()), - payout_proof: Some(json!({ "proof": "attacker" })), - }) - .await - .expect("completed refund replay should resolve"); - assert!(matches!( - repeated_completion, - WalletMutationOutcome::Invalid(_) - )); - let repeated_refund = repository - .find_wallet_refund(&wallet.id, &refund.id) - .await - .expect("completed refund should still load") - .expect("completed refund should still exist"); - assert_eq!( - repeated_refund.gateway_refund_id.as_deref(), - Some("gateway-refund-write-1") - ); - assert_eq!( - repeated_refund.payout_reference.as_deref(), - Some("payout-ref-write-1") - ); - assert_eq!(repeated_refund.payout_proof, completed_refund.payout_proof); - - let refund_to_fail = repository - .create_wallet_refund_request(CreateWalletRefundRequestInput { - wallet_id: wallet.id.clone(), - user_id: "user-write-1".to_string(), - amount_usd: 1.5, - payment_order_id: Some(order.id.clone()), - source_type: None, - source_id: None, - refund_mode: None, - reason: Some("requested again".to_string()), - idempotency_key: Some("idem-refund-write-2".to_string()), - refund_no: "refund-no-write-2".to_string(), - }) - .await - .expect("second refund request should create"); - let CreateWalletRefundRequestOutcome::Created(refund_to_fail) = refund_to_fail else { - panic!("second refund request should be created"); - }; - let before_fail = refund_mutation_snapshot( - &pool, - &RefundMutationFixture { - wallet_id: wallet.id.clone(), - payment_order_id: order.id.clone(), - refund_id: refund_to_fail.id.clone(), - }, - ) - .await; - let failed = repository - .fail_admin_wallet_refund(FailAdminWalletRefundInput { - wallet_id: wallet.id.clone(), - refund_id: refund_to_fail.id.clone(), - reason: "manual failure".to_string(), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("second refund failure should resolve"); - assert!(matches!(&failed, WalletMutationOutcome::Applied(_))); - let failed_refund = match failed { - WalletMutationOutcome::Applied((_, refund, transaction)) => { - assert!(transaction.is_none()); - refund - } - _ => unreachable!(), - }; - assert_eq!(failed_refund.status, "failed"); - let after_fail = refund_mutation_snapshot( - &pool, - &RefundMutationFixture { - wallet_id: wallet.id.clone(), - payment_order_id: order.id.clone(), - refund_id: refund_to_fail.id.clone(), - }, - ) - .await; - assert_eq!(after_fail.wallet_balance, before_fail.wallet_balance); - assert_eq!( - after_fail.wallet_total_refunded, - before_fail.wallet_total_refunded - ); - assert_eq!(after_fail.refund_status, "failed"); - - let batch = repository - .create_admin_redeem_code_batch(CreateAdminRedeemCodeBatchInput { - name: "Write Batch".to_string(), - amount_usd: 3.5, - currency: "USD".to_string(), - balance_bucket: "gift".to_string(), - total_count: 1, - expires_at_unix_secs: None, - description: Some("write smoke".to_string()), - created_by: Some("admin-1".to_string()), - }) - .await - .expect("redeem batch should create"); - assert_eq!(batch.batch.active_count, 1); - let redeem_code = batch.codes[0].code.clone(); - - let redeem = repository - .redeem_wallet_code(RedeemWalletCodeInput { - code: redeem_code, - user_id: "user-write-1".to_string(), - order_no: "redeem-order-write-1".to_string(), - }) - .await - .expect("redeem should apply"); - let RedeemWalletCodeOutcome::Redeemed { - wallet, - order, - amount_usd, - batch_name, - } = redeem - else { - panic!("redeem should succeed"); - }; - assert_eq!(wallet.gift_balance, 3.5); - assert_eq!(order.payment_method, "gift_code"); - assert_eq!(amount_usd, 3.5); - assert_eq!(batch_name, "Write Batch"); - - let disabled_batch = repository - .create_admin_redeem_code_batch(CreateAdminRedeemCodeBatchInput { - name: "Disabled Batch".to_string(), - amount_usd: 1.25, - currency: "USD".to_string(), - balance_bucket: "gift".to_string(), - total_count: 2, - expires_at_unix_secs: None, - description: Some("disable smoke".to_string()), - created_by: Some("admin-1".to_string()), - }) - .await - .expect("disable batch should create"); - let disabled_code = repository - .disable_admin_redeem_code(DisableAdminRedeemCodeInput { - code_id: disabled_batch.codes[0].code_id.clone(), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("redeem code should disable"); - let WalletMutationOutcome::Applied(disabled_code) = disabled_code else { - panic!("redeem code should be disabled"); - }; - assert_eq!(disabled_code.status, "disabled"); - - let disabled_batch = repository - .disable_admin_redeem_code_batch(DisableAdminRedeemCodeBatchInput { - batch_id: disabled_batch.batch.id.clone(), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("redeem batch should disable"); - let WalletMutationOutcome::Applied(disabled_batch) = disabled_batch else { - panic!("redeem batch should be disabled"); - }; - assert_eq!(disabled_batch.status, "disabled"); - assert_eq!(disabled_batch.active_count, 0); - - let deleted_batch = repository - .delete_admin_redeem_code_batch(DeleteAdminRedeemCodeBatchInput { - batch_id: disabled_batch.id, - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("disabled unredeemed batch should delete"); - assert!(matches!(deleted_batch, WalletMutationOutcome::Applied(_))); - - let (wallet, adjustment) = repository - .adjust_wallet_balance(AdjustWalletBalanceInput { - wallet_id: wallet.id.clone(), - amount_usd: -2.0, - balance_type: "gift".to_string(), - operator_id: Some("admin-1".to_string()), - description: Some("trim gift".to_string()), - }) - .await - .expect("adjustment should run") - .expect("wallet should exist"); - assert_eq!(wallet.gift_balance, 1.5); - assert_eq!(adjustment.reason_code, "adjust_admin"); - assert_eq!( - adjustment.balance_after, - wallet.balance + wallet.gift_balance - ); - - let (wallet, order) = repository - .create_manual_wallet_recharge(CreateManualWalletRechargeInput { - wallet_id: wallet.id, - amount_usd: 5.0, - payment_method: "admin_manual".to_string(), - operator_id: Some("admin-1".to_string()), - description: Some("manual topup".to_string()), - order_no: "manual-order-write-1".to_string(), - }) - .await - .expect("manual recharge should run") - .expect("wallet should exist"); - assert_eq!(wallet.balance, 13.5); - assert_eq!(order.status, "credited"); - assert_eq!(order.payment_method, "admin_manual"); - - let credit_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-credit-1".to_string()), - user_id: "user-credit-1".to_string(), - amount_usd: 2.25, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - payment_method: "manual_gateway".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-order-credit-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-no-credit-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("credit order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new credit wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new credit order should not already exist") - } - }; - let credited = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: credit_order.id.clone(), - gateway_order_id: Some("gateway-order-credit-paid-1".to_string()), - pay_amount: Some(2.25), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - gateway_response_patch: Some(json!({ "settled": true })), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("credit order should apply"); - let WalletMutationOutcome::Applied((credited_order, applied)) = credited else { - panic!("credit order should be applied"); - }; - assert!(applied); - assert_eq!(credited_order.status, "credited"); - assert_eq!( - credited_order.gateway_response.as_ref().unwrap()["manual_credit"], - true - ); - let credited_again = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: credit_order.id, - gateway_order_id: None, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - gateway_response_patch: None, - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("credit order should be idempotent"); - assert!(matches!( - credited_again, - WalletMutationOutcome::Applied((_, false)) - )); - - let expiring_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-expire-1".to_string()), - user_id: "user-expire-1".to_string(), - amount_usd: 1.0, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - payment_method: "alipay".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-order-expire-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-no-expire-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("expiring order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new expire wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new expiring order should not already exist") - } - }; - let expired = repository - .expire_admin_payment_order(&expiring_order.id) - .await - .expect("expire should run"); - assert!(matches!(expired, WalletMutationOutcome::Applied((_, true)))); - let expired_again = repository - .expire_admin_payment_order(&expiring_order.id) - .await - .expect("expire should be idempotent"); - assert!(matches!( - expired_again, - WalletMutationOutcome::Applied((_, false)) - )); - - let failing_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-fail-1".to_string()), - user_id: "user-fail-1".to_string(), - amount_usd: 1.0, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - payment_method: "alipay".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-order-fail-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-no-fail-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("failing order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new fail wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new failing order should not already exist") - } - }; - let failed_order = repository - .fail_admin_payment_order(&failing_order.id) - .await - .expect("fail should run"); - let WalletMutationOutcome::Applied(failed_order) = failed_order else { - panic!("payment order should fail"); - }; - assert_eq!(failed_order.status, "failed"); -} - -#[tokio::test] -async fn sqlite_payment_callback_rejects_gateway_identifier_mismatch_without_crediting() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-identifier-mismatch").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-identifier-mismatch".to_string()), - user_id: "user-callback-identifier-mismatch".to_string(), - amount_usd: 15.0, - pay_amount: Some(15.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-order-original".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-no-identifier-mismatch".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new identifier-mismatch order should not already exist") - } - }; - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-identifier-mismatch".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-order-attacker".to_string()), - amount_usd: 15.0, - pay_amount: Some(15.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-identifier-mismatch".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("mismatched callback should resolve"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment gateway order mismatch" - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("wallet-callback-identifier-mismatch") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (0.0, 0.0)); - - let stored_order: (String, String, Option, Option) = sqlx::query_as( - "SELECT status, gateway_order_id, paid_at, credited_at FROM payment_orders WHERE id = ?", - ) - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order should load"); - assert_eq!(stored_order.0, "pending"); - assert_eq!(stored_order.1, "gateway-order-original"); - assert_eq!(stored_order.2, None); - assert_eq!(stored_order.3, None); - - let callback: (String, Option) = sqlx::query_as( - "SELECT status, error_message FROM payment_callbacks WHERE callback_key = ?", - ) - .bind("callback-identifier-mismatch") - .fetch_one(&pool) - .await - .expect("callback should load"); - assert_eq!(callback.0, "failed"); - assert_eq!( - callback.1.as_deref(), - Some("payment gateway order mismatch") - ); -} - -#[tokio::test] -async fn sqlite_payment_callback_rejects_wallet_owner_mismatch_without_crediting() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users(&pool, &["callback-order-owner", "callback-wallet-owner"]).await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-owner-mismatch".to_string()), - user_id: "callback-wallet-owner".to_string(), - amount_usd: 8.0, - pay_amount: Some(8.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-owner-mismatch".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-owner-mismatch".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - // Simulate a corrupted/imported order that points at a different user - // than the wallet selected at checkout. - sqlx::query("UPDATE payment_orders SET user_id = ? WHERE id = ?") - .bind("callback-order-owner") - .bind(&order.id) - .execute(&pool) - .await - .expect("order owner should update for regression setup"); - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-owner-mismatch".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-owner-mismatch".to_string()), - amount_usd: 8.0, - pay_amount: Some(8.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-owner-mismatch".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("owner-mismatch callback should resolve"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment order wallet owner mismatch" - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind(&order.wallet_id) - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (0.0, 0.0)); - let stored_order: (String, Option, Option) = - sqlx::query_as("SELECT status, paid_at, credited_at FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order should load"); - assert_eq!(stored_order, ("pending".to_string(), None, None)); - let callback: (String, Option) = sqlx::query_as( - "SELECT status, error_message FROM payment_callbacks WHERE callback_key = ?", - ) - .bind("callback-owner-mismatch") - .fetch_one(&pool) - .await - .expect("callback should load"); - assert_eq!(callback.0, "failed"); - assert_eq!( - callback.1.as_deref(), - Some("payment order wallet owner mismatch") - ); -} - -#[tokio::test] -async fn sqlite_payment_callback_credits_overdrawn_recharge_balance() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-overdrawn").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-overdrawn".to_string()), - user_id: "user-callback-overdrawn".to_string(), - amount_usd: 5.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-callback-overdrawn".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-callback-overdrawn".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - sqlx::query("UPDATE wallets SET balance = -3.0, total_recharged = 0.0 WHERE id = ?") - .bind(&order.wallet_id) - .execute(&pool) - .await - .expect("wallet should be made overdrawn"); - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-overdrawn".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-callback-overdrawn".to_string()), - amount_usd: 5.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-callback-overdrawn".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("payment callback should process"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - - let wallet = repository - .find(WalletLookupKey::UserId("user-callback-overdrawn")) - .await - .expect("wallet should query") - .expect("wallet should exist"); - assert_eq!(wallet.balance, 2.0); - assert_eq!(wallet.total_recharged, 5.0); - - let callback: (String, Option) = sqlx::query_as( - "SELECT status, error_message FROM payment_callbacks WHERE callback_key = ?", - ) - .bind("callback-overdrawn") - .fetch_one(&pool) - .await - .expect("callback should load"); - assert_eq!(callback.0, "processed"); - assert_eq!(callback.1, None); -} - -#[tokio::test] -async fn sqlite_manual_credit_rejects_invalid_order_and_wallet_values() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-manual-credit-invalid").await; - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-manual-credit-invalid".to_string()), - user_id: "user-manual-credit-invalid".to_string(), - amount_usd: 5.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "manual_gateway".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-manual-credit-invalid".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-manual-credit-invalid".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("credit order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - sqlx::query("UPDATE payment_orders SET amount_usd = ? WHERE id = ?") - .bind(-5.0) - .bind(&order.id) - .execute(&pool) - .await - .expect("invalid order fixture should update"); - let invalid_order = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: order.id.clone(), - gateway_order_id: None, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - gateway_response_patch: None, - operator_id: Some("admin-invalid".to_string()), - }) - .await - .expect("invalid order credit should resolve"); - assert!(matches!( - invalid_order, - WalletMutationOutcome::Invalid(ref error) if error == "payment order amount is invalid" - )); - let order_status: String = sqlx::query_scalar("SELECT status FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order status should query"); - assert_eq!(order_status, "pending"); - - sqlx::query("UPDATE payment_orders SET amount_usd = ? WHERE id = ?") - .bind(5.0) - .bind(&order.id) - .execute(&pool) - .await - .expect("order fixture should restore"); - sqlx::query("UPDATE wallets SET gift_balance = ? WHERE id = ?") - .bind(-1.0) - .bind(&order.wallet_id) - .execute(&pool) - .await - .expect("invalid wallet fixture should update"); - let invalid_wallet = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: order.id.clone(), - gateway_order_id: None, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - gateway_response_patch: None, - operator_id: Some("admin-invalid".to_string()), - }) - .await - .expect("invalid wallet credit should resolve"); - assert!(matches!( - invalid_wallet, - WalletMutationOutcome::Invalid(ref error) if error == "wallet balance is invalid" - )); - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, gift_balance FROM wallets WHERE id = ?") - .bind(&order.wallet_id) - .fetch_one(&pool) - .await - .expect("wallet should query"); - assert_eq!(wallet, (0.0, -1.0)); - let order_status: String = sqlx::query_scalar("SELECT status FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order status should query"); - assert_eq!(order_status, "pending"); -} - -#[tokio::test] -async fn sqlite_payment_callback_rejects_unknown_order_status_without_crediting() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-invalid-state").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-invalid-state".to_string()), - user_id: "user-callback-invalid-state".to_string(), - amount_usd: 11.0, - pay_amount: Some(11.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-invalid-state".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-invalid-state".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - sqlx::query("UPDATE payment_orders SET status = 'cancelled' WHERE id = ?") - .bind(&order.id) - .execute(&pool) - .await - .expect("test order status should update"); - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-invalid-state".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-invalid-state".to_string()), - amount_usd: 11.0, - pay_amount: Some(11.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-invalid-state".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("invalid-state callback should resolve"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment order is not creditable: cancelled" - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("wallet-callback-invalid-state") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (0.0, 0.0)); - let stored_order: (String, Option, Option) = - sqlx::query_as("SELECT status, paid_at, credited_at FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order should load"); - assert_eq!(stored_order, ("cancelled".to_string(), None, None)); -} - -#[tokio::test] -async fn sqlite_payment_callback_recovers_failed_checkout_placeholder_without_losing_credit() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-failed-checkout").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-failed-checkout".to_string()), - user_id: "user-callback-failed-checkout".to_string(), - amount_usd: 12.0, - pay_amount: Some(12.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "order-callback-failed-checkout".to_string(), - gateway_response: json!({ - "gateway": "alipay", - "gateway_order_id": "order-callback-failed-checkout", - "order_kind": "wallet_recharge", - "payment_channel": "alipay", - "pay_amount": 12.0, - "pay_currency": "USD", - "integration_status": "checkout_pending", - "checkout_claim_token": "claim-failed-checkout", - "checkout_claimed_at_unix_secs": 1, - }), - order_no: "order-callback-failed-checkout".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - let claim_token = order - .gateway_response - .as_ref() - .and_then(|response| response.get("checkout_claim_token")) - .and_then(serde_json::Value::as_str) - .expect("checkout claim token should persist") - .to_string(); - - let failed = repository - .fail_wallet_recharge_checkout(FailWalletRechargeCheckoutInput { - order_id: order.id.clone(), - claim_token, - reason: "checkout response timed out after provider acceptance".to_string(), - provider_request_may_have_succeeded: true, - }) - .await - .expect("checkout failure should resolve"); - let WalletMutationOutcome::Applied(failed) = failed else { - panic!("checkout placeholder should become failed"); - }; - assert_eq!(failed.status, "failed"); - assert_eq!( - failed - .gateway_response - .as_ref() - .and_then(|response| response.get("integration_status")) - .and_then(serde_json::Value::as_str), - Some("checkout_uncertain") - ); - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-failed-checkout-recovery".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("provider-order-failed-checkout".to_string()), - amount_usd: 12.0, - pay_amount: Some(12.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-failed-checkout-recovery".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("provider callback should process"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind(&order.wallet_id) - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (12.0, 12.0)); - - let stored_order: (String, Option) = - sqlx::query_as("SELECT status, gateway_order_id FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(&pool) - .await - .expect("order should load"); - assert_eq!( - stored_order, - ( - "credited".to_string(), - Some("provider-order-failed-checkout".to_string()) - ) - ); - - let callback_status: String = - sqlx::query_scalar("SELECT status FROM payment_callbacks WHERE callback_key = ?") - .bind("callback-failed-checkout-recovery") - .fetch_one(&pool) - .await - .expect("callback should load"); - assert_eq!(callback_status, "processed"); -} - -#[tokio::test] -async fn sqlite_payment_callback_rejects_corrupt_stored_credit_amount() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-corrupt-amount").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-corrupt-amount".to_string()), - user_id: "user-callback-corrupt-amount".to_string(), - amount_usd: 11.0, - pay_amount: Some(11.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-corrupt-amount".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-corrupt-amount".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - sqlx::query("UPDATE payment_orders SET amount_usd = -11 WHERE id = ?") - .bind(&order.id) - .execute(&pool) - .await - .expect("test order amount should update"); - - let outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-corrupt-amount".to_string(), - order_no: Some(order.order_no), - gateway_order_id: Some("gateway-corrupt-amount".to_string()), - amount_usd: 11.0, - pay_amount: Some(11.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-corrupt-amount".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("corrupt-amount callback should resolve"); - assert!(matches!( - outcome, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment order amount is invalid" - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("wallet-callback-corrupt-amount") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (0.0, 0.0)); -} - -#[tokio::test] -async fn sqlite_payment_callback_reconstructs_legacy_provider_amount_from_order_terms() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users( - &pool, - &["legacy-cny-callback-user", "legacy-usd-callback-user"], - ) - .await; - - let cny_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("legacy-cny-callback-wallet".to_string()), - user_id: "legacy-cny-callback-user".to_string(), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "legacy-cny-callback-gateway".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "legacy-cny-callback-order".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("CNY recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - sqlx::query("UPDATE payment_orders SET pay_amount = NULL WHERE id = ?") - .bind(&cny_order.id) - .execute(&pool) - .await - .expect("legacy CNY order should drop provider amount"); - - let cny_outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "legacy-cny-callback".to_string(), - order_no: Some(cny_order.order_no), - gateway_order_id: Some("legacy-cny-callback-gateway".to_string()), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payload_hash: "legacy-cny-callback-payload".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("legacy CNY callback should resolve"); - assert!(matches!( - cny_outcome, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - let cny_wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("legacy-cny-callback-wallet") - .fetch_one(&pool) - .await - .expect("legacy CNY wallet should load"); - assert_eq!(cny_wallet, (10.0, 10.0)); - - let usd_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("legacy-usd-callback-wallet".to_string()), - user_id: "legacy-usd-callback-user".to_string(), - amount_usd: 10.0, - pay_amount: Some(10.0), - pay_currency: Some("USD".to_string()), - // Old rows could retain the historical CNY default for USD. - exchange_rate: Some(7.2), - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "legacy-usd-callback-gateway".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "legacy-usd-callback-order".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("USD recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - sqlx::query("UPDATE payment_orders SET pay_amount = NULL WHERE id = ?") - .bind(&usd_order.id) - .execute(&pool) - .await - .expect("legacy USD order should drop provider amount"); - - let wrong_usd_outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - callback_key: "legacy-usd-callback-wrong".to_string(), - order_no: Some(usd_order.order_no.clone()), - gateway_order_id: Some("legacy-usd-callback-gateway".to_string()), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(7.2), - payload_hash: "legacy-usd-callback-wrong-payload".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("wrong USD callback should resolve"); - assert!(matches!( - wrong_usd_outcome, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "callback amount mismatch" - )); - let usd_wallet_before: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("legacy-usd-callback-wallet") - .fetch_one(&pool) - .await - .expect("legacy USD wallet should load before valid callback"); - assert_eq!(usd_wallet_before, (0.0, 0.0)); - - let valid_usd_outcome = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - callback_key: "legacy-usd-callback-valid".to_string(), - order_no: Some(usd_order.order_no), - gateway_order_id: Some("legacy-usd-callback-gateway".to_string()), - amount_usd: 10.0, - pay_amount: Some(10.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(7.2), - payload_hash: "legacy-usd-callback-valid-payload".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("valid USD callback should resolve"); - assert!(matches!( - valid_usd_outcome, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - let usd_wallet_after: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("legacy-usd-callback-wallet") - .fetch_one(&pool) - .await - .expect("legacy USD wallet should load after valid callback"); - assert_eq!(usd_wallet_after, (10.0, 10.0)); -} - -#[tokio::test] -async fn sqlite_payment_callback_requires_provider_namespace_to_match_exactly() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-provider-boundary").await; - - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-provider-boundary".to_string()), - user_id: "user-provider-boundary".to_string(), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-provider-boundary".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-provider-boundary".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - let error = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: None, - payment_channel: None, - callback_key: "callback-provider-boundary".to_string(), - order_no: Some(order.order_no), - gateway_order_id: Some("gateway-provider-boundary".to_string()), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payload_hash: "payload-provider-boundary".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect_err("provider-less official callback must be rejected at repository boundary"); - assert!(matches!( - error, - aether_data_contracts::DataLayerError::InvalidInput(ref detail) - if detail == "official payment callback provider binding mismatch" - )); - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("wallet-provider-boundary") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (0.0, 0.0)); -} - -#[tokio::test] -async fn sqlite_legacy_epay_channel_order_without_provider_is_compatible() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-legacy-epay").await; - - // Rows written before payment_provider/payment_channel were introduced - // stored the selected EPay channel as payment_method and left both new - // columns NULL. - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-legacy-epay".to_string()), - user_id: "user-legacy-epay".to_string(), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payment_method: "alipay".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-legacy-epay".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-legacy-epay".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("legacy order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - let wrong_channel = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("wxpay".to_string()), - callback_key: "callback-legacy-epay-wrong-channel".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-legacy-epay".to_string()), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payload_hash: "payload-legacy-epay-wrong-channel".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("wrong-channel callback should resolve"); - assert!(matches!( - wrong_channel, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment channel mismatch" - )); - - let applied = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-legacy-epay-success".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-legacy-epay".to_string()), - amount_usd: 10.0, - pay_amount: Some(72.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payload_hash: "payload-legacy-epay-success".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("legacy EPay callback should process"); - assert!(matches!( - applied, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind(&order.wallet_id) - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (10.0, 10.0)); -} - -#[tokio::test] -async fn sqlite_payment_callback_validates_placeholder_gateway_binding_and_conflicts() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users( - &pool, - &[ - "user-callback-placeholder-a", - "user-callback-placeholder-b", - "user-callback-placeholder-c", - ], - ) - .await; - - // Order B already owns the real provider transaction id. - let order_b = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-placeholder-b".to_string()), - user_id: "user-callback-placeholder-b".to_string(), - amount_usd: 7.0, - pay_amount: Some(7.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-b".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-b".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("order B should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - // Order A stores its merchant order number as a provider-id placeholder. - let order_a = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-placeholder-a".to_string()), - user_id: "user-callback-placeholder-a".to_string(), - amount_usd: 7.0, - pay_amount: Some(7.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "order-a".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-a".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("order A should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - let conflict = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-placeholder-conflict".to_string(), - order_no: Some(order_a.order_no.clone()), - gateway_order_id: Some("gateway-b".to_string()), - amount_usd: 7.0, - pay_amount: Some(7.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-placeholder-conflict".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("placeholder conflict callback should resolve"); - assert!(matches!( - conflict, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment gateway order belongs to another payment order" - )); - let order_a_state: (String, String) = - sqlx::query_as("SELECT status, gateway_order_id FROM payment_orders WHERE id = ?") - .bind(&order_a.id) - .fetch_one(&pool) - .await - .expect("order A should load"); - assert_eq!( - order_a_state, - ("pending".to_string(), "order-a".to_string()) - ); - let order_b_state: (String, f64) = - sqlx::query_as("SELECT status, amount_usd FROM payment_orders WHERE id = ?") - .bind(&order_b.id) - .fetch_one(&pool) - .await - .expect("order B should load"); - assert_eq!(order_b_state.0, "pending"); - assert_eq!(order_b_state.1, 7.0); - - // A fresh provider id that is not owned by another order is bound and - // credited atomically, replacing the placeholder. - let order_c = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-placeholder-c".to_string()), - user_id: "user-callback-placeholder-c".to_string(), - amount_usd: 5.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "order-c".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-c".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("order C should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - let applied = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-placeholder-success".to_string(), - order_no: Some(order_c.order_no.clone()), - gateway_order_id: Some("gateway-c".to_string()), - amount_usd: 5.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-placeholder-success".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("placeholder binding callback should apply"); - assert!(matches!( - applied, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - let order_c_state: (String, String) = - sqlx::query_as("SELECT status, gateway_order_id FROM payment_orders WHERE id = ?") - .bind(&order_c.id) - .fetch_one(&pool) - .await - .expect("order C should load"); - assert_eq!( - order_c_state, - ("credited".to_string(), "gateway-c".to_string()) - ); - let wallet_c_balance: f64 = sqlx::query_scalar("SELECT balance FROM wallets WHERE id = ?") - .bind("wallet-callback-placeholder-c") - .fetch_one(&pool) - .await - .expect("wallet C should load"); - assert_eq!(wallet_c_balance, 5.0); -} - -#[tokio::test] -async fn sqlite_payment_gateway_order_identifier_is_unique_within_payment_method() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users( - &pool, - &[ - "user-gateway-unique-first", - "user-gateway-unique-second", - "user-gateway-other-method", - "user-gateway-case-distinct", - ], - ) - .await; - - let first = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-gateway-unique-first".to_string()), - user_id: "user-gateway-unique-first".to_string(), - amount_usd: 3.0, - pay_amount: Some(3.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: " EPAY ".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "shared-provider-transaction".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-gateway-unique-first".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("first order should create"); - assert!(matches!( - first, - CreateWalletRechargeOrderOutcome::Created(ref order) if order.payment_method == "epay" - )); - - let duplicate = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-gateway-unique-second".to_string()), - user_id: "user-gateway-unique-second".to_string(), - amount_usd: 3.0, - pay_amount: Some(3.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "shared-provider-transaction".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-gateway-unique-second".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await; - assert!(duplicate.is_err()); - assert!(repository - .find(WalletLookupKey::UserId("user-gateway-unique-second")) - .await - .expect("conflicting order must not leave a wallet") - .is_none()); - - let other_method = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-gateway-other-method".to_string()), - user_id: "user-gateway-other-method".to_string(), - amount_usd: 3.0, - pay_amount: Some(3.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "stripe".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "shared-provider-transaction".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-gateway-other-method".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("another payment method may reuse the identifier"); - assert!(matches!( - other_method, - CreateWalletRechargeOrderOutcome::Created(_) - )); - - let case_distinct_identifier = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-gateway-case-distinct".to_string()), - user_id: "user-gateway-case-distinct".to_string(), - amount_usd: 3.0, - pay_amount: Some(3.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "Shared-Provider-Transaction".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-gateway-case-distinct".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("opaque identifiers that differ by case may coexist"); - assert!(matches!( - case_distinct_identifier, - CreateWalletRechargeOrderOutcome::Created(_) - )); -} - -#[tokio::test] -async fn sqlite_payment_callback_failure_does_not_rebind_existing_identifiers() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-failure-preserve").await; - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-failure-preserve".to_string()), - user_id: "user-callback-failure-preserve".to_string(), - amount_usd: 6.0, - pay_amount: Some(6.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-preserve".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-failure-preserve".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => panic!("wallet should be active"), - CreateWalletRechargeOrderOutcome::Existing(_) => panic!("order should be new"), - }; - - let first_failure = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-failure-preserve".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-preserve".to_string()), - amount_usd: 6.0, - pay_amount: Some(5.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-failure-preserve".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("first failure should resolve"); - assert!(matches!( - first_failure, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "callback amount mismatch" - )); - - let second_failure = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-failure-preserve".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-attacker-retry".to_string()), - amount_usd: 6.0, - pay_amount: Some(6.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-failure-preserve".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("second failure should resolve"); - assert!(matches!( - second_failure, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment gateway order mismatch" - )); - - let callback: (String, String, Option) = sqlx::query_as( - "SELECT gateway_order_id, status, error_message FROM payment_callbacks WHERE callback_key = ?", - ) - .bind("callback-failure-preserve") - .fetch_one(&pool) - .await - .expect("callback should load"); - assert_eq!(callback.0, "gateway-preserve"); - assert_eq!(callback.1, "failed"); - assert_eq!( - callback.2.as_deref(), - Some("payment gateway order mismatch") - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sqlite_payment_callback_registration_is_atomic_under_concurrency() { - let database_path = std::env::temp_dir().join(format!( - "aether-sqlite-payment-callback-race-{}.db", - uuid::Uuid::new_v4() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .foreign_keys(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .busy_timeout(Duration::from_secs(30)); - 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"); - - let repository = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "user-callback-race").await; - let order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-callback-race".to_string()), - user_id: "user-callback-race".to_string(), - amount_usd: 9.0, - pay_amount: Some(9.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-callback-race".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-callback-race".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new callback-race order should not already exist") - } - }; - - let callback_input = ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("alipay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-key-race".to_string(), - order_no: Some(order.order_no.clone()), - gateway_order_id: Some("gateway-callback-race".to_string()), - amount_usd: 9.0, - pay_amount: Some(9.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payload_hash: "payload-hash-race".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }; - let barrier = Arc::new(tokio::sync::Barrier::new(3)); - let first_repository = repository.clone(); - let first_barrier = barrier.clone(); - let first_input = callback_input.clone(); - let first = tokio::spawn(async move { - first_barrier.wait().await; - first_repository.process_payment_callback(first_input).await - }); - let second_repository = repository.clone(); - let second_barrier = barrier.clone(); - let second = tokio::spawn(async move { - second_barrier.wait().await; - second_repository - .process_payment_callback(callback_input) - .await - }); - barrier.wait().await; - - let outcomes = [ - first - .await - .expect("first callback task should join") - .expect("first callback should resolve"), - second - .await - .expect("second callback task should join") - .expect("second callback should resolve"), - ]; - let mut applied = 0; - let mut duplicate = 0; - for outcome in outcomes { - match outcome { - ProcessPaymentCallbackOutcome::Applied { - duplicate: false, .. - } => applied += 1, - ProcessPaymentCallbackOutcome::DuplicateProcessed { .. } - | ProcessPaymentCallbackOutcome::Applied { - duplicate: true, .. - } - | ProcessPaymentCallbackOutcome::AlreadyCredited { - duplicate: true, .. - } => duplicate += 1, - other => panic!("unexpected concurrent callback outcome: {other:?}"), - } - } - assert_eq!(applied, 1, "exactly one callback should apply the credit"); - assert_eq!(duplicate, 1, "the other callback should be a duplicate"); - - let wallet: (f64, f64) = - sqlx::query_as("SELECT balance, total_recharged FROM wallets WHERE id = ?") - .bind("wallet-callback-race") - .fetch_one(&pool) - .await - .expect("wallet should load"); - assert_eq!(wallet, (9.0, 9.0)); - let callback_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM payment_callbacks WHERE callback_key = ?") - .bind("callback-key-race") - .fetch_one(&pool) - .await - .expect("callback count should query"); - assert_eq!(callback_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())); -} - -#[tokio::test] -async fn sqlite_payment_callback_binds_settlement_amount_before_usd_conversion() { - 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 = SqliteWalletReadRepository::new(pool); - ensure_test_user(repository.pool(), "user-fee-callback").await; - let order = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-fee-callback".to_string()), - user_id: "user-fee-callback".to_string(), - amount_usd: 10.0, - pay_amount: Some(73.5), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.0), - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-fee-callback".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-fee-callback".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("recharge order should be created"); - assert!(matches!( - order, - CreateWalletRechargeOrderOutcome::Created(_) - )); - - let mismatched = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-fee-mismatch".to_string(), - order_no: Some("order-fee-callback".to_string()), - gateway_order_id: Some("gateway-fee-callback".to_string()), - amount_usd: 10.0, - pay_amount: Some(73.49), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.0), - payload_hash: "payload-fee-mismatch".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("mismatched callback should resolve"); - assert!(matches!( - mismatched, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "callback amount mismatch" - )); - - let wrong_currency = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-fee-wrong-currency".to_string(), - order_no: Some("order-fee-callback".to_string()), - gateway_order_id: Some("gateway-fee-callback".to_string()), - amount_usd: 10.0, - pay_amount: Some(73.5), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(7.0), - payload_hash: "payload-fee-wrong-currency".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("wrong-currency callback should resolve"); - assert!(matches!( - wrong_currency, - ProcessPaymentCallbackOutcome::Failed { ref error, .. } - if error == "payment currency mismatch" - )); - - let applied = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-fee-applied".to_string(), - order_no: Some("order-fee-callback".to_string()), - gateway_order_id: Some("gateway-fee-callback".to_string()), - // Gateway callbacks derive USD from the fee-inclusive settlement - // amount. The stored order's USD amount remains the net credit. - amount_usd: 10.5, - pay_amount: Some(73.5000005), - pay_currency: Some("cny".to_string()), - // A callback may carry a conflicting provider-side rate. The - // order's checkout-time rate is the settlement proof and must - // remain unchanged. - exchange_rate: Some(99.0), - payload_hash: "payload-fee-applied".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("fee-inclusive callback should process"); - assert!(matches!( - applied, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - - let wallet = repository - .find(WalletLookupKey::UserId("user-fee-callback")) - .await - .expect("wallet should query") - .expect("wallet should exist"); - assert_eq!(wallet.balance, 10.0); - assert_eq!(wallet.total_recharged, 10.0); - let persisted_terms: (Option, Option, Option) = sqlx::query_as( - "SELECT pay_amount, pay_currency, exchange_rate FROM payment_orders WHERE order_no = ?", - ) - .bind("order-fee-callback") - .fetch_one(repository.pool()) - .await - .expect("settlement terms should query"); - assert_eq!( - persisted_terms, - (Some(73.5), Some("CNY".to_string()), Some(7.0)) - ); - - repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some(wallet.id), - user_id: "user-fee-callback".to_string(), - amount_usd: 5.0, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - payment_method: "manual".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-usd-fallback".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-usd-fallback".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("fallback order should be created"); - let fallback = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "manual".to_string(), - payment_provider: None, - payment_channel: None, - callback_key: "callback-usd-fallback".to_string(), - order_no: Some("order-usd-fallback".to_string()), - // The order already has a verified gateway id. A generic retry - // may identify it by order number without repeating that id. - gateway_order_id: None, - amount_usd: 5.0, - pay_amount: None, - pay_currency: None, - exchange_rate: None, - payload_hash: "payload-usd-fallback".to_string(), - payload: json!({ "status": "paid" }), - signature_valid: true, - }) - .await - .expect("USD fallback callback should process"); - assert!(matches!( - fallback, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - let persisted_gateway: (Option, String) = sqlx::query_as( - "SELECT gateway_order_id, gateway_response FROM payment_orders WHERE order_no = ?", - ) - .bind("order-usd-fallback") - .fetch_one(repository.pool()) - .await - .expect("fallback order should remain queryable"); - assert_eq!(persisted_gateway.0.as_deref(), Some("gateway-usd-fallback")); - assert_eq!( - serde_json::from_str::(&persisted_gateway.1) - .expect("gateway response should be JSON") - .get("gateway_order_id"), - Some(&json!("gateway-usd-fallback")) - ); -} - -#[tokio::test] -async fn sqlite_plan_purchase_rejects_missing_user_without_creating_wallet_or_order() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - let input = CreatePlanPurchaseOrderInput { - preferred_wallet_id: Some("missing-user-wallet".to_string()), - user_id: "missing-plan-user".to_string(), - amount_usd: 1.0, - pay_amount: 1.0, - pay_currency: "USD".to_string(), - exchange_rate: 1.0, - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "missing-user-gateway".to_string(), - gateway_response: json!({"checkout": true}), - order_no: "missing-user-order".to_string(), - product_id: "missing-user-plan".to_string(), - product_snapshot: json!({ - "id": "missing-user-plan", - "duration_unit": "month", - "duration_value": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [] - }), - expires_at_unix_secs: 4_102_444_800, - }; - let error = repository - .create_plan_purchase_order(input) - .await - .expect_err("missing user must be rejected"); - assert!(matches!(error, DataLayerError::InvalidInput(message) if message == "user not found")); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallets WHERE id = 'missing-user-wallet'", - ) - .fetch_one(&pool) - .await - .expect("wallet count should query"), - 0 - ); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM payment_orders WHERE order_no = 'missing-user-order'", - ) - .fetch_one(&pool) - .await - .expect("payment order count should query"), - 0 - ); -} - -#[tokio::test] -async fn sqlite_plan_purchase_rejects_preferred_wallet_id_owned_by_another_user() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_users(&pool, &["plan-wallet-owner", "plan-wallet-conflict"]).await; - let existing_wallet = repository - .initialize_auth_user_wallet("plan-wallet-owner", 0.0, false) - .await - .expect("owner wallet initialization should run") - .expect("owner wallet should exist"); - - let result = repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: Some(existing_wallet.id.clone()), - user_id: "plan-wallet-conflict".to_string(), - amount_usd: 1.0, - pay_amount: 1.0, - pay_currency: "USD".to_string(), - exchange_rate: 1.0, - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "gateway-plan-wallet-conflict".to_string(), - gateway_response: json!({"checkout": true}), - order_no: "order-plan-wallet-conflict".to_string(), - product_id: "plan-wallet-conflict-product".to_string(), - product_snapshot: json!({ - "id": "plan-wallet-conflict-product", - "duration_unit": "month", - "duration_value": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [] - }), - expires_at_unix_secs: 4_102_444_800, - }) - .await; - - assert!(matches!( - result, - Err(DataLayerError::InvalidInput(message)) - if message == "wallet identifier already belongs to another owner" - )); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallets WHERE user_id = 'plan-wallet-conflict'", - ) - .fetch_one(&pool) - .await - .expect("conflicting user wallet count should query"), - 0 - ); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM payment_orders WHERE order_no = 'order-plan-wallet-conflict'", - ) - .fetch_one(&pool) - .await - .expect("conflicting plan order count should query"), - 0 - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sqlite_plan_purchase_initializes_one_wallet_under_concurrency() { - let database_path = std::env::temp_dir().join(format!( - "aether-sqlite-plan-wallet-race-{}.db", - uuid::Uuid::new_v4() - )); - let options = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&database_path) - .create_if_missing(true) - .foreign_keys(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .busy_timeout(Duration::from_secs(30)); - 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"); - let repository = SqliteWalletReadRepository::new(pool.clone()); - ensure_test_user(&pool, "plan-wallet-race-user").await; - - let plan_snapshot = json!({ - "id": "plan-wallet-race-product", - "duration_unit": "month", - "duration_value": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [] - }); - let first_input = CreatePlanPurchaseOrderInput { - preferred_wallet_id: Some("plan-wallet-race-first".to_string()), - user_id: "plan-wallet-race-user".to_string(), - amount_usd: 1.0, - pay_amount: 1.0, - pay_currency: "USD".to_string(), - exchange_rate: 1.0, - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "gateway-plan-wallet-race-first".to_string(), - gateway_response: json!({"checkout": true}), - order_no: "order-plan-wallet-race-first".to_string(), - product_id: "plan-wallet-race-product".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }; - let second_input = CreatePlanPurchaseOrderInput { - preferred_wallet_id: Some("plan-wallet-race-second".to_string()), - gateway_order_id: "gateway-plan-wallet-race-second".to_string(), - order_no: "order-plan-wallet-race-second".to_string(), - ..first_input.clone() - }; - 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 - .create_plan_purchase_order(first_input) - .await - }); - let second_repository = repository.clone(); - let second_barrier = barrier.clone(); - let second = tokio::spawn(async move { - second_barrier.wait().await; - second_repository - .create_plan_purchase_order(second_input) - .await - }); - barrier.wait().await; - - let first = first - .await - .expect("first plan task should join") - .expect("first plan task should resolve"); - let second = second - .await - .expect("second plan task should join") - .expect("second plan task should resolve"); - let first_wallet_id = match first { - CreatePlanPurchaseOrderOutcome::Created(order) => order.wallet_id, - other => panic!("first concurrent plan should be created, got {other:?}"), - }; - let second_wallet_id = match second { - CreatePlanPurchaseOrderOutcome::Created(order) => order.wallet_id, - other => panic!("second concurrent plan should be created, got {other:?}"), - }; - assert_eq!(first_wallet_id, second_wallet_id); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallets WHERE user_id = 'plan-wallet-race-user'", - ) - .fetch_one(&pool) - .await - .expect("wallet count should query"), - 1 - ); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM payment_orders WHERE user_id = 'plan-wallet-race-user' AND order_kind = 'plan_purchase'", - ) - .fetch_one(&pool) - .await - .expect("plan order count should query"), - 2 - ); - - pool.close().await; - let _ = std::fs::remove_file(database_path); -} - -#[tokio::test] -async fn sqlite_wallet_recharge_rejects_missing_user_without_creating_wallet_or_order() { - 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 = SqliteWalletReadRepository::new(pool.clone()); - - let error = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("missing-recharge-wallet".to_string()), - user_id: "missing-recharge-user".to_string(), - amount_usd: 2.0, - pay_amount: Some(2.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "missing-recharge-gateway".to_string(), - gateway_response: json!({"checkout": true}), - order_no: "missing-recharge-order".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect_err("missing user must be rejected"); - assert!(matches!( - error, - DataLayerError::InvalidInput(message) if message == "user not found" - )); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM wallets WHERE id = 'missing-recharge-wallet'", - ) - .fetch_one(&pool) - .await - .expect("wallet count should query"), - 0 - ); - assert_eq!( - sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM payment_orders WHERE order_no = 'missing-recharge-order'", - ) - .fetch_one(&pool) - .await - .expect("payment order count should query"), - 0 - ); -} - -#[tokio::test] -async fn sqlite_plan_purchase_blocks_duplicate_pending_active_period_order_and_manual_credit_fulfills( -) { - 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 = SqliteWalletReadRepository::new(pool); - sqlx::query( - "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("user-active-period-1") - .bind("Active Period Buyer") - .bind("active-period@example.com") - .bind("local") - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("user should seed"); - - let _wallet_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-active-period-1".to_string()), - user_id: "user-active-period-1".to_string(), - amount_usd: 1.0, - pay_amount: Some(1.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "bootstrap".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-bootstrap-active-period-1".to_string(), - gateway_response: json!({ "bootstrap": true }), - order_no: "order-bootstrap-active-period-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("wallet should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new bootstrap order should not already exist") - } - }; - - let plan_snapshot = json!({ - "id": "active-period-plan", - "title": "每日额度月卡", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "active_period", - "entitlements": [ - { - "type": "daily_quota", - "daily_quota_usd": 50.0, - "reset_timezone": "Asia/Shanghai", - "allow_wallet_overage": false - } - ] - }); - sqlx::query( - r#" -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, duration_value, - max_active_per_user, purchase_limit_scope, entitlements_json, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind("active-period-plan") - .bind("每日额度月卡") - .bind(100.0_f64) - .bind("CNY") - .bind("month") - .bind(1_i64) - .bind(1_i64) - .bind("active_period") - .bind(plan_snapshot["entitlements"].to_string()) - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("billing plan should seed"); - - let first_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-active-period-1".to_string(), - amount_usd: 13.8, - pay_amount: 100.0, - pay_currency: "CNY".to_string(), - exchange_rate: 7.24637681, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-plan-active-period-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-plan-active-period-1".to_string(), - product_id: "active-period-plan".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("first active period order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("first active period order should be created, got {other:?}"), - }; - let duplicate_pending = repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-active-period-1".to_string(), - amount_usd: 13.8, - pay_amount: 100.0, - pay_currency: "CNY".to_string(), - exchange_rate: 7.24637681, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-plan-active-period-2".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-plan-active-period-2".to_string(), - product_id: "active-period-plan".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("duplicate active period order should resolve"); - assert!(matches!( - duplicate_pending, - CreatePlanPurchaseOrderOutcome::ActivePlanLimitReached - )); - - let credited = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: first_order.id.clone(), - gateway_order_id: Some("gateway-plan-active-period-paid-1".to_string()), - pay_amount: Some(100.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.24637681), - gateway_response_patch: Some(json!({ "settled": true })), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("manual plan credit should run"); - let WalletMutationOutcome::Applied((credited_order, applied)) = credited else { - panic!("manual plan credit should be applied"); - }; - assert!(applied); - assert_eq!(credited_order.status, "credited"); - assert_eq!(credited_order.refundable_amount_usd, 0.0); - - let entitlement_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE user_id = ? AND plan_id = ?", - ) - .bind("user-active-period-1") - .bind("active-period-plan") - .fetch_one(repository.pool()) - .await - .expect("entitlement count should query"); - assert_eq!(entitlement_count, 1); - - let wallet_balance: f64 = sqlx::query_scalar("SELECT balance FROM wallets WHERE id = ?") - .bind("wallet-active-period-1") - .fetch_one(repository.pool()) - .await - .expect("wallet balance should query"); - assert_eq!(wallet_balance, 0.0); - - // A malformed wallet_credit must abort fulfillment rather than silently - // activating the plan without delivering its promised balance. - let malformed_snapshot = json!({ - "id": "active-period-plan", - "duration_unit": "month", - "duration_value": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [{ - "type": "wallet_credit", - "amount_usd": 5.0, - "balance_bucket": "not-a-wallet-bucket" - }] - }); - let valid_legacy_snapshot = json!({ - "id": "active-period-plan", - "duration_unit": "month", - "duration_value": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [{ - "type": "wallet_credit", - "amount_usd": 5.0, - "balance_bucket": "gift" - }] - }); - let malformed_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-active-period-1".to_string(), - amount_usd: 2.0, - pay_amount: 2.0, - pay_currency: "USD".to_string(), - exchange_rate: 1.0, - payment_method: "stripe".to_string(), - payment_provider: Some("stripe".to_string()), - payment_channel: Some("card".to_string()), - gateway_order_id: "gateway-malformed-wallet-credit".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-malformed-wallet-credit".to_string(), - product_id: "active-period-plan".to_string(), - // Create through the validated boundary first; the malformed - // snapshot is installed below to simulate a legacy/corrupt row - // that predates the boundary validator. - product_snapshot: valid_legacy_snapshot, - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("malformed plan order should be persisted for fulfillment test") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("malformed plan order should be created, got {other:?}"), - }; - sqlx::query("UPDATE payment_orders SET product_snapshot = ? WHERE id = ?") - .bind(malformed_snapshot.to_string()) - .bind(&malformed_order.id) - .execute(repository.pool()) - .await - .expect("malformed legacy snapshot should be installed"); - let credit_result = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: malformed_order.id.clone(), - gateway_order_id: Some("gateway-malformed-wallet-credit-paid".to_string()), - pay_amount: Some(2.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - gateway_response_patch: Some(json!({ "settled": true })), - operator_id: Some("admin-1".to_string()), - }) - .await; - assert!(matches!( - credit_result, - Err(DataLayerError::InvalidInput(_)) - )); - let malformed_status: String = - sqlx::query_scalar("SELECT status FROM payment_orders WHERE id = ?") - .bind(&malformed_order.id) - .fetch_one(repository.pool()) - .await - .expect("malformed order status should remain queryable"); - assert_eq!(malformed_status, "pending"); - let malformed_entitlements: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE payment_order_id = ?", - ) - .bind(&malformed_order.id) - .fetch_one(repository.pool()) - .await - .expect("malformed entitlement count should query"); - assert_eq!(malformed_entitlements, 0); - let wallet_balance_after: f64 = sqlx::query_scalar("SELECT balance FROM wallets WHERE id = ?") - .bind("wallet-active-period-1") - .fetch_one(repository.pool()) - .await - .expect("wallet balance after rejected credit should query"); - assert_eq!(wallet_balance_after, 0.0); -} - -#[tokio::test] -async fn sqlite_finds_reusable_pending_plan_purchase_order() { - 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 = SqliteWalletReadRepository::new(pool); - sqlx::query( - "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("user-pending-plan-1") - .bind("Pending Buyer") - .bind("pending-plan@example.com") - .bind("local") - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("user should seed"); - - let _wallet_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-pending-plan-1".to_string()), - user_id: "user-pending-plan-1".to_string(), - amount_usd: 1.0, - pay_amount: Some(1.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "bootstrap".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-bootstrap-pending-plan-1".to_string(), - gateway_response: json!({ "bootstrap": true }), - order_no: "order-bootstrap-pending-plan-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("wallet should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new pending-plan order should not already exist") - } - }; - - let plan_snapshot = json!({ - "id": "pending-plan", - "title": "每日额度月卡", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "active_period", - "entitlements": [ - { - "type": "daily_quota", - "daily_quota_usd": 50.0, - "reset_timezone": "Asia/Shanghai", - "allow_wallet_overage": false - } - ] - }); - let pending_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-pending-plan-1".to_string(), - amount_usd: 13.8, - pay_amount: 100.0, - pay_currency: "CNY".to_string(), - exchange_rate: 7.24637681, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-pending-plan-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-pending-plan-1".to_string(), - product_id: "pending-plan".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("pending plan order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("pending plan order should be created, got {other:?}"), - }; - let now = chrono::Utc::now().timestamp().max(0); - for (id, order_no, status, product_id, user_id, expires_at, created_at) in [ - ( - "expired-pending-plan-order", - "order-expired-pending-plan", - "pending", - "pending-plan", - "user-pending-plan-1", - now - 10, - now + 10, - ), - ( - "credited-pending-plan-order", - "order-credited-pending-plan", - "credited", - "pending-plan", - "user-pending-plan-1", - now + 3_600, - now + 20, - ), - ( - "other-user-pending-plan-order", - "order-other-user-pending-plan", - "pending", - "pending-plan", - "other-user", - now + 3_600, - now + 30, - ), - ] { - sqlx::query( - r#" -INSERT INTO payment_orders ( - 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, - fulfillment_status, gateway_order_id, gateway_response, status, created_at, expires_at -) VALUES (?, ?, ?, ?, 13.8, 100.0, 'CNY', 7.24637681, 0, 0, 'alipay', - 'epay', 'alipay', 'plan_purchase', ?, ?, 'pending', ?, ?, ?, ?, ?) - "#, - ) - .bind(id) - .bind(order_no) - .bind("wallet-pending-plan-1") - .bind(user_id) - .bind(product_id) - .bind(plan_snapshot.to_string()) - .bind(format!("gateway-{id}")) - .bind(json!({ "checkout": id }).to_string()) - .bind(status) - .bind(created_at) - .bind(expires_at) - .execute(repository.pool()) - .await - .expect("extra payment order should seed"); - } - - let found = repository - .find_pending_plan_purchase_order_by_user_id("user-pending-plan-1", "pending-plan") - .await - .expect("pending plan lookup should run") - .expect("pending plan order should be found"); - assert_eq!(found.id, pending_order.id); - assert_eq!( - repository - .find_pending_plan_purchase_order_by_user_id("user-pending-plan-1", "missing-plan") - .await - .expect("missing plan lookup should run"), - None - ); - assert_eq!( - repository - .find_pending_plan_purchase_order_by_user_id("missing-user", "pending-plan") - .await - .expect("missing user lookup should run"), - None - ); -} - -#[tokio::test] -async fn sqlite_plan_purchase_replaces_same_class_entitlements_on_manual_credit() { - 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 = SqliteWalletReadRepository::new(pool); - sqlx::query( - "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("user-upgrade-1") - .bind("Upgrade Buyer") - .bind("upgrade@example.com") - .bind("local") - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("user should seed"); - - let _wallet_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-upgrade-1".to_string()), - user_id: "user-upgrade-1".to_string(), - amount_usd: 1.0, - pay_amount: Some(1.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "bootstrap".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-bootstrap-upgrade-1".to_string(), - gateway_response: json!({ "bootstrap": true }), - order_no: "order-bootstrap-upgrade-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("wallet should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new upgrade order should not already exist") - } - }; - - let low_snapshot = json!({ - "id": "pro-basic", - "title": "Pro Basic", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "active_period", - "entitlements": [{"type": "daily_quota", "daily_quota_usd": 10.0}] - }); - let high_snapshot = json!({ - "id": "pro-plus", - "title": "Pro Plus", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "active_period", - "entitlements": [{"type": "daily_quota", "daily_quota_usd": 50.0}] - }); - for (id, title, snapshot) in [ - ("pro-basic", "Pro Basic", &low_snapshot), - ("pro-plus", "Pro Plus", &high_snapshot), - ] { - sqlx::query( - r#" -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, duration_value, - max_active_per_user, purchase_limit_scope, entitlements_json, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(id) - .bind(title) - .bind(100.0_f64) - .bind("CNY") - .bind("month") - .bind(1_i64) - .bind(1_i64) - .bind("active_period") - .bind(snapshot["entitlements"].to_string()) - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("billing plan should seed"); - } - - let low_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-upgrade-1".to_string(), - amount_usd: 13.8, - pay_amount: 100.0, - pay_currency: "CNY".to_string(), - exchange_rate: 7.24637681, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-pro-basic-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-pro-basic-1".to_string(), - product_id: "pro-basic".to_string(), - product_snapshot: low_snapshot, - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("low order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("low order should be created, got {other:?}"), - }; - let WalletMutationOutcome::Applied((_, true)) = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: low_order.id, - gateway_order_id: Some("gateway-pro-basic-paid-1".to_string()), - pay_amount: Some(100.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.24637681), - gateway_response_patch: Some(json!({ "settled": true })), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("low plan credit should run") - else { - panic!("low plan credit should apply"); - }; - - let high_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-upgrade-1".to_string(), - amount_usd: 13.8, - pay_amount: 100.0, - pay_currency: "CNY".to_string(), - exchange_rate: 7.24637681, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-pro-plus-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-pro-plus-1".to_string(), - product_id: "pro-plus".to_string(), - product_snapshot: high_snapshot, - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("high order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("high order should be created, got {other:?}"), - }; - let WalletMutationOutcome::Applied((_, true)) = repository - .credit_admin_payment_order(CreditAdminPaymentOrderInput { - order_id: high_order.id, - gateway_order_id: Some("gateway-pro-plus-paid-1".to_string()), - pay_amount: Some(100.0), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.24637681), - gateway_response_patch: Some(json!({ "settled": true })), - operator_id: Some("admin-1".to_string()), - }) - .await - .expect("high plan credit should run") - else { - panic!("high plan credit should apply"); - }; - - let low_status: String = sqlx::query_scalar( - "SELECT status FROM user_plan_entitlements WHERE user_id = ? AND plan_id = ?", - ) - .bind("user-upgrade-1") - .bind("pro-basic") - .fetch_one(repository.pool()) - .await - .expect("low entitlement status should query"); - assert_eq!(low_status, "replaced"); - let active_high_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE user_id = ? AND plan_id = ? AND status = 'active'", - ) - .bind("user-upgrade-1") - .bind("pro-plus") - .fetch_one(repository.pool()) - .await - .expect("high entitlement count should query"); - assert_eq!(active_high_count, 1); -} - -#[tokio::test] -async fn sqlite_plan_replacement_stacks_usage_policies_unless_groups_match() { - 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 = SqliteWalletReadRepository::new(pool); - sqlx::query("PRAGMA foreign_keys = OFF") - .execute(repository.pool()) - .await - .expect("foreign keys should be disabled for isolated replacement fixtures"); - - let fixtures = [ - ( - "usage-plain", - json!([{"type": "usage_policy", "policy_id": "weekly", "rules": []}]), - ), - ( - "usage-pro", - json!([{ - "type": "usage_policy", - "replacement_group": "pro-tier", - "rules": [] - }]), - ), - ( - "usage-team", - json!([{ - "type": "usage_policy", - "replacement_group": "team-tier", - "rules": [] - }]), - ), - ( - "daily-legacy", - json!([{"type": "daily_quota", "daily_quota_usd": 10.0}]), - ), - ]; - for (id, entitlements) in fixtures { - sqlx::query( - r#" -INSERT INTO user_plan_entitlements ( - id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, - entitlements_snapshot, created_at, updated_at -) VALUES (?, 'replacement-user', ?, ?, 'active', 1, 4102444800, ?, 1, 1) - "#, - ) - .bind(id) - .bind(format!("plan-{id}")) - .bind(format!("order-{id}")) - .bind(entitlements.to_string()) - .execute(repository.pool()) - .await - .expect("entitlement fixture should seed"); - } - - let mut tx = repository.pool().begin().await.expect("tx should start"); - replace_matching_plan_entitlements_sqlite( - &mut tx, - "replacement-user", - &json!({ - "entitlements": [{"type": "usage_policy", "policy_id": "five-hour", "rules": []}] - }), - 100, - ) - .await - .expect("ungrouped usage policy replacement should run"); - tx.commit().await.expect("tx should commit"); - let active_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE user_id = ? AND status = 'active'", - ) - .bind("replacement-user") - .fetch_one(repository.pool()) - .await - .expect("active entitlement count should query"); - assert_eq!(active_count, 4); - - let mut tx = repository.pool().begin().await.expect("tx should start"); - replace_matching_plan_entitlements_sqlite( - &mut tx, - "replacement-user", - &json!({ - "entitlements": [{ - "type": "usage_policy", - "replacement_group": "pro-tier", - "rules": [] - }] - }), - 200, - ) - .await - .expect("grouped usage policy replacement should run"); - tx.commit().await.expect("tx should commit"); - - let statuses = sqlx::query_as::<_, (String, String)>( - "SELECT id, status FROM user_plan_entitlements ORDER BY id", - ) - .fetch_all(repository.pool()) - .await - .expect("entitlement statuses should query") - .into_iter() - .collect::>(); - assert_eq!( - statuses.get("usage-pro").map(String::as_str), - Some("replaced") - ); - assert_eq!( - statuses.get("usage-plain").map(String::as_str), - Some("active") - ); - assert_eq!( - statuses.get("usage-team").map(String::as_str), - Some("active") - ); - assert_eq!( - statuses.get("daily-legacy").map(String::as_str), - Some("active") - ); - - let mut tx = repository.pool().begin().await.expect("tx should start"); - replace_matching_plan_entitlements_sqlite( - &mut tx, - "replacement-user", - &json!({ - "entitlements": [{"type": "daily_quota", "daily_quota_usd": 50.0}] - }), - 300, - ) - .await - .expect("legacy daily quota replacement should run"); - tx.commit().await.expect("tx should commit"); - let daily_status: String = - sqlx::query_scalar("SELECT status FROM user_plan_entitlements WHERE id = 'daily-legacy'") - .fetch_one(repository.pool()) - .await - .expect("daily entitlement status should query"); - assert_eq!(daily_status, "replaced"); -} - -#[tokio::test] -async fn sqlite_plan_purchase_respects_lifetime_purchase_limit() { - 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 = SqliteWalletReadRepository::new(pool); - sqlx::query( - "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("user-lifetime-1") - .bind("Lifetime Buyer") - .bind("lifetime@example.com") - .bind("local") - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("user should seed"); - - let _wallet_order = match repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-lifetime-1".to_string()), - user_id: "user-lifetime-1".to_string(), - amount_usd: 1.0, - pay_amount: Some(1.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "bootstrap".to_string(), - payment_provider: None, - payment_channel: None, - gateway_order_id: "gateway-bootstrap-lifetime-1".to_string(), - gateway_response: json!({ "bootstrap": true }), - order_no: "order-bootstrap-lifetime-1".to_string(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("wallet should be created") - { - CreateWalletRechargeOrderOutcome::Created(order) => order, - CreateWalletRechargeOrderOutcome::WalletInactive => { - panic!("new wallet should be active") - } - CreateWalletRechargeOrderOutcome::Existing(_) => { - panic!("new lifetime order should not already exist") - } - }; - - let plan_snapshot = json!({ - "id": "first-plan", - "title": "首购特惠包", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "lifetime", - "entitlements": [ - { - "type": "wallet_credit", - "amount_usd": 1.0, - "balance_bucket": "gift" - } - ] - }); - sqlx::query( - r#" -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, duration_value, - max_active_per_user, purchase_limit_scope, entitlements_json, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind("first-plan") - .bind("首购特惠包") - .bind(7.2_f64) - .bind("CNY") - .bind("month") - .bind(1_i64) - .bind(1_i64) - .bind("lifetime") - .bind(plan_snapshot["entitlements"].to_string()) - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("billing plan should seed"); - - let first_order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-lifetime-1".to_string(), - amount_usd: 1.0, - pay_amount: 7.2, - pay_currency: "CNY".to_string(), - exchange_rate: 7.2, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-plan-lifetime-1".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-plan-lifetime-1".to_string(), - product_id: "first-plan".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("first plan order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("first plan order should be created, got {other:?}"), - }; - assert_eq!(first_order.status, "pending"); - - let callback = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: "callback-plan-lifetime-1".to_string(), - order_no: Some("order-plan-lifetime-1".to_string()), - gateway_order_id: Some("gateway-plan-lifetime-1".to_string()), - amount_usd: 1.0, - pay_amount: Some(7.2000005), - pay_currency: Some("cny".to_string()), - exchange_rate: Some(99.0), - payload_hash: "payload-plan-lifetime-1".to_string(), - payload: json!({ "trade_status": "TRADE_SUCCESS" }), - signature_valid: true, - }) - .await - .expect("plan payment callback should process"); - assert!(matches!( - callback, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - let persisted_plan_terms: (f64, String, f64) = sqlx::query_as( - "SELECT pay_amount, pay_currency, exchange_rate FROM payment_orders WHERE order_no = ?", - ) - .bind("order-plan-lifetime-1") - .fetch_one(repository.pool()) - .await - .expect("plan settlement terms should query"); - assert_eq!( - persisted_plan_terms, - (7.2, "CNY".to_string(), 7.2), - "callback data must not overwrite checkout-time settlement terms" - ); - - let entitlement_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE user_id = ? AND plan_id = ?", - ) - .bind("user-lifetime-1") - .bind("first-plan") - .fetch_one(repository.pool()) - .await - .expect("entitlement count should query"); - assert_eq!(entitlement_count, 1); - - let second_order = repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-lifetime-1".to_string(), - amount_usd: 1.0, - pay_amount: 7.2, - pay_currency: "CNY".to_string(), - exchange_rate: 7.2, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "gateway-plan-lifetime-2".to_string(), - gateway_response: json!({ "checkout": true }), - order_no: "order-plan-lifetime-2".to_string(), - product_id: "first-plan".to_string(), - product_snapshot: plan_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("second plan order should resolve"); - assert!(matches!( - second_order, - CreatePlanPurchaseOrderOutcome::ActivePlanLimitReached - )); - - let unlimited_snapshot = json!({ - "id": "unlimited-plan", - "title": "不限购余额包", - "duration_unit": "month", - "duration_value": 1, - "max_active_per_user": 1, - "purchase_limit_scope": "unlimited", - "entitlements": [ - { - "type": "wallet_credit", - "amount_usd": 1.0, - "balance_bucket": "gift" - } - ] - }); - sqlx::query( - r#" -INSERT INTO billing_plans ( - id, title, price_amount, price_currency, duration_unit, duration_value, - max_active_per_user, purchase_limit_scope, entitlements_json, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind("unlimited-plan") - .bind("不限购余额包") - .bind(7.2_f64) - .bind("CNY") - .bind("month") - .bind(1_i64) - .bind(1_i64) - .bind("unlimited") - .bind(unlimited_snapshot["entitlements"].to_string()) - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("unlimited billing plan should seed"); - - for index in 1..=2 { - let order_no = format!("order-plan-unlimited-{index}"); - let gateway_order_id = format!("gateway-plan-unlimited-{index}"); - let order = match repository - .create_plan_purchase_order(CreatePlanPurchaseOrderInput { - preferred_wallet_id: None, - user_id: "user-lifetime-1".to_string(), - amount_usd: 1.0, - pay_amount: 7.2, - pay_currency: "CNY".to_string(), - exchange_rate: 7.2, - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: gateway_order_id.clone(), - gateway_response: json!({ "checkout": true }), - order_no: order_no.clone(), - product_id: "unlimited-plan".to_string(), - product_snapshot: unlimited_snapshot.clone(), - expires_at_unix_secs: 4_102_444_800, - }) - .await - .expect("unlimited plan order should create") - { - CreatePlanPurchaseOrderOutcome::Created(order) => order, - other => panic!("unlimited plan order should be created, got {other:?}"), - }; - assert_eq!(order.status, "pending"); - - let callback = repository - .process_payment_callback(ProcessPaymentCallbackInput { - payment_method: "alipay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - callback_key: format!("callback-plan-unlimited-{index}"), - order_no: Some(order_no), - gateway_order_id: Some(gateway_order_id), - amount_usd: 1.0, - pay_amount: Some(7.2), - pay_currency: Some("CNY".to_string()), - exchange_rate: Some(7.2), - payload_hash: format!("payload-plan-unlimited-{index}"), - payload: json!({ "trade_status": "TRADE_SUCCESS" }), - signature_valid: true, - }) - .await - .expect("unlimited plan payment callback should process"); - assert!(matches!( - callback, - ProcessPaymentCallbackOutcome::Applied { .. } - )); - } - - let unlimited_entitlement_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM user_plan_entitlements WHERE user_id = ? AND plan_id = ?", - ) - .bind("user-lifetime-1") - .bind("unlimited-plan") - .fetch_one(repository.pool()) - .await - .expect("unlimited entitlement count should query"); - assert_eq!(unlimited_entitlement_count, 2); -} - -impl SqliteWalletReadRepository { - fn pool(&self) -> &sqlx::SqlitePool { - &self.pool - } -} - -#[tokio::test] -async fn sqlite_recharge_checkout_update_rejects_expired_order() { - 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 = SqliteWalletReadRepository::new(pool); - - sqlx::query( - "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("user-expired-checkout") - .bind("Expired Checkout") - .bind("expired-checkout@example.com") - .bind("local") - .bind(1_i64) - .bind(1_i64) - .execute(repository.pool()) - .await - .expect("user should seed"); - - let created = repository - .create_wallet_recharge_order(CreateWalletRechargeOrderInput { - preferred_wallet_id: Some("wallet-expired-checkout".to_string()), - user_id: "user-expired-checkout".to_string(), - amount_usd: 3.0, - pay_amount: Some(3.0), - pay_currency: Some("USD".to_string()), - exchange_rate: Some(1.0), - payment_method: "epay".to_string(), - payment_provider: Some("epay".to_string()), - payment_channel: Some("alipay".to_string()), - gateway_order_id: "order-expired-checkout".to_string(), - gateway_response: json!({ - "order_kind": "wallet_recharge", - "integration_status": "checkout_pending" - }), - order_no: "order-expired-checkout".to_string(), - expires_at_unix_secs: 1, - }) - .await - .expect("expired recharge order should be creatable for regression setup"); - let CreateWalletRechargeOrderOutcome::Created(order) = created else { - panic!("expected a newly created recharge order"); - }; - - let result = repository - .update_wallet_recharge_checkout(UpdateWalletRechargeCheckoutInput { - order_id: order.id.clone(), - gateway_order_id: "provider-expired-checkout".to_string(), - gateway_response: json!({ - "order_kind": "wallet_recharge", - "payment_url": "https://pay.example.test/expired" - }), - }) - .await - .expect("expired checkout update should resolve"); - assert!(matches!(result, WalletMutationOutcome::Invalid(_))); - - let persisted: (Option, String) = - sqlx::query_as("SELECT gateway_order_id, status FROM payment_orders WHERE id = ?") - .bind(&order.id) - .fetch_one(repository.pool()) - .await - .expect("expired recharge order should remain queryable"); - assert_eq!(persisted.0.as_deref(), Some("order-expired-checkout")); - assert_eq!(persisted.1, "pending"); -} - -async fn seed_rows(pool: &sqlx::SqlitePool) { - 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) -"#, - ) - .execute(pool) - .await - .expect("user should seed"); - - sqlx::query( - r#" -INSERT INTO wallets ( - id, user_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 -) -"#, - ) - .execute(pool) - .await - .expect("wallet 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 -) -"#, - ) - .execute(pool) - .await - .expect("payment order 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 -) -"#, - ) - .execute(pool) - .await - .expect("callback 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 -) -"#, - ) - .execute(pool) - .await - .expect("refund should seed"); - - sqlx::query( - r#" -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 -) -"#, - ) - .execute(pool) - .await - .expect("transaction 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 -) -"#, - ) - .execute(pool) - .await - .expect("redeem batch should seed"); - - sqlx::query( - r#" -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 -) -"#, - ) - .execute(pool) - .await - .expect("redeem code should seed"); - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES ( - 'daily-1', 'wallet-1', '2000-01-01', 'UTC', 1.25, 2, 10, 20, 3, 4, 11, 11, 11 -) -"#, - ) - .execute(pool) - .await - .expect("daily usage should seed"); -} diff --git a/crates/aether-data/contracts/src/database.rs b/crates/aether-data/contracts/src/database.rs index c64d6ef44..f1c0e2088 100644 --- a/crates/aether-data/contracts/src/database.rs +++ b/crates/aether-data/contracts/src/database.rs @@ -3,8 +3,6 @@ use std::str::FromStr; use crate::DataLayerError; -pub const DEFAULT_SQLITE_DATABASE_URL: &str = "sqlite://./data/aether.db"; - #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)] pub struct PostgresPoolConfig { pub database_url: String, @@ -39,6 +37,7 @@ impl PostgresPoolConfig { "postgres database_url cannot be empty".to_string(), )); } + validate_database_url(&self.database_url)?; if self.min_connections > self.max_connections { return Err(DataLayerError::InvalidConfiguration( "postgres min_connections cannot exceed max_connections".to_string(), @@ -58,16 +57,12 @@ impl PostgresPoolConfig { )] #[serde(rename_all = "snake_case")] pub enum DatabaseDriver { - Sqlite, - Mysql, Postgres, } impl DatabaseDriver { pub const fn as_str(self) -> &'static str { match self { - Self::Sqlite => "sqlite", - Self::Mysql => "mysql", Self::Postgres => "postgres", } } @@ -75,8 +70,6 @@ impl DatabaseDriver { pub fn from_database_url(url: &str) -> Option { let scheme = url.split_once(':')?.0.to_ascii_lowercase(); match scheme.as_str() { - "sqlite" => Some(Self::Sqlite), - "mysql" | "mariadb" => Some(Self::Mysql), "postgres" | "postgresql" => Some(Self::Postgres), _ => None, } @@ -94,16 +87,27 @@ impl FromStr for DatabaseDriver { fn from_str(value: &str) -> Result { match value.trim().to_ascii_lowercase().as_str() { - "sqlite" => Ok(Self::Sqlite), - "mysql" | "mariadb" => Ok(Self::Mysql), "postgres" | "postgresql" => Ok(Self::Postgres), other => Err(DataLayerError::InvalidConfiguration(format!( - "unsupported database driver '{other}'; expected sqlite, mysql, or postgres" + "unsupported database driver '{other}'; expected postgres or postgresql" ))), } } } +fn validate_database_url(url: &str) -> Result<(), DataLayerError> { + if DatabaseDriver::from_database_url(url) != Some(DatabaseDriver::Postgres) { + let scheme = url + .split_once(':') + .map(|(scheme, _)| scheme) + .unwrap_or("missing"); + return Err(DataLayerError::InvalidConfiguration(format!( + "unsupported database URL scheme '{scheme}'; expected postgres or postgresql" + ))); + } + Ok(()) +} + #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)] pub struct SqlPoolConfig { pub min_connections: u32, @@ -141,11 +145,6 @@ impl SqlPoolConfig { "{driver} statement_cache_capacity must be positive" ))); } - if driver == DatabaseDriver::Sqlite && self.require_ssl { - return Err(DataLayerError::InvalidConfiguration( - "sqlite database does not support require_ssl".to_string(), - )); - } Ok(()) } } @@ -172,14 +171,6 @@ impl SqlDatabaseConfig { Ok(config) } - pub fn sqlite_default() -> Self { - Self { - driver: DatabaseDriver::Sqlite, - url: DEFAULT_SQLITE_DATABASE_URL.to_string(), - pool: SqlPoolConfig::default(), - } - } - pub fn validate(&self) -> Result<(), DataLayerError> { if self.url.trim().is_empty() { return Err(DataLayerError::InvalidConfiguration(format!( @@ -187,14 +178,7 @@ impl SqlDatabaseConfig { self.driver ))); } - if let Some(url_driver) = DatabaseDriver::from_database_url(&self.url) { - if url_driver != self.driver { - return Err(DataLayerError::InvalidConfiguration(format!( - "database driver '{}' does not match url scheme '{}'", - self.driver, url_driver - ))); - } - } + validate_database_url(&self.url)?; self.pool.validate(self.driver) } @@ -242,21 +226,10 @@ impl From for SqlDatabaseConfig { #[cfg(test)] mod tests { - use super::{ - DatabaseDriver, PostgresPoolConfig, SqlDatabaseConfig, SqlPoolConfig, - DEFAULT_SQLITE_DATABASE_URL, - }; + use super::{DatabaseDriver, PostgresPoolConfig, SqlDatabaseConfig, SqlPoolConfig}; #[test] fn parses_database_driver_aliases() { - assert_eq!( - "sqlite".parse::().unwrap(), - DatabaseDriver::Sqlite - ); - assert_eq!( - "mariadb".parse::().unwrap(), - DatabaseDriver::Mysql - ); assert_eq!( "postgresql".parse::().unwrap(), DatabaseDriver::Postgres @@ -266,14 +239,6 @@ mod tests { #[test] fn infers_driver_from_database_url_scheme() { - assert_eq!( - DatabaseDriver::from_database_url("sqlite://./data/aether.db"), - Some(DatabaseDriver::Sqlite) - ); - assert_eq!( - DatabaseDriver::from_database_url("mysql://localhost/aether"), - Some(DatabaseDriver::Mysql) - ); assert_eq!( DatabaseDriver::from_database_url("postgres://localhost/aether"), Some(DatabaseDriver::Postgres) @@ -283,14 +248,32 @@ mod tests { #[test] fn validates_driver_url_mismatch() { let config = SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "postgres://localhost/aether".to_string(), + driver: DatabaseDriver::Postgres, + url: "unsupported://localhost/aether".to_string(), pool: SqlPoolConfig::default(), }; assert!(config.validate().is_err()); } + #[test] + fn rejects_removed_database_drivers_and_url_schemes() { + for unsupported in ["mysql", "mariadb", "sqlite"] { + assert!(unsupported.parse::().is_err()); + assert!(serde_json::from_str::(&format!("\"{unsupported}\"")).is_err()); + let url = format!("{unsupported}://localhost/aether"); + assert_eq!(DatabaseDriver::from_database_url(&url), None); + { + let driver = DatabaseDriver::Postgres; + let error = SqlDatabaseConfig::new(driver, &url, SqlPoolConfig::default()) + .expect_err("unsupported URL must not fall back to another driver"); + assert!(error + .to_string() + .contains("unsupported database URL scheme")); + } + } + } + #[test] fn builds_legacy_postgres_config_round_trip() { let postgres = PostgresPoolConfig { @@ -309,12 +292,4 @@ mod tests { assert_eq!(database.driver, DatabaseDriver::Postgres); assert_eq!(database.to_postgres_config().unwrap(), postgres); } - - #[test] - fn sqlite_default_uses_local_database_path() { - let database = SqlDatabaseConfig::sqlite_default(); - - assert_eq!(database.driver, DatabaseDriver::Sqlite); - assert_eq!(database.url, DEFAULT_SQLITE_DATABASE_URL); - } } diff --git a/crates/aether-data/contracts/src/lib.rs b/crates/aether-data/contracts/src/lib.rs index bacd9bc93..c2992874a 100644 --- a/crates/aether-data/contracts/src/lib.rs +++ b/crates/aether-data/contracts/src/lib.rs @@ -3,9 +3,6 @@ mod error; pub mod migration; pub mod repository; -pub use database::{ - DatabaseDriver, PostgresPoolConfig, SqlDatabaseConfig, SqlPoolConfig, - DEFAULT_SQLITE_DATABASE_URL, -}; +pub use database::{DatabaseDriver, PostgresPoolConfig, SqlDatabaseConfig, SqlPoolConfig}; pub use error::DataLayerError; pub use migration::PendingMigrationInfo; diff --git a/crates/aether-data/query/Cargo.toml b/crates/aether-data/query/Cargo.toml index 530501abe..b8dc4ca68 100644 --- a/crates/aether-data/query/Cargo.toml +++ b/crates/aether-data/query/Cargo.toml @@ -10,4 +10,4 @@ description = "Small SQL dialect helpers for Aether data repositories" sqlx.workspace = true [dev-dependencies] -sqlx = { workspace = true, features = ["mysql", "postgres", "sqlite"] } +sqlx = { workspace = true, features = ["postgres"] } diff --git a/crates/aether-data/query/src/lib.rs b/crates/aether-data/query/src/lib.rs index 0efd6a13d..ce12da6e3 100644 --- a/crates/aether-data/query/src/lib.rs +++ b/crates/aether-data/query/src/lib.rs @@ -3,15 +3,12 @@ use sqlx::{Database, Encode, QueryBuilder, Type}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SqlDialect { Postgres, - MySql, - Sqlite, } impl SqlDialect { pub fn quote_ident(self, ident: &str) -> String { let quote = match self { - Self::Postgres | Self::Sqlite => '"', - Self::MySql => '`', + Self::Postgres => '"', }; let escaped = ident.replace(quote, &format!("{quote}{quote}")); format!("{quote}{escaped}{quote}") @@ -30,8 +27,6 @@ impl SqlDialect { pub struct DialectSql<'a> { common: Option<&'a str>, postgres: Option<&'a str>, - mysql: Option<&'a str>, - sqlite: Option<&'a str>, } impl<'a> DialectSql<'a> { @@ -39,17 +34,6 @@ impl<'a> DialectSql<'a> { Self { common: Some(sql), postgres: None, - mysql: None, - sqlite: None, - } - } - - pub const fn dialect(postgres: &'a str, sqlite: &'a str) -> Self { - Self { - common: None, - postgres: Some(postgres), - mysql: None, - sqlite: Some(sqlite), } } @@ -58,21 +42,9 @@ impl<'a> DialectSql<'a> { self } - pub fn with_mysql(mut self, sql: &'a str) -> Self { - self.mysql = Some(sql); - self - } - - pub fn with_sqlite(mut self, sql: &'a str) -> Self { - self.sqlite = Some(sql); - self - } - pub fn sql(self, dialect: SqlDialect) -> &'a str { match dialect { SqlDialect::Postgres => self.postgres.or(self.common), - SqlDialect::MySql => self.mysql.or(self.common), - SqlDialect::Sqlite => self.sqlite.or(self.common), } .expect("dialect SQL expression is missing for selected dialect") } @@ -470,13 +442,6 @@ fn push_ci_contains_predicate<'args, DB>( .push(" ILIKE ") .push_bind(format!("%{trimmed}%")); } - SqlDialect::MySql | SqlDialect::Sqlite => { - builder - .push("LOWER(") - .push(column_sql) - .push(") LIKE ") - .push_bind(format!("%{}%", trimmed.to_ascii_lowercase())); - } } } @@ -522,27 +487,24 @@ pub fn push_order_by( #[cfg(test)] mod tests { use super::*; - use sqlx::{Execute, MySql, Postgres, QueryBuilder, Sqlite}; + use sqlx::{Execute, Postgres, QueryBuilder}; #[test] fn quotes_identifiers_by_dialect() { assert_eq!(SqlDialect::Postgres.quote_ident("trigger"), "\"trigger\""); - assert_eq!(SqlDialect::MySql.quote_ident("trigger"), "`trigger`"); - assert_eq!(SqlDialect::MySql.quote_ident("tri`gger"), "`tri``gger`"); - assert_eq!(SqlDialect::Sqlite.quote_ident("trigger"), "\"trigger\""); assert_eq!( SqlDialect::Postgres.quote_path(&["usage", "id"]), "\"usage\".\"id\"" ); assert_eq!( - SqlDialect::MySql.quote_path(&["usage", "item`id"]), - "`usage`.`item``id`" + SqlDialect::Postgres.quote_path(&["usage", "item\"id"]), + "\"usage\".\"item\"\"id\"" ); } #[test] fn where_clause_pushes_where_then_and() { - let mut builder = QueryBuilder::::new("SELECT * FROM items"); + let mut builder = QueryBuilder::::new("SELECT * FROM items"); let mut where_clause = WhereClause::new(); push_eq( &mut builder, @@ -557,7 +519,7 @@ mod tests { "running".to_string(), ); let query = builder.build(); - assert!(query.sql().contains(" WHERE kind = ? AND status = ?")); + assert!(query.sql().contains(" WHERE kind = $1 AND status = $2")); } #[test] @@ -575,60 +537,26 @@ mod tests { assert!(query.sql().contains(" WHERE task_key ILIKE $1")); } - #[test] - fn ci_contains_uses_lower_like_for_sqlite() { - let mut sqlite_builder = QueryBuilder::::new("SELECT * FROM items"); - let mut sqlite_where = WhereClause::new(); - push_ci_contains( - &mut sqlite_builder, - &mut sqlite_where, - SqlDialect::Sqlite, - "task_key", - " Fetch ", - ); - assert!(sqlite_builder - .build() - .sql() - .contains(" WHERE LOWER(task_key) LIKE ?")); - } - - #[test] - fn ci_contains_uses_lower_like_for_mysql() { - let mut mysql_builder = QueryBuilder::::new("SELECT * FROM items"); - let mut mysql_where = WhereClause::new(); - push_ci_contains( - &mut mysql_builder, - &mut mysql_where, - SqlDialect::MySql, - "task_key", - " Fetch ", - ); - assert!(mysql_builder - .build() - .sql() - .contains(" WHERE LOWER(task_key) LIKE ?")); - } - #[test] fn ci_contains_any_groups_or_predicates() { - let mut builder = QueryBuilder::::new("SELECT * FROM items"); + let mut builder = QueryBuilder::::new("SELECT * FROM items"); let mut where_clause = WhereClause::new(); push_ci_contains_any( &mut builder, &mut where_clause, - SqlDialect::Sqlite, + SqlDialect::Postgres, &["file_name", "COALESCE(display_name, '')"], "Avatar", ); let query = builder.build(); - assert!(query.sql().contains( - " WHERE (LOWER(file_name) LIKE ? OR LOWER(COALESCE(display_name, '')) LIKE ?)" - )); + assert!(query + .sql() + .contains(" WHERE (file_name ILIKE $1 OR COALESCE(display_name, '') ILIKE $2)")); } #[test] fn in_limit_offset_and_order_are_rendered() { - let mut builder = QueryBuilder::::new("SELECT * FROM items"); + let mut builder = QueryBuilder::::new("SELECT * FROM items"); let mut where_clause = WhereClause::new(); push_in( &mut builder, @@ -648,22 +576,18 @@ mod tests { ); push_limit_offset(&mut builder, 10, 5); let query = builder.build(); - assert!(query.sql().contains(" WHERE id IN (?, ?)")); + assert!(query.sql().contains(" WHERE id IN ($1, $2)")); assert!(query.sql().contains(" ORDER BY created_at DESC")); - assert!(query.sql().contains(" LIMIT ? OFFSET ?")); + assert!(query.sql().contains(" LIMIT $3 OFFSET $4")); } #[test] fn select_query_renders_dialect_specific_projection() { let query = SelectQuery::new("providers").select_columns([ SelectColumn::expr("id").alias("provider_id"), - SelectColumn::expr( - DialectSql::dialect( - "CAST(monthly_quota_usd AS DOUBLE PRECISION)", - "CAST(monthly_quota_usd AS REAL)", - ) - .with_mysql("CAST(monthly_quota_usd AS DOUBLE)"), - ) + SelectColumn::expr(DialectSql::common( + "CAST(monthly_quota_usd AS DOUBLE PRECISION)", + )) .alias("monthly_quota_usd"), ]); @@ -671,14 +595,6 @@ mod tests { query.render(SqlDialect::Postgres), "SELECT id AS \"provider_id\", CAST(monthly_quota_usd AS DOUBLE PRECISION) AS \"monthly_quota_usd\" FROM providers" ); - assert_eq!( - query.render(SqlDialect::MySql), - "SELECT id AS `provider_id`, CAST(monthly_quota_usd AS DOUBLE) AS `monthly_quota_usd` FROM providers" - ); - assert_eq!( - query.render(SqlDialect::Sqlite), - "SELECT id AS \"provider_id\", CAST(monthly_quota_usd AS REAL) AS \"monthly_quota_usd\" FROM providers" - ); } #[test] diff --git a/crates/aether-data/runtime/Cargo.toml b/crates/aether-data/runtime/Cargo.toml index 8cee17c9f..100f1b135 100644 --- a/crates/aether-data/runtime/Cargo.toml +++ b/crates/aether-data/runtime/Cargo.toml @@ -10,16 +10,12 @@ build = "build.rs" [features] default = ["postgres"] postgres = ["dep:aether-data-postgres", "sqlx/postgres"] -mysql = ["dep:aether-data-mysql", "sqlx/mysql"] -sqlite = ["dep:aether-data-sqlite", "sqlx/sqlite"] -all-drivers = ["postgres", "mysql", "sqlite"] +all-drivers = ["postgres"] [dependencies] aether-ai-formats.workspace = true aether-data-contracts.workspace = true aether-data-postgres = { workspace = true, optional = true } -aether-data-mysql = { workspace = true, optional = true } -aether-data-sqlite = { workspace = true, optional = true } aether-cache.workspace = true aether-wallet.workspace = true async-trait.workspace = true diff --git a/crates/aether-data/runtime/README.md b/crates/aether-data/runtime/README.md index 8d7cd9cf5..6d3ca1292 100644 --- a/crates/aether-data/runtime/README.md +++ b/crates/aether-data/runtime/README.md @@ -16,21 +16,21 @@ task crates live in the sibling `../contracts` crate (`aether-data-contracts`). | `src/database.rs` | Logical SQL driver selection and shared pool configuration. | | `src/config.rs` | Data-layer config for SQL drivers and repository wiring. | | `src/maintenance.rs` | Maintenance DTOs and aggregation summaries used by backend dispatch and runtime maintenance entrypoints. | -| `src/driver/{postgres,mysql,sqlite}.rs` | Thin compatibility facades for the selected adapter crates. Low-level pools, transactions, and leases live in `aether-data-postgres`, `aether-data-mysql`, and `aether-data-sqlite`. | +| `src/driver/postgres.rs` | Thin compatibility facades for the selected adapter crates. Low-level pools, transactions, and leases live in `aether-data-postgres`. | | `src/repository` | Compatibility modules that re-export contracts and selected SQL adapters, plus concrete in-memory implementations. Driver-specific request-path SQL lives in the adapter crates. | | `src/backend` | Composition root. Builds concrete driver backends and exposes app-facing read/write/worker/lock/lease handles. | | `src/backend/{maintenance,stats,wallet,system}.rs` | Backend-owned maintenance, aggregation, wallet ledger, and system config workflows that are not normal request-path repositories. | | `src/lifecycle/migrate.rs` and `src/lifecycle/migrate/*` | Compatibility entry points and Postgres snapshot bootstrap adapter. | | `src/lifecycle/backfill.rs` | Backfill entry points and backfill discovery. | | `src/lifecycle/export.rs` | Cross-database export/import workflows. | -| `../adapters/{postgres,mysql,sqlite}/migrations` | Executable `sqlx` migrations owned and embedded by each adapter crate. | +| `../adapters/postgres/migrations` | Executable `sqlx` migrations owned and embedded by each adapter crate. | | `schema` | Schema maintenance workspace for logical definitions, driver fragments, and generated output. | | `schema/logical` | Human-maintained logical table definitions used by `aether-data-schema`. | -| `schema/drivers/{postgres,mysql,sqlite}` | Human-maintained driver fragments that compose back into executable SQL while generation is being promoted. | +| `schema/drivers/postgres` | Human-maintained driver fragments that compose back into executable SQL while generation is being promoted. | | `schema/bootstrap/postgres` | Human-maintained source fragments for the Postgres bootstrap snapshot. `build.rs` composes them into the runtime embedded artifact during crate builds. | | `schema/generated` | Machine-written SQL generated from logical schema for audit and drift detection only. | | `schema/overrides` | Rare driver-specific SQL escape hatch. Keep README-only until a real override is needed. | -| `backfills/{postgres,mysql,sqlite}` | Executable backfill SQL grouped by driver. | +| `backfills/postgres` | Executable backfill SQL grouped by driver. | ## Layering Rules @@ -39,11 +39,10 @@ The crate is easiest to read as five layers: 1. Contracts: DTOs, input structs, repository traits, and `DataLayerError`. Prefer `aether-data-contracts` for anything that another crate needs to compile against. -2. Driver primitives: `aether-data-postgres`, `aether-data-mysql`, and - `aether-data-sqlite` connect to infrastructure and expose pools, runners, +2. Driver primitives: `aether-data-postgres` connect to infrastructure and expose pools, runners, and executable migrations; the matching `src/driver/*.rs` files only preserve the existing import paths. -3. Repository implementations: `aether-data-{postgres,mysql,sqlite}` translate +3. Repository implementations: `aether-data-postgres` translate contract types to driver-specific SQL. `src/repository/` keeps the compatibility import path and owns the in-memory implementation where one exists. @@ -68,8 +67,6 @@ src/repository// mod.rs # exports trait/type names and concrete implementations types.rs # implementation-local DTOs when they are not already in contracts postgres.rs # Postgres implementation - mysql.rs # MySQL implementation - sqlite.rs # SQLite implementation memory.rs # tests/dev in-memory implementation ``` @@ -78,49 +75,31 @@ new generic `sql.rs` modules for driver-specific code. ## SQL Driver Policy -The project supports three SQL drivers at the repository/backend boundary: -Postgres, MySQL, and SQLite. That does not mean every raw SQL file is shared. -The portable contract is the Rust shape and behavior; the physical SQL stays -driver-specific where syntax, indexes, JSON support, timestamps, locking, or -upsert semantics differ. +The repository/backend boundary supports PostgreSQL only. Repository contracts +remain independent of SQLx, while physical SQL, pools, and migrations belong to +`aether-data-postgres`. -Use logical types in design docs and reviews: - -| Logical type | Postgres | MySQL | SQLite | -|---|---|---|---| -| `json` | `json` or `jsonb` | `json` or text JSON | text JSON | -| `bool` | `boolean` | `boolean` / `tinyint(1)` | integer | -| `time_unix` | `bigint` or legacy timestamp | `bigint` | integer | -| `money_decimal` | `numeric` / legacy double | `double` | real | - -`jsonb` is acceptable only in Postgres SQL. MySQL and SQLite migrations must not -contain `jsonb`; this is guarded by migration tests. Prefer `serde_json::Value` -or typed Rust structs at the repository boundary so callers do not depend on the -physical storage type. +| Logical type | PostgreSQL | +|---|---| +| `json` | `json` or `jsonb` | +| `bool` | `boolean` | +| `time_unix` | `bigint` or timestamp | +| `money_decimal` | `numeric` / legacy double | ### Feature Matrix -`aether-data` enables only PostgreSQL by default so local checks do not compile -all SQLx drivers: +`aether-data` enables PostgreSQL by default. The `all-drivers` feature selects +the same PostgreSQL backend; a no-driver build remains available for pure +in-memory repository tests. ```bash cargo check -p aether-data -cargo check -p aether-data --no-default-features --features mysql -cargo check -p aether-data --no-default-features --features sqlite -cargo check -p aether-data --no-default-features --features all-drivers +cargo check -p aether-data --no-default-features +cargo check -p aether-data --no-default-features --features postgres ``` -Services selecting MySQL or SQLite in `Cargo.toml` must also disable the default -Postgres feature: - -```toml -aether-data = { workspace = true, default-features = false, features = ["mysql"] } -``` - -The gateway explicitly enables `all-drivers` for deployment compatibility. New -services should select only the driver they deploy with. A configured driver -that is not enabled in the build returns an explicit configuration error rather -than silently constructing an empty backend. +Database URLs must use `postgres:` or `postgresql:`. Unsupported drivers and +URL schemes fail explicitly instead of falling back to another backend. ## Schema Maintenance @@ -147,11 +126,11 @@ normal driver fragments. `compose_schema.sh check` also verifies that required baseline/portable table-creation SQL is represented in `schema/logical`. This is the guardrail -that keeps table structure from drifting back into three manually maintained +that keeps table structure from drifting back into multiple manually maintained definitions. For executable fragments that have not been promoted to generated output yet, -edit fragments under `schema/drivers/{postgres,mysql,sqlite}` directly, run +edit fragments under `schema/drivers/postgres` directly, run `compose`, then run `check`. Do not edit baseline executable SQL and fragments independently. @@ -172,7 +151,7 @@ When adding a table: ## Known Cleanup Targets -These are intentionally staged to keep the multi-database refactor reviewable: +These are intentionally staged to keep the data-layer refactor reviewable: 1. Retire the compatibility re-export paths once downstream crates use `aether-data-contracts` for contracts and `aether-data` only as the runtime 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 deleted file mode 100644 index 5a81e59c2..000000000 --- a/crates/aether-data/runtime/backfills/mysql/20260422120000_backfill_api_key_usage_stats.sql +++ /dev/null @@ -1,36 +0,0 @@ -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 deleted file mode 100644 index c0e41fef5..000000000 --- a/crates/aether-data/runtime/backfills/mysql/20260505120000_rebuild_global_model_usage_count.sql +++ /dev/null @@ -1,15 +0,0 @@ -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 deleted file mode 100644 index 88626dfd7..000000000 --- a/crates/aether-data/runtime/backfills/mysql/20260517012000_sync_legacy_enabled_active_flags.sql +++ /dev/null @@ -1,11 +0,0 @@ -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 deleted file mode 100644 index bdb4d9190..000000000 --- a/crates/aether-data/runtime/backfills/mysql/20260716010000_rebuild_canonical_usage_token_totals.sql +++ /dev/null @@ -1,229 +0,0 @@ -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 deleted file mode 100644 index 733ea7188..000000000 --- a/crates/aether-data/runtime/backfills/mysql/README.md +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index b73cb7f8d..000000000 --- a/crates/aether-data/runtime/backfills/sqlite/20260422120000_backfill_api_key_usage_stats.sql +++ /dev/null @@ -1,28 +0,0 @@ -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 deleted file mode 100644 index f67d6cfaf..000000000 --- a/crates/aether-data/runtime/backfills/sqlite/20260505120000_rebuild_global_model_usage_count.sql +++ /dev/null @@ -1,9 +0,0 @@ -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 deleted file mode 100644 index 88626dfd7..000000000 --- a/crates/aether-data/runtime/backfills/sqlite/20260517012000_sync_legacy_enabled_active_flags.sql +++ /dev/null @@ -1,11 +0,0 @@ -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 deleted file mode 100644 index 75caf9057..000000000 --- a/crates/aether-data/runtime/backfills/sqlite/20260716010000_rebuild_canonical_usage_token_totals.sql +++ /dev/null @@ -1,225 +0,0 @@ -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 deleted file mode 100644 index bd3728e7e..000000000 --- a/crates/aether-data/runtime/backfills/sqlite/README.md +++ /dev/null @@ -1,2 +0,0 @@ -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/README.md b/crates/aether-data/runtime/schema/README.md index 59ca6beb8..f0f47c4f3 100644 --- a/crates/aether-data/runtime/schema/README.md +++ b/crates/aether-data/runtime/schema/README.md @@ -1,7 +1,7 @@ # Aether Schema Source This directory is the schema maintenance workspace. Executable migrations are -owned by the Postgres, MySQL, and SQLite adapters under `../../adapters/`. +owned by the Postgres adapter under `../../adapters/`. The Postgres bootstrap snapshot is compiled from the source fragments here during `aether-data` builds, so there is no checked-in bootstrap artifact. @@ -14,7 +14,7 @@ bash crates/aether-data/runtime/schema/compose_schema.sh check ``` - `generate` renders `logical/*.toml` through `aether-data-schema` into - `generated/{postgres,mysql,sqlite}`. This is a build output, not another SQL + `generated/postgres`. This is a build output, not another SQL source to maintain. - `compose` rewrites the executable SQL from the manifest order. - `check` verifies generated output is current, confirms the bootstrap source @@ -30,15 +30,15 @@ The schema workspace has three normal source areas: | Path | Role | Edit policy | |---|---|---| | `logical/*.toml` | Long-term logical table model shared by all SQL drivers. | Edit first for portable table-shape changes. | -| `drivers/{postgres,mysql,sqlite}/` | Current maintenance fragments for executable SQL. | Edit only for deployment compatibility, ordering, or generator gaps. | +| `drivers/postgres/` | Current maintenance fragments for executable SQL. | Edit only for deployment compatibility, ordering, or generator gaps. | | `bootstrap/postgres/` | Source fragments for the Postgres empty-database bootstrap snapshot. | Edit here when the bootstrap snapshot changes, then rebuild `aether-data` so `build.rs` regenerates the embedded snapshot. | Everything else is output: | Path | Role | Edit policy | |---|---|---| -| `generated/{postgres,mysql,sqlite}/` | Machine-written SQL emitted from `logical/*.toml` for audit and drift detection. | Do not edit; regenerate with `compose_schema.sh generate`. | -| `../../adapters/{postgres,mysql,sqlite}/migrations/` | Runtime SQL artifacts embedded by each database adapter. | Regenerate through `compose_schema.sh compose`; do not edit independently. | +| `generated/postgres/` | Machine-written SQL emitted from `logical/*.toml` for audit and drift detection. | Do not edit; regenerate with `compose_schema.sh generate`. | +| `../../adapters/postgres/migrations/` | Runtime SQL artifacts embedded by each database adapter. | Regenerate through `compose_schema.sh compose`; do not edit independently. | `generated/**` is deliberately checked in so reviews and CI can see exactly what the logical schema compiler emits for each driver. It is not a fourth SQL @@ -59,8 +59,7 @@ cargo run -p aether-data-schema --bin aether-schema -- print --driver postgres ## Logical Schema `logical/*.toml` is the long-term source for table definitions. It covers the -clean baseline table set and the portable MySQL/SQLite table-creation -migrations. The generator emits driver-specific SQL under `generated/`; those +PostgreSQL baseline table set. The generator emits driver-specific SQL under `generated/`; those files include a directory README plus `Do not edit` headers and should only change through `compose_schema.sh generate`. @@ -82,7 +81,7 @@ The migration path is incremental: `overrides/` only when they cannot live cleanly in a driver fragment. 6. Once a domain matches, move its baseline maintenance to generated output. -The existing `drivers/postgres`, `drivers/mysql`, and `drivers/sqlite` +The existing `drivers/postgres` fragment trees remain authoritative for executable migrations until a generated fragment is deliberately promoted. @@ -95,8 +94,6 @@ by logical schema or the normal driver fragments. Keep it small and explicit. |---|---|---| | Postgres baseline | `../../adapters/postgres/migrations/20260403000000_baseline.sql` | `drivers/postgres/baseline/manifest.txt` | | Postgres empty-database snapshot | `aether-data` build output (`OUT_DIR/empty_database_snapshot.sql`) | `bootstrap/postgres/manifest.txt` | -| MySQL baseline | `../../adapters/mysql/migrations/20260403000000_baseline.sql` | `drivers/mysql/baseline/manifest.txt` | -| SQLite baseline | `../../adapters/sqlite/migrations/20260403000000_baseline.sql` | `drivers/sqlite/baseline/manifest.txt` | Driver baseline source manifests are kept as a small set of numbered SQL fragments. Postgres uses execution-phase fragments so the pg_dump ordering @@ -110,9 +107,5 @@ remains byte-for-byte stable when composed: - `006_footer.sql` - `100_*` extension files for empty-database snapshot-only additions -MySQL and SQLite use similarly numbered domain fragments (`001_identity.sql` -through `006_usage.sql`) because their baselines are shorter and already -organized by domain. - The Rust migration tests compose these manifests too, so fragment drift is caught during `cargo test -p aether-data split_baseline_sources_match_executable_migrations`. diff --git a/crates/aether-data/runtime/schema/compose_schema.sh b/crates/aether-data/runtime/schema/compose_schema.sh index a369f4614..0a9a71f18 100644 --- a/crates/aether-data/runtime/schema/compose_schema.sh +++ b/crates/aether-data/runtime/schema/compose_schema.sh @@ -8,8 +8,6 @@ driver_schema_root="${schema_root}/drivers" bootstrap_schema_root="${schema_root}/bootstrap/postgres" adapters_root="${root}/../adapters" postgres_migrations_root="${adapters_root}/postgres/migrations" -mysql_migrations_root="${adapters_root}/mysql/migrations" -sqlite_migrations_root="${adapters_root}/sqlite/migrations" usage() { cat <<'USAGE' @@ -34,10 +32,6 @@ check_logical_generated() { local args=() local path args+=(--require-tables-from "${postgres_migrations_root}/20260403000000_baseline.sql") - for path in "${mysql_migrations_root}/"*.sql "${sqlite_migrations_root}/"*.sql; do - [[ -f "${path}" ]] || continue - args+=(--require-tables-from "${path}") - done (cd "${workspace_root}" && cargo run -q -p aether-data-schema --bin aether-schema -- check "${args[@]}") printf 'ok generated logical schema\n' } @@ -58,12 +52,6 @@ output_path() { postgres/baseline) printf '%s/20260403000000_baseline.sql' "${postgres_migrations_root}" ;; - mysql/baseline) - printf '%s/20260403000000_baseline.sql' "${mysql_migrations_root}" - ;; - sqlite/baseline) - printf '%s/20260403000000_baseline.sql' "${sqlite_migrations_root}" - ;; *) printf 'unknown schema target: %s\n' "${target}" >&2 exit 2 @@ -185,52 +173,15 @@ check_bootstrap_sources() { printf 'ok bootstrap/postgres source\n' } -split_linear_baseline() { - local driver="$1" - local source - if [[ "${driver}" == "mysql" ]]; then - source="${mysql_migrations_root}/20260403000000_baseline.sql" - else - source="${sqlite_migrations_root}/20260403000000_baseline.sql" - fi - local target="${driver}/baseline" - write_manifest "${target}" \ - "001_identity.sql" \ - "002_provider_catalog.sql" \ - "003_auth_config.sql" \ - "004_proxy_nodes.sql" \ - "005_wallet_billing.sql" \ - "006_usage.sql" - - if [[ "${driver}" == "mysql" ]]; then - split_fragment "${source}" "${target}" 001_identity.sql 1 121 - split_fragment "${source}" "${target}" 002_provider_catalog.sql 122 417 - split_fragment "${source}" "${target}" 003_auth_config.sql 418 487 - split_fragment "${source}" "${target}" 004_proxy_nodes.sql 488 527 - split_fragment "${source}" "${target}" 005_wallet_billing.sql 528 711 - split_fragment "${source}" "${target}" 006_usage.sql 712 EOF - else - split_fragment "${source}" "${target}" 001_identity.sql 1 117 - split_fragment "${source}" "${target}" 002_provider_catalog.sql 118 417 - split_fragment "${source}" "${target}" 003_auth_config.sql 418 485 - split_fragment "${source}" "${target}" 004_proxy_nodes.sql 486 525 - split_fragment "${source}" "${target}" 005_wallet_billing.sql 526 728 - split_fragment "${source}" "${target}" 006_usage.sql 729 EOF - fi -} targets=( "postgres/baseline" - "mysql/baseline" - "sqlite/baseline" ) cmd="${1:-}" case "${cmd}" in split) split_postgres_baseline - split_linear_baseline mysql - split_linear_baseline sqlite ;; generate) generate_logical_schema diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/001_identity.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/001_identity.sql deleted file mode 100644 index 66546f5f4..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/001_identity.sql +++ /dev/null @@ -1,167 +0,0 @@ -CREATE TABLE IF NOT EXISTS users ( - id VARCHAR(64) PRIMARY KEY, - external_id VARCHAR(255), - email VARCHAR(320), - username VARCHAR(255), - password_hash VARCHAR(255), - role VARCHAR(64), - auth_source VARCHAR(64) NOT NULL DEFAULT 'local', - email_verified TINYINT(1) NOT NULL DEFAULT 0, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_deleted TINYINT(1) NOT NULL DEFAULT 0, - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - model_capability_settings TEXT, - rate_limit INT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - last_login_at BIGINT, - ldap_dn VARCHAR(1024), - ldap_username VARCHAR(255), - UNIQUE KEY users_email_key (email), - UNIQUE KEY users_username_key (username) -); - -CREATE TABLE IF NOT EXISTS api_keys ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - key_hash VARCHAR(255) NOT NULL, - key_encrypted TEXT, - name VARCHAR(255), - key_prefix VARCHAR(64), - status VARCHAR(64) NOT NULL DEFAULT 'active', - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - rate_limit INT DEFAULT 100, - concurrent_limit INT, - force_capabilities TEXT, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_locked TINYINT(1) NOT NULL DEFAULT 0, - is_standalone TINYINT(1) NOT NULL DEFAULT 0, - auto_delete_on_expiry TINYINT(1) NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - metadata TEXT, - expires_at BIGINT, - last_used_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY api_keys_key_hash_key (key_hash), - KEY api_keys_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS audit_logs ( - id VARCHAR(64) PRIMARY KEY, - event_type VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - description TEXT NOT NULL, - ip_address VARCHAR(64), - user_agent VARCHAR(512), - request_id VARCHAR(128), - event_metadata TEXT, - status_code INT, - error_message TEXT, - created_at BIGINT NOT NULL, - KEY audit_logs_created_at_idx (created_at), - KEY audit_logs_event_type_idx (event_type), - KEY audit_logs_request_id_idx (request_id), - KEY audit_logs_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS announcements ( - id VARCHAR(64) PRIMARY KEY, - title VARCHAR(200) NOT NULL, - content TEXT NOT NULL, - `type` VARCHAR(32) NOT NULL DEFAULT 'info', - priority INT NOT NULL DEFAULT 0, - author_id VARCHAR(64), - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_pinned TINYINT(1) NOT NULL DEFAULT 0, - start_time BIGINT, - end_time BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY announcements_author_id_idx (author_id), - KEY announcements_created_at_idx (created_at), - KEY announcements_is_active_idx (is_active) -); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - announcement_id VARCHAR(64) NOT NULL, - read_at BIGINT NOT NULL, - UNIQUE KEY uq_user_announcement (user_id, announcement_id), - KEY announcement_reads_announcement_id_idx (announcement_id), - KEY announcement_reads_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS management_tokens ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - description TEXT, - token_hash VARCHAR(255) NOT NULL, - token_prefix VARCHAR(64), - allowed_ips TEXT, - expires_at BIGINT, - last_used_at BIGINT, - last_used_ip VARCHAR(255), - usage_count BIGINT NOT NULL DEFAULT 0, - is_active TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY management_tokens_token_hash_key (token_hash), - UNIQUE KEY uq_management_tokens_user_name (user_id, name), - KEY management_tokens_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS user_preferences ( - id VARCHAR(36) PRIMARY KEY, - user_id VARCHAR(36) NOT NULL UNIQUE, - avatar_url VARCHAR(500), - bio TEXT, - default_provider_id VARCHAR(36), - theme VARCHAR(20) NOT NULL DEFAULT 'light', - language VARCHAR(10) NOT NULL DEFAULT 'zh-CN', - timezone VARCHAR(50) NOT NULL DEFAULT 'Asia/Shanghai', - email_notifications BOOLEAN NOT NULL DEFAULT TRUE, - usage_alerts BOOLEAN NOT NULL DEFAULT TRUE, - announcement_notifications BOOLEAN NOT NULL DEFAULT TRUE, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - INDEX user_preferences_default_provider_id_idx (default_provider_id), - INDEX user_preferences_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS user_sessions ( - id VARCHAR(36) PRIMARY KEY, - user_id VARCHAR(36) NOT NULL, - client_device_id VARCHAR(128) NOT NULL, - device_label VARCHAR(120), - device_type VARCHAR(20) NOT NULL DEFAULT 'unknown', - browser_name VARCHAR(50), - browser_version VARCHAR(50), - os_name VARCHAR(50), - os_version VARCHAR(50), - device_model VARCHAR(100), - ip_address VARCHAR(45), - user_agent VARCHAR(1000), - client_hints TEXT, - refresh_token_hash VARCHAR(64) NOT NULL, - prev_refresh_token_hash VARCHAR(64), - rotated_at BIGINT, - last_seen_at BIGINT NOT NULL, - expires_at BIGINT NOT NULL, - revoked_at BIGINT, - revoke_reason VARCHAR(100), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - INDEX user_sessions_user_active_idx (user_id, revoked_at, expires_at), - INDEX user_sessions_user_device_idx (user_id, client_device_id) -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/002_provider_catalog.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/002_provider_catalog.sql deleted file mode 100644 index a4777197d..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/002_provider_catalog.sql +++ /dev/null @@ -1,326 +0,0 @@ -CREATE TABLE IF NOT EXISTS billing_rules ( - id VARCHAR(64) PRIMARY KEY, - global_model_id VARCHAR(64), - model_id VARCHAR(64), - name VARCHAR(255) NOT NULL, - task_type VARCHAR(64) NOT NULL DEFAULT 'chat', - expression TEXT NOT NULL, - variables TEXT NOT NULL, - dimension_mappings TEXT NOT NULL, - is_enabled TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY billing_rules_global_model_task_idx (global_model_id, task_type, is_enabled), - KEY billing_rules_model_task_idx (model_id, task_type, is_enabled) -); - -CREATE TABLE IF NOT EXISTS dimension_collectors ( - id VARCHAR(64) PRIMARY KEY, - api_format VARCHAR(64) NOT NULL, - task_type VARCHAR(64) NOT NULL, - dimension_name VARCHAR(128) NOT NULL, - source_type VARCHAR(64) NOT NULL, - source_path VARCHAR(255), - value_type VARCHAR(64) NOT NULL DEFAULT 'float', - transform_expression TEXT, - default_value VARCHAR(255), - priority INT NOT NULL DEFAULT 0, - is_enabled TINYINT(1) NOT NULL DEFAULT 1, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY dimension_collectors_enabled_idx ( - api_format, - task_type, - dimension_name, - priority, - is_enabled - ) -); - -CREATE TABLE IF NOT EXISTS providers ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - description TEXT, - website VARCHAR(500), - provider_type VARCHAR(64) NOT NULL, - billing_type VARCHAR(64), - monthly_quota_usd DOUBLE, - monthly_used_usd DOUBLE, - quota_reset_day INT, - quota_last_reset_at BIGINT, - quota_expires_at BIGINT, - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - priority BIGINT NOT NULL DEFAULT 0, - provider_priority INT NOT NULL DEFAULT 100, - keep_priority_on_conversion TINYINT(1) NOT NULL DEFAULT 0, - enable_format_conversion TINYINT(1) NOT NULL DEFAULT 1, - concurrent_limit INT, - max_retries INT, - proxy TEXT, - request_timeout DOUBLE, - stream_first_byte_timeout DOUBLE, - config TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY providers_name_key (name) -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - api_key TEXT, - encrypted_key TEXT, - auth_type VARCHAR(32) NOT NULL DEFAULT 'api_key', - auth_config TEXT, - note TEXT, - internal_priority INT NOT NULL DEFAULT 50, - capabilities TEXT, - api_formats TEXT, - auth_type_by_format TEXT, - allow_auth_channel_mismatch_formats TEXT, - rate_multipliers TEXT, - global_priority_by_format TEXT, - allowed_models TEXT, - expires_at BIGINT, - cache_ttl_minutes INT NOT NULL DEFAULT 5, - max_probe_interval_minutes INT NOT NULL DEFAULT 32, - proxy TEXT, - fingerprint TEXT, - concurrent_limit INT, - learned_rpm_limit INT, - concurrent_429_count INT NOT NULL DEFAULT 0, - rpm_429_count INT NOT NULL DEFAULT 0, - last_429_at BIGINT, - last_429_type VARCHAR(64), - adjustment_history TEXT, - utilization_samples TEXT, - last_probe_increase_at BIGINT, - last_rpm_peak INT, - request_count BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - success_count BIGINT NOT NULL DEFAULT 0, - error_count BIGINT NOT NULL DEFAULT 0, - total_response_time_ms BIGINT NOT NULL DEFAULT 0, - last_used_at BIGINT, - auto_fetch_models TINYINT(1) NOT NULL DEFAULT 0, - last_models_fetch_at BIGINT, - last_models_fetch_error TEXT, - locked_models TEXT, - model_include_patterns TEXT, - model_exclude_patterns TEXT, - upstream_metadata TEXT, - oauth_invalid_at BIGINT, - oauth_invalid_reason VARCHAR(255), - status_snapshot TEXT, - health_by_format TEXT, - circuit_breaker_by_format TEXT, - status VARCHAR(64) NOT NULL DEFAULT 'active', - is_active TINYINT(1) NOT NULL DEFAULT 1, - weight BIGINT NOT NULL DEFAULT 1, - rpm_limit BIGINT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY provider_api_keys_provider_id_idx (provider_id), - KEY idx_provider_api_keys_provider_default_sort (provider_id, internal_priority, name, id) -); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id VARCHAR(192) PRIMARY KEY, - pool_kind VARCHAR(64) NOT NULL, - pool_id VARCHAR(64) NOT NULL, - member_kind VARCHAR(64) NOT NULL, - member_id VARCHAR(64) NOT NULL, - capability VARCHAR(64) NOT NULL, - scope_kind VARCHAR(64) NOT NULL, - scope_id VARCHAR(128), - score DOUBLE NOT NULL DEFAULT 0, - hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown', - score_version BIGINT NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at BIGINT, - last_scheduled_at BIGINT, - last_success_at BIGINT, - last_failure_at BIGINT, - failure_count BIGINT NOT NULL DEFAULT 0, - last_probe_attempt_at BIGINT, - last_probe_success_at BIGINT, - last_probe_failure_at BIGINT, - probe_failure_count BIGINT NOT NULL DEFAULT 0, - probe_status VARCHAR(64) NOT NULL DEFAULT 'never', - updated_at BIGINT NOT NULL, - KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC), - KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id), - KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at), - KEY pool_member_scores_updated_at_idx (updated_at) -); - -CREATE TABLE IF NOT EXISTS gemini_file_mappings ( - id VARCHAR(64) PRIMARY KEY, - file_name VARCHAR(512) NOT NULL, - key_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - display_name VARCHAR(512), - mime_type VARCHAR(255), - source_hash VARCHAR(128), - created_at BIGINT NOT NULL, - expires_at BIGINT NOT NULL, - UNIQUE KEY gemini_file_mappings_file_name_key (file_name), - KEY gemini_file_mappings_key_id_idx (key_id), - KEY gemini_file_mappings_user_id_idx (user_id), - KEY gemini_file_mappings_expires_at_idx (expires_at), - KEY gemini_file_mappings_source_hash_idx (source_hash) -); - -CREATE TABLE IF NOT EXISTS request_candidates ( - id VARCHAR(64) PRIMARY KEY, - request_id VARCHAR(128) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - username VARCHAR(255), - api_key_name VARCHAR(255), - candidate_index INT NOT NULL, - retry_index INT NOT NULL DEFAULT 0, - provider_id VARCHAR(64), - endpoint_id VARCHAR(64), - key_id VARCHAR(64), - status VARCHAR(32) NOT NULL, - skip_reason TEXT, - is_cached TINYINT(1) NOT NULL DEFAULT 0, - status_code INT, - error_type VARCHAR(128), - error_message TEXT, - latency_ms INT, - concurrent_requests INT, - extra_data TEXT, - required_capabilities TEXT, - created_at BIGINT NOT NULL, - started_at BIGINT, - finished_at BIGINT, - UNIQUE KEY uq_request_candidate_with_retry (request_id, candidate_index, retry_index), - KEY request_candidates_request_id_idx (request_id), - KEY request_candidates_provider_id_idx (provider_id), - KEY request_candidates_endpoint_id_idx (endpoint_id), - KEY request_candidates_status_idx (status), - KEY request_candidates_created_at_idx (created_at), - KEY request_candidates_endpoint_status_created_idx (endpoint_id, status, created_at) -); - -CREATE TABLE IF NOT EXISTS video_tasks ( - id VARCHAR(64) PRIMARY KEY, - short_id VARCHAR(32), - request_id VARCHAR(128) NOT NULL, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - username VARCHAR(255), - api_key_name VARCHAR(255), - external_task_id VARCHAR(255), - provider_id VARCHAR(64), - endpoint_id VARCHAR(64), - key_id VARCHAR(64), - client_api_format VARCHAR(128), - provider_api_format VARCHAR(128), - format_converted TINYINT(1) NOT NULL DEFAULT 0, - model VARCHAR(255), - prompt TEXT, - original_request_body TEXT, - duration_seconds INT, - resolution VARCHAR(64), - aspect_ratio VARCHAR(32), - size VARCHAR(64), - status VARCHAR(32) NOT NULL DEFAULT 'pending', - progress_percent INT NOT NULL DEFAULT 0, - progress_message TEXT, - retry_count INT NOT NULL DEFAULT 0, - poll_interval_seconds INT NOT NULL DEFAULT 10, - next_poll_at BIGINT, - poll_count INT NOT NULL DEFAULT 0, - max_poll_count INT NOT NULL DEFAULT 360, - created_at BIGINT NOT NULL, - submitted_at BIGINT, - completed_at BIGINT, - updated_at BIGINT NOT NULL, - error_code VARCHAR(128), - error_message TEXT, - video_url TEXT, - request_metadata TEXT, - UNIQUE KEY video_tasks_short_id_key (short_id), - UNIQUE KEY video_tasks_request_id_key (request_id), - KEY video_tasks_external_id_idx (external_task_id), - KEY video_tasks_next_poll_idx (next_poll_at), - KEY video_tasks_user_status_idx (user_id, status), - KEY video_tasks_api_key_id_idx (api_key_id), - KEY video_tasks_provider_id_idx (provider_id), - KEY video_tasks_endpoint_id_idx (endpoint_id), - KEY video_tasks_key_id_idx (key_id) -); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - name VARCHAR(255) NOT NULL, - base_url TEXT NOT NULL, - api_format VARCHAR(128), - api_family VARCHAR(128), - endpoint_kind VARCHAR(128), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - health_score DOUBLE NOT NULL DEFAULT 1.0, - weight BIGINT NOT NULL DEFAULT 1, - header_rules TEXT, - body_rules TEXT, - max_retries INT, - custom_path TEXT, - metadata TEXT, - config TEXT, - format_acceptance_config TEXT, - proxy TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY provider_endpoints_provider_id_idx (provider_id) -); - -CREATE TABLE IF NOT EXISTS models ( - id VARCHAR(64) PRIMARY KEY, - provider_id VARCHAR(64) NOT NULL, - global_model_id VARCHAR(64), - provider_model_name VARCHAR(255) NOT NULL, - global_model_name VARCHAR(255), - api_format VARCHAR(128), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - is_available TINYINT(1) NOT NULL DEFAULT 1, - price_per_request DOUBLE, - tiered_pricing TEXT, - supports_vision TINYINT(1), - supports_function_calling TINYINT(1), - supports_streaming TINYINT(1), - supports_extended_thinking TINYINT(1), - supports_image_generation TINYINT(1), - provider_model_mappings TEXT, - config TEXT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY models_provider_id_idx (provider_id) -); - -CREATE TABLE IF NOT EXISTS global_models ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - display_name VARCHAR(255), - enabled TINYINT(1) NOT NULL DEFAULT 1, - is_active TINYINT(1) NOT NULL DEFAULT 1, - default_price_per_request DOUBLE, - default_tiered_pricing TEXT, - supported_capabilities TEXT, - usage_count BIGINT NOT NULL DEFAULT 0, - config TEXT, - metadata TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY global_models_name_key (name) -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/003_auth_config.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/003_auth_config.sql deleted file mode 100644 index 73e37f7e5..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/003_auth_config.sql +++ /dev/null @@ -1,71 +0,0 @@ -CREATE TABLE IF NOT EXISTS system_configs ( - id VARCHAR(64) PRIMARY KEY, - `key` VARCHAR(255) NOT NULL, - value TEXT NOT NULL, - description TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY system_configs_key_key (`key`) -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - id VARCHAR(64) PRIMARY KEY, - module_type VARCHAR(128) NOT NULL, - enabled TINYINT(1) NOT NULL DEFAULT 1, - config TEXT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY auth_modules_module_type_key (module_type) -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - provider_type VARCHAR(64) PRIMARY KEY, - display_name VARCHAR(255) NOT NULL, - client_id TEXT NOT NULL, - client_secret_encrypted TEXT, - authorization_url_override VARCHAR(500), - token_url_override VARCHAR(500), - userinfo_url_override VARCHAR(500), - scopes TEXT, - redirect_uri VARCHAR(500) NOT NULL, - frontend_callback_url VARCHAR(500) NOT NULL, - attribute_mapping TEXT, - extra_config TEXT, - is_enabled TINYINT(1) NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - id BIGINT PRIMARY KEY AUTO_INCREMENT, - server_url VARCHAR(255) NOT NULL, - bind_dn TEXT NOT NULL, - bind_password_encrypted TEXT, - base_dn TEXT NOT NULL, - user_search_filter VARCHAR(512) NOT NULL DEFAULT '(uid={username})', - username_attr VARCHAR(50) NOT NULL DEFAULT 'uid', - email_attr VARCHAR(50) NOT NULL DEFAULT 'mail', - display_name_attr VARCHAR(50) NOT NULL DEFAULT 'cn', - is_enabled TINYINT(1) NOT NULL DEFAULT 0, - is_exclusive TINYINT(1) NOT NULL DEFAULT 0, - use_starttls TINYINT(1) NOT NULL DEFAULT 0, - connect_timeout INT NOT NULL DEFAULT 10, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - provider_type VARCHAR(64) NOT NULL, - provider_user_id VARCHAR(255) NOT NULL, - provider_username VARCHAR(255), - provider_email VARCHAR(255), - extra_data TEXT, - linked_at BIGINT NOT NULL, - last_login_at BIGINT, - UNIQUE KEY uq_user_oauth_links_provider_user (provider_type, provider_user_id), - UNIQUE KEY uq_user_oauth_links_user_provider (user_id, provider_type), - KEY user_oauth_links_provider_type_idx (provider_type), - KEY user_oauth_links_user_id_idx (user_id) -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/004_proxy_nodes.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/004_proxy_nodes.sql deleted file mode 100644 index 18904c145..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/004_proxy_nodes.sql +++ /dev/null @@ -1,39 +0,0 @@ -CREATE TABLE IF NOT EXISTS proxy_nodes ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - ip VARCHAR(512) NOT NULL, - port INT NOT NULL, - region VARCHAR(100), - status VARCHAR(32) NOT NULL DEFAULT 'online', - registered_by VARCHAR(64), - last_heartbeat_at BIGINT, - heartbeat_interval INT NOT NULL DEFAULT 30, - active_connections INT NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - avg_latency_ms DOUBLE, - is_manual TINYINT(1) NOT NULL DEFAULT 0, - proxy_url VARCHAR(500), - proxy_username VARCHAR(255), - proxy_password TEXT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - remote_config TEXT, - config_version INT NOT NULL DEFAULT 0, - hardware_info TEXT, - estimated_max_concurrency INT, - tunnel_mode TINYINT(1) NOT NULL DEFAULT 0, - tunnel_connected TINYINT(1) NOT NULL DEFAULT 0, - tunnel_connected_at BIGINT, - failed_requests BIGINT NOT NULL DEFAULT 0, - dns_failures BIGINT NOT NULL DEFAULT 0, - stream_errors BIGINT NOT NULL DEFAULT 0, - proxy_metadata TEXT -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - id BIGINT PRIMARY KEY AUTO_INCREMENT, - node_id VARCHAR(64) NOT NULL, - event_type VARCHAR(64) NOT NULL, - detail VARCHAR(500), - created_at BIGINT NOT NULL -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/005_wallet_billing.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/005_wallet_billing.sql deleted file mode 100644 index 52113ed9d..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/005_wallet_billing.sql +++ /dev/null @@ -1,183 +0,0 @@ -CREATE TABLE IF NOT EXISTS wallets ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64), - api_key_id VARCHAR(64), - balance DOUBLE NOT NULL DEFAULT 0, - gift_balance DOUBLE NOT NULL DEFAULT 0, - limit_mode VARCHAR(64) NOT NULL DEFAULT 'finite', - currency VARCHAR(16) NOT NULL DEFAULT 'USD', - status VARCHAR(64) NOT NULL DEFAULT 'active', - total_recharged DOUBLE NOT NULL DEFAULT 0, - total_consumed DOUBLE NOT NULL DEFAULT 0, - total_refunded DOUBLE NOT NULL DEFAULT 0, - total_adjusted DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY wallets_user_id_key (user_id), - UNIQUE KEY wallets_api_key_id_key (api_key_id), - KEY wallets_api_key_id_idx (api_key_id), - KEY wallets_user_id_idx (user_id) -); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - id VARCHAR(64) PRIMARY KEY, - wallet_id VARCHAR(64) NOT NULL, - category VARCHAR(64) NOT NULL, - reason_code VARCHAR(64) NOT NULL, - amount DOUBLE NOT NULL, - balance_before DOUBLE NOT NULL, - balance_after DOUBLE NOT NULL, - recharge_balance_before DOUBLE NOT NULL, - recharge_balance_after DOUBLE NOT NULL, - gift_balance_before DOUBLE NOT NULL, - gift_balance_after DOUBLE NOT NULL, - link_type VARCHAR(64), - link_id VARCHAR(128), - operator_id VARCHAR(64), - description TEXT, - created_at BIGINT NOT NULL, - KEY idx_wallet_tx_wallet_created (wallet_id, created_at), - KEY idx_wallet_tx_category_created (category, created_at), - KEY idx_wallet_tx_reason_created (reason_code, created_at), - KEY idx_wallet_tx_link (link_type, link_id), - KEY ix_wallet_transactions_operator_id (operator_id) -); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - id VARCHAR(64) PRIMARY KEY, - wallet_id VARCHAR(64) NOT NULL, - billing_date VARCHAR(16) NOT NULL, - billing_timezone VARCHAR(64) NOT NULL, - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - first_finalized_at BIGINT, - last_finalized_at BIGINT, - aggregated_at BIGINT NOT NULL, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_wallet_daily_usage_wallet_date (wallet_id, billing_timezone, billing_date) -); - -CREATE TABLE IF NOT EXISTS payment_orders ( - id VARCHAR(64) PRIMARY KEY, - order_no VARCHAR(128) NOT NULL, - wallet_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - amount_usd DOUBLE NOT NULL, - pay_amount DOUBLE, - pay_currency VARCHAR(16), - exchange_rate DOUBLE, - refunded_amount_usd DOUBLE NOT NULL DEFAULT 0, - refundable_amount_usd DOUBLE NOT NULL DEFAULT 0, - payment_method VARCHAR(64) NOT NULL, - gateway_order_id VARCHAR(128), - gateway_response TEXT, - status VARCHAR(64) NOT NULL DEFAULT 'pending', - created_at BIGINT NOT NULL, - paid_at BIGINT, - credited_at BIGINT, - expires_at BIGINT, - UNIQUE KEY uq_payment_orders_order_no (order_no), - KEY idx_payment_orders_wallet_created (wallet_id, created_at), - KEY idx_payment_orders_user_created (user_id, created_at), - KEY idx_payment_orders_status (status), - KEY idx_payment_orders_gateway_order_id (gateway_order_id) -); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - id VARCHAR(64) PRIMARY KEY, - payment_order_id VARCHAR(64), - payment_method VARCHAR(64) NOT NULL, - callback_key VARCHAR(128) NOT NULL, - order_no VARCHAR(128), - gateway_order_id VARCHAR(128), - payload_hash VARCHAR(128), - signature_valid TINYINT(1) NOT NULL DEFAULT 0, - status VARCHAR(64) NOT NULL DEFAULT 'received', - payload TEXT, - error_message TEXT, - created_at BIGINT NOT NULL, - processed_at BIGINT, - UNIQUE KEY uq_payment_callbacks_callback_key (callback_key), - KEY idx_payment_callbacks_order (order_no), - KEY idx_payment_callbacks_gateway_order (gateway_order_id), - KEY idx_payment_callbacks_created (created_at), - KEY ix_payment_callbacks_payment_order_id (payment_order_id) -); - -CREATE TABLE IF NOT EXISTS refund_requests ( - id VARCHAR(64) PRIMARY KEY, - refund_no VARCHAR(128) NOT NULL, - wallet_id VARCHAR(64) NOT NULL, - user_id VARCHAR(64), - payment_order_id VARCHAR(64), - source_type VARCHAR(64) NOT NULL, - source_id VARCHAR(128), - refund_mode VARCHAR(64) NOT NULL, - amount_usd DOUBLE NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'pending_approval', - reason TEXT, - requested_by VARCHAR(64), - approved_by VARCHAR(64), - processed_by VARCHAR(64), - gateway_refund_id VARCHAR(128), - payout_method VARCHAR(64), - payout_reference VARCHAR(255), - payout_proof TEXT, - failure_reason TEXT, - idempotency_key VARCHAR(128), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - processed_at BIGINT, - completed_at BIGINT, - UNIQUE KEY uq_refund_requests_refund_no (refund_no), - UNIQUE KEY uq_refund_requests_idempotency_key (idempotency_key), - KEY idx_refund_wallet_created (wallet_id, created_at), - KEY idx_refund_user_created (user_id, created_at), - KEY idx_refund_status (status), - KEY ix_refund_requests_payment_order_id (payment_order_id), - KEY ix_refund_requests_requested_by (requested_by), - KEY ix_refund_requests_approved_by (approved_by), - KEY ix_refund_requests_processed_by (processed_by) -); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - id VARCHAR(64) PRIMARY KEY, - name VARCHAR(255) NOT NULL, - amount_usd DOUBLE NOT NULL, - currency VARCHAR(16) NOT NULL DEFAULT 'USD', - balance_bucket VARCHAR(64) NOT NULL DEFAULT 'gift', - total_count INT NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'active', - description TEXT, - created_by VARCHAR(64), - expires_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY idx_redeem_code_batches_status (status, created_at) -); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - id VARCHAR(64) PRIMARY KEY, - batch_id VARCHAR(64) NOT NULL, - code_hash VARCHAR(128) NOT NULL, - code_prefix VARCHAR(16) NOT NULL, - code_suffix VARCHAR(16) NOT NULL, - status VARCHAR(64) NOT NULL DEFAULT 'active', - redeemed_by_user_id VARCHAR(64), - redeemed_wallet_id VARCHAR(64), - redeemed_payment_order_id VARCHAR(64), - redeemed_at BIGINT, - disabled_by VARCHAR(64), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_redeem_codes_code_hash (code_hash), - KEY idx_redeem_codes_batch_created (batch_id, created_at), - KEY idx_redeem_codes_status (status, updated_at), - KEY idx_redeem_codes_redeemed_user (redeemed_by_user_id, redeemed_at), - KEY idx_redeem_codes_redeemed_order (redeemed_payment_order_id) -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/006_usage.sql b/crates/aether-data/runtime/schema/drivers/mysql/baseline/006_usage.sql deleted file mode 100644 index a46b26aa5..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/006_usage.sql +++ /dev/null @@ -1,265 +0,0 @@ -CREATE TABLE IF NOT EXISTS `usage` ( - request_id VARCHAR(128) PRIMARY KEY, - id VARCHAR(128), - user_id VARCHAR(64), - api_key_id VARCHAR(64), - provider_name VARCHAR(255) NOT NULL DEFAULT 'unknown', - model VARCHAR(255) NOT NULL DEFAULT 'unknown', - target_model VARCHAR(255), - provider_id VARCHAR(64), - provider_endpoint_id VARCHAR(64), - provider_api_key_id VARCHAR(64), - request_type VARCHAR(64), - api_format VARCHAR(64), - api_family VARCHAR(64), - endpoint_kind VARCHAR(64), - endpoint_api_format VARCHAR(64), - provider_api_family VARCHAR(64), - provider_endpoint_kind VARCHAR(64), - has_format_conversion TINYINT(1) NOT NULL DEFAULT 0, - is_stream TINYINT(1) NOT NULL DEFAULT 0, - upstream_is_stream TINYINT(1), - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_input_tokens 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, - cache_creation_cost_usd DOUBLE NOT NULL DEFAULT 0, - cache_read_cost_usd DOUBLE NOT NULL DEFAULT 0, - output_price_per_1m DOUBLE, - status_code INT, - error_message TEXT, - error_category VARCHAR(255), - response_time_ms BIGINT, - first_byte_time_ms BIGINT, - wallet_id VARCHAR(64), - status VARCHAR(64) NOT NULL DEFAULT 'completed', - billing_status VARCHAR(64) NOT NULL DEFAULT 'pending', - total_cost_usd DOUBLE NOT NULL DEFAULT 0, - actual_total_cost_usd DOUBLE NOT NULL DEFAULT 0, - request_metadata TEXT, - candidate_id VARCHAR(128), - candidate_index BIGINT, - key_name VARCHAR(255), - planner_kind VARCHAR(64), - route_family VARCHAR(128), - route_kind VARCHAR(128), - execution_path VARCHAR(128), - local_execution_runtime_miss_reason VARCHAR(255), - wallet_balance_before DOUBLE, - wallet_balance_after DOUBLE, - wallet_recharge_balance_before DOUBLE, - wallet_recharge_balance_after DOUBLE, - wallet_gift_balance_before DOUBLE, - wallet_gift_balance_after DOUBLE, - finalized_at BIGINT, - created_at_unix_ms BIGINT NOT NULL DEFAULT 0, - updated_at_unix_secs BIGINT NOT NULL DEFAULT 0, - KEY usage_api_key_id_idx (api_key_id), - KEY usage_billing_status_idx (billing_status), - KEY usage_created_at_idx (created_at_unix_ms), - KEY usage_provider_api_key_id_idx (provider_api_key_id), - KEY usage_provider_id_idx (provider_id), - KEY usage_request_id_idx (request_id), - KEY usage_user_id_idx (user_id), - KEY usage_wallet_id_idx (wallet_id) -); - -CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( - request_id VARCHAR(128) PRIMARY KEY, - billing_status VARCHAR(64) NOT NULL, - wallet_id VARCHAR(64), - wallet_balance_before DOUBLE, - wallet_balance_after DOUBLE, - wallet_recharge_balance_before DOUBLE, - wallet_recharge_balance_after DOUBLE, - wallet_gift_balance_before DOUBLE, - wallet_gift_balance_after DOUBLE, - provider_monthly_used_usd DOUBLE, - finalized_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - KEY usage_settlement_snapshots_billing_status_idx (billing_status), - KEY usage_settlement_snapshots_wallet_id_idx (wallet_id) -); - - -CREATE TABLE IF NOT EXISTS stats_hourly ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - is_complete TINYINT(1) NOT NULL DEFAULT 0, - aggregated_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_hour (hour_utc) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - user_id VARCHAR(64) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_user (hour_utc, user_id) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - user_id VARCHAR(64) NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_user_model (hour_utc, user_id, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_model (hour_utc, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - id VARCHAR(64) PRIMARY KEY, - hour_utc BIGINT NOT NULL, - provider_name VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_hourly_provider (hour_utc, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - fallback_count BIGINT NOT NULL DEFAULT 0, - unique_models BIGINT NOT NULL DEFAULT 0, - unique_providers BIGINT NOT NULL DEFAULT 0, - is_complete TINYINT(1) NOT NULL DEFAULT 0, - aggregated_at BIGINT, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_date (`date`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - model VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - avg_response_time_ms DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_model (`date`, model) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - provider_name VARCHAR(255) NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_provider (`date`, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - id VARCHAR(64) PRIMARY KEY, - api_key_id VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - api_key_name VARCHAR(255), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_api_key (`date`, api_key_id) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - id VARCHAR(64) PRIMARY KEY, - `date` BIGINT NOT NULL, - error_category VARCHAR(255) NOT NULL, - provider_name VARCHAR(255), - model VARCHAR(255), - count BIGINT NOT NULL DEFAULT 0, - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_daily_error (`date`, error_category, provider_name, model) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - id VARCHAR(64) PRIMARY KEY, - user_id VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - total_requests BIGINT NOT NULL DEFAULT 0, - success_requests BIGINT NOT NULL DEFAULT 0, - error_requests BIGINT NOT NULL DEFAULT 0, - input_tokens BIGINT NOT NULL DEFAULT 0, - output_tokens BIGINT NOT NULL DEFAULT 0, - cache_creation_tokens BIGINT NOT NULL DEFAULT 0, - cache_read_tokens BIGINT NOT NULL DEFAULT 0, - total_cost DOUBLE NOT NULL DEFAULT 0, - username VARCHAR(255), - created_at BIGINT NOT NULL, - updated_at BIGINT NOT NULL, - UNIQUE KEY uq_stats_user_daily (`date`, user_id) -); diff --git a/crates/aether-data/runtime/schema/drivers/mysql/baseline/manifest.txt b/crates/aether-data/runtime/schema/drivers/mysql/baseline/manifest.txt deleted file mode 100644 index 6caec84c6..000000000 --- a/crates/aether-data/runtime/schema/drivers/mysql/baseline/manifest.txt +++ /dev/null @@ -1,6 +0,0 @@ -001_identity.sql -002_provider_catalog.sql -003_auth_config.sql -004_proxy_nodes.sql -005_wallet_billing.sql -006_usage.sql diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/001_identity.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/001_identity.sql deleted file mode 100644 index d67ba57cb..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/001_identity.sql +++ /dev/null @@ -1,167 +0,0 @@ -CREATE TABLE IF NOT EXISTS users ( - id TEXT PRIMARY KEY, - external_id TEXT, - email TEXT UNIQUE, - username TEXT UNIQUE, - password_hash TEXT, - role TEXT, - auth_source TEXT NOT NULL DEFAULT 'local', - email_verified INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - is_deleted INTEGER NOT NULL DEFAULT 0, - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - model_capability_settings TEXT, - rate_limit INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - last_login_at INTEGER, - ldap_dn TEXT, - ldap_username TEXT -); - -CREATE TABLE IF NOT EXISTS api_keys ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - key_hash TEXT NOT NULL UNIQUE, - key_encrypted TEXT, - name TEXT, - key_prefix TEXT, - status TEXT NOT NULL DEFAULT 'active', - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - rate_limit INTEGER DEFAULT 100, - concurrent_limit INTEGER, - force_capabilities TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_locked INTEGER NOT NULL DEFAULT 0, - is_standalone INTEGER NOT NULL DEFAULT 0, - auto_delete_on_expiry INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - metadata TEXT, - expires_at INTEGER, - last_used_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS api_keys_user_id_idx ON api_keys (user_id); - -CREATE TABLE IF NOT EXISTS audit_logs ( - id TEXT PRIMARY KEY, - event_type TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - description TEXT NOT NULL, - ip_address TEXT, - user_agent TEXT, - request_id TEXT, - event_metadata TEXT, - status_code INTEGER, - error_message TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS audit_logs_created_at_idx ON audit_logs (created_at); -CREATE INDEX IF NOT EXISTS audit_logs_event_type_idx ON audit_logs (event_type); -CREATE INDEX IF NOT EXISTS audit_logs_request_id_idx ON audit_logs (request_id); -CREATE INDEX IF NOT EXISTS audit_logs_user_id_idx ON audit_logs (user_id); - -CREATE TABLE IF NOT EXISTS announcements ( - id TEXT PRIMARY KEY, - title TEXT NOT NULL, - content TEXT NOT NULL, - type TEXT NOT NULL DEFAULT 'info', - priority INTEGER NOT NULL DEFAULT 0, - author_id TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_pinned INTEGER NOT NULL DEFAULT 0, - start_time INTEGER, - end_time INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS announcements_author_id_idx ON announcements (author_id); -CREATE INDEX IF NOT EXISTS announcements_created_at_idx ON announcements (created_at); -CREATE INDEX IF NOT EXISTS announcements_is_active_idx ON announcements (is_active); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - announcement_id TEXT NOT NULL, - read_at INTEGER NOT NULL, - UNIQUE (user_id, announcement_id) -); -CREATE INDEX IF NOT EXISTS announcement_reads_announcement_id_idx ON announcement_reads (announcement_id); -CREATE INDEX IF NOT EXISTS announcement_reads_user_id_idx ON announcement_reads (user_id); - -CREATE TABLE IF NOT EXISTS management_tokens ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - name TEXT NOT NULL, - description TEXT, - token_hash TEXT NOT NULL UNIQUE, - token_prefix TEXT, - allowed_ips TEXT, - expires_at INTEGER, - last_used_at INTEGER, - last_used_ip TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (user_id, name) -); -CREATE INDEX IF NOT EXISTS management_tokens_user_id_idx ON management_tokens (user_id); - -CREATE TABLE IF NOT EXISTS user_preferences ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL UNIQUE, - avatar_url TEXT, - bio TEXT, - default_provider_id TEXT, - theme TEXT NOT NULL DEFAULT 'light', - language TEXT NOT NULL DEFAULT 'zh-CN', - timezone TEXT NOT NULL DEFAULT 'Asia/Shanghai', - email_notifications INTEGER NOT NULL DEFAULT 1, - usage_alerts INTEGER NOT NULL DEFAULT 1, - announcement_notifications INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS user_preferences_default_provider_id_idx - ON user_preferences (default_provider_id); -CREATE INDEX IF NOT EXISTS user_preferences_user_id_idx - ON user_preferences (user_id); - -CREATE TABLE IF NOT EXISTS user_sessions ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - client_device_id TEXT NOT NULL, - device_label TEXT, - device_type TEXT NOT NULL DEFAULT 'unknown', - browser_name TEXT, - browser_version TEXT, - os_name TEXT, - os_version TEXT, - device_model TEXT, - ip_address TEXT, - user_agent TEXT, - client_hints TEXT, - refresh_token_hash TEXT NOT NULL, - prev_refresh_token_hash TEXT, - rotated_at INTEGER, - last_seen_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - revoked_at INTEGER, - revoke_reason TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS user_sessions_user_active_idx - ON user_sessions (user_id, revoked_at, expires_at); -CREATE INDEX IF NOT EXISTS user_sessions_user_device_idx - ON user_sessions (user_id, client_device_id); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/002_provider_catalog.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/002_provider_catalog.sql deleted file mode 100644 index 59bbba96e..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/002_provider_catalog.sql +++ /dev/null @@ -1,330 +0,0 @@ -CREATE TABLE IF NOT EXISTS billing_rules ( - id TEXT PRIMARY KEY, - global_model_id TEXT, - model_id TEXT, - name TEXT NOT NULL, - task_type TEXT NOT NULL DEFAULT 'chat', - expression TEXT NOT NULL, - variables TEXT NOT NULL DEFAULT '{}', - dimension_mappings TEXT NOT NULL DEFAULT '{}', - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CHECK ( - (global_model_id IS NOT NULL AND model_id IS NULL) - OR (global_model_id IS NULL AND model_id IS NOT NULL) - ) -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_billing_rules_global_model_task - ON billing_rules (global_model_id, task_type) - WHERE is_enabled = 1 AND global_model_id IS NOT NULL; -CREATE UNIQUE INDEX IF NOT EXISTS uq_billing_rules_model_task - ON billing_rules (model_id, task_type) - WHERE is_enabled = 1 AND model_id IS NOT NULL; - -CREATE TABLE IF NOT EXISTS dimension_collectors ( - id TEXT PRIMARY KEY, - api_format TEXT NOT NULL, - task_type TEXT NOT NULL, - dimension_name TEXT NOT NULL, - source_type TEXT NOT NULL, - source_path TEXT, - value_type TEXT NOT NULL DEFAULT 'float', - transform_expression TEXT, - default_value TEXT, - priority INTEGER NOT NULL DEFAULT 0, - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CHECK ( - (source_type = 'computed' AND source_path IS NULL AND transform_expression IS NOT NULL) - OR (source_type <> 'computed' AND source_path IS NOT NULL) - ) -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_dimension_collectors_enabled - ON dimension_collectors (api_format, task_type, dimension_name, priority) - WHERE is_enabled = 1; - -CREATE TABLE IF NOT EXISTS providers ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - description TEXT, - website TEXT, - provider_type TEXT NOT NULL, - billing_type TEXT, - monthly_quota_usd REAL, - monthly_used_usd REAL, - quota_reset_day INTEGER, - quota_last_reset_at INTEGER, - quota_expires_at INTEGER, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - priority INTEGER NOT NULL DEFAULT 0, - provider_priority INTEGER NOT NULL DEFAULT 100, - keep_priority_on_conversion INTEGER NOT NULL DEFAULT 0, - enable_format_conversion INTEGER NOT NULL DEFAULT 1, - concurrent_limit INTEGER, - max_retries INTEGER, - proxy TEXT, - request_timeout REAL, - stream_first_byte_timeout REAL, - config TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - api_key TEXT, - encrypted_key TEXT, - auth_type TEXT NOT NULL DEFAULT 'api_key', - auth_config TEXT, - note TEXT, - internal_priority INTEGER NOT NULL DEFAULT 50, - capabilities TEXT, - api_formats TEXT, - auth_type_by_format TEXT, - allow_auth_channel_mismatch_formats TEXT, - rate_multipliers TEXT, - global_priority_by_format TEXT, - allowed_models TEXT, - expires_at INTEGER, - cache_ttl_minutes INTEGER NOT NULL DEFAULT 5, - max_probe_interval_minutes INTEGER NOT NULL DEFAULT 32, - proxy TEXT, - fingerprint TEXT, - concurrent_limit INTEGER, - learned_rpm_limit INTEGER, - concurrent_429_count INTEGER NOT NULL DEFAULT 0, - rpm_429_count INTEGER NOT NULL DEFAULT 0, - last_429_at INTEGER, - last_429_type TEXT, - adjustment_history TEXT, - utilization_samples TEXT, - last_probe_increase_at INTEGER, - last_rpm_peak INTEGER, - request_count INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - success_count INTEGER NOT NULL DEFAULT 0, - error_count INTEGER NOT NULL DEFAULT 0, - total_response_time_ms INTEGER NOT NULL DEFAULT 0, - last_used_at INTEGER, - auto_fetch_models INTEGER NOT NULL DEFAULT 0, - last_models_fetch_at INTEGER, - last_models_fetch_error TEXT, - locked_models TEXT, - model_include_patterns TEXT, - model_exclude_patterns TEXT, - upstream_metadata TEXT, - oauth_invalid_at INTEGER, - oauth_invalid_reason TEXT, - status_snapshot TEXT, - health_by_format TEXT, - circuit_breaker_by_format TEXT, - status TEXT NOT NULL DEFAULT 'active', - is_active INTEGER NOT NULL DEFAULT 1, - weight INTEGER NOT NULL DEFAULT 1, - rpm_limit INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id); -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id TEXT PRIMARY KEY, - pool_kind TEXT NOT NULL, - pool_id TEXT NOT NULL, - member_kind TEXT NOT NULL, - member_id TEXT NOT NULL, - capability TEXT NOT NULL, - scope_kind TEXT NOT NULL, - scope_id TEXT, - score REAL NOT NULL DEFAULT 0, - hard_state TEXT NOT NULL DEFAULT 'unknown', - score_version INTEGER NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at INTEGER, - last_scheduled_at INTEGER, - last_success_at INTEGER, - last_failure_at INTEGER, - failure_count INTEGER NOT NULL DEFAULT 0, - last_probe_attempt_at INTEGER, - last_probe_success_at INTEGER, - last_probe_failure_at INTEGER, - probe_failure_count INTEGER NOT NULL DEFAULT 0, - probe_status TEXT NOT NULL DEFAULT 'never', - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC); -CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id); -CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at); -CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at); - -CREATE TABLE IF NOT EXISTS gemini_file_mappings ( - id TEXT PRIMARY KEY, - file_name TEXT NOT NULL UNIQUE, - key_id TEXT NOT NULL, - user_id TEXT, - display_name TEXT, - mime_type TEXT, - source_hash TEXT, - created_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_key_id_idx ON gemini_file_mappings (key_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_user_id_idx ON gemini_file_mappings (user_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_expires_at_idx ON gemini_file_mappings (expires_at); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_source_hash_idx ON gemini_file_mappings (source_hash); - -CREATE TABLE IF NOT EXISTS request_candidates ( - id TEXT PRIMARY KEY, - request_id TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - candidate_index INTEGER NOT NULL, - retry_index INTEGER NOT NULL DEFAULT 0, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - status TEXT NOT NULL, - skip_reason TEXT, - is_cached INTEGER NOT NULL DEFAULT 0, - status_code INTEGER, - error_type TEXT, - error_message TEXT, - latency_ms INTEGER, - concurrent_requests INTEGER, - extra_data TEXT, - required_capabilities TEXT, - created_at INTEGER NOT NULL, - started_at INTEGER, - finished_at INTEGER, - UNIQUE (request_id, candidate_index, retry_index) -); -CREATE INDEX IF NOT EXISTS request_candidates_request_id_idx ON request_candidates (request_id); -CREATE INDEX IF NOT EXISTS request_candidates_provider_id_idx ON request_candidates (provider_id); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_id_idx ON request_candidates (endpoint_id); -CREATE INDEX IF NOT EXISTS request_candidates_status_idx ON request_candidates (status); -CREATE INDEX IF NOT EXISTS request_candidates_created_at_idx ON request_candidates (created_at); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_status_created_idx ON request_candidates (endpoint_id, status, created_at); - -CREATE TABLE IF NOT EXISTS video_tasks ( - id TEXT PRIMARY KEY, - short_id TEXT UNIQUE, - request_id TEXT NOT NULL UNIQUE, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - external_task_id TEXT, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - client_api_format TEXT, - provider_api_format TEXT, - format_converted INTEGER NOT NULL DEFAULT 0, - model TEXT, - prompt TEXT, - original_request_body TEXT, - duration_seconds INTEGER, - resolution TEXT, - aspect_ratio TEXT, - size TEXT, - status TEXT NOT NULL DEFAULT 'pending', - progress_percent INTEGER NOT NULL DEFAULT 0, - progress_message TEXT, - retry_count INTEGER NOT NULL DEFAULT 0, - poll_interval_seconds INTEGER NOT NULL DEFAULT 10, - next_poll_at INTEGER, - poll_count INTEGER NOT NULL DEFAULT 0, - max_poll_count INTEGER NOT NULL DEFAULT 360, - created_at INTEGER NOT NULL, - submitted_at INTEGER, - completed_at INTEGER, - updated_at INTEGER NOT NULL, - error_code TEXT, - error_message TEXT, - video_url TEXT, - request_metadata TEXT -); -CREATE INDEX IF NOT EXISTS video_tasks_external_id_idx ON video_tasks (external_task_id); -CREATE INDEX IF NOT EXISTS video_tasks_next_poll_idx ON video_tasks (next_poll_at); -CREATE INDEX IF NOT EXISTS video_tasks_request_id_idx ON video_tasks (request_id); -CREATE INDEX IF NOT EXISTS video_tasks_user_status_idx ON video_tasks (user_id, status); -CREATE INDEX IF NOT EXISTS video_tasks_api_key_id_idx ON video_tasks (api_key_id); -CREATE INDEX IF NOT EXISTS video_tasks_provider_id_idx ON video_tasks (provider_id); -CREATE INDEX IF NOT EXISTS video_tasks_endpoint_id_idx ON video_tasks (endpoint_id); -CREATE INDEX IF NOT EXISTS video_tasks_key_id_idx ON video_tasks (key_id); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - base_url TEXT NOT NULL, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - health_score REAL NOT NULL DEFAULT 1.0, - weight INTEGER NOT NULL DEFAULT 1, - header_rules TEXT, - body_rules TEXT, - max_retries INTEGER, - custom_path TEXT, - metadata TEXT, - config TEXT, - format_acceptance_config TEXT, - proxy TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_endpoints_provider_id_idx ON provider_endpoints (provider_id); - -CREATE TABLE IF NOT EXISTS models ( - id TEXT PRIMARY KEY, - provider_id TEXT NOT NULL, - global_model_id TEXT, - provider_model_name TEXT NOT NULL, - global_model_name TEXT, - api_format TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - is_available INTEGER NOT NULL DEFAULT 1, - price_per_request REAL, - tiered_pricing TEXT, - supports_vision INTEGER, - supports_function_calling INTEGER, - supports_streaming INTEGER, - supports_extended_thinking INTEGER, - supports_image_generation INTEGER, - provider_model_mappings TEXT, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS models_provider_id_idx ON models (provider_id); - -CREATE TABLE IF NOT EXISTS global_models ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - display_name TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - default_price_per_request REAL, - default_tiered_pricing TEXT, - supported_capabilities TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/003_auth_config.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/003_auth_config.sql deleted file mode 100644 index 54cb2a37d..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/003_auth_config.sql +++ /dev/null @@ -1,71 +0,0 @@ -CREATE TABLE IF NOT EXISTS system_configs ( - id TEXT PRIMARY KEY, - key TEXT NOT NULL UNIQUE, - value TEXT NOT NULL, - description TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - id TEXT PRIMARY KEY, - module_type TEXT NOT NULL UNIQUE, - enabled INTEGER NOT NULL DEFAULT 1, - config TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - provider_type TEXT PRIMARY KEY, - display_name TEXT NOT NULL, - client_id TEXT NOT NULL, - client_secret_encrypted TEXT, - authorization_url_override TEXT, - token_url_override TEXT, - userinfo_url_override TEXT, - scopes TEXT, - redirect_uri TEXT NOT NULL, - frontend_callback_url TEXT NOT NULL, - attribute_mapping TEXT, - extra_config TEXT, - is_enabled INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - server_url TEXT NOT NULL, - bind_dn TEXT NOT NULL, - bind_password_encrypted TEXT, - base_dn TEXT NOT NULL, - user_search_filter TEXT DEFAULT '(uid={username})' NOT NULL, - username_attr TEXT DEFAULT 'uid' NOT NULL, - email_attr TEXT DEFAULT 'mail' NOT NULL, - display_name_attr TEXT DEFAULT 'cn' NOT NULL, - is_enabled INTEGER NOT NULL DEFAULT 0, - is_exclusive INTEGER NOT NULL DEFAULT 0, - use_starttls INTEGER NOT NULL DEFAULT 0, - connect_timeout INTEGER NOT NULL DEFAULT 10, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - provider_type TEXT NOT NULL, - provider_user_id TEXT NOT NULL, - provider_username TEXT, - provider_email TEXT, - extra_data TEXT, - linked_at INTEGER NOT NULL, - last_login_at INTEGER -); -CREATE UNIQUE INDEX IF NOT EXISTS uq_user_oauth_links_provider_user - ON user_oauth_links (provider_type, provider_user_id); -CREATE UNIQUE INDEX IF NOT EXISTS uq_user_oauth_links_user_provider - ON user_oauth_links (user_id, provider_type); -CREATE INDEX IF NOT EXISTS user_oauth_links_provider_type_idx ON user_oauth_links (provider_type); -CREATE INDEX IF NOT EXISTS user_oauth_links_user_id_idx ON user_oauth_links (user_id); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/004_proxy_nodes.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/004_proxy_nodes.sql deleted file mode 100644 index ac95e3135..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/004_proxy_nodes.sql +++ /dev/null @@ -1,39 +0,0 @@ -CREATE TABLE IF NOT EXISTS proxy_nodes ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - ip TEXT NOT NULL, - port INTEGER NOT NULL, - region TEXT, - status TEXT NOT NULL DEFAULT 'online', - registered_by TEXT, - last_heartbeat_at INTEGER, - heartbeat_interval INTEGER NOT NULL DEFAULT 30, - active_connections INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - avg_latency_ms REAL, - is_manual INTEGER NOT NULL DEFAULT 0, - proxy_url TEXT, - proxy_username TEXT, - proxy_password TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - remote_config TEXT, - config_version INTEGER NOT NULL DEFAULT 0, - hardware_info TEXT, - estimated_max_concurrency INTEGER, - tunnel_mode INTEGER NOT NULL DEFAULT 0, - tunnel_connected INTEGER NOT NULL DEFAULT 0, - tunnel_connected_at INTEGER, - failed_requests INTEGER NOT NULL DEFAULT 0, - dns_failures INTEGER NOT NULL DEFAULT 0, - stream_errors INTEGER NOT NULL DEFAULT 0, - proxy_metadata TEXT -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - node_id TEXT NOT NULL, - event_type TEXT NOT NULL, - detail TEXT, - created_at INTEGER NOT NULL -); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/005_wallet_billing.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/005_wallet_billing.sql deleted file mode 100644 index ccb46d7bc..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/005_wallet_billing.sql +++ /dev/null @@ -1,202 +0,0 @@ -CREATE TABLE IF NOT EXISTS wallets ( - id TEXT PRIMARY KEY, - user_id TEXT UNIQUE, - api_key_id TEXT UNIQUE, - balance REAL NOT NULL DEFAULT 0, - gift_balance REAL NOT NULL DEFAULT 0, - limit_mode TEXT NOT NULL DEFAULT 'finite', - currency TEXT NOT NULL DEFAULT 'USD', - status TEXT NOT NULL DEFAULT 'active', - total_recharged REAL NOT NULL DEFAULT 0, - total_consumed REAL NOT NULL DEFAULT 0, - total_refunded REAL NOT NULL DEFAULT 0, - total_adjusted REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS wallets_api_key_id_idx ON wallets (api_key_id); -CREATE INDEX IF NOT EXISTS wallets_user_id_idx ON wallets (user_id); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - id TEXT PRIMARY KEY, - wallet_id TEXT NOT NULL, - category TEXT NOT NULL, - reason_code TEXT NOT NULL, - amount REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - recharge_balance_before REAL NOT NULL, - recharge_balance_after REAL NOT NULL, - gift_balance_before REAL NOT NULL, - gift_balance_after REAL NOT NULL, - link_type TEXT, - link_id TEXT, - operator_id TEXT, - description TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_wallet_created - ON wallet_transactions (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_category_created - ON wallet_transactions (category, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_reason_created - ON wallet_transactions (reason_code, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_link - ON wallet_transactions (link_type, link_id); -CREATE INDEX IF NOT EXISTS ix_wallet_transactions_operator_id - ON wallet_transactions (operator_id); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - id TEXT PRIMARY KEY, - wallet_id TEXT NOT NULL, - billing_date TEXT NOT NULL, - billing_timezone TEXT NOT NULL, - total_cost_usd REAL NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - first_finalized_at INTEGER, - last_finalized_at INTEGER, - aggregated_at INTEGER NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_daily_usage_wallet_date - ON wallet_daily_usage_ledgers (wallet_id, billing_timezone, billing_date); - -CREATE TABLE IF NOT EXISTS payment_orders ( - id TEXT PRIMARY KEY, - order_no TEXT NOT NULL UNIQUE, - wallet_id TEXT NOT NULL, - user_id TEXT, - amount_usd REAL NOT NULL, - pay_amount REAL, - pay_currency TEXT, - exchange_rate REAL, - refunded_amount_usd REAL NOT NULL DEFAULT 0, - refundable_amount_usd REAL NOT NULL DEFAULT 0, - payment_method TEXT NOT NULL, - gateway_order_id TEXT, - gateway_response TEXT, - status TEXT NOT NULL DEFAULT 'pending', - created_at INTEGER NOT NULL, - paid_at INTEGER, - credited_at INTEGER, - expires_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_payment_orders_wallet_created - ON payment_orders (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_user_created - ON payment_orders (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_status - ON payment_orders (status); -CREATE INDEX IF NOT EXISTS idx_payment_orders_gateway_order_id - ON payment_orders (gateway_order_id); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - id TEXT PRIMARY KEY, - payment_order_id TEXT, - payment_method TEXT NOT NULL, - callback_key TEXT NOT NULL UNIQUE, - order_no TEXT, - gateway_order_id TEXT, - payload_hash TEXT, - signature_valid INTEGER NOT NULL DEFAULT 0, - status TEXT NOT NULL DEFAULT 'received', - payload TEXT, - error_message TEXT, - created_at INTEGER NOT NULL, - processed_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_order - ON payment_callbacks (order_no); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_gateway_order - ON payment_callbacks (gateway_order_id); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_created - ON payment_callbacks (created_at); -CREATE INDEX IF NOT EXISTS ix_payment_callbacks_payment_order_id - ON payment_callbacks (payment_order_id); - -CREATE TABLE IF NOT EXISTS refund_requests ( - id TEXT PRIMARY KEY, - refund_no TEXT NOT NULL UNIQUE, - wallet_id TEXT NOT NULL, - user_id TEXT, - payment_order_id TEXT, - source_type TEXT NOT NULL, - source_id TEXT, - refund_mode TEXT NOT NULL, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending_approval', - reason TEXT, - requested_by TEXT, - approved_by TEXT, - processed_by TEXT, - gateway_refund_id TEXT, - payout_method TEXT, - payout_reference TEXT, - payout_proof TEXT, - failure_reason TEXT, - idempotency_key TEXT UNIQUE, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - processed_at INTEGER, - completed_at INTEGER -); -CREATE INDEX IF NOT EXISTS idx_refund_wallet_created - ON refund_requests (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_user_created - ON refund_requests (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_status - ON refund_requests (status); -CREATE INDEX IF NOT EXISTS ix_refund_requests_payment_order_id - ON refund_requests (payment_order_id); -CREATE INDEX IF NOT EXISTS ix_refund_requests_requested_by - ON refund_requests (requested_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_approved_by - ON refund_requests (approved_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_processed_by - ON refund_requests (processed_by); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - amount_usd REAL NOT NULL, - currency TEXT NOT NULL DEFAULT 'USD', - balance_bucket TEXT NOT NULL DEFAULT 'gift', - total_count INTEGER NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - description TEXT, - created_by TEXT, - expires_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_redeem_code_batches_status - ON redeem_code_batches (status, created_at); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - id TEXT PRIMARY KEY, - batch_id TEXT NOT NULL, - code_hash TEXT NOT NULL UNIQUE, - code_prefix TEXT NOT NULL, - code_suffix TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - redeemed_by_user_id TEXT, - redeemed_wallet_id TEXT, - redeemed_payment_order_id TEXT, - redeemed_at INTEGER, - disabled_by TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_batch_created - ON redeem_codes (batch_id, created_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_status - ON redeem_codes (status, updated_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_user - ON redeem_codes (redeemed_by_user_id, redeemed_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_order - ON redeem_codes (redeemed_payment_order_id); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/006_usage.sql b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/006_usage.sql deleted file mode 100644 index 045b841f6..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/006_usage.sql +++ /dev/null @@ -1,265 +0,0 @@ -CREATE TABLE IF NOT EXISTS "usage" ( - request_id TEXT PRIMARY KEY, - id TEXT, - user_id TEXT, - api_key_id TEXT, - provider_name TEXT NOT NULL DEFAULT 'unknown', - model TEXT NOT NULL DEFAULT 'unknown', - target_model TEXT, - provider_id TEXT, - provider_endpoint_id TEXT, - provider_api_key_id TEXT, - request_type TEXT, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - endpoint_api_format TEXT, - provider_api_family TEXT, - provider_endpoint_kind TEXT, - has_format_conversion INTEGER NOT NULL DEFAULT 0, - is_stream INTEGER NOT NULL DEFAULT 0, - upstream_is_stream INTEGER, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_5m_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_1h_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_cost_usd REAL NOT NULL DEFAULT 0, - cache_read_cost_usd REAL NOT NULL DEFAULT 0, - output_price_per_1m REAL, - status_code INTEGER, - error_message TEXT, - error_category TEXT, - response_time_ms INTEGER, - first_byte_time_ms INTEGER, - wallet_id TEXT, - status TEXT NOT NULL DEFAULT 'completed', - billing_status TEXT NOT NULL DEFAULT 'pending', - total_cost_usd REAL NOT NULL DEFAULT 0, - actual_total_cost_usd REAL NOT NULL DEFAULT 0, - request_metadata TEXT, - 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, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - wallet_gift_balance_before REAL, - wallet_gift_balance_after REAL, - finalized_at INTEGER, - created_at_unix_ms INTEGER NOT NULL DEFAULT 0, - updated_at_unix_secs INTEGER NOT NULL DEFAULT 0 -); -CREATE INDEX IF NOT EXISTS usage_api_key_id_idx ON "usage" (api_key_id); -CREATE INDEX IF NOT EXISTS usage_billing_status_idx ON "usage" (billing_status); -CREATE INDEX IF NOT EXISTS usage_created_at_idx ON "usage" (created_at_unix_ms); -CREATE INDEX IF NOT EXISTS usage_provider_api_key_id_idx ON "usage" (provider_api_key_id); -CREATE INDEX IF NOT EXISTS usage_provider_id_idx ON "usage" (provider_id); -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_settlement_snapshots ( - request_id TEXT PRIMARY KEY, - billing_status TEXT NOT NULL, - wallet_id TEXT, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - wallet_gift_balance_before REAL, - wallet_gift_balance_after REAL, - provider_monthly_used_usd REAL, - 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 TABLE IF NOT EXISTS stats_hourly ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL UNIQUE, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - aggregated_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, user_id) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, user_id, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, model) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - id TEXT PRIMARY KEY, - hour_utc INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL UNIQUE, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - fallback_count INTEGER NOT NULL DEFAULT 0, - unique_models INTEGER NOT NULL DEFAULT 0, - unique_providers INTEGER NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - aggregated_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", model) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - id TEXT PRIMARY KEY, - api_key_id TEXT NOT NULL, - "date" INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - api_key_name TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", api_key_id) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - id TEXT PRIMARY KEY, - "date" INTEGER NOT NULL, - error_category TEXT NOT NULL, - provider_name TEXT, - model TEXT, - count INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", error_category, provider_name, model) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - "date" INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - username TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE ("date", user_id) -); diff --git a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/manifest.txt b/crates/aether-data/runtime/schema/drivers/sqlite/baseline/manifest.txt deleted file mode 100644 index 6caec84c6..000000000 --- a/crates/aether-data/runtime/schema/drivers/sqlite/baseline/manifest.txt +++ /dev/null @@ -1,6 +0,0 @@ -001_identity.sql -002_provider_catalog.sql -003_auth_config.sql -004_proxy_nodes.sql -005_wallet_billing.sql -006_usage.sql diff --git a/crates/aether-data/runtime/schema/generated/README.md b/crates/aether-data/runtime/schema/generated/README.md index b64881807..d9ad3013b 100644 --- a/crates/aether-data/runtime/schema/generated/README.md +++ b/crates/aether-data/runtime/schema/generated/README.md @@ -9,4 +9,4 @@ Do not edit files in this directory by hand. Update `../logical/*.toml`, then ru bash crates/aether-data/runtime/schema/compose_schema.sh generate ``` -Runtime migrations are not loaded from this directory. The executable SQL lives under `crates/aether-data/adapters/{postgres,mysql,sqlite}/migrations`, and the Postgres bootstrap snapshot is generated at build time from `crates/aether-data/runtime/schema/bootstrap/postgres` into the crate build output until a generated fragment is deliberately promoted into the driver-specific schema manifests. +Runtime migrations are not loaded from this directory. The executable SQL lives under `crates/aether-data/adapters/postgres/migrations`, and the Postgres bootstrap snapshot is generated at build time from `crates/aether-data/runtime/schema/bootstrap/postgres` into the crate build output until a generated fragment is deliberately promoted into the driver-specific schema manifests. diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql deleted file mode 100644 index fd7a01c27..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/001_identity.sql +++ /dev/null @@ -1,222 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS users ( - `id` VARCHAR(64) NOT NULL, - `external_id` VARCHAR(255), - `email` VARCHAR(320), - `username` VARCHAR(255), - `password_hash` VARCHAR(255), - `role` VARCHAR(64), - `auth_source` VARCHAR(64) NOT NULL DEFAULT 'local', - `email_verified` TINYINT(1) NOT NULL DEFAULT 0, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `is_deleted` TINYINT(1) NOT NULL DEFAULT 0, - `security_version` BIGINT NOT NULL DEFAULT 0, - `allowed_models` JSON, - `allowed_models_mode` VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - `allowed_providers` JSON, - `allowed_providers_mode` VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - `allowed_api_formats` JSON, - `allowed_api_formats_mode` VARCHAR(32) NOT NULL DEFAULT 'unrestricted', - `model_capability_settings` JSON, - `feature_settings` JSON, - `rate_limit` INT, - `rate_limit_mode` VARCHAR(32) NOT NULL DEFAULT 'system', - `privacy_policy_accepted_version` VARCHAR(64), - `privacy_policy_accepted_at` BIGINT, - `metadata` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `last_login_at` BIGINT, - `ldap_dn` VARCHAR(1024), - `ldap_username` VARCHAR(255), - PRIMARY KEY (`id`), - UNIQUE KEY users_email_key (`email`), - UNIQUE KEY users_username_key (`username`) -); - -CREATE TABLE IF NOT EXISTS user_groups ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(100) NOT NULL, - `normalized_name` VARCHAR(100) NOT NULL, - `description` LONGTEXT, - `priority` INT NOT NULL DEFAULT 0, - `allowed_providers` JSON, - `allowed_providers_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', - `allowed_api_formats` JSON, - `allowed_api_formats_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', - `allowed_models` JSON, - `allowed_models_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', - `rate_limit` INT, - `rate_limit_mode` VARCHAR(32) NOT NULL DEFAULT 'inherit', - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY user_groups_normalized_name_key (`normalized_name`), - KEY user_groups_priority_name_idx (`priority`, `name`, `id`) -); - -CREATE TABLE IF NOT EXISTS user_group_members ( - `group_id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`group_id`, `user_id`), - KEY user_group_members_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS api_keys ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `key_hash` VARCHAR(255) NOT NULL, - `key_encrypted` LONGTEXT, - `name` VARCHAR(255), - `key_prefix` VARCHAR(64), - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `allowed_models` JSON, - `allowed_providers` JSON, - `allowed_api_formats` JSON, - `ip_rules` JSON, - `rate_limit` INT DEFAULT 100, - `concurrent_limit` INT, - `force_capabilities` JSON, - `feature_settings` JSON, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `is_locked` TINYINT(1) NOT NULL DEFAULT 0, - `is_standalone` TINYINT(1) NOT NULL DEFAULT 0, - `auto_delete_on_expiry` TINYINT(1) NOT NULL DEFAULT 0, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `total_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `metadata` JSON, - `expires_at` BIGINT, - `last_used_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY api_keys_key_hash_key (`key_hash`), - KEY api_keys_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS audit_logs ( - `id` VARCHAR(64) NOT NULL, - `event_type` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64), - `api_key_id` VARCHAR(64), - `description` LONGTEXT NOT NULL, - `ip_address` VARCHAR(64), - `user_agent` VARCHAR(512), - `request_id` VARCHAR(128), - `event_metadata` JSON, - `status_code` INT, - `error_message` LONGTEXT, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY audit_logs_created_at_idx (`created_at`), - KEY audit_logs_event_type_idx (`event_type`), - KEY audit_logs_request_id_idx (`request_id`), - KEY audit_logs_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS announcements ( - `id` VARCHAR(64) NOT NULL, - `title` VARCHAR(200) NOT NULL, - `content` LONGTEXT NOT NULL, - `type` VARCHAR(32) NOT NULL DEFAULT 'info', - `priority` INT NOT NULL DEFAULT 0, - `author_id` VARCHAR(64), - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `is_pinned` TINYINT(1) NOT NULL DEFAULT 0, - `requires_ack` TINYINT(1) NOT NULL DEFAULT 0, - `start_time` BIGINT, - `end_time` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY announcements_author_id_idx (`author_id`), - KEY announcements_created_at_idx (`created_at`), - KEY announcements_is_active_idx (`is_active`) -); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `announcement_id` VARCHAR(64) NOT NULL, - `read_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_user_announcement (`user_id`, `announcement_id`), - KEY announcement_reads_announcement_id_idx (`announcement_id`), - KEY announcement_reads_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS management_tokens ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `description` LONGTEXT, - `token_hash` VARCHAR(255) NOT NULL, - `token_prefix` VARCHAR(64), - `allowed_ips` JSON, - `permissions` JSON, - `expires_at` BIGINT, - `last_used_at` BIGINT, - `last_used_ip` VARCHAR(255), - `usage_count` BIGINT NOT NULL DEFAULT 0, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY management_tokens_token_hash_key (`token_hash`), - UNIQUE KEY uq_management_tokens_user_name (`user_id`, `name`), - KEY management_tokens_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS user_preferences ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `avatar_url` VARCHAR(500), - `bio` LONGTEXT, - `default_provider_id` VARCHAR(64), - `theme` VARCHAR(20) NOT NULL DEFAULT 'light', - `language` VARCHAR(10) NOT NULL DEFAULT 'zh-CN', - `timezone` VARCHAR(50) NOT NULL DEFAULT 'Asia/Shanghai', - `email_notifications` TINYINT(1) NOT NULL DEFAULT 1, - `usage_alerts` TINYINT(1) NOT NULL DEFAULT 1, - `announcement_notifications` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY user_preferences_user_id_key (`user_id`), - KEY user_preferences_default_provider_id_idx (`default_provider_id`), - KEY user_preferences_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS user_sessions ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `security_version` BIGINT NOT NULL DEFAULT 0, - `client_device_id` VARCHAR(128) NOT NULL, - `device_label` VARCHAR(120), - `device_type` VARCHAR(20) NOT NULL DEFAULT 'unknown', - `browser_name` VARCHAR(50), - `browser_version` VARCHAR(50), - `os_name` VARCHAR(50), - `os_version` VARCHAR(50), - `device_model` VARCHAR(100), - `ip_address` VARCHAR(45), - `user_agent` VARCHAR(1000), - `client_hints` JSON, - `refresh_token_hash` VARCHAR(64) NOT NULL, - `prev_refresh_token_hash` VARCHAR(64), - `rotated_at` BIGINT, - `last_seen_at` BIGINT NOT NULL, - `expires_at` BIGINT NOT NULL, - `revoked_at` BIGINT, - `revoke_reason` VARCHAR(100), - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY user_sessions_user_active_idx (`user_id`, `revoked_at`, `expires_at`), - KEY user_sessions_user_device_idx (`user_id`, `client_device_id`) -); - diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/002_provider_catalog.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/002_provider_catalog.sql deleted file mode 100644 index f926876ae..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/002_provider_catalog.sql +++ /dev/null @@ -1,431 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS billing_rules ( - `id` VARCHAR(64) NOT NULL, - `global_model_id` VARCHAR(64), - `model_id` VARCHAR(64), - `name` VARCHAR(255) NOT NULL, - `task_type` VARCHAR(64) NOT NULL DEFAULT 'chat', - `expression` LONGTEXT NOT NULL, - `variables` JSON NOT NULL, - `dimension_mappings` JSON NOT NULL, - `is_enabled` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY billing_rules_global_model_task_idx (`global_model_id`, `task_type`, `is_enabled`), - KEY billing_rules_model_task_idx (`model_id`, `task_type`, `is_enabled`) -); - -CREATE TABLE IF NOT EXISTS dimension_collectors ( - `id` VARCHAR(64) NOT NULL, - `api_format` VARCHAR(64) NOT NULL, - `task_type` VARCHAR(64) NOT NULL, - `dimension_name` VARCHAR(128) NOT NULL, - `source_type` VARCHAR(64) NOT NULL, - `source_path` VARCHAR(255), - `value_type` VARCHAR(64) NOT NULL DEFAULT 'float', - `transform_expression` LONGTEXT, - `default_value` VARCHAR(255), - `priority` INT NOT NULL DEFAULT 0, - `is_enabled` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY dimension_collectors_enabled_idx (`api_format`, `task_type`, `dimension_name`, `priority`, `is_enabled`) -); - -CREATE TABLE IF NOT EXISTS providers ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `description` LONGTEXT, - `website` VARCHAR(500), - `provider_type` VARCHAR(64) NOT NULL, - `billing_type` VARCHAR(64), - `monthly_quota_usd` DOUBLE, - `monthly_used_usd` DOUBLE, - `quota_reset_day` INT, - `quota_last_reset_at` BIGINT, - `quota_expires_at` BIGINT, - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `priority` BIGINT NOT NULL DEFAULT 0, - `provider_priority` INT NOT NULL DEFAULT 100, - `keep_priority_on_conversion` TINYINT(1) NOT NULL DEFAULT 0, - `enable_format_conversion` TINYINT(1) NOT NULL DEFAULT 1, - `concurrent_limit` INT, - `max_retries` INT, - `proxy` JSON, - `request_timeout` DOUBLE, - `stream_first_byte_timeout` DOUBLE, - `config` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY providers_name_key (`name`) -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - `id` VARCHAR(64) NOT NULL, - `provider_id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `api_key` LONGTEXT, - `encrypted_key` LONGTEXT, - `auth_type` VARCHAR(32) NOT NULL DEFAULT 'api_key', - `auth_config` JSON, - `note` LONGTEXT, - `internal_priority` INT NOT NULL DEFAULT 50, - `capabilities` JSON, - `api_formats` JSON, - `auth_type_by_format` JSON, - `allow_auth_channel_mismatch_formats` JSON, - `rate_multipliers` JSON, - `global_priority_by_format` JSON, - `allowed_models` JSON, - `expires_at` BIGINT, - `cache_ttl_minutes` INT NOT NULL DEFAULT 5, - `max_probe_interval_minutes` INT NOT NULL DEFAULT 32, - `proxy` JSON, - `fingerprint` JSON, - `concurrent_limit` INT, - `learned_rpm_limit` INT, - `concurrent_429_count` INT NOT NULL DEFAULT 0, - `rpm_429_count` INT NOT NULL DEFAULT 0, - `last_429_at` BIGINT, - `last_429_type` VARCHAR(64), - `adjustment_history` JSON, - `utilization_samples` JSON, - `last_probe_increase_at` BIGINT, - `last_rpm_peak` INT, - `request_count` BIGINT NOT NULL DEFAULT 0, - `total_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `success_count` BIGINT NOT NULL DEFAULT 0, - `error_count` BIGINT NOT NULL DEFAULT 0, - `total_response_time_ms` BIGINT NOT NULL DEFAULT 0, - `last_used_at` BIGINT, - `last_error_at` BIGINT, - `last_error_msg` LONGTEXT, - `auto_fetch_models` TINYINT(1) NOT NULL DEFAULT 0, - `last_models_fetch_at` BIGINT, - `last_models_fetch_error` LONGTEXT, - `locked_models` JSON, - `model_include_patterns` JSON, - `model_exclude_patterns` JSON, - `upstream_metadata` JSON, - `oauth_invalid_at` BIGINT, - `oauth_invalid_reason` VARCHAR(255), - `status_snapshot` JSON, - `health_by_format` JSON, - `circuit_breaker_by_format` JSON, - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `weight` BIGINT NOT NULL DEFAULT 1, - `rpm_limit` BIGINT, - `metadata` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY provider_api_keys_provider_id_idx (`provider_id`), - KEY idx_provider_api_keys_provider_default_sort (`provider_id`, `internal_priority`, `name`, `id`) -); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - `id` VARCHAR(192) NOT NULL, - `pool_kind` VARCHAR(64) NOT NULL, - `pool_id` VARCHAR(64) NOT NULL, - `member_kind` VARCHAR(64) NOT NULL, - `member_id` VARCHAR(64) NOT NULL, - `capability` VARCHAR(64) NOT NULL, - `scope_kind` VARCHAR(64) NOT NULL, - `scope_id` VARCHAR(128), - `score` DOUBLE NOT NULL DEFAULT 0, - `hard_state` VARCHAR(64) NOT NULL DEFAULT 'unknown', - `score_version` BIGINT NOT NULL DEFAULT 1, - `score_reason` JSON NOT NULL, - `last_ranked_at` BIGINT, - `last_scheduled_at` BIGINT, - `last_success_at` BIGINT, - `last_failure_at` BIGINT, - `failure_count` BIGINT NOT NULL DEFAULT 0, - `last_probe_attempt_at` BIGINT, - `last_probe_success_at` BIGINT, - `last_probe_failure_at` BIGINT, - `probe_failure_count` BIGINT NOT NULL DEFAULT 0, - `probe_status` VARCHAR(64) NOT NULL DEFAULT 'never', - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY pool_member_scores_rank_idx (`pool_kind`, `pool_id`, `capability`, `scope_kind`, `scope_id`, `hard_state`, `score`), - KEY pool_member_scores_member_idx (`pool_kind`, `pool_id`, `member_kind`, `member_id`), - KEY pool_member_scores_probe_idx (`pool_kind`, `pool_id`, `probe_status`, `last_probe_success_at`), - KEY pool_member_scores_updated_at_idx (`updated_at`) -); - -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 gemini_file_mappings ( - `id` VARCHAR(64) NOT NULL, - `file_name` VARCHAR(512) NOT NULL, - `key_id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64), - `display_name` VARCHAR(512), - `mime_type` VARCHAR(255), - `source_hash` VARCHAR(128), - `created_at` BIGINT NOT NULL, - `expires_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY gemini_file_mappings_file_name_key (`file_name`), - KEY gemini_file_mappings_key_id_idx (`key_id`), - KEY gemini_file_mappings_user_id_idx (`user_id`), - KEY gemini_file_mappings_expires_at_idx (`expires_at`), - KEY gemini_file_mappings_source_hash_idx (`source_hash`) -); - -CREATE TABLE IF NOT EXISTS request_candidates ( - `id` VARCHAR(64) NOT NULL, - `request_id` VARCHAR(128) NOT NULL, - `user_id` VARCHAR(64), - `api_key_id` VARCHAR(64), - `username` VARCHAR(255), - `api_key_name` VARCHAR(255), - `candidate_index` INT NOT NULL, - `retry_index` INT NOT NULL DEFAULT 0, - `provider_id` VARCHAR(64), - `endpoint_id` VARCHAR(64), - `key_id` VARCHAR(64), - `status` VARCHAR(32) NOT NULL, - `skip_reason` LONGTEXT, - `is_cached` TINYINT(1) NOT NULL DEFAULT 0, - `status_code` INT, - `error_type` VARCHAR(128), - `error_message` LONGTEXT, - `latency_ms` INT, - `concurrent_requests` INT, - `extra_data` JSON, - `required_capabilities` JSON, - `created_at` BIGINT NOT NULL, - `started_at` BIGINT, - `finished_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY uq_request_candidate_with_retry (`request_id`, `candidate_index`, `retry_index`), - KEY request_candidates_request_id_idx (`request_id`), - KEY request_candidates_provider_id_idx (`provider_id`), - KEY request_candidates_endpoint_id_idx (`endpoint_id`), - KEY request_candidates_status_idx (`status`), - KEY request_candidates_created_at_idx (`created_at`), - KEY request_candidates_endpoint_status_created_idx (`endpoint_id`, `status`, `created_at`) -); - -CREATE TABLE IF NOT EXISTS video_tasks ( - `id` VARCHAR(64) NOT NULL, - `short_id` VARCHAR(32), - `request_id` VARCHAR(128) NOT NULL, - `user_id` VARCHAR(64), - `api_key_id` VARCHAR(64), - `username` VARCHAR(255), - `api_key_name` VARCHAR(255), - `external_task_id` VARCHAR(255), - `provider_id` VARCHAR(64), - `endpoint_id` VARCHAR(64), - `key_id` VARCHAR(64), - `client_api_format` VARCHAR(128), - `provider_api_format` VARCHAR(128), - `format_converted` TINYINT(1) NOT NULL DEFAULT 0, - `model` VARCHAR(255), - `prompt` LONGTEXT, - `original_request_body` JSON, - `converted_request_body` JSON, - `duration_seconds` INT, - `resolution` VARCHAR(64), - `aspect_ratio` VARCHAR(32), - `size` VARCHAR(64), - `status` VARCHAR(32) NOT NULL DEFAULT 'pending', - `progress_percent` INT NOT NULL DEFAULT 0, - `progress_message` LONGTEXT, - `retry_count` INT NOT NULL DEFAULT 0, - `max_retries` INT NOT NULL DEFAULT 3, - `poll_interval_seconds` INT NOT NULL DEFAULT 10, - `next_poll_at` BIGINT, - `poll_count` INT NOT NULL DEFAULT 0, - `max_poll_count` INT NOT NULL DEFAULT 360, - `created_at` BIGINT NOT NULL, - `submitted_at` BIGINT, - `completed_at` BIGINT, - `updated_at` BIGINT NOT NULL, - `error_code` VARCHAR(128), - `error_message` LONGTEXT, - `video_url` LONGTEXT, - `video_urls` JSON, - `thumbnail_url` LONGTEXT, - `video_size_bytes` BIGINT, - `video_expires_at` BIGINT, - `stored_video_path` VARCHAR(500), - `storage_provider` VARCHAR(50), - `remixed_from_task_id` VARCHAR(64), - `webhook_url` VARCHAR(500), - `webhook_sent` TINYINT(1) NOT NULL DEFAULT 0, - `webhook_sent_at` BIGINT, - `request_metadata` JSON, - `video_duration_seconds` DOUBLE, - PRIMARY KEY (`id`), - UNIQUE KEY video_tasks_short_id_key (`short_id`), - UNIQUE KEY video_tasks_request_id_key (`request_id`), - KEY video_tasks_external_id_idx (`external_task_id`), - KEY video_tasks_next_poll_idx (`next_poll_at`), - KEY video_tasks_due_poll_idx (`status`, `next_poll_at`, `updated_at`), - KEY video_tasks_user_status_idx (`user_id`, `status`), - KEY video_tasks_api_key_id_idx (`api_key_id`), - KEY video_tasks_provider_id_idx (`provider_id`), - KEY video_tasks_endpoint_id_idx (`endpoint_id`), - KEY video_tasks_key_id_idx (`key_id`) -); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - `id` VARCHAR(64) NOT NULL, - `provider_id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `base_url` LONGTEXT NOT NULL, - `api_format` VARCHAR(128), - `api_family` VARCHAR(128), - `endpoint_kind` VARCHAR(128), - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `health_score` DOUBLE NOT NULL DEFAULT 1.0, - `weight` BIGINT NOT NULL DEFAULT 1, - `header_rules` JSON, - `body_rules` JSON, - `max_retries` INT, - `custom_path` LONGTEXT, - `metadata` JSON, - `config` JSON, - `format_acceptance_config` JSON, - `proxy` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY provider_endpoints_provider_id_idx (`provider_id`) -); - -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`) -); - -CREATE TABLE IF NOT EXISTS models ( - `id` VARCHAR(64) NOT NULL, - `provider_id` VARCHAR(64) NOT NULL, - `global_model_id` VARCHAR(64), - `provider_model_name` VARCHAR(255) NOT NULL, - `global_model_name` VARCHAR(255), - `api_format` VARCHAR(128), - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `is_available` TINYINT(1) NOT NULL DEFAULT 1, - `price_per_request` DOUBLE, - `tiered_pricing` JSON, - `supports_vision` TINYINT(1), - `supports_function_calling` TINYINT(1), - `supports_streaming` TINYINT(1), - `supports_extended_thinking` TINYINT(1), - `supports_image_generation` TINYINT(1), - `provider_model_mappings` JSON, - `config` JSON, - `metadata` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY models_provider_id_idx (`provider_id`) -); - -CREATE TABLE IF NOT EXISTS global_models ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `display_name` VARCHAR(255), - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_active` TINYINT(1) NOT NULL DEFAULT 1, - `default_price_per_request` DOUBLE, - `default_tiered_pricing` JSON, - `supported_capabilities` JSON, - `usage_count` BIGINT NOT NULL DEFAULT 0, - `config` JSON, - `metadata` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY global_models_name_key (`name`) -); - -CREATE TABLE IF NOT EXISTS routing_groups ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `description` LONGTEXT, - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `is_system_default` TINYINT(1) NOT NULL DEFAULT 0, - `sort_order` BIGINT NOT NULL DEFAULT 0, - `config_json` JSON NOT NULL, - `version` BIGINT NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `published_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY routing_groups_name_key (`name`), - KEY routing_groups_system_default_idx (`is_system_default`, `enabled`), - KEY routing_groups_enabled_sort_idx (`enabled`, `sort_order`, `name`, `id`) -); - -CREATE TABLE IF NOT EXISTS routing_group_bindings ( - `id` VARCHAR(64) NOT NULL, - `group_id` VARCHAR(64) NOT NULL, - `subject_type` VARCHAR(32) NOT NULL, - `subject_id` VARCHAR(64) NOT NULL, - `is_default` TINYINT(1) NOT NULL DEFAULT 0, - `allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY routing_group_bindings_group_id_idx (`group_id`), - KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`) -); - -CREATE TABLE IF NOT EXISTS routing_group_versions ( - `id` VARCHAR(64) NOT NULL, - `group_id` VARCHAR(64) NOT NULL, - `version` BIGINT NOT NULL, - `config_json` JSON NOT NULL, - `created_at` BIGINT NOT NULL, - `created_by` VARCHAR(64), - PRIMARY KEY (`id`), - UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`), - KEY routing_group_versions_group_id_idx (`group_id`) -); - diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/003_auth_config.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/003_auth_config.sql deleted file mode 100644 index 7ddf3f6a8..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/003_auth_config.sql +++ /dev/null @@ -1,83 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS system_configs ( - `id` VARCHAR(64) NOT NULL, - `key` VARCHAR(255) NOT NULL, - `value` LONGTEXT NOT NULL, - `description` LONGTEXT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY system_configs_key_key (`key`) -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - `id` VARCHAR(64) NOT NULL, - `module_type` VARCHAR(128) NOT NULL, - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `config` JSON NOT NULL, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY auth_modules_module_type_key (`module_type`) -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - `provider_type` VARCHAR(64) NOT NULL, - `display_name` VARCHAR(255) NOT NULL, - `client_id` LONGTEXT NOT NULL, - `client_secret_encrypted` LONGTEXT, - `authorization_url_override` VARCHAR(500), - `token_url_override` VARCHAR(500), - `userinfo_url_override` VARCHAR(500), - `scopes` JSON, - `redirect_uri` VARCHAR(500) NOT NULL, - `frontend_callback_url` VARCHAR(500) NOT NULL, - `attribute_mapping` JSON, - `extra_config` JSON, - `icon_url` VARCHAR(500), - `is_enabled` TINYINT(1) NOT NULL DEFAULT 0, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`provider_type`) -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - `id` BIGINT NOT NULL AUTO_INCREMENT, - `singleton_key` INT NOT NULL DEFAULT 1, - `server_url` VARCHAR(255) NOT NULL, - `bind_dn` LONGTEXT NOT NULL, - `bind_password_encrypted` LONGTEXT, - `base_dn` LONGTEXT NOT NULL, - `user_search_filter` VARCHAR(512) NOT NULL DEFAULT '(uid={username})', - `username_attr` VARCHAR(50) NOT NULL DEFAULT 'uid', - `email_attr` VARCHAR(50) NOT NULL DEFAULT 'mail', - `display_name_attr` VARCHAR(50) NOT NULL DEFAULT 'cn', - `is_enabled` TINYINT(1) NOT NULL DEFAULT 0, - `is_exclusive` TINYINT(1) NOT NULL DEFAULT 0, - `use_starttls` TINYINT(1) NOT NULL DEFAULT 0, - `connect_timeout` INT NOT NULL DEFAULT 10, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY ldap_configs_singleton_key_key (`singleton_key`) -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `provider_type` VARCHAR(64) NOT NULL, - `provider_user_id` VARCHAR(255) NOT NULL, - `provider_username` VARCHAR(255), - `provider_email` VARCHAR(255), - `extra_data` JSON, - `linked_at` BIGINT NOT NULL, - `last_login_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY uq_user_oauth_links_provider_user (`provider_type`, `provider_user_id`), - UNIQUE KEY uq_user_oauth_links_user_provider (`user_id`, `provider_type`), - KEY user_oauth_links_provider_type_idx (`provider_type`), - KEY user_oauth_links_user_id_idx (`user_id`) -); - diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/004_proxy_nodes.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/004_proxy_nodes.sql deleted file mode 100644 index cd104a852..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/004_proxy_nodes.sql +++ /dev/null @@ -1,88 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS proxy_nodes ( - `id` VARCHAR(64) NOT NULL, - `tunnel_generation` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `ip` VARCHAR(512) NOT NULL, - `port` INT NOT NULL, - `region` VARCHAR(100), - `status` VARCHAR(32) NOT NULL DEFAULT 'online', - `registered_by` VARCHAR(64), - `last_heartbeat_at` BIGINT, - `heartbeat_interval` INT NOT NULL DEFAULT 30, - `active_connections` INT NOT NULL DEFAULT 0, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `avg_latency_ms` DOUBLE, - `is_manual` TINYINT(1) NOT NULL DEFAULT 0, - `proxy_url` VARCHAR(500), - `proxy_username` VARCHAR(255), - `proxy_password` TEXT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `remote_config` JSON, - `config_version` INT NOT NULL DEFAULT 0, - `hardware_info` JSON, - `estimated_max_concurrency` INT, - `tunnel_mode` TINYINT(1) NOT NULL DEFAULT 0, - `tunnel_connected` TINYINT(1) NOT NULL DEFAULT 0, - `tunnel_connected_at` BIGINT, - `failed_requests` BIGINT NOT NULL DEFAULT 0, - `dns_failures` BIGINT NOT NULL DEFAULT 0, - `stream_errors` BIGINT NOT NULL DEFAULT 0, - `proxy_metadata` JSON, - PRIMARY KEY (`id`), - UNIQUE KEY uq_proxy_node_ip_port (`ip`, `port`) -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - `id` BIGINT NOT NULL AUTO_INCREMENT, - `node_id` VARCHAR(64) NOT NULL, - `event_type` VARCHAR(64) NOT NULL, - `detail` VARCHAR(500), - `event_metadata` JSON, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`id`) -); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m ( - `node_id` VARCHAR(64) NOT NULL, - `bucket_start_unix_secs` BIGINT NOT NULL, - `samples` BIGINT NOT NULL DEFAULT 0, - `uptime_samples` BIGINT NOT NULL DEFAULT 0, - `active_connections_sum` BIGINT NOT NULL DEFAULT 0, - `active_connections_max` BIGINT NOT NULL DEFAULT 0, - `heartbeat_rtt_ms_sum` BIGINT NOT NULL DEFAULT 0, - `heartbeat_rtt_ms_max` BIGINT NOT NULL DEFAULT 0, - `connect_errors_delta` BIGINT NOT NULL DEFAULT 0, - `disconnects_delta` BIGINT NOT NULL DEFAULT 0, - `error_events_delta` BIGINT NOT NULL DEFAULT 0, - `ws_in_bytes_delta` BIGINT NOT NULL DEFAULT 0, - `ws_out_bytes_delta` BIGINT NOT NULL DEFAULT 0, - `ws_in_frames_delta` BIGINT NOT NULL DEFAULT 0, - `ws_out_frames_delta` BIGINT NOT NULL DEFAULT 0, - PRIMARY KEY (`node_id`, `bucket_start_unix_secs`), - KEY idx_proxy_node_metrics_1m_bucket_start (`bucket_start_unix_secs`) -); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h ( - `node_id` VARCHAR(64) NOT NULL, - `bucket_start_unix_secs` BIGINT NOT NULL, - `samples` BIGINT NOT NULL DEFAULT 0, - `uptime_samples` BIGINT NOT NULL DEFAULT 0, - `active_connections_sum` BIGINT NOT NULL DEFAULT 0, - `active_connections_max` BIGINT NOT NULL DEFAULT 0, - `heartbeat_rtt_ms_sum` BIGINT NOT NULL DEFAULT 0, - `heartbeat_rtt_ms_max` BIGINT NOT NULL DEFAULT 0, - `connect_errors_delta` BIGINT NOT NULL DEFAULT 0, - `disconnects_delta` BIGINT NOT NULL DEFAULT 0, - `error_events_delta` BIGINT NOT NULL DEFAULT 0, - `ws_in_bytes_delta` BIGINT NOT NULL DEFAULT 0, - `ws_out_bytes_delta` BIGINT NOT NULL DEFAULT 0, - `ws_in_frames_delta` BIGINT NOT NULL DEFAULT 0, - `ws_out_frames_delta` BIGINT NOT NULL DEFAULT 0, - PRIMARY KEY (`node_id`, `bucket_start_unix_secs`), - KEY idx_proxy_node_metrics_1h_bucket_start (`bucket_start_unix_secs`) -); - diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/005_wallet_billing.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/005_wallet_billing.sql deleted file mode 100644 index 6b46a309c..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/005_wallet_billing.sql +++ /dev/null @@ -1,330 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS wallets ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64), - `api_key_id` VARCHAR(64), - `balance` DOUBLE NOT NULL DEFAULT 0, - `gift_balance` DOUBLE NOT NULL DEFAULT 0, - `limit_mode` VARCHAR(64) NOT NULL DEFAULT 'finite', - `currency` VARCHAR(16) NOT NULL DEFAULT 'USD', - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `total_recharged` DOUBLE NOT NULL DEFAULT 0, - `total_consumed` DOUBLE NOT NULL DEFAULT 0, - `total_refunded` DOUBLE NOT NULL DEFAULT 0, - `total_adjusted` DOUBLE NOT NULL DEFAULT 0, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY wallets_user_id_key (`user_id`), - UNIQUE KEY wallets_api_key_id_key (`api_key_id`), - KEY wallets_api_key_id_idx (`api_key_id`), - KEY wallets_user_id_idx (`user_id`) -); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - `id` VARCHAR(64) NOT NULL, - `wallet_id` VARCHAR(64) NOT NULL, - `category` VARCHAR(64) NOT NULL, - `reason_code` VARCHAR(64) NOT NULL, - `amount` DOUBLE NOT NULL, - `balance_before` DOUBLE NOT NULL, - `balance_after` DOUBLE NOT NULL, - `recharge_balance_before` DOUBLE NOT NULL, - `recharge_balance_after` DOUBLE NOT NULL, - `gift_balance_before` DOUBLE NOT NULL, - `gift_balance_after` DOUBLE NOT NULL, - `link_type` VARCHAR(64), - `link_id` VARCHAR(128), - `operator_id` VARCHAR(64), - `description` LONGTEXT, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_wallet_tx_wallet_created (`wallet_id`, `created_at`), - KEY idx_wallet_tx_category_created (`category`, `created_at`), - KEY idx_wallet_tx_reason_created (`reason_code`, `created_at`), - KEY idx_wallet_tx_link (`link_type`, `link_id`), - KEY ix_wallet_transactions_operator_id (`operator_id`) -); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - `id` VARCHAR(64) NOT NULL, - `wallet_id` VARCHAR(64) NOT NULL, - `billing_date` VARCHAR(16) NOT NULL, - `billing_timezone` VARCHAR(64) NOT NULL, - `total_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, - `first_finalized_at` BIGINT, - `last_finalized_at` BIGINT, - `aggregated_at` BIGINT NOT NULL, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_wallet_daily_usage_wallet_date (`wallet_id`, `billing_timezone`, `billing_date`) -); - -CREATE TABLE IF NOT EXISTS payment_orders ( - `id` VARCHAR(64) NOT NULL, - `order_no` VARCHAR(128) NOT NULL, - `wallet_id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64), - `amount_usd` DOUBLE NOT NULL, - `pay_amount` DOUBLE, - `pay_currency` VARCHAR(16), - `exchange_rate` DOUBLE, - `refunded_amount_usd` DOUBLE NOT NULL DEFAULT 0, - `refundable_amount_usd` DOUBLE NOT NULL DEFAULT 0, - `payment_method` VARCHAR(64) NOT NULL, - `payment_provider` VARCHAR(64), - `payment_channel` VARCHAR(64), - `order_kind` VARCHAR(64) NOT NULL DEFAULT 'wallet_recharge', - `product_id` VARCHAR(64), - `product_snapshot` JSON, - `fulfillment_status` VARCHAR(64) NOT NULL DEFAULT 'pending', - `fulfillment_error` LONGTEXT, - `gateway_order_id` VARCHAR(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin, - `gateway_response` JSON, - `status` VARCHAR(64) NOT NULL DEFAULT 'pending', - `created_at` BIGINT NOT NULL, - `paid_at` BIGINT, - `credited_at` BIGINT, - `expires_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY uq_payment_orders_order_no (`order_no`), - KEY idx_payment_orders_wallet_created (`wallet_id`, `created_at`), - KEY idx_payment_orders_user_created (`user_id`, `created_at`), - KEY idx_payment_orders_status (`status`), - KEY idx_payment_orders_gateway_order_id (`gateway_order_id`), - UNIQUE KEY uq_payment_orders_payment_method_gateway_order_id (`payment_method`, `gateway_order_id`), - KEY idx_payment_orders_kind_status (`order_kind`, `status`), - KEY idx_payment_orders_product (`product_id`) -); - -CREATE TABLE IF NOT EXISTS user_invite_codes ( - `user_id` VARCHAR(64) NOT NULL, - `invite_code` VARCHAR(64) NOT NULL, - `active` TINYINT(1) NOT NULL DEFAULT 1, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`user_id`), - UNIQUE KEY user_invite_codes_invite_code_key (`invite_code`), - CONSTRAINT user_invite_codes_user_id_fkey FOREIGN KEY (`user_id`) REFERENCES users (`id`) ON DELETE CASCADE -); - -CREATE TABLE IF NOT EXISTS user_referrals ( - `id` VARCHAR(64) NOT NULL, - `inviter_user_id` VARCHAR(64) NOT NULL, - `invitee_user_id` VARCHAR(64) NOT NULL, - `invite_code_snapshot` VARCHAR(64) NOT NULL, - `source_json` JSON, - `first_paid_order_id` VARCHAR(64), - `first_paid_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY user_referrals_invitee_user_id_key (`invitee_user_id`), - KEY idx_user_referrals_inviter (`inviter_user_id`, `created_at`), - KEY idx_user_referrals_created (`created_at`), - KEY idx_user_referrals_invite_code (`invite_code_snapshot`), - CONSTRAINT user_referrals_first_paid_order_fkey FOREIGN KEY (`first_paid_order_id`) REFERENCES payment_orders (`id`) ON DELETE SET NULL -); - -CREATE TABLE IF NOT EXISTS referral_rewards ( - `id` VARCHAR(64) NOT NULL, - `referral_id` VARCHAR(64) NOT NULL, - `inviter_user_id` VARCHAR(64) NOT NULL, - `invitee_user_id` VARCHAR(64) NOT NULL, - `reward_type` VARCHAR(32) NOT NULL, - `trigger_point` VARCHAR(64) NOT NULL, - `source_order_id` VARCHAR(64), - `idempotency_key` VARCHAR(128) NOT NULL, - `amount_usd` DOUBLE NOT NULL, - `status` VARCHAR(32) NOT NULL DEFAULT 'pending', - `wallet_transaction_id` VARCHAR(64), - `reversed_amount_usd` DOUBLE NOT NULL DEFAULT 0, - `pending_reversal_amount_usd` DOUBLE NOT NULL DEFAULT 0, - `failure_reason` LONGTEXT, - `admin_operator_id` VARCHAR(64), - `admin_note` LONGTEXT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY referral_rewards_idempotency_key_key (`idempotency_key`), - KEY idx_referral_rewards_inviter_status (`inviter_user_id`, `status`, `created_at`), - KEY idx_referral_rewards_inviter_created (`inviter_user_id`, `created_at`), - KEY idx_referral_rewards_created (`created_at`), - KEY idx_referral_rewards_source_order (`source_order_id`), - CONSTRAINT referral_rewards_referral_id_fkey FOREIGN KEY (`referral_id`) REFERENCES user_referrals (`id`) ON DELETE CASCADE, - CONSTRAINT referral_rewards_source_order_fkey FOREIGN KEY (`source_order_id`) REFERENCES payment_orders (`id`) ON DELETE SET NULL -); - -CREATE TABLE IF NOT EXISTS payment_gateway_configs ( - `provider` VARCHAR(64) NOT NULL, - `enabled` TINYINT(1) NOT NULL DEFAULT 0, - `endpoint_url` VARCHAR(512) NOT NULL, - `callback_base_url` VARCHAR(512), - `merchant_id` VARCHAR(128) NOT NULL, - `merchant_key_encrypted` LONGTEXT, - `pay_currency` VARCHAR(16) NOT NULL DEFAULT 'CNY', - `usd_exchange_rate` DOUBLE NOT NULL DEFAULT 7.2, - `min_recharge_usd` DOUBLE NOT NULL DEFAULT 1, - `channels_json` JSON, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`provider`) -); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - `id` VARCHAR(64) NOT NULL, - `payment_order_id` VARCHAR(64), - `payment_method` VARCHAR(64) NOT NULL, - `callback_key` VARCHAR(128) NOT NULL, - `order_no` VARCHAR(128), - `gateway_order_id` VARCHAR(128), - `payload_hash` VARCHAR(128), - `signature_valid` TINYINT(1) NOT NULL DEFAULT 0, - `status` VARCHAR(64) NOT NULL DEFAULT 'received', - `payload` JSON, - `error_message` LONGTEXT, - `created_at` BIGINT NOT NULL, - `processed_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY uq_payment_callbacks_callback_key (`callback_key`), - KEY idx_payment_callbacks_order (`order_no`), - KEY idx_payment_callbacks_gateway_order (`gateway_order_id`), - KEY idx_payment_callbacks_created (`created_at`), - KEY ix_payment_callbacks_payment_order_id (`payment_order_id`) -); - -CREATE TABLE IF NOT EXISTS billing_plans ( - `id` VARCHAR(64) NOT NULL, - `title` VARCHAR(128) NOT NULL, - `description` LONGTEXT, - `price_amount` DOUBLE NOT NULL, - `price_currency` VARCHAR(16) NOT NULL DEFAULT 'CNY', - `duration_unit` VARCHAR(32) NOT NULL, - `duration_value` BIGINT NOT NULL, - `enabled` TINYINT(1) NOT NULL DEFAULT 1, - `sort_order` BIGINT NOT NULL DEFAULT 0, - `max_active_per_user` BIGINT NOT NULL DEFAULT 1, - `purchase_limit_scope` VARCHAR(32) NOT NULL DEFAULT 'active_period', - `entitlements_json` JSON NOT NULL, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_billing_plans_enabled_sort (`enabled`, `sort_order`) -); - -CREATE TABLE IF NOT EXISTS user_plan_entitlements ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `plan_id` VARCHAR(64) NOT NULL, - `payment_order_id` VARCHAR(64) NOT NULL, - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `starts_at` BIGINT NOT NULL, - `expires_at` BIGINT NOT NULL, - `entitlements_snapshot` JSON NOT NULL, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_user_plan_entitlements_user_active (`user_id`, `status`, `expires_at`), - KEY idx_user_plan_entitlements_order (`payment_order_id`) -); - -CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers ( - `id` VARCHAR(64) NOT NULL, - `user_entitlement_id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `request_id` VARCHAR(128) NOT NULL, - `amount_usd` DOUBLE NOT NULL, - `balance_before` DOUBLE NOT NULL, - `balance_after` DOUBLE NOT NULL, - `usage_date` VARCHAR(16) NOT NULL, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_entitlement_usage_request (`user_entitlement_id`, `request_id`), - KEY idx_entitlement_usage_user_date (`user_id`, `usage_date`), - KEY idx_entitlement_usage_entitlement_date (`user_entitlement_id`, `usage_date`) -); - -CREATE TABLE IF NOT EXISTS refund_requests ( - `id` VARCHAR(64) NOT NULL, - `refund_no` VARCHAR(128) NOT NULL, - `wallet_id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64), - `payment_order_id` VARCHAR(64), - `source_type` VARCHAR(64) NOT NULL, - `source_id` VARCHAR(128), - `refund_mode` VARCHAR(64) NOT NULL, - `amount_usd` DOUBLE NOT NULL, - `status` VARCHAR(64) NOT NULL DEFAULT 'pending_approval', - `reason` LONGTEXT, - `requested_by` VARCHAR(64), - `approved_by` VARCHAR(64), - `processed_by` VARCHAR(64), - `gateway_refund_id` VARCHAR(128), - `payout_method` VARCHAR(64), - `payout_reference` VARCHAR(255), - `payout_proof` JSON, - `failure_reason` LONGTEXT, - `idempotency_key` VARCHAR(128), - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `processed_at` BIGINT, - `completed_at` BIGINT, - PRIMARY KEY (`id`), - UNIQUE KEY uq_refund_requests_refund_no (`refund_no`), - UNIQUE KEY uq_refund_requests_idempotency_key (`idempotency_key`), - KEY idx_refund_wallet_created (`wallet_id`, `created_at`), - KEY idx_refund_user_created (`user_id`, `created_at`), - KEY idx_refund_status (`status`), - KEY ix_refund_requests_payment_order_id (`payment_order_id`), - KEY ix_refund_requests_requested_by (`requested_by`), - KEY ix_refund_requests_approved_by (`approved_by`), - KEY ix_refund_requests_processed_by (`processed_by`) -); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - `id` VARCHAR(64) NOT NULL, - `name` VARCHAR(255) NOT NULL, - `amount_usd` DOUBLE NOT NULL, - `currency` VARCHAR(16) NOT NULL DEFAULT 'USD', - `balance_bucket` VARCHAR(64) NOT NULL DEFAULT 'gift', - `total_count` INT NOT NULL, - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `description` LONGTEXT, - `created_by` VARCHAR(64), - `expires_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_redeem_code_batches_status (`status`, `created_at`) -); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - `id` VARCHAR(64) NOT NULL, - `batch_id` VARCHAR(64) NOT NULL, - `code_hash` VARCHAR(128) NOT NULL, - `code_prefix` VARCHAR(16) NOT NULL, - `code_suffix` VARCHAR(16) NOT NULL, - `status` VARCHAR(64) NOT NULL DEFAULT 'active', - `redeemed_by_user_id` VARCHAR(64), - `redeemed_wallet_id` VARCHAR(64), - `redeemed_payment_order_id` VARCHAR(64), - `redeemed_at` BIGINT, - `disabled_by` VARCHAR(64), - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_redeem_codes_code_hash (`code_hash`), - KEY idx_redeem_codes_batch_created (`batch_id`, `created_at`), - KEY idx_redeem_codes_status (`status`, `updated_at`), - KEY idx_redeem_codes_redeemed_user (`redeemed_by_user_id`, `redeemed_at`), - KEY idx_redeem_codes_redeemed_order (`redeemed_payment_order_id`) -); - 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 deleted file mode 100644 index 15f272948..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/006_usage.sql +++ /dev/null @@ -1,282 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS `usage` ( - `request_id` VARCHAR(128) NOT NULL, - `id` VARCHAR(128), - `user_id` VARCHAR(64), - `api_key_id` VARCHAR(64), - `provider_name` VARCHAR(255) NOT NULL DEFAULT 'unknown', - `model` VARCHAR(255) NOT NULL DEFAULT 'unknown', - `target_model` VARCHAR(255), - `provider_id` VARCHAR(64), - `provider_endpoint_id` VARCHAR(64), - `provider_api_key_id` VARCHAR(64), - `request_type` VARCHAR(64), - `api_format` VARCHAR(64), - `api_family` VARCHAR(64), - `endpoint_kind` VARCHAR(64), - `endpoint_api_format` VARCHAR(64), - `provider_api_family` VARCHAR(64), - `provider_endpoint_kind` VARCHAR(64), - `has_format_conversion` TINYINT(1) NOT NULL DEFAULT 0, - `is_stream` TINYINT(1) NOT NULL DEFAULT 0, - `upstream_is_stream` TINYINT(1), - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `input_output_total_tokens` BIGINT NOT NULL DEFAULT 0, - `total_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, - `input_context_tokens` BIGINT NOT NULL DEFAULT 0, - `input_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `output_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `cache_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `cache_creation_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `cache_creation_cost_usd_5m` DOUBLE NOT NULL DEFAULT 0, - `cache_creation_cost_usd_1h` DOUBLE NOT NULL DEFAULT 0, - `cache_read_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `request_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_input_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_output_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_cache_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_cache_creation_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_cache_creation_cost_usd_5m` DOUBLE NOT NULL DEFAULT 0, - `actual_cache_creation_cost_usd_1h` DOUBLE NOT NULL DEFAULT 0, - `actual_cache_read_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_request_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `rate_multiplier` DOUBLE NOT NULL DEFAULT 1, - `input_price_per_1m` DOUBLE, - `output_price_per_1m` DOUBLE, - `cache_creation_price_per_1m` DOUBLE, - `cache_creation_price_per_1m_5m` DOUBLE, - `cache_creation_price_per_1m_1h` DOUBLE, - `cache_read_price_per_1m` DOUBLE, - `price_per_request` DOUBLE, - `status_code` INT, - `error_message` LONGTEXT, - `error_category` VARCHAR(255), - `response_time_ms` BIGINT, - `first_byte_time_ms` BIGINT, - `wallet_id` VARCHAR(64), - `status` VARCHAR(64) NOT NULL DEFAULT 'completed', - `billing_status` VARCHAR(64) NOT NULL DEFAULT 'pending', - `total_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `actual_total_cost_usd` DOUBLE NOT NULL DEFAULT 0, - `request_headers` JSON, - `request_body` JSON, - `provider_request_headers` JSON, - `provider_request_body` JSON, - `response_headers` JSON, - `response_body` JSON, - `client_response_headers` JSON, - `client_response_body` JSON, - `request_body_compressed` LONGBLOB, - `provider_request_body_compressed` LONGBLOB, - `response_body_compressed` LONGBLOB, - `client_response_body_compressed` LONGBLOB, - `request_metadata` JSON, - `created_at` BIGINT, - `candidate_id` VARCHAR(128), - `candidate_index` BIGINT, - `key_name` VARCHAR(255), - `username` VARCHAR(255), - `api_key_name` VARCHAR(255), - `planner_kind` VARCHAR(64), - `route_family` VARCHAR(128), - `route_kind` VARCHAR(128), - `execution_path` VARCHAR(128), - `local_execution_runtime_miss_reason` VARCHAR(255), - `wallet_balance_before` DOUBLE, - `wallet_balance_after` DOUBLE, - `wallet_recharge_balance_before` DOUBLE, - `wallet_recharge_balance_after` DOUBLE, - `wallet_gift_balance_before` DOUBLE, - `wallet_gift_balance_after` DOUBLE, - `finalized_at` BIGINT, - `created_at_unix_ms` BIGINT NOT NULL DEFAULT 0, - `updated_at_unix_secs` BIGINT NOT NULL DEFAULT 0, - PRIMARY KEY (`request_id`), - KEY usage_api_key_id_idx (`api_key_id`), - KEY usage_billing_status_idx (`billing_status`), - KEY usage_created_at_idx (`created_at_unix_ms`), - KEY usage_provider_api_key_id_idx (`provider_api_key_id`), - KEY usage_provider_id_idx (`provider_id`), - KEY usage_request_id_idx (`request_id`), - KEY usage_user_id_idx (`user_id`), - 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, - `kind` VARCHAR(64) NOT NULL, - `target_id` TEXT NOT NULL, - `target_tunnel_generation` VARCHAR(64), - `request_count_delta` BIGINT NOT NULL DEFAULT 0, - `total_requests_delta` BIGINT NOT NULL DEFAULT 0, - `success_count_delta` BIGINT NOT NULL DEFAULT 0, - `error_count_delta` BIGINT NOT NULL DEFAULT 0, - `dns_failures_delta` BIGINT NOT NULL DEFAULT 0, - `stream_errors_delta` BIGINT NOT NULL DEFAULT 0, - `total_tokens_delta` BIGINT NOT NULL DEFAULT 0, - `total_cost_usd_delta` DOUBLE NOT NULL DEFAULT 0, - `total_response_time_ms_delta` BIGINT NOT NULL DEFAULT 0, - `last_used_at_unix_secs` BIGINT, - `last_used_ip` TEXT, - `candidate_last_used_at_unix_secs` BIGINT, - `removed_last_used_at_unix_secs` BIGINT, - `usage_created_at_unix_secs` BIGINT, - `created_at` BIGINT NOT NULL, - `processed_at` BIGINT, - PRIMARY KEY (`id`), - KEY ix_usage_counter_deltas_unprocessed (`created_at`, `id`), - KEY ix_usage_counter_deltas_processed (`processed_at`, `created_at`, `id`), - KEY ix_usage_counter_deltas_request_kind (`request_id`, `kind`, `target_id`) -); - -CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( - `request_id` VARCHAR(128) NOT NULL, - `billing_status` VARCHAR(64) NOT NULL, - `wallet_id` VARCHAR(64), - `wallet_balance_before` DOUBLE, - `wallet_balance_after` DOUBLE, - `wallet_recharge_balance_before` DOUBLE, - `wallet_recharge_balance_after` DOUBLE, - `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 ix_usage_settlement_snapshots_schema_version (`settlement_snapshot_schema_version`), - KEY ix_usage_settlement_snapshots_pricing_source (`billing_pricing_source`) -); - -CREATE TABLE IF NOT EXISTS usage_cost_reservations ( - `request_id` VARCHAR(128) NOT NULL, - `subject_id` VARCHAR(128) NOT NULL, - `reservation_token` VARCHAR(128) NOT NULL, - `admitted_at` BIGINT NOT NULL, - `reserved_cost_units` BIGINT NOT NULL, - `actual_cost_units` BIGINT, - `state` VARCHAR(20) NOT NULL, - `reservation_expires_at` BIGINT NOT NULL, - `retain_until` BIGINT NOT NULL, - `finalized_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`reservation_token`), - KEY usage_cost_reservations_request_id_idx (`request_id`), - KEY usage_cost_reservations_subject_admitted_at_idx (`subject_id`, `admitted_at`), - KEY usage_cost_reservations_reservation_expires_at_idx (`reservation_expires_at`), - KEY usage_cost_reservations_retain_until_token_idx (`retain_until`, `reservation_token`), - CONSTRAINT usage_cost_reservations_subject_id_fkey FOREIGN KEY (`subject_id`) REFERENCES users (`id`) ON DELETE CASCADE -); - -CREATE TABLE IF NOT EXISTS usage_request_admissions ( - `request_id` VARCHAR(128) NOT NULL, - `subject_id` VARCHAR(128) NOT NULL, - `event_token` VARCHAR(128) NOT NULL, - `admitted_at` BIGINT NOT NULL, - `retain_until` BIGINT NOT NULL, - `state` VARCHAR(20) NOT NULL, - `released_at` BIGINT, - `created_at` BIGINT NOT NULL, - PRIMARY KEY (`event_token`), - KEY usage_request_admissions_subject_admitted_at_idx (`subject_id`, `admitted_at`), - KEY usage_request_admissions_retain_until_token_idx (`retain_until`, `event_token`), - CONSTRAINT usage_request_admissions_subject_id_fkey FOREIGN KEY (`subject_id`) REFERENCES users (`id`) ON DELETE CASCADE -); - 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 deleted file mode 100644 index ae5d2aee2..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/007_stats.sql +++ /dev/null @@ -1,599 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS stats_hourly ( - `id` VARCHAR(64) NOT NULL, - `hour_utc` BIGINT NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `success_requests` BIGINT NOT NULL DEFAULT 0, - `error_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_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, - `avg_response_time_ms` DOUBLE NOT NULL DEFAULT 0, - `is_complete` TINYINT(1) NOT NULL DEFAULT 0, - `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`) -); - -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 stats_hourly_user ( - `id` VARCHAR(64) NOT NULL, - `hour_utc` BIGINT NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `success_requests` BIGINT NOT NULL DEFAULT 0, - `error_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `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`) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - `id` VARCHAR(64) NOT NULL, - `hour_utc` BIGINT NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `model` VARCHAR(255) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `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`) -); - -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`) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - `id` VARCHAR(64) NOT NULL, - `hour_utc` BIGINT NOT NULL, - `model` VARCHAR(255) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `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`) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - `id` VARCHAR(64) NOT NULL, - `hour_utc` BIGINT NOT NULL, - `provider_name` VARCHAR(255) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_stats_hourly_provider (`hour_utc`, `provider_name`) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - `id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `success_requests` BIGINT NOT NULL DEFAULT 0, - `error_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_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, - `input_cost` DOUBLE NOT NULL DEFAULT 0, - `output_cost` DOUBLE NOT NULL DEFAULT 0, - `cache_creation_cost` DOUBLE NOT NULL DEFAULT 0, - `cache_read_cost` DOUBLE NOT NULL DEFAULT 0, - `avg_response_time_ms` DOUBLE NOT NULL DEFAULT 0, - `fallback_count` BIGINT NOT NULL DEFAULT 0, - `unique_models` BIGINT NOT NULL DEFAULT 0, - `unique_providers` BIGINT NOT NULL DEFAULT 0, - `is_complete` TINYINT(1) NOT NULL DEFAULT 0, - `aggregated_at` BIGINT, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - `p50_response_time_ms` BIGINT, - `p90_response_time_ms` BIGINT, - `p99_response_time_ms` BIGINT, - `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`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - `id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `model` VARCHAR(255) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `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`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - `id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `provider_name` VARCHAR(255) NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_stats_daily_provider (`date`, `provider_name`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - `id` VARCHAR(64) NOT NULL, - `api_key_id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `success_requests` BIGINT NOT NULL DEFAULT 0, - `error_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `api_key_name` VARCHAR(255), - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_stats_daily_api_key (`date`, `api_key_id`) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - `id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `error_category` VARCHAR(255) NOT NULL, - `provider_name` VARCHAR(255), - `model` VARCHAR(255), - `count` BIGINT NOT NULL DEFAULT 0, - `created_at` BIGINT NOT NULL, - `updated_at` BIGINT NOT NULL, - PRIMARY KEY (`id`), - UNIQUE KEY uq_stats_daily_error (`date`, `error_category`, `provider_name`, `model`) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - `id` VARCHAR(64) NOT NULL, - `user_id` VARCHAR(64) NOT NULL, - `date` BIGINT NOT NULL, - `total_requests` BIGINT NOT NULL DEFAULT 0, - `success_requests` BIGINT NOT NULL DEFAULT 0, - `error_requests` BIGINT NOT NULL DEFAULT 0, - `input_tokens` BIGINT NOT NULL DEFAULT 0, - `output_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_creation_tokens` BIGINT NOT NULL DEFAULT 0, - `cache_read_tokens` BIGINT NOT NULL DEFAULT 0, - `total_cost` DOUBLE NOT NULL DEFAULT 0, - `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/mysql/baseline/008_background_tasks.sql b/crates/aether-data/runtime/schema/generated/mysql/baseline/008_background_tasks.sql deleted file mode 100644 index 465f2cea7..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/008_background_tasks.sql +++ /dev/null @@ -1,42 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS background_task_runs ( - `id` VARCHAR(64) NOT NULL, - `task_key` VARCHAR(200) NOT NULL, - `kind` VARCHAR(32) NOT NULL, - `trigger` VARCHAR(64) NOT NULL, - `status` VARCHAR(32) NOT NULL, - `attempt` INT NOT NULL DEFAULT 0, - `max_attempts` INT NOT NULL DEFAULT 0, - `owner_instance` VARCHAR(200), - `progress_percent` INT NOT NULL DEFAULT 0, - `progress_message` TEXT, - `payload_json` JSON, - `result_json` JSON, - `error_message` TEXT, - `cancel_requested` TINYINT(1) NOT NULL DEFAULT 0, - `created_by` VARCHAR(200), - `created_at_unix_secs` BIGINT NOT NULL, - `started_at_unix_secs` BIGINT, - `finished_at_unix_secs` BIGINT, - `updated_at_unix_secs` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_background_task_runs_task_key (`task_key`), - KEY idx_background_task_runs_status (`status`), - KEY idx_background_task_runs_kind (`kind`), - KEY idx_background_task_runs_created_at (`created_at_unix_secs`) -); - -CREATE TABLE IF NOT EXISTS background_task_events ( - `id` VARCHAR(64) NOT NULL, - `run_id` VARCHAR(64) NOT NULL, - `event_type` VARCHAR(64) NOT NULL, - `message` TEXT NOT NULL, - `payload_json` JSON, - `created_at_unix_secs` BIGINT NOT NULL, - PRIMARY KEY (`id`), - KEY idx_background_task_events_run_id (`run_id`, `created_at_unix_secs`), - CONSTRAINT fk_background_task_events_run FOREIGN KEY (`run_id`) REFERENCES background_task_runs (`id`) ON DELETE CASCADE -); - diff --git a/crates/aether-data/runtime/schema/generated/mysql/baseline/manifest.txt b/crates/aether-data/runtime/schema/generated/mysql/baseline/manifest.txt deleted file mode 100644 index 32ce47e5f..000000000 --- a/crates/aether-data/runtime/schema/generated/mysql/baseline/manifest.txt +++ /dev/null @@ -1,11 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -001_identity.sql -002_provider_catalog.sql -003_auth_config.sql -004_proxy_nodes.sql -005_wallet_billing.sql -006_usage.sql -007_stats.sql -008_background_tasks.sql diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql deleted file mode 100644 index 8a9eb62e6..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/001_identity.sql +++ /dev/null @@ -1,213 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS users ( - id TEXT PRIMARY KEY NOT NULL, - external_id TEXT, - email TEXT, - username TEXT, - password_hash TEXT, - role TEXT, - auth_source TEXT NOT NULL DEFAULT 'local', - email_verified INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - is_deleted INTEGER NOT NULL DEFAULT 0, - security_version INTEGER NOT NULL DEFAULT 0, - allowed_models TEXT, - allowed_models_mode TEXT NOT NULL DEFAULT 'unrestricted', - allowed_providers TEXT, - allowed_providers_mode TEXT NOT NULL DEFAULT 'unrestricted', - allowed_api_formats TEXT, - allowed_api_formats_mode TEXT NOT NULL DEFAULT 'unrestricted', - model_capability_settings TEXT, - feature_settings TEXT, - rate_limit INTEGER, - rate_limit_mode TEXT NOT NULL DEFAULT 'system', - privacy_policy_accepted_version TEXT, - privacy_policy_accepted_at INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - last_login_at INTEGER, - ldap_dn TEXT, - ldap_username TEXT, - UNIQUE (email), - UNIQUE (username) -); - -CREATE TABLE IF NOT EXISTS user_groups ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - normalized_name TEXT NOT NULL, - description TEXT, - priority INTEGER NOT NULL DEFAULT 0, - allowed_providers TEXT, - allowed_providers_mode TEXT NOT NULL DEFAULT 'inherit', - allowed_api_formats TEXT, - allowed_api_formats_mode TEXT NOT NULL DEFAULT 'inherit', - allowed_models TEXT, - allowed_models_mode TEXT NOT NULL DEFAULT 'inherit', - rate_limit INTEGER, - rate_limit_mode TEXT NOT NULL DEFAULT 'inherit', - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (normalized_name) -); -CREATE INDEX IF NOT EXISTS user_groups_priority_name_idx ON user_groups (priority, name, id); - -CREATE TABLE IF NOT EXISTS user_group_members ( - group_id TEXT NOT NULL, - user_id TEXT NOT NULL, - created_at INTEGER NOT NULL, - PRIMARY KEY (group_id, user_id) -); -CREATE INDEX IF NOT EXISTS user_group_members_user_id_idx ON user_group_members (user_id); - -CREATE TABLE IF NOT EXISTS api_keys ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - key_hash TEXT NOT NULL, - key_encrypted TEXT, - name TEXT, - key_prefix TEXT, - status TEXT NOT NULL DEFAULT 'active', - allowed_models TEXT, - allowed_providers TEXT, - allowed_api_formats TEXT, - ip_rules TEXT, - rate_limit INTEGER DEFAULT 100, - concurrent_limit INTEGER, - force_capabilities TEXT, - feature_settings TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_locked INTEGER NOT NULL DEFAULT 0, - is_standalone INTEGER NOT NULL DEFAULT 0, - auto_delete_on_expiry INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - metadata TEXT, - expires_at INTEGER, - last_used_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (key_hash) -); -CREATE INDEX IF NOT EXISTS api_keys_user_id_idx ON api_keys (user_id); - -CREATE TABLE IF NOT EXISTS audit_logs ( - id TEXT PRIMARY KEY NOT NULL, - event_type TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - description TEXT NOT NULL, - ip_address TEXT, - user_agent TEXT, - request_id TEXT, - event_metadata TEXT, - status_code INTEGER, - error_message TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS audit_logs_created_at_idx ON audit_logs (created_at); -CREATE INDEX IF NOT EXISTS audit_logs_event_type_idx ON audit_logs (event_type); -CREATE INDEX IF NOT EXISTS audit_logs_request_id_idx ON audit_logs (request_id); -CREATE INDEX IF NOT EXISTS audit_logs_user_id_idx ON audit_logs (user_id); - -CREATE TABLE IF NOT EXISTS announcements ( - id TEXT PRIMARY KEY NOT NULL, - title TEXT NOT NULL, - content TEXT NOT NULL, - type TEXT NOT NULL DEFAULT 'info', - priority INTEGER NOT NULL DEFAULT 0, - author_id TEXT, - is_active INTEGER NOT NULL DEFAULT 1, - is_pinned INTEGER NOT NULL DEFAULT 0, - requires_ack INTEGER NOT NULL DEFAULT 0, - start_time INTEGER, - end_time INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS announcements_author_id_idx ON announcements (author_id); -CREATE INDEX IF NOT EXISTS announcements_created_at_idx ON announcements (created_at); -CREATE INDEX IF NOT EXISTS announcements_is_active_idx ON announcements (is_active); - -CREATE TABLE IF NOT EXISTS announcement_reads ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - announcement_id TEXT NOT NULL, - read_at INTEGER NOT NULL, - UNIQUE (user_id, announcement_id) -); -CREATE INDEX IF NOT EXISTS announcement_reads_announcement_id_idx ON announcement_reads (announcement_id); -CREATE INDEX IF NOT EXISTS announcement_reads_user_id_idx ON announcement_reads (user_id); - -CREATE TABLE IF NOT EXISTS management_tokens ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - name TEXT NOT NULL, - description TEXT, - token_hash TEXT NOT NULL, - token_prefix TEXT, - allowed_ips TEXT, - permissions TEXT, - expires_at INTEGER, - last_used_at INTEGER, - last_used_ip TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - is_active INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (token_hash), - UNIQUE (user_id, name) -); -CREATE INDEX IF NOT EXISTS management_tokens_user_id_idx ON management_tokens (user_id); - -CREATE TABLE IF NOT EXISTS user_preferences ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - avatar_url TEXT, - bio TEXT, - default_provider_id TEXT, - theme TEXT NOT NULL DEFAULT 'light', - language TEXT NOT NULL DEFAULT 'zh-CN', - timezone TEXT NOT NULL DEFAULT 'Asia/Shanghai', - email_notifications INTEGER NOT NULL DEFAULT 1, - usage_alerts INTEGER NOT NULL DEFAULT 1, - announcement_notifications INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (user_id) -); -CREATE INDEX IF NOT EXISTS user_preferences_default_provider_id_idx ON user_preferences (default_provider_id); -CREATE INDEX IF NOT EXISTS user_preferences_user_id_idx ON user_preferences (user_id); - -CREATE TABLE IF NOT EXISTS user_sessions ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - security_version INTEGER NOT NULL DEFAULT 0, - client_device_id TEXT NOT NULL, - device_label TEXT, - device_type TEXT NOT NULL DEFAULT 'unknown', - browser_name TEXT, - browser_version TEXT, - os_name TEXT, - os_version TEXT, - device_model TEXT, - ip_address TEXT, - user_agent TEXT, - client_hints TEXT, - refresh_token_hash TEXT NOT NULL, - prev_refresh_token_hash TEXT, - rotated_at INTEGER, - last_seen_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - revoked_at INTEGER, - revoke_reason TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS user_sessions_user_active_idx ON user_sessions (user_id, revoked_at, expires_at); -CREATE INDEX IF NOT EXISTS user_sessions_user_device_idx ON user_sessions (user_id, client_device_id); - diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/002_provider_catalog.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/002_provider_catalog.sql deleted file mode 100644 index a09bebaf1..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/002_provider_catalog.sql +++ /dev/null @@ -1,415 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS billing_rules ( - id TEXT PRIMARY KEY NOT NULL, - global_model_id TEXT, - model_id TEXT, - name TEXT NOT NULL, - task_type TEXT NOT NULL DEFAULT 'chat', - expression TEXT NOT NULL, - variables TEXT NOT NULL, - dimension_mappings TEXT NOT NULL, - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -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 TABLE IF NOT EXISTS dimension_collectors ( - id TEXT PRIMARY KEY NOT NULL, - api_format TEXT NOT NULL, - task_type TEXT NOT NULL, - dimension_name TEXT NOT NULL, - source_type TEXT NOT NULL, - source_path TEXT, - value_type TEXT NOT NULL DEFAULT 'float', - transform_expression TEXT, - default_value TEXT, - priority INTEGER NOT NULL DEFAULT 0, - is_enabled INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS dimension_collectors_enabled_idx ON dimension_collectors (api_format, task_type, dimension_name, priority, is_enabled); - -CREATE TABLE IF NOT EXISTS providers ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - description TEXT, - website TEXT, - provider_type TEXT NOT NULL, - billing_type TEXT, - monthly_quota_usd REAL, - monthly_used_usd REAL, - quota_reset_day INTEGER, - quota_last_reset_at INTEGER, - quota_expires_at INTEGER, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - priority INTEGER NOT NULL DEFAULT 0, - provider_priority INTEGER NOT NULL DEFAULT 100, - keep_priority_on_conversion INTEGER NOT NULL DEFAULT 0, - enable_format_conversion INTEGER NOT NULL DEFAULT 1, - concurrent_limit INTEGER, - max_retries INTEGER, - proxy TEXT, - request_timeout REAL, - stream_first_byte_timeout REAL, - config TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (name) -); - -CREATE TABLE IF NOT EXISTS provider_api_keys ( - id TEXT PRIMARY KEY NOT NULL, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - api_key TEXT, - encrypted_key TEXT, - auth_type TEXT NOT NULL DEFAULT 'api_key', - auth_config TEXT, - note TEXT, - internal_priority INTEGER NOT NULL DEFAULT 50, - capabilities TEXT, - api_formats TEXT, - auth_type_by_format TEXT, - allow_auth_channel_mismatch_formats TEXT, - rate_multipliers TEXT, - global_priority_by_format TEXT, - allowed_models TEXT, - expires_at INTEGER, - cache_ttl_minutes INTEGER NOT NULL DEFAULT 5, - max_probe_interval_minutes INTEGER NOT NULL DEFAULT 32, - proxy TEXT, - fingerprint TEXT, - concurrent_limit INTEGER, - learned_rpm_limit INTEGER, - concurrent_429_count INTEGER NOT NULL DEFAULT 0, - rpm_429_count INTEGER NOT NULL DEFAULT 0, - last_429_at INTEGER, - last_429_type TEXT, - adjustment_history TEXT, - utilization_samples TEXT, - last_probe_increase_at INTEGER, - last_rpm_peak INTEGER, - request_count INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - total_cost_usd REAL NOT NULL DEFAULT 0, - success_count INTEGER NOT NULL DEFAULT 0, - error_count INTEGER NOT NULL DEFAULT 0, - total_response_time_ms INTEGER NOT NULL DEFAULT 0, - last_used_at INTEGER, - last_error_at INTEGER, - last_error_msg TEXT, - auto_fetch_models INTEGER NOT NULL DEFAULT 0, - last_models_fetch_at INTEGER, - last_models_fetch_error TEXT, - locked_models TEXT, - model_include_patterns TEXT, - model_exclude_patterns TEXT, - upstream_metadata TEXT, - oauth_invalid_at INTEGER, - oauth_invalid_reason TEXT, - status_snapshot TEXT, - health_by_format TEXT, - circuit_breaker_by_format TEXT, - status TEXT NOT NULL DEFAULT 'active', - is_active INTEGER NOT NULL DEFAULT 1, - weight INTEGER NOT NULL DEFAULT 1, - rpm_limit INTEGER, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id); -CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id); - -CREATE TABLE IF NOT EXISTS pool_member_scores ( - id TEXT PRIMARY KEY NOT NULL, - pool_kind TEXT NOT NULL, - pool_id TEXT NOT NULL, - member_kind TEXT NOT NULL, - member_id TEXT NOT NULL, - capability TEXT NOT NULL, - scope_kind TEXT NOT NULL, - scope_id TEXT, - score REAL NOT NULL DEFAULT 0, - hard_state TEXT NOT NULL DEFAULT 'unknown', - score_version INTEGER NOT NULL DEFAULT 1, - score_reason TEXT NOT NULL, - last_ranked_at INTEGER, - last_scheduled_at INTEGER, - last_success_at INTEGER, - last_failure_at INTEGER, - failure_count INTEGER NOT NULL DEFAULT 0, - last_probe_attempt_at INTEGER, - last_probe_success_at INTEGER, - last_probe_failure_at INTEGER, - probe_failure_count INTEGER NOT NULL DEFAULT 0, - probe_status TEXT NOT NULL DEFAULT 'never', - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score); -CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id); -CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at); -CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at); - -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 gemini_file_mappings ( - id TEXT PRIMARY KEY NOT NULL, - file_name TEXT NOT NULL, - key_id TEXT NOT NULL, - user_id TEXT, - display_name TEXT, - mime_type TEXT, - source_hash TEXT, - created_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - UNIQUE (file_name) -); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_key_id_idx ON gemini_file_mappings (key_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_user_id_idx ON gemini_file_mappings (user_id); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_expires_at_idx ON gemini_file_mappings (expires_at); -CREATE INDEX IF NOT EXISTS gemini_file_mappings_source_hash_idx ON gemini_file_mappings (source_hash); - -CREATE TABLE IF NOT EXISTS request_candidates ( - id TEXT PRIMARY KEY NOT NULL, - request_id TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - candidate_index INTEGER NOT NULL, - retry_index INTEGER NOT NULL DEFAULT 0, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - status TEXT NOT NULL, - skip_reason TEXT, - is_cached INTEGER NOT NULL DEFAULT 0, - status_code INTEGER, - error_type TEXT, - error_message TEXT, - latency_ms INTEGER, - concurrent_requests INTEGER, - extra_data TEXT, - required_capabilities TEXT, - created_at INTEGER NOT NULL, - started_at INTEGER, - finished_at INTEGER, - UNIQUE (request_id, candidate_index, retry_index) -); -CREATE INDEX IF NOT EXISTS request_candidates_request_id_idx ON request_candidates (request_id); -CREATE INDEX IF NOT EXISTS request_candidates_provider_id_idx ON request_candidates (provider_id); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_id_idx ON request_candidates (endpoint_id); -CREATE INDEX IF NOT EXISTS request_candidates_status_idx ON request_candidates (status); -CREATE INDEX IF NOT EXISTS request_candidates_created_at_idx ON request_candidates (created_at); -CREATE INDEX IF NOT EXISTS request_candidates_endpoint_status_created_idx ON request_candidates (endpoint_id, status, created_at); - -CREATE TABLE IF NOT EXISTS video_tasks ( - id TEXT PRIMARY KEY NOT NULL, - short_id TEXT, - request_id TEXT NOT NULL, - user_id TEXT, - api_key_id TEXT, - username TEXT, - api_key_name TEXT, - external_task_id TEXT, - provider_id TEXT, - endpoint_id TEXT, - key_id TEXT, - client_api_format TEXT, - provider_api_format TEXT, - format_converted INTEGER NOT NULL DEFAULT 0, - model TEXT, - prompt TEXT, - original_request_body TEXT, - converted_request_body TEXT, - duration_seconds INTEGER, - resolution TEXT, - aspect_ratio TEXT, - size TEXT, - status TEXT NOT NULL DEFAULT 'pending', - progress_percent INTEGER NOT NULL DEFAULT 0, - progress_message TEXT, - retry_count INTEGER NOT NULL DEFAULT 0, - max_retries INTEGER NOT NULL DEFAULT 3, - poll_interval_seconds INTEGER NOT NULL DEFAULT 10, - next_poll_at INTEGER, - poll_count INTEGER NOT NULL DEFAULT 0, - max_poll_count INTEGER NOT NULL DEFAULT 360, - created_at INTEGER NOT NULL, - submitted_at INTEGER, - completed_at INTEGER, - updated_at INTEGER NOT NULL, - error_code TEXT, - error_message TEXT, - video_url TEXT, - video_urls TEXT, - thumbnail_url TEXT, - video_size_bytes INTEGER, - video_expires_at INTEGER, - stored_video_path TEXT, - storage_provider TEXT, - remixed_from_task_id TEXT, - webhook_url TEXT, - webhook_sent INTEGER NOT NULL DEFAULT 0, - webhook_sent_at INTEGER, - request_metadata TEXT, - video_duration_seconds REAL, - UNIQUE (short_id), - UNIQUE (request_id) -); -CREATE INDEX IF NOT EXISTS video_tasks_external_id_idx ON video_tasks (external_task_id); -CREATE INDEX IF NOT EXISTS video_tasks_next_poll_idx ON video_tasks (next_poll_at); -CREATE INDEX IF NOT EXISTS video_tasks_due_poll_idx ON video_tasks (status, next_poll_at, updated_at); -CREATE INDEX IF NOT EXISTS video_tasks_user_status_idx ON video_tasks (user_id, status); -CREATE INDEX IF NOT EXISTS video_tasks_api_key_id_idx ON video_tasks (api_key_id); -CREATE INDEX IF NOT EXISTS video_tasks_provider_id_idx ON video_tasks (provider_id); -CREATE INDEX IF NOT EXISTS video_tasks_endpoint_id_idx ON video_tasks (endpoint_id); -CREATE INDEX IF NOT EXISTS video_tasks_key_id_idx ON video_tasks (key_id); - -CREATE TABLE IF NOT EXISTS provider_endpoints ( - id TEXT PRIMARY KEY NOT NULL, - provider_id TEXT NOT NULL, - name TEXT NOT NULL, - base_url TEXT NOT NULL, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - health_score REAL NOT NULL DEFAULT 1.0, - weight INTEGER NOT NULL DEFAULT 1, - header_rules TEXT, - body_rules TEXT, - max_retries INTEGER, - custom_path TEXT, - metadata TEXT, - config TEXT, - format_acceptance_config TEXT, - proxy TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS provider_endpoints_provider_id_idx ON provider_endpoints (provider_id); - -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); - -CREATE TABLE IF NOT EXISTS models ( - id TEXT PRIMARY KEY NOT NULL, - provider_id TEXT NOT NULL, - global_model_id TEXT, - provider_model_name TEXT NOT NULL, - global_model_name TEXT, - api_format TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - is_available INTEGER NOT NULL DEFAULT 1, - price_per_request REAL, - tiered_pricing TEXT, - supports_vision INTEGER, - supports_function_calling INTEGER, - supports_streaming INTEGER, - supports_extended_thinking INTEGER, - supports_image_generation INTEGER, - provider_model_mappings TEXT, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS models_provider_id_idx ON models (provider_id); - -CREATE TABLE IF NOT EXISTS global_models ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - display_name TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_active INTEGER NOT NULL DEFAULT 1, - default_price_per_request REAL, - default_tiered_pricing TEXT, - supported_capabilities TEXT, - usage_count INTEGER NOT NULL DEFAULT 0, - config TEXT, - metadata TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (name) -); - -CREATE TABLE IF NOT EXISTS routing_groups ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - description TEXT, - enabled INTEGER NOT NULL DEFAULT 1, - is_system_default INTEGER NOT NULL DEFAULT 0, - sort_order INTEGER NOT NULL DEFAULT 0, - config_json TEXT NOT NULL, - version INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - published_at INTEGER, - UNIQUE (name) -); -CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx ON routing_groups (is_system_default, enabled); -CREATE INDEX IF NOT EXISTS routing_groups_enabled_sort_idx ON routing_groups (enabled, sort_order, name, id); - -CREATE TABLE IF NOT EXISTS routing_group_bindings ( - id TEXT PRIMARY KEY NOT NULL, - group_id TEXT NOT NULL, - subject_type TEXT NOT NULL, - subject_id TEXT NOT NULL, - is_default INTEGER NOT NULL DEFAULT 0, - allow_explicit_select INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx ON routing_group_bindings (group_id); -CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx ON routing_group_bindings (subject_type, subject_id); - -CREATE TABLE IF NOT EXISTS routing_group_versions ( - id TEXT PRIMARY KEY NOT NULL, - group_id TEXT NOT NULL, - version INTEGER NOT NULL, - config_json TEXT NOT NULL, - created_at INTEGER NOT NULL, - created_by TEXT, - UNIQUE (group_id, version) -); -CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx ON routing_group_versions (group_id); - diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/003_auth_config.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/003_auth_config.sql deleted file mode 100644 index 7a0f3218b..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/003_auth_config.sql +++ /dev/null @@ -1,78 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS system_configs ( - id TEXT PRIMARY KEY NOT NULL, - key TEXT NOT NULL, - value TEXT NOT NULL, - description TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (key) -); - -CREATE TABLE IF NOT EXISTS auth_modules ( - id TEXT PRIMARY KEY NOT NULL, - module_type TEXT NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - config TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (module_type) -); - -CREATE TABLE IF NOT EXISTS oauth_providers ( - provider_type TEXT PRIMARY KEY NOT NULL, - display_name TEXT NOT NULL, - client_id TEXT NOT NULL, - client_secret_encrypted TEXT, - authorization_url_override TEXT, - token_url_override TEXT, - userinfo_url_override TEXT, - scopes TEXT, - redirect_uri TEXT NOT NULL, - frontend_callback_url TEXT NOT NULL, - attribute_mapping TEXT, - extra_config TEXT, - icon_url TEXT, - is_enabled INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS ldap_configs ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - singleton_key INTEGER NOT NULL DEFAULT 1, - server_url TEXT NOT NULL, - bind_dn TEXT NOT NULL, - bind_password_encrypted TEXT, - base_dn TEXT NOT NULL, - user_search_filter TEXT NOT NULL DEFAULT '(uid={username})', - username_attr TEXT NOT NULL DEFAULT 'uid', - email_attr TEXT NOT NULL DEFAULT 'mail', - display_name_attr TEXT NOT NULL DEFAULT 'cn', - is_enabled INTEGER NOT NULL DEFAULT 0, - is_exclusive INTEGER NOT NULL DEFAULT 0, - use_starttls INTEGER NOT NULL DEFAULT 0, - connect_timeout INTEGER NOT NULL DEFAULT 10, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (singleton_key) -); - -CREATE TABLE IF NOT EXISTS user_oauth_links ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - provider_type TEXT NOT NULL, - provider_user_id TEXT NOT NULL, - provider_username TEXT, - provider_email TEXT, - extra_data TEXT, - linked_at INTEGER NOT NULL, - last_login_at INTEGER, - UNIQUE (provider_type, provider_user_id), - UNIQUE (user_id, provider_type) -); -CREATE INDEX IF NOT EXISTS user_oauth_links_provider_type_idx ON user_oauth_links (provider_type); -CREATE INDEX IF NOT EXISTS user_oauth_links_user_id_idx ON user_oauth_links (user_id); - diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/004_proxy_nodes.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/004_proxy_nodes.sql deleted file mode 100644 index 2e7294ba9..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/004_proxy_nodes.sql +++ /dev/null @@ -1,86 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS proxy_nodes ( - id TEXT PRIMARY KEY NOT NULL, - tunnel_generation TEXT NOT NULL, - name TEXT NOT NULL, - ip TEXT NOT NULL, - port INTEGER NOT NULL, - region TEXT, - status TEXT NOT NULL DEFAULT 'online', - registered_by TEXT, - last_heartbeat_at INTEGER, - heartbeat_interval INTEGER NOT NULL DEFAULT 30, - active_connections INTEGER NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - avg_latency_ms REAL, - is_manual INTEGER NOT NULL DEFAULT 0, - proxy_url TEXT, - proxy_username TEXT, - proxy_password TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - remote_config TEXT, - config_version INTEGER NOT NULL DEFAULT 0, - hardware_info TEXT, - estimated_max_concurrency INTEGER, - tunnel_mode INTEGER NOT NULL DEFAULT 0, - tunnel_connected INTEGER NOT NULL DEFAULT 0, - tunnel_connected_at INTEGER, - failed_requests INTEGER NOT NULL DEFAULT 0, - dns_failures INTEGER NOT NULL DEFAULT 0, - stream_errors INTEGER NOT NULL DEFAULT 0, - proxy_metadata TEXT, - UNIQUE (ip, port) -); - -CREATE TABLE IF NOT EXISTS proxy_node_events ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - node_id TEXT NOT NULL, - event_type TEXT NOT NULL, - detail TEXT, - event_metadata TEXT, - created_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1m ( - node_id TEXT NOT NULL, - bucket_start_unix_secs INTEGER NOT NULL, - samples INTEGER NOT NULL DEFAULT 0, - uptime_samples INTEGER NOT NULL DEFAULT 0, - active_connections_sum INTEGER NOT NULL DEFAULT 0, - active_connections_max INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0, - connect_errors_delta INTEGER NOT NULL DEFAULT 0, - disconnects_delta INTEGER NOT NULL DEFAULT 0, - error_events_delta INTEGER NOT NULL DEFAULT 0, - ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_in_frames_delta INTEGER NOT NULL DEFAULT 0, - ws_out_frames_delta INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs) -); -CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1m_bucket_start ON proxy_node_metrics_1m (bucket_start_unix_secs); - -CREATE TABLE IF NOT EXISTS proxy_node_metrics_1h ( - node_id TEXT NOT NULL, - bucket_start_unix_secs INTEGER NOT NULL, - samples INTEGER NOT NULL DEFAULT 0, - uptime_samples INTEGER NOT NULL DEFAULT 0, - active_connections_sum INTEGER NOT NULL DEFAULT 0, - active_connections_max INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_sum INTEGER NOT NULL DEFAULT 0, - heartbeat_rtt_ms_max INTEGER NOT NULL DEFAULT 0, - connect_errors_delta INTEGER NOT NULL DEFAULT 0, - disconnects_delta INTEGER NOT NULL DEFAULT 0, - error_events_delta INTEGER NOT NULL DEFAULT 0, - ws_in_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_out_bytes_delta INTEGER NOT NULL DEFAULT 0, - ws_in_frames_delta INTEGER NOT NULL DEFAULT 0, - ws_out_frames_delta INTEGER NOT NULL DEFAULT 0, - PRIMARY KEY (node_id, bucket_start_unix_secs) -); -CREATE INDEX IF NOT EXISTS idx_proxy_node_metrics_1h_bucket_start ON proxy_node_metrics_1h (bucket_start_unix_secs); - diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/005_wallet_billing.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/005_wallet_billing.sql deleted file mode 100644 index c1f09090b..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/005_wallet_billing.sql +++ /dev/null @@ -1,315 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS wallets ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT, - api_key_id TEXT, - balance REAL NOT NULL DEFAULT 0, - gift_balance REAL NOT NULL DEFAULT 0, - limit_mode TEXT NOT NULL DEFAULT 'finite', - currency TEXT NOT NULL DEFAULT 'USD', - status TEXT NOT NULL DEFAULT 'active', - total_recharged REAL NOT NULL DEFAULT 0, - total_consumed REAL NOT NULL DEFAULT 0, - total_refunded REAL NOT NULL DEFAULT 0, - total_adjusted REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (user_id), - UNIQUE (api_key_id) -); -CREATE INDEX IF NOT EXISTS wallets_api_key_id_idx ON wallets (api_key_id); -CREATE INDEX IF NOT EXISTS wallets_user_id_idx ON wallets (user_id); - -CREATE TABLE IF NOT EXISTS wallet_transactions ( - id TEXT PRIMARY KEY NOT NULL, - wallet_id TEXT NOT NULL, - category TEXT NOT NULL, - reason_code TEXT NOT NULL, - amount REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - recharge_balance_before REAL NOT NULL, - recharge_balance_after REAL NOT NULL, - gift_balance_before REAL NOT NULL, - gift_balance_after REAL NOT NULL, - link_type TEXT, - link_id TEXT, - operator_id TEXT, - description TEXT, - created_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_wallet_created ON wallet_transactions (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_category_created ON wallet_transactions (category, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_reason_created ON wallet_transactions (reason_code, created_at); -CREATE INDEX IF NOT EXISTS idx_wallet_tx_link ON wallet_transactions (link_type, link_id); -CREATE INDEX IF NOT EXISTS ix_wallet_transactions_operator_id ON wallet_transactions (operator_id); - -CREATE TABLE IF NOT EXISTS wallet_daily_usage_ledgers ( - id TEXT PRIMARY KEY NOT NULL, - wallet_id TEXT NOT NULL, - billing_date TEXT NOT NULL, - billing_timezone TEXT NOT NULL, - total_cost_usd REAL NOT NULL DEFAULT 0, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - first_finalized_at INTEGER, - last_finalized_at INTEGER, - aggregated_at INTEGER NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_wallet_daily_usage_wallet_date ON wallet_daily_usage_ledgers (wallet_id, billing_timezone, billing_date); - -CREATE TABLE IF NOT EXISTS payment_orders ( - id TEXT PRIMARY KEY NOT NULL, - order_no TEXT NOT NULL, - wallet_id TEXT NOT NULL, - user_id TEXT, - amount_usd REAL NOT NULL, - pay_amount REAL, - pay_currency TEXT, - exchange_rate REAL, - refunded_amount_usd REAL NOT NULL DEFAULT 0, - refundable_amount_usd REAL NOT NULL DEFAULT 0, - payment_method TEXT NOT NULL, - payment_provider TEXT, - payment_channel TEXT, - order_kind TEXT NOT NULL DEFAULT 'wallet_recharge', - product_id TEXT, - product_snapshot TEXT, - fulfillment_status TEXT NOT NULL DEFAULT 'pending', - fulfillment_error TEXT, - gateway_order_id TEXT, - gateway_response TEXT, - status TEXT NOT NULL DEFAULT 'pending', - created_at INTEGER NOT NULL, - paid_at INTEGER, - credited_at INTEGER, - expires_at INTEGER, - UNIQUE (order_no) -); -CREATE INDEX IF NOT EXISTS idx_payment_orders_wallet_created ON payment_orders (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_user_created ON payment_orders (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_payment_orders_status ON payment_orders (status); -CREATE INDEX IF NOT EXISTS idx_payment_orders_gateway_order_id ON payment_orders (gateway_order_id); -CREATE UNIQUE INDEX IF NOT EXISTS uq_payment_orders_payment_method_gateway_order_id ON payment_orders (payment_method, gateway_order_id); -CREATE INDEX IF NOT EXISTS idx_payment_orders_kind_status ON payment_orders (order_kind, status); -CREATE INDEX IF NOT EXISTS idx_payment_orders_product ON payment_orders (product_id); - -CREATE TABLE IF NOT EXISTS user_invite_codes ( - user_id TEXT PRIMARY KEY NOT NULL, - invite_code TEXT NOT NULL, - active INTEGER NOT NULL DEFAULT 1, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (invite_code), - CONSTRAINT user_invite_codes_user_id_fkey FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE -); - -CREATE TABLE IF NOT EXISTS user_referrals ( - id TEXT PRIMARY KEY NOT NULL, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL, - invite_code_snapshot TEXT NOT NULL, - source_json TEXT, - first_paid_order_id TEXT, - first_paid_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (invitee_user_id), - CONSTRAINT user_referrals_first_paid_order_fkey FOREIGN KEY (first_paid_order_id) REFERENCES payment_orders (id) ON DELETE SET NULL -); -CREATE INDEX IF NOT EXISTS idx_user_referrals_inviter ON user_referrals (inviter_user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_user_referrals_created ON user_referrals (created_at); -CREATE INDEX IF NOT EXISTS idx_user_referrals_invite_code ON user_referrals (invite_code_snapshot); - -CREATE TABLE IF NOT EXISTS referral_rewards ( - id TEXT PRIMARY KEY NOT NULL, - referral_id TEXT NOT NULL, - inviter_user_id TEXT NOT NULL, - invitee_user_id TEXT NOT NULL, - reward_type TEXT NOT NULL, - trigger_point TEXT NOT NULL, - source_order_id TEXT, - idempotency_key TEXT NOT NULL, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - wallet_transaction_id TEXT, - reversed_amount_usd REAL NOT NULL DEFAULT 0, - pending_reversal_amount_usd REAL NOT NULL DEFAULT 0, - failure_reason TEXT, - admin_operator_id TEXT, - admin_note TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (idempotency_key), - CONSTRAINT referral_rewards_referral_id_fkey FOREIGN KEY (referral_id) REFERENCES user_referrals (id) ON DELETE CASCADE, - CONSTRAINT referral_rewards_source_order_fkey FOREIGN KEY (source_order_id) REFERENCES payment_orders (id) ON DELETE SET NULL -); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_inviter_status ON referral_rewards (inviter_user_id, status, created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_inviter_created ON referral_rewards (inviter_user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_created ON referral_rewards (created_at); -CREATE INDEX IF NOT EXISTS idx_referral_rewards_source_order ON referral_rewards (source_order_id); - -CREATE TABLE IF NOT EXISTS payment_gateway_configs ( - provider TEXT PRIMARY KEY NOT NULL, - enabled INTEGER NOT NULL DEFAULT 0, - endpoint_url TEXT NOT NULL, - callback_base_url TEXT, - merchant_id TEXT NOT NULL, - merchant_key_encrypted TEXT, - pay_currency TEXT NOT NULL DEFAULT 'CNY', - usd_exchange_rate REAL NOT NULL DEFAULT 7.2, - min_recharge_usd REAL NOT NULL DEFAULT 1, - channels_json TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -CREATE TABLE IF NOT EXISTS payment_callbacks ( - id TEXT PRIMARY KEY NOT NULL, - payment_order_id TEXT, - payment_method TEXT NOT NULL, - callback_key TEXT NOT NULL, - order_no TEXT, - gateway_order_id TEXT, - payload_hash TEXT, - signature_valid INTEGER NOT NULL DEFAULT 0, - status TEXT NOT NULL DEFAULT 'received', - payload TEXT, - error_message TEXT, - created_at INTEGER NOT NULL, - processed_at INTEGER, - UNIQUE (callback_key) -); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_order ON payment_callbacks (order_no); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_gateway_order ON payment_callbacks (gateway_order_id); -CREATE INDEX IF NOT EXISTS idx_payment_callbacks_created ON payment_callbacks (created_at); -CREATE INDEX IF NOT EXISTS ix_payment_callbacks_payment_order_id ON payment_callbacks (payment_order_id); - -CREATE TABLE IF NOT EXISTS billing_plans ( - id TEXT PRIMARY KEY NOT NULL, - title TEXT NOT NULL, - description TEXT, - price_amount REAL NOT NULL, - price_currency TEXT NOT NULL DEFAULT 'CNY', - duration_unit TEXT NOT NULL, - duration_value INTEGER NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - sort_order INTEGER NOT NULL DEFAULT 0, - max_active_per_user INTEGER NOT NULL DEFAULT 1, - purchase_limit_scope TEXT NOT NULL DEFAULT 'active_period', - entitlements_json TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_billing_plans_enabled_sort ON billing_plans (enabled, sort_order); - -CREATE TABLE IF NOT EXISTS user_plan_entitlements ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - plan_id TEXT NOT NULL, - payment_order_id TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - starts_at INTEGER NOT NULL, - expires_at INTEGER NOT NULL, - entitlements_snapshot TEXT NOT NULL, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_user_plan_entitlements_user_active ON user_plan_entitlements (user_id, status, expires_at); -CREATE INDEX IF NOT EXISTS idx_user_plan_entitlements_order ON user_plan_entitlements (payment_order_id); - -CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers ( - id TEXT PRIMARY KEY NOT NULL, - user_entitlement_id TEXT NOT NULL, - user_id TEXT NOT NULL, - request_id TEXT NOT NULL, - amount_usd REAL NOT NULL, - balance_before REAL NOT NULL, - balance_after REAL NOT NULL, - usage_date TEXT NOT NULL, - created_at INTEGER NOT NULL, - UNIQUE (user_entitlement_id, request_id) -); -CREATE INDEX IF NOT EXISTS idx_entitlement_usage_user_date ON entitlement_usage_ledgers (user_id, usage_date); -CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date ON entitlement_usage_ledgers (user_entitlement_id, usage_date); - -CREATE TABLE IF NOT EXISTS refund_requests ( - id TEXT PRIMARY KEY NOT NULL, - refund_no TEXT NOT NULL, - wallet_id TEXT NOT NULL, - user_id TEXT, - payment_order_id TEXT, - source_type TEXT NOT NULL, - source_id TEXT, - refund_mode TEXT NOT NULL, - amount_usd REAL NOT NULL, - status TEXT NOT NULL DEFAULT 'pending_approval', - reason TEXT, - requested_by TEXT, - approved_by TEXT, - processed_by TEXT, - gateway_refund_id TEXT, - payout_method TEXT, - payout_reference TEXT, - payout_proof TEXT, - failure_reason TEXT, - idempotency_key TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - processed_at INTEGER, - completed_at INTEGER, - UNIQUE (refund_no), - UNIQUE (idempotency_key) -); -CREATE INDEX IF NOT EXISTS idx_refund_wallet_created ON refund_requests (wallet_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_user_created ON refund_requests (user_id, created_at); -CREATE INDEX IF NOT EXISTS idx_refund_status ON refund_requests (status); -CREATE INDEX IF NOT EXISTS ix_refund_requests_payment_order_id ON refund_requests (payment_order_id); -CREATE INDEX IF NOT EXISTS ix_refund_requests_requested_by ON refund_requests (requested_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_approved_by ON refund_requests (approved_by); -CREATE INDEX IF NOT EXISTS ix_refund_requests_processed_by ON refund_requests (processed_by); - -CREATE TABLE IF NOT EXISTS redeem_code_batches ( - id TEXT PRIMARY KEY NOT NULL, - name TEXT NOT NULL, - amount_usd REAL NOT NULL, - currency TEXT NOT NULL DEFAULT 'USD', - balance_bucket TEXT NOT NULL DEFAULT 'gift', - total_count INTEGER NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - description TEXT, - created_by TEXT, - expires_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_redeem_code_batches_status ON redeem_code_batches (status, created_at); - -CREATE TABLE IF NOT EXISTS redeem_codes ( - id TEXT PRIMARY KEY NOT NULL, - batch_id TEXT NOT NULL, - code_hash TEXT NOT NULL, - code_prefix TEXT NOT NULL, - code_suffix TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'active', - redeemed_by_user_id TEXT, - redeemed_wallet_id TEXT, - redeemed_payment_order_id TEXT, - redeemed_at INTEGER, - disabled_by TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (code_hash) -); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_batch_created ON redeem_codes (batch_id, created_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_status ON redeem_codes (status, updated_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_user ON redeem_codes (redeemed_by_user_id, redeemed_at); -CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_order ON redeem_codes (redeemed_payment_order_id); - 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 deleted file mode 100644 index 890d6915b..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/006_usage.sql +++ /dev/null @@ -1,274 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS "usage" ( - request_id TEXT PRIMARY KEY NOT NULL, - id TEXT, - user_id TEXT, - api_key_id TEXT, - provider_name TEXT NOT NULL DEFAULT 'unknown', - model TEXT NOT NULL DEFAULT 'unknown', - target_model TEXT, - provider_id TEXT, - provider_endpoint_id TEXT, - provider_api_key_id TEXT, - request_type TEXT, - api_format TEXT, - api_family TEXT, - endpoint_kind TEXT, - endpoint_api_format TEXT, - provider_api_family TEXT, - provider_endpoint_kind TEXT, - has_format_conversion INTEGER NOT NULL DEFAULT 0, - is_stream INTEGER NOT NULL DEFAULT 0, - upstream_is_stream INTEGER, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - input_output_total_tokens INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_input_tokens_5m INTEGER NOT NULL DEFAULT 0, - cache_creation_input_tokens_1h INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_5m_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_ephemeral_1h_input_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_input_tokens INTEGER NOT NULL DEFAULT 0, - input_context_tokens INTEGER NOT NULL DEFAULT 0, - input_cost_usd REAL NOT NULL DEFAULT 0, - output_cost_usd REAL NOT NULL DEFAULT 0, - cache_cost_usd REAL NOT NULL DEFAULT 0, - cache_creation_cost_usd REAL NOT NULL DEFAULT 0, - cache_creation_cost_usd_5m REAL NOT NULL DEFAULT 0, - cache_creation_cost_usd_1h REAL NOT NULL DEFAULT 0, - cache_read_cost_usd REAL NOT NULL DEFAULT 0, - request_cost_usd REAL NOT NULL DEFAULT 0, - actual_input_cost_usd REAL NOT NULL DEFAULT 0, - actual_output_cost_usd REAL NOT NULL DEFAULT 0, - actual_cache_cost_usd REAL NOT NULL DEFAULT 0, - actual_cache_creation_cost_usd REAL NOT NULL DEFAULT 0, - actual_cache_creation_cost_usd_5m REAL NOT NULL DEFAULT 0, - actual_cache_creation_cost_usd_1h REAL NOT NULL DEFAULT 0, - actual_cache_read_cost_usd REAL NOT NULL DEFAULT 0, - actual_request_cost_usd REAL NOT NULL DEFAULT 0, - rate_multiplier REAL NOT NULL DEFAULT 1, - input_price_per_1m REAL, - output_price_per_1m REAL, - cache_creation_price_per_1m REAL, - cache_creation_price_per_1m_5m REAL, - cache_creation_price_per_1m_1h REAL, - cache_read_price_per_1m REAL, - price_per_request REAL, - status_code INTEGER, - error_message TEXT, - error_category TEXT, - response_time_ms INTEGER, - first_byte_time_ms INTEGER, - wallet_id TEXT, - status TEXT NOT NULL DEFAULT 'completed', - billing_status TEXT NOT NULL DEFAULT 'pending', - total_cost_usd REAL NOT NULL DEFAULT 0, - actual_total_cost_usd REAL NOT NULL DEFAULT 0, - request_headers TEXT, - request_body TEXT, - provider_request_headers TEXT, - provider_request_body TEXT, - response_headers TEXT, - response_body TEXT, - client_response_headers TEXT, - client_response_body TEXT, - request_body_compressed BLOB, - provider_request_body_compressed BLOB, - response_body_compressed BLOB, - client_response_body_compressed BLOB, - request_metadata TEXT, - created_at INTEGER, - candidate_id TEXT, - candidate_index INTEGER, - key_name TEXT, - username TEXT, - api_key_name TEXT, - planner_kind TEXT, - route_family TEXT, - route_kind TEXT, - execution_path TEXT, - local_execution_runtime_miss_reason TEXT, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - wallet_gift_balance_before REAL, - wallet_gift_balance_after REAL, - finalized_at INTEGER, - created_at_unix_ms INTEGER NOT NULL DEFAULT 0, - updated_at_unix_secs INTEGER NOT NULL DEFAULT 0 -); -CREATE INDEX IF NOT EXISTS usage_api_key_id_idx ON "usage" (api_key_id); -CREATE INDEX IF NOT EXISTS usage_billing_status_idx ON "usage" (billing_status); -CREATE INDEX IF NOT EXISTS usage_created_at_idx ON "usage" (created_at_unix_ms); -CREATE INDEX IF NOT EXISTS usage_provider_api_key_id_idx ON "usage" (provider_api_key_id); -CREATE INDEX IF NOT EXISTS usage_provider_id_idx ON "usage" (provider_id); -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, - kind TEXT NOT NULL, - target_id TEXT NOT NULL, - target_tunnel_generation TEXT, - request_count_delta INTEGER NOT NULL DEFAULT 0, - total_requests_delta INTEGER NOT NULL DEFAULT 0, - success_count_delta INTEGER NOT NULL DEFAULT 0, - error_count_delta INTEGER NOT NULL DEFAULT 0, - dns_failures_delta INTEGER NOT NULL DEFAULT 0, - stream_errors_delta INTEGER NOT NULL DEFAULT 0, - total_tokens_delta INTEGER NOT NULL DEFAULT 0, - total_cost_usd_delta REAL NOT NULL DEFAULT 0, - total_response_time_ms_delta INTEGER NOT NULL DEFAULT 0, - last_used_at_unix_secs INTEGER, - last_used_ip TEXT, - candidate_last_used_at_unix_secs INTEGER, - removed_last_used_at_unix_secs INTEGER, - usage_created_at_unix_secs INTEGER, - created_at INTEGER NOT NULL, - processed_at INTEGER -); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed ON usage_counter_deltas (created_at, id); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed ON usage_counter_deltas (processed_at, created_at, id); -CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind ON usage_counter_deltas (request_id, kind, target_id); - -CREATE TABLE IF NOT EXISTS usage_settlement_snapshots ( - request_id TEXT PRIMARY KEY NOT NULL, - billing_status TEXT NOT NULL, - wallet_id TEXT, - wallet_balance_before REAL, - wallet_balance_after REAL, - wallet_recharge_balance_before REAL, - wallet_recharge_balance_after REAL, - 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); - -CREATE TABLE IF NOT EXISTS usage_cost_reservations ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - reservation_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - reserved_cost_units INTEGER NOT NULL, - actual_cost_units INTEGER, - state TEXT NOT NULL, - reservation_expires_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - finalized_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - CONSTRAINT usage_cost_reservations_subject_id_fkey FOREIGN KEY (subject_id) REFERENCES users (id) ON DELETE CASCADE -); -CREATE INDEX IF NOT EXISTS usage_cost_reservations_request_id_idx ON usage_cost_reservations (request_id); -CREATE INDEX IF NOT EXISTS usage_cost_reservations_subject_admitted_at_idx ON usage_cost_reservations (subject_id, admitted_at); -CREATE INDEX IF NOT EXISTS usage_cost_reservations_reservation_expires_at_idx ON usage_cost_reservations (reservation_expires_at); -CREATE INDEX IF NOT EXISTS usage_cost_reservations_retain_until_token_idx ON usage_cost_reservations (retain_until, reservation_token); - -CREATE TABLE IF NOT EXISTS usage_request_admissions ( - request_id TEXT NOT NULL, - subject_id TEXT NOT NULL, - event_token TEXT PRIMARY KEY NOT NULL, - admitted_at INTEGER NOT NULL, - retain_until INTEGER NOT NULL, - state TEXT NOT NULL, - released_at INTEGER, - created_at INTEGER NOT NULL, - CONSTRAINT usage_request_admissions_subject_id_fkey FOREIGN KEY (subject_id) REFERENCES users (id) ON DELETE CASCADE -); -CREATE INDEX IF NOT EXISTS usage_request_admissions_subject_admitted_at_idx ON usage_request_admissions (subject_id, admitted_at); -CREATE INDEX IF NOT EXISTS usage_request_admissions_retain_until_token_idx ON usage_request_admissions (retain_until, event_token); - 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 deleted file mode 100644 index 60af13fe3..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/007_stats.sql +++ /dev/null @@ -1,572 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS stats_hourly ( - id TEXT PRIMARY KEY NOT NULL, - hour_utc INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - 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) -); - -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 stats_hourly_user ( - id TEXT PRIMARY KEY NOT NULL, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - 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) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_user_model ( - id TEXT PRIMARY KEY NOT NULL, - hour_utc INTEGER NOT NULL, - user_id TEXT NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - 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) -); - -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); - -CREATE TABLE IF NOT EXISTS stats_hourly_model ( - id TEXT PRIMARY KEY NOT NULL, - hour_utc INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - 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) -); - -CREATE TABLE IF NOT EXISTS stats_hourly_provider ( - id TEXT PRIMARY KEY NOT NULL, - hour_utc INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (hour_utc, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily ( - id TEXT PRIMARY KEY NOT NULL, - date INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_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, - input_cost REAL NOT NULL DEFAULT 0, - output_cost REAL NOT NULL DEFAULT 0, - cache_creation_cost REAL NOT NULL DEFAULT 0, - cache_read_cost REAL NOT NULL DEFAULT 0, - avg_response_time_ms REAL NOT NULL DEFAULT 0, - fallback_count INTEGER NOT NULL DEFAULT 0, - unique_models INTEGER NOT NULL DEFAULT 0, - unique_providers INTEGER NOT NULL DEFAULT 0, - is_complete INTEGER NOT NULL DEFAULT 0, - aggregated_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - p50_response_time_ms INTEGER, - p90_response_time_ms INTEGER, - p99_response_time_ms INTEGER, - 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) -); - -CREATE TABLE IF NOT EXISTS stats_daily_model ( - id TEXT PRIMARY KEY NOT NULL, - date INTEGER NOT NULL, - model TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - 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) -); - -CREATE TABLE IF NOT EXISTS stats_daily_provider ( - id TEXT PRIMARY KEY NOT NULL, - date INTEGER NOT NULL, - provider_name TEXT NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (date, provider_name) -); - -CREATE TABLE IF NOT EXISTS stats_daily_api_key ( - id TEXT PRIMARY KEY NOT NULL, - api_key_id TEXT NOT NULL, - date INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - api_key_name TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (date, api_key_id) -); - -CREATE TABLE IF NOT EXISTS stats_daily_error ( - id TEXT PRIMARY KEY NOT NULL, - date INTEGER NOT NULL, - error_category TEXT NOT NULL, - provider_name TEXT, - model TEXT, - count INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (date, error_category, provider_name, model) -); - -CREATE TABLE IF NOT EXISTS stats_user_daily ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL, - date INTEGER NOT NULL, - total_requests INTEGER NOT NULL DEFAULT 0, - success_requests INTEGER NOT NULL DEFAULT 0, - error_requests INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cache_creation_tokens INTEGER NOT NULL DEFAULT 0, - cache_read_tokens INTEGER NOT NULL DEFAULT 0, - total_cost REAL NOT NULL DEFAULT 0, - 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/generated/sqlite/baseline/008_background_tasks.sql b/crates/aether-data/runtime/schema/generated/sqlite/baseline/008_background_tasks.sql deleted file mode 100644 index 1035c86d2..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/008_background_tasks.sql +++ /dev/null @@ -1,40 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -CREATE TABLE IF NOT EXISTS background_task_runs ( - id TEXT PRIMARY KEY NOT NULL, - task_key TEXT NOT NULL, - kind TEXT NOT NULL, - trigger TEXT NOT NULL, - status TEXT NOT NULL, - attempt INTEGER NOT NULL DEFAULT 0, - max_attempts INTEGER NOT NULL DEFAULT 0, - owner_instance TEXT, - progress_percent INTEGER NOT NULL DEFAULT 0, - progress_message TEXT, - payload_json TEXT, - result_json TEXT, - error_message TEXT, - cancel_requested INTEGER NOT NULL DEFAULT 0, - created_by TEXT, - created_at_unix_secs INTEGER NOT NULL, - started_at_unix_secs INTEGER, - finished_at_unix_secs INTEGER, - updated_at_unix_secs INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_task_key ON background_task_runs (task_key); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_status ON background_task_runs (status); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_kind ON background_task_runs (kind); -CREATE INDEX IF NOT EXISTS idx_background_task_runs_created_at ON background_task_runs (created_at_unix_secs); - -CREATE TABLE IF NOT EXISTS background_task_events ( - id TEXT PRIMARY KEY NOT NULL, - run_id TEXT NOT NULL, - event_type TEXT NOT NULL, - message TEXT NOT NULL, - payload_json TEXT, - created_at_unix_secs INTEGER NOT NULL, - CONSTRAINT fk_background_task_events_run FOREIGN KEY (run_id) REFERENCES background_task_runs (id) ON DELETE CASCADE -); -CREATE INDEX IF NOT EXISTS idx_background_task_events_run_id ON background_task_events (run_id, created_at_unix_secs); - diff --git a/crates/aether-data/runtime/schema/generated/sqlite/baseline/manifest.txt b/crates/aether-data/runtime/schema/generated/sqlite/baseline/manifest.txt deleted file mode 100644 index 32ce47e5f..000000000 --- a/crates/aether-data/runtime/schema/generated/sqlite/baseline/manifest.txt +++ /dev/null @@ -1,11 +0,0 @@ --- Generated by aether-data-schema from crates/aether-data/runtime/schema/logical/*.toml. --- Do not edit generated files directly; edit logical schema or explicit overrides instead. - -001_identity.sql -002_provider_catalog.sql -003_auth_config.sql -004_proxy_nodes.sql -005_wallet_billing.sql -006_usage.sql -007_stats.sql -008_background_tasks.sql diff --git a/crates/aether-data/runtime/schema/logical/005_wallet_billing.toml b/crates/aether-data/runtime/schema/logical/005_wallet_billing.toml index d2f6be018..99bda2f71 100644 --- a/crates/aether-data/runtime/schema/logical/005_wallet_billing.toml +++ b/crates/aether-data/runtime/schema/logical/005_wallet_billing.toml @@ -380,9 +380,6 @@ type = "text" length = 128 nullable = true - [table.payment_orders.columns.driver.mysql] -type = "VARCHAR(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin" - [[table.payment_orders.columns]] name = "gateway_response" type = "json" diff --git a/crates/aether-data/runtime/src/backend/maintenance.rs b/crates/aether-data/runtime/src/backend/maintenance.rs index fa465c6d0..3382f7db3 100644 --- a/crates/aether-data/runtime/src/backend/maintenance.rs +++ b/crates/aether-data/runtime/src/backend/maintenance.rs @@ -1,9 +1,5 @@ -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; use super::{summarize_pool, DataBackends, SqlBackendRef}; use crate::maintenance::{ @@ -315,15 +311,7 @@ impl<'a> SqlBackendRef<'a> { Self::Postgres(postgres) => { warm_pool(postgres.pool(), postgres.config().min_connections).await } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - warm_pool(mysql.pool(), mysql.config().pool.min_connections).await - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - warm_pool(sqlite.pool(), sqlite.config().pool.min_connections).await - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -335,11 +323,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.run_table_maintenance(table_names).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.run_table_maintenance(table_names).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.run_table_maintenance(table_names).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -351,17 +335,7 @@ impl<'a> SqlBackendRef<'a> { crate::lifecycle::migrate::run_migrations(postgres.pool()).await?; Ok(true) } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - crate::lifecycle::migrate::run_mysql_migrations(mysql.pool()).await?; - Ok(true) - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - crate::lifecycle::migrate::run_sqlite_migrations(sqlite.pool()).await?; - Ok(true) - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -373,17 +347,7 @@ impl<'a> SqlBackendRef<'a> { crate::lifecycle::backfill::run_backfills(postgres.pool()).await?; Ok(true) } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - crate::lifecycle::backfill::run_mysql_backfills(mysql.pool()).await?; - Ok(true) - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - crate::lifecycle::backfill::run_sqlite_backfills(sqlite.pool()).await?; - Ok(true) - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -396,15 +360,7 @@ impl<'a> SqlBackendRef<'a> { Self::Postgres(postgres) => Ok(Some( crate::lifecycle::migrate::pending_migrations(postgres.pool()).await?, )), - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => Ok(Some( - crate::lifecycle::migrate::pending_mysql_migrations(mysql.pool()).await?, - )), - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => Ok(Some( - crate::lifecycle::migrate::pending_sqlite_migrations(sqlite.pool()).await?, - )), - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -417,16 +373,7 @@ impl<'a> SqlBackendRef<'a> { Self::Postgres(postgres) => Ok(Some( crate::lifecycle::migrate::prepare_database_for_startup(postgres.pool()).await?, )), - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => Ok(Some( - crate::lifecycle::migrate::prepare_mysql_database_for_startup(mysql.pool()).await?, - )), - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => Ok(Some( - crate::lifecycle::migrate::prepare_sqlite_database_for_startup(sqlite.pool()) - .await?, - )), - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -439,15 +386,7 @@ impl<'a> SqlBackendRef<'a> { Self::Postgres(postgres) => Ok(Some( crate::lifecycle::backfill::pending_backfills(postgres.pool()).await?, )), - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => Ok(Some( - crate::lifecycle::backfill::pending_mysql_backfills(mysql.pool()).await?, - )), - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => Ok(Some( - crate::lifecycle::backfill::pending_sqlite_backfills(sqlite.pool()).await?, - )), - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -461,21 +400,7 @@ impl<'a> SqlBackendRef<'a> { postgres.pool().num_idle(), postgres.config().max_connections, ), - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => summarize_pool( - crate::database::DatabaseDriver::Mysql, - usize::try_from(mysql.pool().size()).unwrap_or(usize::MAX), - mysql.pool().num_idle(), - mysql.config().pool.max_connections, - ), - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => summarize_pool( - crate::database::DatabaseDriver::Sqlite, - usize::try_from(sqlite.pool().size()).unwrap_or(usize::MAX), - sqlite.pool().num_idle(), - sqlite.config().pool.max_connections, - ), - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -487,11 +412,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.aggregate_wallet_daily_usage(input).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.aggregate_wallet_daily_usage(input).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.aggregate_wallet_daily_usage(input).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -503,11 +424,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.aggregate_stats_hourly(input).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.aggregate_stats_hourly(input).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.aggregate_stats_hourly(input).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -519,11 +436,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.aggregate_stats_daily(input).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.aggregate_stats_daily(input).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.aggregate_stats_daily(input).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -535,11 +448,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.find_system_config_value(key).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.find_system_config_value(key).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.find_system_config_value(key).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -557,19 +466,7 @@ impl<'a> SqlBackendRef<'a> { .compare_and_set_system_config_string_value(key, expected, replacement) .await } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - mysql - .compare_and_set_system_config_string_value(key, expected, replacement) - .await - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - sqlite - .compare_and_set_system_config_string_value(key, expected, replacement) - .await - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -580,11 +477,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.list_system_config_entries().await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.list_system_config_entries().await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.list_system_config_entries().await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -602,19 +495,7 @@ impl<'a> SqlBackendRef<'a> { .upsert_system_config_entry(key, value, description) .await } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - mysql - .upsert_system_config_entry(key, value, description) - .await - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - sqlite - .upsert_system_config_entry(key, value, description) - .await - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -623,11 +504,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.delete_system_config_value(key).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.delete_system_config_value(key).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.delete_system_config_value(key).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -636,11 +513,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.read_admin_system_stats().await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.read_admin_system_stats().await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.read_admin_system_stats().await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -652,11 +525,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.purge_admin_system_data(target).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.purge_admin_system_data(target).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.purge_admin_system_data(target).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -667,11 +536,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.export_admin_system_usage_aggregates().await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.export_admin_system_usage_aggregates().await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.export_admin_system_usage_aggregates().await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -695,29 +560,7 @@ impl<'a> SqlBackendRef<'a> { ) .await } - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => { - mysql - .import_admin_system_usage_aggregates( - snapshot, - user_id_map, - api_key_id_map, - mode, - ) - .await - } - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => { - sqlite - .import_admin_system_usage_aggregates( - snapshot, - user_id_map, - api_key_id_map, - mode, - ) - .await - } - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } @@ -729,11 +572,7 @@ impl<'a> SqlBackendRef<'a> { match self { #[cfg(feature = "postgres")] Self::Postgres(postgres) => postgres.purge_admin_request_bodies_batch(batch_size).await, - #[cfg(feature = "mysql")] - Self::Mysql(mysql) => mysql.purge_admin_request_bodies_batch(batch_size).await, - #[cfg(feature = "sqlite")] - Self::Sqlite(sqlite) => sqlite.purge_admin_request_bodies_batch(batch_size).await, - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Self::Disabled(_) => unreachable!("a SQL backend cannot exist without a driver"), } } diff --git a/crates/aether-data/runtime/src/backend/maintenance/mysql.rs b/crates/aether-data/runtime/src/backend/maintenance/mysql.rs deleted file mode 100644 index 87ecbeeb0..000000000 --- a/crates/aether-data/runtime/src/backend/maintenance/mysql.rs +++ /dev/null @@ -1,67 +0,0 @@ -use crate::backend::MysqlBackend; -use crate::error::SqlResultExt; -use crate::{DataLayerError, DatabaseMaintenanceSummary}; - -use super::maintenance_identifier; - -impl MysqlBackend { - pub async fn run_table_maintenance( - &self, - table_names: &[&str], - ) -> Result { - let mut summary = DatabaseMaintenanceSummary::default(); - for table_name in table_names { - let table_name = maintenance_identifier(table_name)?; - summary.attempted += 1; - let statement = format!("ANALYZE TABLE `{table_name}`"); - if sqlx::query_as::<_, (String, String, String, String)>(&statement) - .fetch_all(self.pool()) - .await - .map_sql_err() - .is_ok_and(|rows| mysql_analyze_succeeded(&rows)) - { - summary.succeeded += 1; - } - } - 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/maintenance/sqlite.rs b/crates/aether-data/runtime/src/backend/maintenance/sqlite.rs deleted file mode 100644 index b880e83c0..000000000 --- a/crates/aether-data/runtime/src/backend/maintenance/sqlite.rs +++ /dev/null @@ -1,34 +0,0 @@ -use crate::backend::SqliteBackend; -use crate::error::SqlResultExt; -use crate::{DataLayerError, DatabaseMaintenanceSummary}; - -use super::maintenance_identifier; - -impl SqliteBackend { - pub async fn run_table_maintenance( - &self, - table_names: &[&str], - ) -> Result { - let mut summary = DatabaseMaintenanceSummary::default(); - for table_name in table_names { - let table_name = maintenance_identifier(table_name)?; - summary.attempted += 1; - let statement = format!("ANALYZE \"{table_name}\""); - if sqlx::raw_sql(&statement) - .execute(self.pool()) - .await - .map_sql_err() - .is_ok() - { - summary.succeeded += 1; - } - } - if summary.succeeded > 0 { - sqlx::raw_sql("PRAGMA optimize") - .execute(self.pool()) - .await - .map_sql_err()?; - } - Ok(summary) - } -} diff --git a/crates/aether-data/runtime/src/backend/mod.rs b/crates/aether-data/runtime/src/backend/mod.rs index af170ed4b..503b39ecc 100644 --- a/crates/aether-data/runtime/src/backend/mod.rs +++ b/crates/aether-data/runtime/src/backend/mod.rs @@ -8,17 +8,11 @@ mod leases; mod maintenance; -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; mod read; mod referrals; -#[cfg(feature = "sqlite")] -mod sqlite; mod stats; -#[cfg(any(feature = "mysql", feature = "sqlite"))] -mod stats_common; mod system; mod transactions; mod wallet; @@ -26,8 +20,6 @@ mod write; use crate::maintenance::DatabasePoolSummary; pub use leases::DataLeaseBackends; -#[cfg(feature = "mysql")] -pub use mysql::MysqlBackend; #[cfg(feature = "postgres")] pub use postgres::PostgresBackend; pub use read::DataReadRepositories; @@ -36,8 +28,6 @@ pub use referrals::{ ReferralRelationshipListQuery, ReferralRelationshipRecord, ReferralRewardConfig, ReferralRewardListQuery, ReferralRewardRecord, ReferralUserDashboard, }; -#[cfg(feature = "sqlite")] -pub use sqlite::SqliteBackend; pub use transactions::DataTransactionBackends; pub use write::DataWriteRepositories; @@ -48,14 +38,10 @@ use crate::{DataLayerConfig, DataLayerError}; enum SqlBackendRef<'a> { #[cfg(feature = "postgres")] Postgres(&'a PostgresBackend), - #[cfg(feature = "mysql")] - Mysql(&'a MysqlBackend), - #[cfg(feature = "sqlite")] - Sqlite(&'a SqliteBackend), // Keep the reference lifetime represented when this crate is built without // any SQL driver features. The no-driver build still exposes the // maintenance facade, but has no concrete backend variant to carry `'a`. - #[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))] + #[cfg(not(feature = "postgres"))] Disabled(std::marker::PhantomData<&'a ()>), } @@ -64,10 +50,6 @@ pub struct DataBackends { config: DataLayerConfig, #[cfg(feature = "postgres")] postgres: Option, - #[cfg(feature = "mysql")] - mysql: Option, - #[cfg(feature = "sqlite")] - sqlite: Option, leases: DataLeaseBackends, read: DataReadRepositories, transactions: DataTransactionBackends, @@ -102,18 +84,6 @@ fn ensure_driver_enabled(driver: DatabaseDriver) -> Result<(), DataLayerError> { DatabaseDriver::Postgres => Err(DataLayerError::InvalidInput( "PostgreSQL driver is not enabled for this aether-data build".to_string(), )), - #[cfg(feature = "mysql")] - DatabaseDriver::Mysql => Ok(()), - #[cfg(not(feature = "mysql"))] - DatabaseDriver::Mysql => Err(DataLayerError::InvalidInput( - "MySQL driver is not enabled for this aether-data build".to_string(), - )), - #[cfg(feature = "sqlite")] - DatabaseDriver::Sqlite => Ok(()), - #[cfg(not(feature = "sqlite"))] - DatabaseDriver::Sqlite => Err(DataLayerError::InvalidInput( - "SQLite driver is not enabled for this aether-data build".to_string(), - )), } } @@ -123,14 +93,6 @@ impl DataBackends { if let Some(postgres) = self.postgres.as_ref() { return Some(SqlBackendRef::Postgres(postgres)); } - #[cfg(feature = "mysql")] - if let Some(mysql) = self.mysql.as_ref() { - return Some(SqlBackendRef::Mysql(mysql)); - } - #[cfg(feature = "sqlite")] - if let Some(sqlite) = self.sqlite.as_ref() { - return Some(SqlBackendRef::Sqlite(sqlite)); - } None } @@ -148,20 +110,6 @@ impl DataBackends { ), _ => None, }; - #[cfg(feature = "mysql")] - let mysql = match database.clone() { - Some(database) if database.driver == DatabaseDriver::Mysql => { - Some(MysqlBackend::from_config(database)?) - } - _ => None, - }; - #[cfg(feature = "sqlite")] - let sqlite = match database.clone() { - Some(database) if database.driver == DatabaseDriver::Sqlite => { - Some(SqliteBackend::from_config(database)?) - } - _ => None, - }; #[cfg(feature = "postgres")] let leases = DataLeaseBackends::from_postgres(postgres.as_ref())?; #[cfg(not(feature = "postgres"))] @@ -169,10 +117,6 @@ impl DataBackends { let read = DataReadRepositories::from_backends( #[cfg(feature = "postgres")] postgres.as_ref(), - #[cfg(feature = "mysql")] - mysql.as_ref(), - #[cfg(feature = "sqlite")] - sqlite.as_ref(), ); #[cfg(feature = "postgres")] let transactions = DataTransactionBackends::from_postgres(postgres.as_ref()); @@ -181,20 +125,12 @@ impl DataBackends { let write = DataWriteRepositories::from_backends( #[cfg(feature = "postgres")] postgres.as_ref(), - #[cfg(feature = "mysql")] - mysql.as_ref(), - #[cfg(feature = "sqlite")] - sqlite.as_ref(), ); Ok(Self { config, #[cfg(feature = "postgres")] postgres, - #[cfg(feature = "mysql")] - mysql, - #[cfg(feature = "sqlite")] - sqlite, leases, read, transactions, @@ -217,16 +153,6 @@ impl DataBackends { .map(|database| database.driver) } - #[cfg(feature = "mysql")] - pub fn mysql(&self) -> Option<&MysqlBackend> { - self.mysql.as_ref() - } - - #[cfg(feature = "sqlite")] - pub fn sqlite(&self) -> Option<&SqliteBackend> { - self.sqlite.as_ref() - } - pub fn read(&self) -> &DataReadRepositories { &self.read } @@ -256,31 +182,7 @@ mod tests { use super::DataBackends; #[cfg(feature = "postgres")] use crate::driver::postgres::PostgresPoolConfig; - use crate::{DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; - - #[test] - #[cfg(not(feature = "mysql"))] - fn rejects_mysql_when_driver_is_not_enabled() { - let error = DataBackends::from_config(DataLayerConfig::from_database(SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "mysql://user:pass@localhost/aether".to_string(), - pool: SqlPoolConfig::default(), - })) - .expect_err("disabled mysql should fail explicitly"); - assert!(error.to_string().contains("MySQL driver is not enabled")); - } - - #[test] - #[cfg(not(feature = "sqlite"))] - fn rejects_sqlite_when_driver_is_not_enabled() { - let error = DataBackends::from_config(DataLayerConfig::from_database(SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite://./data/aether.db".to_string(), - pool: SqlPoolConfig::default(), - })) - .expect_err("disabled sqlite should fail explicitly"); - assert!(error.to_string().contains("SQLite driver is not enabled")); - } + use crate::DataLayerConfig; #[test] fn builds_empty_backends_from_default_config() { @@ -290,10 +192,6 @@ mod tests { assert!(!backends.has_runtime_backends()); #[cfg(feature = "postgres")] assert!(backends.postgres().is_none()); - #[cfg(feature = "mysql")] - assert!(backends.mysql().is_none()); - #[cfg(feature = "sqlite")] - assert!(backends.sqlite().is_none()); #[cfg(feature = "postgres")] assert!(backends.leases().postgres().is_none()); assert!(backends.read().auth_api_keys().is_none()); @@ -336,10 +234,6 @@ mod tests { assert!(backends.has_runtime_backends()); #[cfg(feature = "postgres")] assert!(backends.postgres().is_some()); - #[cfg(feature = "mysql")] - assert!(backends.mysql().is_none()); - #[cfg(feature = "sqlite")] - assert!(backends.sqlite().is_none()); #[cfg(feature = "postgres")] assert!(backends.leases().postgres().is_some()); assert!(backends.read().auth_api_keys().is_some()); @@ -371,113 +265,4 @@ mod tests { assert!(backends.write().wallets().is_some()); assert!(backends.config().effective_database().is_some()); } - - #[tokio::test] - #[cfg(feature = "mysql")] - async fn builds_mysql_backend_from_database_config_with_first_core_repository() { - let backends = DataBackends::from_config(DataLayerConfig { - database: Some(SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "mysql://user:pass@localhost:3306/aether".to_string(), - pool: SqlPoolConfig::default(), - }), - postgres: None, - }) - .expect("mysql backend should build"); - - assert!(backends.has_runtime_backends()); - #[cfg(feature = "postgres")] - assert!(backends.postgres().is_none()); - #[cfg(feature = "mysql")] - assert!(backends.mysql().is_some()); - #[cfg(feature = "sqlite")] - assert!(backends.sqlite().is_none()); - assert!(backends.read().has_any()); - assert!(backends.read().announcements().is_some()); - assert!(backends.read().auth_api_keys().is_some()); - assert!(backends.read().auth_modules().is_some()); - assert!(backends.read().billing().is_some()); - assert!(backends.read().gemini_file_mappings().is_some()); - assert!(backends.read().global_models().is_some()); - assert!(backends.read().management_tokens().is_some()); - assert!(backends.read().minimal_candidate_selection().is_some()); - assert!(backends.read().oauth_providers().is_some()); - assert!(backends.read().provider_catalog().is_some()); - assert!(backends.read().provider_quotas().is_some()); - assert!(backends.read().proxy_nodes().is_some()); - assert!(backends.read().request_candidates().is_some()); - assert!(backends.read().users().is_some()); - assert!(backends.read().video_tasks().is_some()); - assert!(backends.has_stats_hourly_aggregation_backend()); - assert!(backends.has_stats_daily_aggregation_backend()); - assert!(backends.write().has_any()); - assert!(backends.write().announcements().is_some()); - assert!(backends.write().auth_api_keys().is_some()); - assert!(backends.write().auth_modules().is_some()); - assert!(backends.write().gemini_file_mappings().is_some()); - assert!(backends.write().global_models().is_some()); - assert!(backends.write().management_tokens().is_some()); - assert!(backends.write().oauth_providers().is_some()); - assert!(backends.write().proxy_nodes().is_some()); - assert!(backends.write().provider_catalog().is_some()); - assert!(backends.write().provider_quotas().is_some()); - assert!(backends.write().request_candidates().is_some()); - assert!(backends.write().video_tasks().is_some()); - assert!(backends.write().wallets().is_some()); - assert!(backends.config().effective_database().is_some()); - } - - #[tokio::test] - #[cfg(feature = "sqlite")] - async fn builds_sqlite_backend_from_database_config_with_first_core_repository() { - let backends = DataBackends::from_config(DataLayerConfig { - database: Some(SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite://./data/aether.db".to_string(), - pool: SqlPoolConfig::default(), - }), - postgres: None, - }) - .expect("sqlite backend should build"); - - assert!(backends.has_runtime_backends()); - #[cfg(feature = "postgres")] - assert!(backends.postgres().is_none()); - #[cfg(feature = "mysql")] - assert!(backends.mysql().is_none()); - #[cfg(feature = "sqlite")] - assert!(backends.sqlite().is_some()); - assert!(backends.read().has_any()); - assert!(backends.read().announcements().is_some()); - assert!(backends.read().auth_api_keys().is_some()); - assert!(backends.read().auth_modules().is_some()); - assert!(backends.read().billing().is_some()); - assert!(backends.read().gemini_file_mappings().is_some()); - assert!(backends.read().global_models().is_some()); - assert!(backends.read().management_tokens().is_some()); - assert!(backends.read().oauth_providers().is_some()); - assert!(backends.read().provider_catalog().is_some()); - assert!(backends.read().provider_quotas().is_some()); - assert!(backends.read().proxy_nodes().is_some()); - assert!(backends.read().request_candidates().is_some()); - assert!(backends.read().users().is_some()); - assert!(backends.read().video_tasks().is_some()); - assert!(backends.has_stats_hourly_aggregation_backend()); - assert!(backends.has_stats_daily_aggregation_backend()); - assert!(backends.write().has_any()); - assert!(backends.write().announcements().is_some()); - assert!(backends.write().auth_api_keys().is_some()); - assert!(backends.write().auth_modules().is_some()); - assert!(backends.write().gemini_file_mappings().is_some()); - assert!(backends.write().global_models().is_some()); - assert!(backends.write().management_tokens().is_some()); - assert!(backends.write().oauth_providers().is_some()); - assert!(backends.write().proxy_nodes().is_some()); - assert!(backends.write().provider_catalog().is_some()); - assert!(backends.write().provider_quotas().is_some()); - assert!(backends.write().request_candidates().is_some()); - assert!(backends.write().video_tasks().is_some()); - assert!(backends.write().wallets().is_some()); - assert!(backends.config().effective_database().is_some()); - } } diff --git a/crates/aether-data/runtime/src/backend/mysql.rs b/crates/aether-data/runtime/src/backend/mysql.rs deleted file mode 100644 index 9ee96bc02..000000000 --- a/crates/aether-data/runtime/src/backend/mysql.rs +++ /dev/null @@ -1,703 +0,0 @@ -use std::sync::Arc; - -use crate::database::SqlDatabaseConfig; -use crate::driver::mysql::{MysqlPool, MysqlPoolFactory}; -use crate::repository::announcements::{ - AnnouncementReadRepository, AnnouncementWriteRepository, MysqlAnnouncementRepository, -}; -use crate::repository::audit::{AuditLogReadRepository, MysqlAuditLogReadRepository}; -use crate::repository::auth::{ - AuthApiKeyReadRepository, AuthApiKeyWriteRepository, MysqlAuthApiKeyReadRepository, -}; -use crate::repository::auth_modules::{ - AuthModuleReadRepository, AuthModuleWriteRepository, MysqlAuthModuleReadRepository, - MysqlAuthModuleRepository, -}; -use crate::repository::background_tasks::{ - BackgroundTaskReadRepository, BackgroundTaskWriteRepository, MysqlBackgroundTaskRepository, -}; -use crate::repository::billing::{BillingReadRepository, MysqlBillingReadRepository}; -use crate::repository::candidate_selection::{ - MinimalCandidateSelectionReadRepository, MysqlMinimalCandidateSelectionReadRepository, -}; -use crate::repository::candidates::{ - MysqlRequestCandidateRepository, RequestCandidateReadRepository, - RequestCandidateWriteRepository, -}; -use crate::repository::gemini_file_mappings::{ - GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository, - MysqlGeminiFileMappingRepository, -}; -use crate::repository::global_models::{ - GlobalModelReadRepository, GlobalModelWriteRepository, MysqlGlobalModelReadRepository, -}; -use crate::repository::management_tokens::{ - ManagementTokenReadRepository, ManagementTokenWriteRepository, MysqlManagementTokenRepository, -}; -use crate::repository::oauth_providers::{ - MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository, -}; -use crate::repository::pool_scores::{ - MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository, -}; -use crate::repository::provider_catalog::{ - MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, -}; -use crate::repository::proxy_nodes::{ - MysqlProxyNodeReadRepository, ProxyNodeReadRepository, ProxyNodeWriteRepository, -}; -use crate::repository::quota::{ - MysqlProviderQuotaRepository, ProviderQuotaReadRepository, ProviderQuotaWriteRepository, -}; -use crate::repository::routing_profiles::{ - MysqlRoutingGroupRepository, RoutingGroupReadRepository, RoutingGroupWriteRepository, -}; -use crate::repository::settlement::{MysqlSettlementRepository, SettlementWriteRepository}; -use crate::repository::usage::{ - MysqlUsageReadRepository, MysqlUsageWriteRepository, UsageReadRepository, UsageWriteRepository, -}; -use crate::repository::users::{MysqlUserReadRepository, UserReadRepository}; -use crate::repository::video_tasks::{ - MysqlVideoTaskRepository, VideoTaskReadRepository, VideoTaskWriteRepository, -}; -use crate::repository::wallet::{ - MysqlWalletReadRepository, WalletReadRepository, WalletWriteRepository, -}; -use crate::DataLayerError; - -#[derive(Debug, Clone)] -pub struct MysqlBackend { - config: SqlDatabaseConfig, - pool: MysqlPool, -} - -impl MysqlBackend { - pub fn from_config(config: SqlDatabaseConfig) -> Result { - let factory = MysqlPoolFactory::new(config.clone())?; - let pool = factory.connect_lazy()?; - - Ok(Self { config, pool }) - } - - pub fn config(&self) -> &SqlDatabaseConfig { - &self.config - } - - pub fn pool(&self) -> &MysqlPool { - &self.pool - } - - pub fn pool_clone(&self) -> MysqlPool { - self.pool.clone() - } - - pub fn auth_api_key_read_repository(&self) -> Arc { - Arc::new(MysqlAuthApiKeyReadRepository::new(self.pool_clone())) - } - - pub fn announcement_read_repository(&self) -> Arc { - Arc::new(MysqlAnnouncementRepository::new(self.pool_clone())) - } - - pub fn audit_log_read_repository(&self) -> Arc { - Arc::new(MysqlAuditLogReadRepository::new(self.pool_clone())) - } - - pub fn announcement_write_repository(&self) -> Arc { - Arc::new(MysqlAnnouncementRepository::new(self.pool_clone())) - } - - pub fn auth_api_key_write_repository(&self) -> Arc { - Arc::new(MysqlAuthApiKeyReadRepository::new(self.pool_clone())) - } - - pub fn management_token_read_repository(&self) -> Arc { - Arc::new(MysqlManagementTokenRepository::new(self.pool_clone())) - } - - pub fn management_token_write_repository(&self) -> Arc { - Arc::new(MysqlManagementTokenRepository::new(self.pool_clone())) - } - - pub fn auth_module_read_repository(&self) -> Arc { - Arc::new(MysqlAuthModuleReadRepository::new(self.pool_clone())) - } - - pub fn auth_module_write_repository(&self) -> Arc { - Arc::new(MysqlAuthModuleRepository::new(self.pool_clone())) - } - - pub fn billing_read_repository(&self) -> Arc { - Arc::new(MysqlBillingReadRepository::new(self.pool_clone())) - } - - pub fn background_task_read_repository(&self) -> Arc { - Arc::new(MysqlBackgroundTaskRepository::new(self.pool_clone())) - } - - pub fn background_task_write_repository(&self) -> Arc { - Arc::new(MysqlBackgroundTaskRepository::new(self.pool_clone())) - } - - pub fn request_candidate_read_repository(&self) -> Arc { - Arc::new(MysqlRequestCandidateRepository::new(self.pool_clone())) - } - - pub fn request_candidate_write_repository(&self) -> Arc { - Arc::new(MysqlRequestCandidateRepository::new(self.pool_clone())) - } - - pub fn minimal_candidate_selection_read_repository( - &self, - ) -> Arc { - Arc::new(MysqlMinimalCandidateSelectionReadRepository::new( - self.pool_clone(), - )) - } - - pub fn gemini_file_mapping_read_repository(&self) -> Arc { - Arc::new(MysqlGeminiFileMappingRepository::new(self.pool_clone())) - } - - pub fn gemini_file_mapping_write_repository( - &self, - ) -> Arc { - Arc::new(MysqlGeminiFileMappingRepository::new(self.pool_clone())) - } - - pub fn global_model_read_repository(&self) -> Arc { - Arc::new(MysqlGlobalModelReadRepository::new(self.pool_clone())) - } - - pub fn global_model_write_repository(&self) -> Arc { - Arc::new(MysqlGlobalModelReadRepository::new(self.pool_clone())) - } - - pub fn oauth_provider_read_repository(&self) -> Arc { - Arc::new(MysqlOAuthProviderRepository::new(self.pool_clone())) - } - - pub fn oauth_provider_write_repository(&self) -> Arc { - Arc::new(MysqlOAuthProviderRepository::new(self.pool_clone())) - } - - pub fn provider_catalog_read_repository(&self) -> Arc { - Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone())) - } - - pub fn provider_catalog_write_repository(&self) -> Arc { - Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone())) - } - - pub fn pool_score_read_repository(&self) -> Arc { - Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone())) - } - - pub fn pool_score_write_repository(&self) -> Arc { - Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone())) - } - - pub fn routing_group_read_repository(&self) -> Arc { - Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone())) - } - - pub fn routing_group_write_repository(&self) -> Arc { - Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone())) - } - - pub fn proxy_node_read_repository(&self) -> Arc { - Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone())) - } - - pub fn proxy_node_write_repository(&self) -> Arc { - Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone())) - } - - pub fn provider_quota_read_repository(&self) -> Arc { - Arc::new(MysqlProviderQuotaRepository::new(self.pool_clone())) - } - - pub fn provider_quota_write_repository(&self) -> Arc { - Arc::new(MysqlProviderQuotaRepository::new(self.pool_clone())) - } - - pub fn settlement_write_repository(&self) -> Arc { - Arc::new(MysqlSettlementRepository::new(self.pool_clone())) - } - - pub fn usage_write_repository(&self) -> Arc { - Arc::new(MysqlUsageWriteRepository::new(self.pool_clone())) - } - - pub fn usage_read_repository(&self) -> Arc { - Arc::new(MysqlUsageReadRepository::new(self.pool_clone())) - } - - pub fn user_read_repository(&self) -> Arc { - Arc::new(MysqlUserReadRepository::new(self.pool_clone())) - } - - pub fn video_task_read_repository(&self) -> Arc { - Arc::new(MysqlVideoTaskRepository::new(self.pool_clone())) - } - - pub fn video_task_write_repository(&self) -> Arc { - Arc::new(MysqlVideoTaskRepository::new(self.pool_clone())) - } - - pub fn wallet_read_repository(&self) -> Arc { - Arc::new(MysqlWalletReadRepository::new(self.pool_clone())) - } - - pub fn wallet_write_repository(&self) -> Arc { - Arc::new(MysqlWalletReadRepository::new(self.pool_clone())) - } -} - -#[cfg(test)] -mod tests { - use super::MysqlBackend; - use crate::lifecycle::migrate::run_mysql_migrations; - use crate::{ - DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, StatsDailyAggregationInput, - StatsHourlyAggregationInput, WalletDailyUsageAggregationInput, - }; - - #[tokio::test] - async fn backend_retains_config_and_pool() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Mysql, - url: "mysql://user:pass@localhost:3306/aether".to_string(), - pool: SqlPoolConfig::default(), - }; - - let backend = MysqlBackend::from_config(config.clone()).expect("backend should build"); - - assert_eq!(backend.config(), &config); - let _pool = backend.pool(); - let _pool_clone = backend.pool_clone(); - } - - #[tokio::test] - async fn mysql_wallet_daily_usage_aggregation_uses_settlement_wallets_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 wallet daily usage aggregation smoke 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 backend = MysqlBackend::from_config(config).expect("backend should build"); - run_mysql_migrations(backend.pool()) - .await - .expect("mysql migrations should run"); - - let suffix = format!( - "{}-{}", - std::process::id(), - chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default() - ); - let wallet_id = format!("wallet-daily-{suffix}"); - let stale_wallet_id = format!("wallet-daily-stale-{suffix}"); - let timezone = format!("Test/WalletDaily/{suffix}"); - let request_one = format!("request-daily-1-{suffix}"); - let request_two = format!("request-daily-2-{suffix}"); - let request_zero = format!("request-daily-zero-{suffix}"); - let request_outside = format!("request-daily-outside-{suffix}"); - let stale_ledger_id = format!("stale-ledger-{suffix}"); - let unique_offset = chrono::Utc::now() - .timestamp_nanos_opt() - .unwrap_or_default() - .rem_euclid(10_000_000); - let window_start = 4_100_000_000_i64 + unique_offset * 1_000; - let window_end = window_start + 200; - let first_finalized_at = window_start; - let last_finalized_at = window_start + 100; - let zero_finalized_at = window_start + 150; - let outside_finalized_at = window_end; - let seed_created_at = window_start - 100; - let aggregated_at = window_end + 100; - - sqlx::query( - r#" -INSERT INTO wallets (id, user_id, balance, gift_balance, limit_mode, created_at, updated_at) -VALUES - (?, ?, 10.0, 2.0, 'finite', 1, 1), - (?, ?, 0.0, 0.0, 'finite', 1, 1) -"#, - ) - .bind(&wallet_id) - .bind(format!("user-{wallet_id}")) - .bind(&stale_wallet_id) - .bind(format!("user-{stale_wallet_id}")) - .execute(backend.pool()) - .await - .expect("wallets should seed"); - - sqlx::query( - r#" -INSERT INTO `usage` ( - request_id, wallet_id, provider_name, model, status, billing_status, - total_cost_usd, input_tokens, output_tokens, cache_creation_input_tokens, - cache_read_input_tokens, finalized_at, created_at_unix_ms, updated_at_unix_secs -) VALUES - (?, 'wrong-wallet', 'provider', 'model', 'completed', 'pending', - 1.25, 10, 20, 3, 4, ?, ?, ?), - (?, NULL, 'provider', 'model', 'completed', 'pending', - 2.00, 5, 7, 1, 2, ?, ?, ?), - (?, NULL, 'provider', 'model', 'completed', 'pending', - 0.00, 100, 100, 0, 0, ?, ?, ?), - (?, NULL, 'provider', 'model', 'completed', 'pending', - 9.00, 50, 50, 0, 0, ?, ?, ?) -"#, - ) - .bind(&request_one) - .bind(seed_created_at) - .bind(seed_created_at * 1000) - .bind(seed_created_at) - .bind(&request_two) - .bind(seed_created_at + 1) - .bind((seed_created_at + 1) * 1000) - .bind(seed_created_at + 1) - .bind(&request_zero) - .bind(seed_created_at + 2) - .bind((seed_created_at + 2) * 1000) - .bind(seed_created_at + 2) - .bind(&request_outside) - .bind(seed_created_at + 3) - .bind((seed_created_at + 3) * 1000) - .bind(seed_created_at + 3) - .execute(backend.pool()) - .await - .expect("usage should seed"); - - sqlx::query( - r#" -INSERT INTO usage_settlement_snapshots ( - request_id, billing_status, wallet_id, finalized_at, created_at, updated_at -) VALUES - (?, 'settled', ?, ?, ?, ?), - (?, 'settled', ?, ?, ?, ?), - (?, 'settled', ?, ?, ?, ?), - (?, 'settled', ?, ?, ?, ?) -"#, - ) - .bind(&request_one) - .bind(&wallet_id) - .bind(first_finalized_at) - .bind(first_finalized_at) - .bind(first_finalized_at) - .bind(&request_two) - .bind(&wallet_id) - .bind(last_finalized_at) - .bind(last_finalized_at) - .bind(last_finalized_at) - .bind(&request_zero) - .bind(&wallet_id) - .bind(zero_finalized_at) - .bind(zero_finalized_at) - .bind(zero_finalized_at) - .bind(&request_outside) - .bind(&wallet_id) - .bind(outside_finalized_at) - .bind(outside_finalized_at) - .bind(outside_finalized_at) - .execute(backend.pool()) - .await - .expect("settlement snapshots should seed"); - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES (?, ?, '2026-05-03', ?, 7.0, 3, 1, 1, 0, 0, ?, ?, ?) -"#, - ) - .bind(&stale_ledger_id) - .bind(&stale_wallet_id) - .bind(&timezone) - .bind(seed_created_at) - .bind(seed_created_at) - .bind(seed_created_at) - .execute(backend.pool()) - .await - .expect("stale ledger should seed"); - - let summary = backend - .aggregate_wallet_daily_usage(&WalletDailyUsageAggregationInput { - billing_date: "2026-05-03".to_string(), - billing_timezone: timezone.clone(), - window_start_unix_secs: window_start as u64, - window_end_unix_secs: window_end as u64, - aggregated_at_unix_secs: aggregated_at as u64, - }) - .await - .expect("wallet daily usage aggregation should run"); - - assert_eq!(summary.aggregated_wallets, 1); - assert_eq!(summary.deleted_stale_ledgers, 1); - - let ledger = sqlx::query_as::< - _, - ( - String, - f64, - i64, - i64, - i64, - i64, - i64, - Option, - Option, - i64, - ), - >( - r#" -SELECT - wallet_id, - total_cost_usd, - total_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - first_finalized_at, - last_finalized_at, - aggregated_at -FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_date = '2026-05-03' - AND billing_timezone = ? -"#, - ) - .bind(&wallet_id) - .bind(&timezone) - .fetch_one(backend.pool()) - .await - .expect("aggregated ledger should load"); - - assert_eq!(ledger.0, wallet_id); - assert!((ledger.1 - 3.25).abs() < f64::EPSILON); - assert_eq!(ledger.2, 2); - assert_eq!(ledger.3, 15); - assert_eq!(ledger.4, 27); - assert_eq!(ledger.5, 4); - assert_eq!(ledger.6, 6); - assert_eq!(ledger.7, Some(first_finalized_at)); - assert_eq!(ledger.8, Some(last_finalized_at)); - assert_eq!(ledger.9, aggregated_at); - - let stale_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM wallet_daily_usage_ledgers WHERE id = ?") - .bind(&stale_ledger_id) - .fetch_one(backend.pool()) - .await - .expect("stale ledger count should load"); - assert_eq!(stale_count, 0); - } - - #[tokio::test] - async fn mysql_stats_aggregation_runs_after_mysql_migrations_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 stats aggregation smoke 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 backend = MysqlBackend::from_config(config).expect("backend should build"); - run_mysql_migrations(backend.pool()) - .await - .expect("mysql migrations should run"); - - 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-%'", - ] { - sqlx::query(sql) - .execute(backend.pool()) - .await - .expect("stats smoke cleanup should run"); - } - - sqlx::query( - r#" -INSERT INTO `usage` ( - 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, 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', '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"); - - let target_hour = chrono::DateTime::::from_timestamp(3600, 0) - .expect("target hour should be valid"); - let aggregated_at = chrono::DateTime::::from_timestamp(7200, 0) - .expect("aggregation time should be valid"); - let hourly = backend - .aggregate_stats_hourly(&StatsHourlyAggregationInput { - target_hour_utc: target_hour, - aggregated_at, - }) - .await - .expect("hourly stats aggregation should run") - .expect("hourly bucket should aggregate"); - assert_eq!(hourly.hour_utc, target_hour); - assert_eq!(hourly.total_requests, 2); - assert_eq!(hourly.user_rows, 2); - assert_eq!(hourly.user_model_rows, 2); - assert_eq!(hourly.model_rows, 2); - assert_eq!(hourly.provider_rows, 2); - - let hourly_row = sqlx::query_as::<_, (i64, i64, i64, i64, f64)>( - r#" -SELECT total_requests, success_requests, error_requests, input_tokens, total_cost -FROM stats_hourly -WHERE hour_utc = 3600 -"#, - ) - .fetch_one(backend.pool()) - .await - .expect("hourly stats row should load"); - assert_eq!(hourly_row.0, 2); - assert_eq!(hourly_row.1, 1); - assert_eq!(hourly_row.2, 1); - assert_eq!(hourly_row.3, 15); - assert!((hourly_row.4 - 0.50).abs() < f64::EPSILON); - - let second_hourly = backend - .aggregate_stats_hourly(&StatsHourlyAggregationInput { - target_hour_utc: target_hour, - aggregated_at, - }) - .await - .expect("second hourly aggregation should run"); - assert!(second_hourly.is_none()); - - let target_day = chrono::DateTime::::from_timestamp(0, 0) - .expect("target day should be valid"); - let daily = backend - .aggregate_stats_daily(&StatsDailyAggregationInput { - target_day_utc: target_day, - aggregated_at, - }) - .await - .expect("daily stats aggregation should run") - .expect("daily bucket should aggregate"); - assert_eq!(daily.day_start_utc, target_day); - assert_eq!(daily.total_requests, 2); - assert_eq!(daily.model_rows, 2); - assert_eq!(daily.provider_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)>( - r#" -SELECT total_requests, success_requests, error_requests, unique_models -FROM stats_daily -WHERE `date` = 0 -"#, - ) - .fetch_one(backend.pool()) - .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/read.rs b/crates/aether-data/runtime/src/backend/read.rs index da0ad98ad..cdfd272c5 100644 --- a/crates/aether-data/runtime/src/backend/read.rs +++ b/crates/aether-data/runtime/src/backend/read.rs @@ -1,12 +1,8 @@ use std::fmt; use std::sync::Arc; -#[cfg(feature = "mysql")] -use super::MysqlBackend; #[cfg(feature = "postgres")] use super::PostgresBackend; -#[cfg(feature = "sqlite")] -use super::SqliteBackend; use crate::repository::announcements::AnnouncementReadRepository; use crate::repository::audit::AuditLogReadRepository; use crate::repository::auth::AuthApiKeyReadRepository; @@ -91,22 +87,12 @@ impl fmt::Debug for DataReadRepositories { impl DataReadRepositories { pub(crate) fn from_backends( #[cfg(feature = "postgres")] postgres: Option<&PostgresBackend>, - #[cfg(feature = "mysql")] mysql: Option<&MysqlBackend>, - #[cfg(feature = "sqlite")] sqlite: Option<&SqliteBackend>, ) -> Self { let mut repositories = Self::default(); #[cfg(feature = "postgres")] if let Some(postgres) = postgres { repositories.install_postgres(postgres); } - #[cfg(feature = "mysql")] - if let Some(mysql) = mysql { - repositories.install_mysql(mysql); - } - #[cfg(feature = "sqlite")] - if let Some(sqlite) = sqlite { - repositories.install_sqlite(sqlite); - } repositories } @@ -183,157 +169,10 @@ impl DataReadRepositories { } } - #[cfg(feature = "mysql")] - fn install_mysql(&mut self, backend: &MysqlBackend) { - if self.announcements.is_none() { - self.announcements = Some(MysqlBackend::announcement_read_repository(backend)); - } - if self.audit_logs.is_none() { - self.audit_logs = Some(MysqlBackend::audit_log_read_repository(backend)); - } - if self.auth_api_keys.is_none() { - self.auth_api_keys = Some(MysqlBackend::auth_api_key_read_repository(backend)); - } - if self.auth_modules.is_none() { - self.auth_modules = Some(MysqlBackend::auth_module_read_repository(backend)); - } - if self.background_tasks.is_none() { - self.background_tasks = Some(MysqlBackend::background_task_read_repository(backend)); - } - if self.billing.is_none() { - self.billing = Some(MysqlBackend::billing_read_repository(backend)); - } - if self.gemini_file_mappings.is_none() { - self.gemini_file_mappings = - Some(MysqlBackend::gemini_file_mapping_read_repository(backend)); - } - if self.global_models.is_none() { - self.global_models = Some(MysqlBackend::global_model_read_repository(backend)); - } - if self.management_tokens.is_none() { - self.management_tokens = Some(MysqlBackend::management_token_read_repository(backend)); - } - if self.oauth_providers.is_none() { - self.oauth_providers = Some(MysqlBackend::oauth_provider_read_repository(backend)); - } - if self.pool_scores.is_none() { - self.pool_scores = Some(MysqlBackend::pool_score_read_repository(backend)); - } - if self.proxy_nodes.is_none() { - self.proxy_nodes = Some(MysqlBackend::proxy_node_read_repository(backend)); - } - if self.minimal_candidate_selection.is_none() { - self.minimal_candidate_selection = Some( - MysqlBackend::minimal_candidate_selection_read_repository(backend), - ); - } - if self.request_candidates.is_none() { - self.request_candidates = - Some(MysqlBackend::request_candidate_read_repository(backend)); - } - if self.provider_catalog.is_none() { - self.provider_catalog = Some(MysqlBackend::provider_catalog_read_repository(backend)); - } - if self.provider_quotas.is_none() { - self.provider_quotas = Some(MysqlBackend::provider_quota_read_repository(backend)); - } - if self.routing_groups.is_none() { - self.routing_groups = Some(MysqlBackend::routing_group_read_repository(backend)); - } - if self.usage.is_none() { - self.usage = Some(MysqlBackend::usage_read_repository(backend)); - } - if self.users.is_none() { - self.users = Some(MysqlBackend::user_read_repository(backend)); - } - if self.video_tasks.is_none() { - self.video_tasks = Some(MysqlBackend::video_task_read_repository(backend)); - } - if self.wallets.is_none() { - self.wallets = Some(MysqlBackend::wallet_read_repository(backend)); - } - } - - #[cfg(feature = "sqlite")] - fn install_sqlite(&mut self, backend: &SqliteBackend) { - if self.announcements.is_none() { - self.announcements = Some(SqliteBackend::announcement_read_repository(backend)); - } - if self.audit_logs.is_none() { - self.audit_logs = Some(SqliteBackend::audit_log_read_repository(backend)); - } - if self.auth_api_keys.is_none() { - self.auth_api_keys = Some(SqliteBackend::auth_api_key_read_repository(backend)); - } - if self.auth_modules.is_none() { - self.auth_modules = Some(SqliteBackend::auth_module_read_repository(backend)); - } - if self.background_tasks.is_none() { - self.background_tasks = Some(SqliteBackend::background_task_read_repository(backend)); - } - if self.billing.is_none() { - self.billing = Some(SqliteBackend::billing_read_repository(backend)); - } - if self.gemini_file_mappings.is_none() { - self.gemini_file_mappings = - Some(SqliteBackend::gemini_file_mapping_read_repository(backend)); - } - if self.global_models.is_none() { - self.global_models = Some(SqliteBackend::global_model_read_repository(backend)); - } - if self.management_tokens.is_none() { - self.management_tokens = Some(SqliteBackend::management_token_read_repository(backend)); - } - if self.oauth_providers.is_none() { - self.oauth_providers = Some(SqliteBackend::oauth_provider_read_repository(backend)); - } - if self.pool_scores.is_none() { - self.pool_scores = Some(SqliteBackend::pool_score_read_repository(backend)); - } - if self.proxy_nodes.is_none() { - self.proxy_nodes = Some(SqliteBackend::proxy_node_read_repository(backend)); - } - if self.minimal_candidate_selection.is_none() { - self.minimal_candidate_selection = Some( - SqliteBackend::minimal_candidate_selection_read_repository(backend), - ); - } - if self.request_candidates.is_none() { - self.request_candidates = - Some(SqliteBackend::request_candidate_read_repository(backend)); - } - if self.provider_catalog.is_none() { - self.provider_catalog = Some(SqliteBackend::provider_catalog_read_repository(backend)); - } - if self.provider_quotas.is_none() { - self.provider_quotas = Some(SqliteBackend::provider_quota_read_repository(backend)); - } - if self.routing_groups.is_none() { - self.routing_groups = Some(SqliteBackend::routing_group_read_repository(backend)); - } - if self.usage.is_none() { - self.usage = Some(SqliteBackend::usage_read_repository(backend)); - } - if self.users.is_none() { - self.users = Some(SqliteBackend::user_read_repository(backend)); - } - if self.video_tasks.is_none() { - self.video_tasks = Some(SqliteBackend::video_task_read_repository(backend)); - } - if self.wallets.is_none() { - self.wallets = Some(SqliteBackend::wallet_read_repository(backend)); - } - } #[cfg(test)] #[cfg(feature = "postgres")] pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self { - Self::from_backends( - postgres, - #[cfg(feature = "mysql")] - None, - #[cfg(feature = "sqlite")] - None, - ) + Self::from_backends(postgres) } pub fn auth_api_keys(&self) -> Option> { diff --git a/crates/aether-data/runtime/src/backend/referrals.rs b/crates/aether-data/runtime/src/backend/referrals.rs index c641e4e12..5b1f5f49c 100644 --- a/crates/aether-data/runtime/src/backend/referrals.rs +++ b/crates/aether-data/runtime/src/backend/referrals.rs @@ -21,8 +21,6 @@ const REFERRAL_RECONCILIATION_LIMIT: usize = 200; // The list tests intentionally build one page larger than the historical // in-memory fetch cap. Keep the fixture cap test-only now that production // queries paginate directly in SQL. -#[cfg(all(test, feature = "sqlite"))] -const REFERRAL_FETCH_LIMIT: usize = 5_000; #[derive(Debug, Clone, Serialize)] pub struct ReferralUserDashboard { @@ -223,11 +221,6 @@ macro_rules! reward_from_row { }}; } -#[cfg(any(feature = "mysql", feature = "sqlite"))] -fn now_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - fn row_unix_secs(row: &R, column: &str) -> Result where R: Row, @@ -266,7 +259,7 @@ fn generate_invite_code() -> String { /// Build a case-insensitive SQL `LIKE` pattern while treating user input as a /// literal substring. `!` is used as the escape character because it is -/// accepted consistently by PostgreSQL, MySQL, and SQLite. +/// accepted by PostgreSQL. fn referral_like_pattern(value: Option<&str>) -> String { let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else { return String::new(); @@ -581,44 +574,6 @@ WHERE id = $1 .rows_affected(); return Ok(affected > 0); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let affected = sqlx::query( - r#" -UPDATE users -SET privacy_policy_accepted_version = ?, - privacy_policy_accepted_at = ? -WHERE id = ? -"#, - ) - .bind(version) - .bind(now_unix_secs() as i64) - .bind(user_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let affected = sqlx::query( - r#" -UPDATE users -SET privacy_policy_accepted_version = ?, - privacy_policy_accepted_at = ? -WHERE id = ? -"#, - ) - .bind(version) - .bind(now_unix_secs() as i64) - .bind(user_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } Ok(false) } @@ -738,120 +693,6 @@ WHERE ($1::TEXT IS NULL OR inviter_user_id = $1) }); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - (SELECT COUNT(*) FROM user_referrals - WHERE (? IS NULL OR inviter_user_id = ?)) AS total_invites, - (SELECT COUNT(*) FROM user_referrals - WHERE (? IS NULL OR inviter_user_id = ?) - AND first_paid_order_id IS NOT NULL) - AS effective_invites, - CAST(COALESCE(SUM(CASE - WHEN status = 'applied' AND amount_usd > 0 THEN amount_usd ELSE 0 END), 0) - AS DOUBLE) AS paid_reward_usd, - CAST(COALESCE(SUM(CASE - WHEN status IN ('pending', 'failed', 'applying') AND amount_usd > 0 THEN amount_usd ELSE 0 END), 0) - AS DOUBLE) AS pending_reward_usd, - CAST(COALESCE(SUM(CASE - WHEN reversed_amount_usd > 0 THEN reversed_amount_usd ELSE 0 END), 0) - AS DOUBLE) AS reversed_reward_usd -FROM referral_rewards -WHERE (? IS NULL OR inviter_user_id = ?) -"#, - ) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return Ok(ReferralAdminStats { - total_invites: referral_stats_count( - row.try_get::("total_invites") - .map_err(DataLayerError::sql)?, - ), - effective_invites: referral_stats_count( - row.try_get::("effective_invites") - .map_err(DataLayerError::sql)?, - ), - paid_reward_usd: referral_stats_amount( - row.try_get::("paid_reward_usd") - .map_err(DataLayerError::sql)?, - ), - pending_reward_usd: referral_stats_amount( - row.try_get::("pending_reward_usd") - .map_err(DataLayerError::sql)?, - ), - reversed_reward_usd: referral_stats_amount( - row.try_get::("reversed_reward_usd") - .map_err(DataLayerError::sql)?, - ), - }); - } - - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - (SELECT COUNT(*) FROM user_referrals - WHERE (? IS NULL OR inviter_user_id = ?)) AS total_invites, - (SELECT COUNT(*) FROM user_referrals - WHERE (? IS NULL OR inviter_user_id = ?) - AND first_paid_order_id IS NOT NULL) - AS effective_invites, - CAST(COALESCE(SUM(CASE - WHEN status = 'applied' AND amount_usd > 0 THEN amount_usd ELSE 0 END), 0) - AS REAL) AS paid_reward_usd, - CAST(COALESCE(SUM(CASE - WHEN status IN ('pending', 'failed', 'applying') AND amount_usd > 0 THEN amount_usd ELSE 0 END), 0) - AS REAL) AS pending_reward_usd, - CAST(COALESCE(SUM(CASE - WHEN reversed_amount_usd > 0 THEN reversed_amount_usd ELSE 0 END), 0) - AS REAL) AS reversed_reward_usd -FROM referral_rewards -WHERE (? IS NULL OR inviter_user_id = ?) -"#, - ) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .bind(inviter_user_id) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return Ok(ReferralAdminStats { - total_invites: referral_stats_count( - row.try_get::("total_invites") - .map_err(DataLayerError::sql)?, - ), - effective_invites: referral_stats_count( - row.try_get::("effective_invites") - .map_err(DataLayerError::sql)?, - ), - paid_reward_usd: referral_stats_amount( - row.try_get::("paid_reward_usd") - .map_err(DataLayerError::sql)?, - ), - pending_reward_usd: referral_stats_amount( - row.try_get::("pending_reward_usd") - .map_err(DataLayerError::sql)?, - ), - reversed_reward_usd: referral_stats_amount( - row.try_get::("reversed_reward_usd") - .map_err(DataLayerError::sql)?, - ), - }); - } - Ok(ReferralAdminStats::default()) } @@ -1329,44 +1170,6 @@ ON CONFLICT DO NOTHING .map_err(DataLayerError::postgres)? .rows_affected(); } - #[cfg(feature = "mysql")] - if inserted == 0 { - if let Some(backend) = backends.mysql() { - inserted = sqlx::query( - r#" -INSERT IGNORE INTO user_invite_codes (user_id, invite_code, active, created_at, updated_at) -VALUES (?, ?, TRUE, ?, ?) -"#, - ) - .bind(user_id) - .bind(&code) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - } - } - #[cfg(feature = "sqlite")] - if inserted == 0 { - if let Some(backend) = backends.sqlite() { - inserted = sqlx::query( - r#" -INSERT OR IGNORE INTO user_invite_codes (user_id, invite_code, active, created_at, updated_at) -VALUES (?, ?, 1, ?, ?) -"#, - ) - .bind(user_id) - .bind(&code) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - } - } if inserted > 0 { return Ok(Some(code)); } @@ -1400,38 +1203,6 @@ VALUES (?, ?, 1, ?, ?) }) .transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - "SELECT invite_code FROM user_invite_codes WHERE user_id = ? AND active = TRUE", - ) - .bind(user_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row - .map(|row| { - row.try_get::("invite_code") - .map_err(DataLayerError::sql) - }) - .transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - "SELECT invite_code FROM user_invite_codes WHERE user_id = ? AND active = 1", - ) - .bind(user_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row - .map(|row| { - row.try_get::("invite_code") - .map_err(DataLayerError::sql) - }) - .transpose(); - } Ok(None) } @@ -1458,38 +1229,6 @@ VALUES (?, ?, 1, ?, ?) }) .transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - "SELECT user_id FROM user_invite_codes WHERE invite_code = ? AND active = TRUE", - ) - .bind(invite_code) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row - .map(|row| { - row.try_get::("user_id") - .map_err(DataLayerError::sql) - }) - .transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - "SELECT user_id FROM user_invite_codes WHERE invite_code = ? AND active = 1", - ) - .bind(invite_code) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row - .map(|row| { - row.try_get::("user_id") - .map_err(DataLayerError::sql) - }) - .transpose(); - } Ok(None) } @@ -1526,52 +1265,6 @@ ON CONFLICT (invitee_user_id) DO NOTHING .rows_affected(); return Ok(affected > 0); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let affected = sqlx::query( - r#" -INSERT IGNORE INTO user_referrals ( - id, inviter_user_id, invitee_user_id, invite_code_snapshot, source_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(referral_id) - .bind(inviter_user_id) - .bind(invitee_user_id) - .bind(invite_code) - .bind(source_json) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let affected = sqlx::query( - r#" -INSERT OR IGNORE INTO user_referrals ( - id, inviter_user_id, invitee_user_id, invite_code_snapshot, source_json, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(referral_id) - .bind(inviter_user_id) - .bind(invitee_user_id) - .bind(invite_code) - .bind(source_json) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } Ok(false) } @@ -1649,152 +1342,6 @@ LIMIT $5 OFFSET $6 .collect::, _>>()?; return Ok((items, total.max(0) as u64)); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let count = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE (? = '' OR LOWER(COALESCE(inviter.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.inviter_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(COALESCE(invitee.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.invitee_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(r.invite_code_snapshot) LIKE ? ESCAPE '!') - AND (? < 0 OR (? = 1 AND r.first_paid_order_id IS NOT NULL) OR (? = 0 AND r.first_paid_order_id IS NULL)) -"#, - ) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invite_code_pattern) - .bind(&invite_code_pattern) - .bind(first_paid) - .bind(first_paid) - .bind(first_paid) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let total = count - .try_get::("total") - .map_err(DataLayerError::sql)?; - let rows = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE (? = '' OR LOWER(COALESCE(inviter.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.inviter_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(COALESCE(invitee.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.invitee_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(r.invite_code_snapshot) LIKE ? ESCAPE '!') - AND (? < 0 OR (? = 1 AND r.first_paid_order_id IS NOT NULL) OR (? = 0 AND r.first_paid_order_id IS NULL)) -ORDER BY r.created_at DESC, r.id DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invite_code_pattern) - .bind(&invite_code_pattern) - .bind(first_paid) - .bind(first_paid) - .bind(first_paid) - .bind(limit) - .bind(offset) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let items = rows - .iter() - .map(|row| relationship_from_row!(row)) - .collect::, _>>()?; - return Ok((items, total.max(0) as u64)); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let count = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE (? = '' OR LOWER(COALESCE(inviter.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.inviter_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(COALESCE(invitee.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.invitee_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(r.invite_code_snapshot) LIKE ? ESCAPE '!') - AND (? < 0 OR (? = 1 AND r.first_paid_order_id IS NOT NULL) OR (? = 0 AND r.first_paid_order_id IS NULL)) -"#, - ) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invite_code_pattern) - .bind(&invite_code_pattern) - .bind(first_paid) - .bind(first_paid) - .bind(first_paid) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let total = count - .try_get::("total") - .map_err(DataLayerError::sql)?; - let rows = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE (? = '' OR LOWER(COALESCE(inviter.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.inviter_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(COALESCE(invitee.username, '')) LIKE ? ESCAPE '!' OR LOWER(r.invitee_user_id) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(r.invite_code_snapshot) LIKE ? ESCAPE '!') - AND (? < 0 OR (? = 1 AND r.first_paid_order_id IS NOT NULL) OR (? = 0 AND r.first_paid_order_id IS NULL)) -ORDER BY r.created_at DESC, r.id DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&inviter_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invitee_pattern) - .bind(&invite_code_pattern) - .bind(&invite_code_pattern) - .bind(first_paid) - .bind(first_paid) - .bind(first_paid) - .bind(limit) - .bind(offset) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let items = rows - .iter() - .map(|row| relationship_from_row!(row)) - .collect::, _>>()?; - return Ok((items, total.max(0) as u64)); - } Ok((Vec::new(), 0)) } @@ -1829,54 +1376,6 @@ LIMIT 1 .map_err(DataLayerError::postgres)?; return row.map(|row| relationship_from_row!(&row)).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE r.id = ? -LIMIT 1 -"#, - ) - .bind(referral_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| relationship_from_row!(&row)).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -LEFT JOIN users inviter ON inviter.id = r.inviter_user_id -LEFT JOIN users invitee ON invitee.id = r.invitee_user_id -WHERE r.id = ? -LIMIT 1 -"#, - ) - .bind(referral_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| relationship_from_row!(&row)).transpose(); - } Ok(None) } @@ -1913,58 +1412,6 @@ LIMIT 1 .map_err(DataLayerError::postgres)?; return row.map(|row| relationship_from_row!(&row)).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -JOIN users inviter ON inviter.id = r.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -JOIN users invitee ON invitee.id = r.invitee_user_id - AND invitee.is_active = 1 AND invitee.is_deleted = 0 -WHERE r.invitee_user_id = ? -LIMIT 1 -"#, - ) - .bind(invitee_user_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| relationship_from_row!(&row)).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - r.id, r.inviter_user_id, inviter.username AS inviter_username, - r.invitee_user_id, invitee.username AS invitee_username, - r.invite_code_snapshot, r.first_paid_order_id, - r.first_paid_at AS first_paid_at_unix_secs, - r.source_json AS source_json, - r.created_at AS created_at_unix_secs -FROM user_referrals r -JOIN users inviter ON inviter.id = r.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -JOIN users invitee ON invitee.id = r.invitee_user_id - AND invitee.is_active = 1 AND invitee.is_deleted = 0 -WHERE r.invitee_user_id = ? -LIMIT 1 -"#, - ) - .bind(invitee_user_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| relationship_from_row!(&row)).transpose(); - } Ok(None) } @@ -2032,122 +1479,6 @@ LIMIT $4 OFFSET $5 .collect::, _>>()?; return Ok((items, total.max(0) as u64)); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let count = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM referral_rewards -WHERE (? = '' OR LOWER(COALESCE(source_order_id, '')) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(reward_type) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(status) LIKE ? ESCAPE '!') -"#, - ) - .bind(&order_pattern) - .bind(&order_pattern) - .bind(&reward_type_pattern) - .bind(&reward_type_pattern) - .bind(&status_pattern) - .bind(&status_pattern) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let total = count - .try_get::("total") - .map_err(DataLayerError::sql)?; - let rows = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, CAST(amount_usd AS DOUBLE) AS amount_usd, - status, wallet_transaction_id, idempotency_key, - CAST(reversed_amount_usd AS DOUBLE) AS reversed_amount_usd, - CAST(pending_reversal_amount_usd AS DOUBLE) AS pending_reversal_amount_usd, - admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE (? = '' OR LOWER(COALESCE(source_order_id, '')) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(reward_type) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(status) LIKE ? ESCAPE '!') -ORDER BY created_at DESC, id DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(&order_pattern) - .bind(&order_pattern) - .bind(&reward_type_pattern) - .bind(&reward_type_pattern) - .bind(&status_pattern) - .bind(&status_pattern) - .bind(limit) - .bind(offset) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let items = rows - .iter() - .map(|row| reward_from_row!(row)) - .collect::, _>>()?; - return Ok((items, total.max(0) as u64)); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let count = sqlx::query( - r#" -SELECT COUNT(*) AS total -FROM referral_rewards -WHERE (? = '' OR LOWER(COALESCE(source_order_id, '')) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(reward_type) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(status) LIKE ? ESCAPE '!') -"#, - ) - .bind(&order_pattern) - .bind(&order_pattern) - .bind(&reward_type_pattern) - .bind(&reward_type_pattern) - .bind(&status_pattern) - .bind(&status_pattern) - .fetch_one(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let total = count - .try_get::("total") - .map_err(DataLayerError::sql)?; - let rows = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, CAST(amount_usd AS DOUBLE PRECISION) AS amount_usd, - status, wallet_transaction_id, idempotency_key, - CAST(reversed_amount_usd AS DOUBLE PRECISION) AS reversed_amount_usd, - CAST(pending_reversal_amount_usd AS DOUBLE PRECISION) AS pending_reversal_amount_usd, - admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE (? = '' OR LOWER(COALESCE(source_order_id, '')) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(reward_type) LIKE ? ESCAPE '!') - AND (? = '' OR LOWER(status) LIKE ? ESCAPE '!') -ORDER BY created_at DESC, id DESC -LIMIT ? OFFSET ? -"#, - ) - .bind(&order_pattern) - .bind(&order_pattern) - .bind(&reward_type_pattern) - .bind(&reward_type_pattern) - .bind(&status_pattern) - .bind(&status_pattern) - .bind(limit) - .bind(offset) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - let items = rows - .iter() - .map(|row| reward_from_row!(row)) - .collect::, _>>()?; - return Ok((items, total.max(0) as u64)); - } Ok((Vec::new(), 0)) } @@ -2172,40 +1503,6 @@ LIMIT $1 .map_err(DataLayerError::postgres)?; return rows.iter().map(|row| Ok(row_string!(row, "id"))).collect(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let rows = sqlx::query( - r#" -SELECT id -FROM referral_rewards -WHERE status = 'applying' -ORDER BY updated_at ASC, created_at ASC, id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| Ok(row_string!(row, "id"))).collect(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let rows = sqlx::query( - r#" -SELECT id -FROM referral_rewards -WHERE status = 'applying' -ORDER BY updated_at ASC, created_at ASC, id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| Ok(row_string!(row, "id"))).collect(); - } Ok(Vec::new()) } @@ -2240,56 +1537,6 @@ LIMIT $1 .map(|row| Ok(row_string!(row, "idempotency_key"))) .collect(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let rows = sqlx::query( - r#" -SELECT rw.idempotency_key -FROM referral_rewards rw -JOIN wallets ON wallets.user_id = rw.inviter_user_id - AND wallets.status = 'active' -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.status IN ('pending', 'failed') - AND rw.amount_usd > 0 -ORDER BY rw.created_at ASC, rw.id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows - .iter() - .map(|row| Ok(row_string!(row, "idempotency_key"))) - .collect(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let rows = sqlx::query( - r#" -SELECT rw.idempotency_key -FROM referral_rewards rw -JOIN wallets ON wallets.user_id = rw.inviter_user_id - AND wallets.status = 'active' -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.status IN ('pending', 'failed') - AND rw.amount_usd > 0 -ORDER BY rw.created_at ASC, rw.id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows - .iter() - .map(|row| Ok(row_string!(row, "idempotency_key"))) - .collect(); - } Ok(Vec::new()) } @@ -2384,160 +1631,6 @@ LIMIT $1 .map_err(DataLayerError::postgres)?; return rows.iter().map(|row| reward_from_row!(row)).collect(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let rows = sqlx::query( - r#" -SELECT - rw.id, rw.referral_id, rw.inviter_user_id, rw.invitee_user_id, - rw.reward_type, rw.source_order_id, rw.trigger_point, rw.amount_usd, - rw.status, rw.wallet_transaction_id, rw.idempotency_key, - rw.reversed_amount_usd, rw.pending_reversal_amount_usd, - rw.admin_operator_id, rw.admin_note, - rw.created_at AS created_at_unix_secs, - rw.updated_at AS updated_at_unix_secs -FROM referral_rewards rw -JOIN ( - SELECT - po0.id, - po0.amount_usd, - po0.credited_at, - po0.paid_at, - po0.created_at, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po0.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po0.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd - FROM payment_orders po0 -) po ON po.id = rw.source_order_id -JOIN wallets wallet ON wallet.user_id = rw.inviter_user_id - AND wallet.status = 'active' -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.status IN ('applied', 'reversed') - AND po.refunded_amount_usd > 0 - AND ( - rw.pending_reversal_amount_usd > 0.00000001 - OR ( - po.amount_usd > 0 - AND rw.amount_usd > 0 - AND rw.reversed_amount_usd + 0.00000001 < - rw.amount_usd * CASE - WHEN po.refunded_amount_usd >= po.amount_usd THEN 1.0 - ELSE po.refunded_amount_usd / po.amount_usd - END - ) - ) -ORDER BY COALESCE(po.credited_at, po.paid_at, po.created_at) ASC, - rw.created_at ASC, rw.id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| reward_from_row!(row)).collect(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let rows = sqlx::query( - r#" -SELECT - rw.id, rw.referral_id, rw.inviter_user_id, rw.invitee_user_id, - rw.reward_type, rw.source_order_id, rw.trigger_point, rw.amount_usd, - rw.status, rw.wallet_transaction_id, rw.idempotency_key, - rw.reversed_amount_usd, rw.pending_reversal_amount_usd, - rw.admin_operator_id, rw.admin_note, - rw.created_at AS created_at_unix_secs, - rw.updated_at AS updated_at_unix_secs -FROM referral_rewards rw -JOIN ( - SELECT - po0.id, - po0.amount_usd, - po0.credited_at, - po0.paid_at, - po0.created_at, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po0.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po0.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po0.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd - FROM payment_orders po0 -) po ON po.id = rw.source_order_id -JOIN wallets wallet ON wallet.user_id = rw.inviter_user_id - AND wallet.status = 'active' -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.status IN ('applied', 'reversed') - AND po.refunded_amount_usd > 0 - AND ( - rw.pending_reversal_amount_usd > 0.00000001 - OR ( - po.amount_usd > 0 - AND rw.amount_usd > 0 - AND rw.reversed_amount_usd + 0.00000001 < - rw.amount_usd * CASE - WHEN po.refunded_amount_usd >= po.amount_usd THEN 1.0 - ELSE po.refunded_amount_usd / po.amount_usd - END - ) - ) -ORDER BY COALESCE(po.credited_at, po.paid_at, po.created_at) ASC, - rw.created_at ASC, rw.id ASC -LIMIT ? -"#, - ) - .bind(REFERRAL_RECONCILIATION_LIMIT as i64) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| reward_from_row!(row)).collect(); - } Ok(Vec::new()) } @@ -2572,52 +1665,6 @@ LIMIT 1 .map_err(DataLayerError::postgres)?; return row.map(|row| reward_from_row!(&row)).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, CAST(amount_usd AS DOUBLE) AS amount_usd, - status, wallet_transaction_id, idempotency_key, - CAST(reversed_amount_usd AS DOUBLE) AS reversed_amount_usd, - CAST(pending_reversal_amount_usd AS DOUBLE) AS pending_reversal_amount_usd, - admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(reward_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| reward_from_row!(&row)).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, CAST(amount_usd AS DOUBLE PRECISION) AS amount_usd, - status, wallet_transaction_id, idempotency_key, - CAST(reversed_amount_usd AS DOUBLE PRECISION) AS reversed_amount_usd, - CAST(pending_reversal_amount_usd AS DOUBLE PRECISION) AS pending_reversal_amount_usd, - admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE id = ? -LIMIT 1 -"#, - ) - .bind(reward_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| reward_from_row!(&row)).transpose(); - } Ok(None) } @@ -2652,46 +1699,6 @@ LIMIT 1 .map_err(DataLayerError::postgres)?; return row.map(|row| reward_from_row!(&row)).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, wallet_transaction_id, idempotency_key, - reversed_amount_usd, pending_reversal_amount_usd, admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE idempotency_key = ? -LIMIT 1 -"#, - ) - .bind(idempotency_key) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| reward_from_row!(&row)).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, wallet_transaction_id, idempotency_key, - reversed_amount_usd, pending_reversal_amount_usd, admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE idempotency_key = ? -LIMIT 1 -"#, - ) - .bind(idempotency_key) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(|row| reward_from_row!(&row)).transpose(); - } Ok(None) } @@ -2727,48 +1734,6 @@ ORDER BY created_at ASC .map_err(DataLayerError::postgres)?; return rows.iter().map(|row| reward_from_row!(row)).collect(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let rows = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, wallet_transaction_id, idempotency_key, - reversed_amount_usd, pending_reversal_amount_usd, admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE source_order_id = ? - AND status IN ('applied', 'reversed') -ORDER BY created_at ASC -"#, - ) - .bind(order_id) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| reward_from_row!(row)).collect(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let rows = sqlx::query( - r#" -SELECT - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, wallet_transaction_id, idempotency_key, - reversed_amount_usd, pending_reversal_amount_usd, admin_operator_id, admin_note, - created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs -FROM referral_rewards -WHERE source_order_id = ? - AND status IN ('applied', 'reversed') -ORDER BY created_at ASC -"#, - ) - .bind(order_id) - .fetch_all(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return rows.iter().map(|row| reward_from_row!(row)).collect(); - } Ok(Vec::new()) } @@ -2812,62 +1777,6 @@ ON CONFLICT (idempotency_key) DO NOTHING .rows_affected(); return Ok(affected > 0); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let affected = sqlx::query( - r#" -INSERT IGNORE INTO referral_rewards ( - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, idempotency_key, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?) -"#, - ) - .bind(&reward_id) - .bind(&relationship.id) - .bind(&relationship.inviter_user_id) - .bind(&relationship.invitee_user_id) - .bind(reward_type) - .bind(source_order_id) - .bind(trigger_point) - .bind(amount_usd) - .bind(idempotency_key) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let affected = sqlx::query( - r#" -INSERT OR IGNORE INTO referral_rewards ( - id, referral_id, inviter_user_id, invitee_user_id, reward_type, source_order_id, - trigger_point, amount_usd, status, idempotency_key, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?) -"#, - ) - .bind(&reward_id) - .bind(&relationship.id) - .bind(&relationship.inviter_user_id) - .bind(&relationship.invitee_user_id) - .bind(reward_type) - .bind(source_order_id) - .bind(trigger_point) - .bind(amount_usd) - .bind(idempotency_key) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } Ok(false) } @@ -2894,36 +1803,6 @@ WHERE id = $1 .map_err(DataLayerError::postgres)?; return row.map(payment_order_context_from_row).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT id, user_id, amount_usd, payment_method, status, order_kind -FROM payment_orders -WHERE id = ? -"#, - ) - .bind(order_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(payment_order_context_from_row).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT id, user_id, amount_usd, payment_method, status, order_kind -FROM payment_orders -WHERE id = ? -"#, - ) - .bind(order_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(payment_order_context_from_row).transpose(); - } Ok(None) } @@ -2977,88 +1856,6 @@ WHERE po.id = $1 .map_err(DataLayerError::postgres)?; return row.map(payment_order_refund_context_from_row).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT po.amount_usd, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd -FROM payment_orders po -WHERE po.id = ? -"#, - ) - .bind(order_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(payment_order_refund_context_from_row).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT po.amount_usd, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd -FROM payment_orders po -WHERE po.id = ? -"#, - ) - .bind(order_id) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(payment_order_refund_context_from_row).transpose(); - } Ok(None) } @@ -3089,48 +1886,6 @@ WHERE id = $1 AND first_paid_order_id IS NULL .rows_affected(); return Ok(affected > 0); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let affected = sqlx::query( - r#" -UPDATE user_referrals -SET first_paid_order_id = ?, - first_paid_at = ?, - updated_at = ? -WHERE id = ? AND first_paid_order_id IS NULL -"#, - ) - .bind(order_id) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .bind(referral_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let affected = sqlx::query( - r#" -UPDATE user_referrals -SET first_paid_order_id = ?, - first_paid_at = ?, - updated_at = ? -WHERE id = ? AND first_paid_order_id IS NULL -"#, - ) - .bind(order_id) - .bind(now_unix_secs() as i64) - .bind(now_unix_secs() as i64) - .bind(referral_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } Ok(false) } @@ -3166,52 +1921,6 @@ WHERE id = $1 AND status IN ('pending', 'failed') .rows_affected(); return Ok(affected > 0); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let affected = sqlx::query( - r#" -UPDATE referral_rewards -SET status = ?, - admin_operator_id = COALESCE(?, admin_operator_id), - admin_note = COALESCE(?, admin_note), - updated_at = ? -WHERE id = ? AND status IN ('pending', 'failed') -"#, - ) - .bind(status) - .bind(operator_id) - .bind(note) - .bind(now_unix_secs() as i64) - .bind(reward_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let affected = sqlx::query( - r#" -UPDATE referral_rewards -SET status = ?, - admin_operator_id = COALESCE(?, admin_operator_id), - admin_note = COALESCE(?, admin_note), - updated_at = ? -WHERE id = ? AND status IN ('pending', 'failed') -"#, - ) - .bind(status) - .bind(operator_id) - .bind(note) - .bind(now_unix_secs() as i64) - .bind(reward_id) - .execute(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - return Ok(affected > 0); - } Ok(false) } @@ -3360,24 +2069,6 @@ WHERE id = $1 AND status = 'applying' tx.commit().await.map_err(DataLayerError::postgres)?; return Ok(recovery); } - #[cfg(feature = "mysql")] - if let Some(backend) = self.backends.and_then(DataBackends::mysql) { - return self - .recover_applying_referral_reward_mysql_numeric_time( - &backend.pool_clone(), - reward_id, - ) - .await; - } - #[cfg(feature = "sqlite")] - if let Some(backend) = self.backends.and_then(DataBackends::sqlite) { - return self - .recover_applying_referral_reward_sqlite_numeric_time( - &backend.pool_clone(), - reward_id, - ) - .await; - } Ok(ReferralApplyingRecovery::Unchanged) } @@ -3431,50 +2122,6 @@ WHERE rw.idempotency_key = $1 .map_err(DataLayerError::postgres)?; return row.map(credit_target_from_row).transpose(); } - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - let row = sqlx::query( - r#" -SELECT - rw.id, rw.inviter_user_id, rw.invitee_user_id, rw.amount_usd, rw.reward_type, - rw.trigger_point, wallets.id AS wallet_id -FROM referral_rewards rw -JOIN wallets ON wallets.user_id = rw.inviter_user_id -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.idempotency_key = ? - AND rw.status IN ('pending', 'failed') - AND wallets.status = 'active' -"#, - ) - .bind(idempotency_key) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(credit_target_from_row).transpose(); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - let row = sqlx::query( - r#" -SELECT - rw.id, rw.inviter_user_id, rw.invitee_user_id, rw.amount_usd, rw.reward_type, - rw.trigger_point, wallets.id AS wallet_id -FROM referral_rewards rw -JOIN wallets ON wallets.user_id = rw.inviter_user_id -JOIN users inviter ON inviter.id = rw.inviter_user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 -WHERE rw.idempotency_key = ? - AND rw.status IN ('pending', 'failed') - AND wallets.status = 'active' -"#, - ) - .bind(idempotency_key) - .fetch_optional(&backend.pool_clone()) - .await - .map_err(DataLayerError::sql)?; - return row.map(credit_target_from_row).transpose(); - } Ok(None) } @@ -3644,28 +2291,6 @@ WHERE id = $1 tx.commit().await.map_err(DataLayerError::postgres)?; return Ok(()); } - #[cfg(feature = "mysql")] - if let Some(backend) = self.backends.and_then(DataBackends::mysql) { - self.credit_referral_reward_mysql_numeric_time( - &backend.pool_clone(), - target, - operator_id, - note, - ) - .await?; - return Ok(()); - } - #[cfg(feature = "sqlite")] - if let Some(backend) = self.backends.and_then(DataBackends::sqlite) { - self.credit_referral_reward_sqlite_numeric_time( - &backend.pool_clone(), - target, - operator_id, - note, - ) - .await?; - return Ok(()); - } Ok(()) } @@ -3957,383 +2582,6 @@ WHERE id = $1 tx.commit().await.map_err(DataLayerError::postgres)?; return Ok(()); } - #[cfg(any(feature = "mysql", feature = "sqlite"))] - { - // MySQL/SQLite refunds use integer timestamps in the wallet tables. - return self - .apply_referral_reward_reversal_numeric_time(reward) - .await; - } - #[cfg(not(any(feature = "mysql", feature = "sqlite")))] - Ok(()) - } -} - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -macro_rules! referral_applying_recovery_numeric_method { - ($name:ident, $pool_ty:ty) => { - async fn $name( - &self, - pool: &$pool_ty, - reward_id: &str, - ) -> Result { - let mut tx = pool.begin().await.map_err(DataLayerError::sql)?; - - // Both drivers begin deferred transactions. This harmless write - // takes the write/row lock before the transaction fact is read. - sqlx::query( - "UPDATE referral_rewards SET updated_at = updated_at WHERE id = ? AND status = 'applying'", - ) - .bind(reward_id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let reward = sqlx::query( - "SELECT id, inviter_user_id, amount_usd FROM referral_rewards WHERE id = ? AND status = 'applying'", - ) - .bind(reward_id) - .fetch_optional(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - if reward.is_none() { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(ReferralApplyingRecovery::Unchanged); - } - - let wallet_transactions = sqlx::query( - r#" -SELECT tx.id, - tx.amount, - tx.balance_before, - tx.balance_after, - tx.recharge_balance_before, - tx.recharge_balance_after, - tx.gift_balance_before, - tx.gift_balance_after -FROM wallet_transactions tx -JOIN wallets wallet ON wallet.id = tx.wallet_id -WHERE tx.category = 'adjust' - AND tx.reason_code = 'referral_reward' - AND tx.link_type = 'referral_reward' - AND tx.link_id = ? - AND wallet.user_id = (SELECT inviter_user_id FROM referral_rewards WHERE id = ?) - AND tx.amount > 0 -ORDER BY tx.created_at ASC, tx.id ASC -LIMIT 32 -"#, - ) - .bind(reward_id) - .bind(reward_id) - .fetch_all(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let reward_amount = reward - .as_ref() - .and_then(|row| row.try_get::("amount_usd").ok()) - .unwrap_or(0.0); - let has_wallet_transaction = !wallet_transactions.is_empty(); - let valid_wallet_transaction_ids = wallet_transactions - .into_iter() - .filter_map(|row| { - let amount = row.try_get::("amount").ok()?; - let balance_before = row.try_get::("balance_before").ok()?; - let balance_after = row.try_get::("balance_after").ok()?; - let recharge_balance_before = row - .try_get::("recharge_balance_before") - .ok()?; - let recharge_balance_after = row - .try_get::("recharge_balance_after") - .ok()?; - let gift_balance_before = row.try_get::("gift_balance_before").ok()?; - let gift_balance_after = row.try_get::("gift_balance_after").ok()?; - if !referral_credit_transaction_fact_valid( - reward_amount, - amount, - balance_before, - balance_after, - recharge_balance_before, - recharge_balance_after, - gift_balance_before, - gift_balance_after, - ) { - return None; - } - row.try_get::("id").ok() - }) - .collect::>(); - // Multiple valid facts for one reward indicate a possible - // duplicate credit. Do not mark the reward applied by selecting - // an arbitrary transaction. - let wallet_transaction_id = (valid_wallet_transaction_ids.len() == 1) - .then(|| valid_wallet_transaction_ids[0].clone()); - let recovery = if !reward_amount.is_finite() || reward_amount <= 0.0 { - // A malformed durable amount must never enter the normal - // failed-reward retry path. Leave it for operator repair, - // just like an ambiguous wallet snapshot. - ReferralApplyingRecovery::Unchanged - } else if wallet_transaction_id.is_some() { - ReferralApplyingRecovery::Applied - } else if has_wallet_transaction { - // A matching transaction with an invalid snapshot is - // ambiguous: retrying it as failed could credit twice. - // Leave the reward applying until the historical fact is - // repaired by an operator. - ReferralApplyingRecovery::Unchanged - } else { - ReferralApplyingRecovery::Failed - }; - if recovery == ReferralApplyingRecovery::Unchanged { - // `applying` rows are processed in a bounded queue. Bump the - // retry timestamp for ambiguous facts so one permanently - // malformed row cannot occupy the oldest page forever. - let rotated_at = now_unix_secs() as i64; - sqlx::query( - "UPDATE referral_rewards SET updated_at = CASE WHEN updated_at >= ? THEN updated_at + 1 ELSE ? END WHERE id = ? AND status = 'applying'", - ) - .bind(rotated_at) - .bind(rotated_at) - .bind(reward_id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(recovery); - } - let status = match recovery { - ReferralApplyingRecovery::Applied => "applied", - ReferralApplyingRecovery::Failed => "failed", - ReferralApplyingRecovery::Unchanged => unreachable!(), - }; - sqlx::query( - r#" -UPDATE referral_rewards -SET status = ?, - wallet_transaction_id = ?, - updated_at = ? -WHERE id = ? AND status = 'applying' -"#, - ) - .bind(status) - .bind(wallet_transaction_id.as_deref()) - .bind(now_unix_secs() as i64) - .bind(reward_id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - Ok(recovery) - } - }; -} - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -macro_rules! referral_credit_numeric_method { - ($name:ident, $pool_ty:ty, $wallet_sql:expr) => { - async fn $name( - &self, - pool: &$pool_ty, - target: ReferralCreditTarget, - operator_id: Option<&str>, - note: Option<&str>, - ) -> Result<(), DataLayerError> { - let mut tx = pool.begin().await.map_err(DataLayerError::sql)?; - let claimed = sqlx::query( - r#" -UPDATE referral_rewards -SET status = 'applying', - admin_operator_id = COALESCE(?, admin_operator_id), - admin_note = COALESCE(?, admin_note), - updated_at = ? -WHERE id = ? AND status IN ('pending', 'failed') -"#, - ) - .bind(operator_id) - .bind(note) - .bind(now_unix_secs() as i64) - .bind(&target.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)? - .rows_affected(); - if claimed == 0 { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - } - let wallet = sqlx::query($wallet_sql) - .bind(&target.wallet_id) - .fetch_optional(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let Some(wallet) = wallet else { - sqlx::query( - r#" -UPDATE referral_rewards -SET status = 'failed', - admin_operator_id = COALESCE(?, admin_operator_id), - admin_note = COALESCE(?, admin_note), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(operator_id) - .bind(note.or(Some("邀请人钱包不存在"))) - .bind(now_unix_secs() as i64) - .bind(&target.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - }; - let balance = row_f64!(wallet, "balance"); - let gift_before = row_f64!(wallet, "gift_balance"); - let total_adjusted_before = row_f64!(wallet, "total_adjusted"); - if !referral_wallet_values_valid(balance, gift_before) - || !total_adjusted_before.is_finite() - { - return Err(DataLayerError::InvalidInput( - "inviter wallet balance is invalid".to_string(), - )); - } - let total_before = balance + gift_before; - let gift_after = gift_before + target.amount_usd; - let total_after = balance + gift_after; - let total_adjusted_after = total_adjusted_before + target.amount_usd; - if !gift_after.is_finite() - || !total_before.is_finite() - || !total_after.is_finite() - || !total_adjusted_after.is_finite() - { - return Err(DataLayerError::InvalidInput( - "inviter wallet balance overflowed".to_string(), - )); - } - let tx_id = uuid::Uuid::new_v4().to_string(); - let description = note - .map(ToOwned::to_owned) - .unwrap_or_else(|| reward_description(&target)); - sqlx::query( - r#" -UPDATE wallets -SET gift_balance = ?, - total_adjusted = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(gift_after) - .bind(total_adjusted_after) - .bind(now_unix_secs() as i64) - .bind(&target.wallet_id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - sqlx::query( - r#" -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, - link_type, link_id, operator_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward', ?, ?, ?, ?, ?, ?, ?, - 'referral_reward', ?, ?, ?, ?) -"#, - ) - .bind(&tx_id) - .bind(&target.wallet_id) - .bind(target.amount_usd) - .bind(total_before) - .bind(total_after) - .bind(balance) - .bind(balance) - .bind(gift_before) - .bind(gift_after) - .bind(&target.id) - .bind(operator_id) - .bind(&description) - .bind(now_unix_secs() as i64) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - sqlx::query( - r#" -UPDATE referral_rewards -SET status = 'applied', - wallet_transaction_id = ?, - admin_operator_id = COALESCE(?, admin_operator_id), - admin_note = COALESCE(?, admin_note), - updated_at = ? -WHERE id = ? -"#, - ) - .bind(&tx_id) - .bind(operator_id) - .bind(note) - .bind(now_unix_secs() as i64) - .bind(&target.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - Ok(()) - } - }; -} - -impl ReferralDataState<'_> { - #[cfg(feature = "mysql")] - referral_applying_recovery_numeric_method!( - recover_applying_referral_reward_mysql_numeric_time, - sqlx::MySqlPool - ); - #[cfg(feature = "sqlite")] - referral_applying_recovery_numeric_method!( - recover_applying_referral_reward_sqlite_numeric_time, - sqlx::SqlitePool - ); - - #[cfg(feature = "mysql")] - referral_credit_numeric_method!( - credit_referral_reward_mysql_numeric_time, - sqlx::MySqlPool, - "SELECT wallets.balance, wallets.gift_balance, wallets.total_adjusted - FROM wallets - JOIN users inviter ON inviter.id = wallets.user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 - WHERE wallets.id = ? AND wallets.status = 'active' - FOR UPDATE" - ); - #[cfg(feature = "sqlite")] - referral_credit_numeric_method!( - credit_referral_reward_sqlite_numeric_time, - sqlx::SqlitePool, - "SELECT wallets.balance, wallets.gift_balance, wallets.total_adjusted - FROM wallets - JOIN users inviter ON inviter.id = wallets.user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 - WHERE wallets.id = ? AND wallets.status = 'active'" - ); - - #[cfg(any(feature = "mysql", feature = "sqlite"))] - async fn apply_referral_reward_reversal_numeric_time( - &self, - reward: &ReferralRewardRecord, - ) -> Result<(), DataLayerError> { - let Some(backends) = self.backends.as_ref() else { - return Ok(()); - }; - #[cfg(feature = "mysql")] - if let Some(backend) = backends.mysql() { - return apply_referral_reward_reversal_for_mysql_pool(&backend.pool_clone(), reward) - .await; - } - #[cfg(feature = "sqlite")] - if let Some(backend) = backends.sqlite() { - return apply_referral_reward_reversal_for_sqlite_pool(&backend.pool_clone(), reward) - .await; - } Ok(()) } } @@ -4390,328 +2638,6 @@ where }) } -#[cfg(any(feature = "mysql", feature = "sqlite"))] -macro_rules! referral_reversal_numeric_fn { - ($name:ident, $pool_ty:ty, $wallet_sql:expr, $order_sql:expr) => { - async fn $name( - pool: &$pool_ty, - reward: &ReferralRewardRecord, - ) -> Result<(), DataLayerError> { - let mut tx = pool.begin().await.map_err(DataLayerError::sql)?; - sqlx::query("UPDATE referral_rewards SET updated_at = updated_at WHERE id = ?") - .bind(&reward.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let reward_row = sqlx::query( - r#" -SELECT status, inviter_user_id, source_order_id, amount_usd, - reversed_amount_usd, pending_reversal_amount_usd -FROM referral_rewards -WHERE id = ? -"#, - ) - .bind(&reward.id) - .fetch_optional(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let Some(reward_row) = reward_row else { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - }; - let reward_status = row_string!(reward_row, "status"); - if !matches!(reward_status.as_str(), "applied" | "reversed") { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - } - let Some(source_order_id) = row_optional_string!(reward_row, "source_order_id") else { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - }; - let inviter_user_id = row_string!(reward_row, "inviter_user_id"); - let reward_amount = row_f64!(reward_row, "amount_usd"); - let current_reversed = row_f64!(reward_row, "reversed_amount_usd"); - let current_pending = row_f64!(reward_row, "pending_reversal_amount_usd"); - - // Match the wallet refund lock order. The payment order is read - // only after its row lock so the target reflects the cumulative - // refund that actually won the race with this transaction. - let wallet = sqlx::query($wallet_sql) - .bind(&inviter_user_id) - .fetch_optional(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let order_row = sqlx::query($order_sql) - .bind(&source_order_id) - .fetch_optional(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - let Some(order_row) = order_row else { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - }; - let refund_context = payment_order_refund_context_from_row(order_row)?; - if !referral_refund_context_valid(&refund_context) { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - } - let target_reversal_amount_usd = referral_reversal_target( - reward_amount, - refund_context.amount_usd, - refund_context.refunded_amount_usd, - ); - if !referral_reversal_inputs_valid( - reward_amount, - target_reversal_amount_usd, - current_reversed, - current_pending, - ) { - return Err(DataLayerError::InvalidInput( - "referral reversal state is invalid".to_string(), - )); - } - let amount_usd = referral_reversal_due_bounded( - target_reversal_amount_usd, - reward_amount, - current_reversed, - current_pending, - ); - if amount_usd <= 0.0 { - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - } - let Some(wallet) = wallet else { - // Preserve the debt when the inviter wallet is temporarily - // absent/inactive; the periodic reconciliation pass will - // retry after the wallet becomes available. - sqlx::query( - r#" -UPDATE referral_rewards -SET pending_reversal_amount_usd = ?, - status = CASE - WHEN status = 'reversed' THEN 'applied' - ELSE status - END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(referral_pending_reversal_capped( - reward_amount, - current_reversed, - current_pending, - amount_usd, - )) - .bind(now_unix_secs() as i64) - .bind(&reward.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - return Ok(()); - }; - let wallet_id = row_string!(wallet, "id"); - let balance = row_f64!(wallet, "balance"); - let gift_before = row_f64!(wallet, "gift_balance"); - let total_adjusted_before = row_f64!(wallet, "total_adjusted"); - if !referral_wallet_values_valid(balance, gift_before) - || !total_adjusted_before.is_finite() - { - return Err(DataLayerError::InvalidInput( - "inviter wallet balance is invalid".to_string(), - )); - } - let actual_reverse = gift_before.max(0.0).min(amount_usd); - let pending_reverse = (amount_usd - actual_reverse).max(0.0); - let gift_after = gift_before - actual_reverse; - let total_before = balance + gift_before; - let total_after = balance + gift_after; - let total_adjusted_after = total_adjusted_before - actual_reverse; - if !actual_reverse.is_finite() - || !pending_reverse.is_finite() - || !gift_after.is_finite() - || !total_before.is_finite() - || !total_after.is_finite() - || !total_adjusted_after.is_finite() - || !referral_reversal_state_valid( - reward_amount, - current_reversed, - current_pending, - actual_reverse, - pending_reverse, - ) - { - return Err(DataLayerError::InvalidInput( - "inviter wallet balance overflowed".to_string(), - )); - } - if actual_reverse > 0.0 { - sqlx::query( - r#" -UPDATE wallets -SET gift_balance = ?, - total_adjusted = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(gift_after) - .bind(total_adjusted_after) - .bind(now_unix_secs() as i64) - .bind(&wallet_id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - sqlx::query( - r#" -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, - link_type, link_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward_reversal', ?, ?, ?, ?, ?, ?, ?, - 'referral_reward', ?, '邀请返利退款冲回', ?) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(&wallet_id) - .bind(-actual_reverse) - .bind(total_before) - .bind(total_after) - .bind(balance) - .bind(balance) - .bind(gift_before) - .bind(gift_after) - .bind(&reward.id) - .bind(now_unix_secs() as i64) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - } - sqlx::query( - r#" -UPDATE referral_rewards -SET reversed_amount_usd = reversed_amount_usd + ?, - pending_reversal_amount_usd = ?, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(actual_reverse) - .bind(pending_reverse) - .bind(now_unix_secs() as i64) - .bind(&reward.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - sqlx::query( - r#" -UPDATE referral_rewards -SET status = CASE - WHEN pending_reversal_amount_usd > 0.00000001 AND status = 'reversed' THEN 'applied' - WHEN pending_reversal_amount_usd <= 0.00000001 - AND reversed_amount_usd >= amount_usd - AND status IN ('applied', 'reversed') THEN 'reversed' - ELSE status - END, - updated_at = ? -WHERE id = ? -"#, - ) - .bind(now_unix_secs() as i64) - .bind(&reward.id) - .execute(&mut *tx) - .await - .map_err(DataLayerError::sql)?; - tx.commit().await.map_err(DataLayerError::sql)?; - Ok(()) - } - }; -} - -#[cfg(feature = "mysql")] -referral_reversal_numeric_fn!( - apply_referral_reward_reversal_for_mysql_pool, - sqlx::MySqlPool, - "SELECT wallets.id, wallets.balance, wallets.gift_balance, wallets.total_adjusted - FROM wallets - JOIN users inviter ON inviter.id = wallets.user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 - WHERE wallets.user_id = ? AND wallets.status = 'active' - FOR UPDATE", - "SELECT po.amount_usd, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd - FROM payment_orders po - WHERE po.id = ? - FOR UPDATE" -); -#[cfg(feature = "sqlite")] -referral_reversal_numeric_fn!( - apply_referral_reward_reversal_for_sqlite_pool, - sqlx::SqlitePool, - "SELECT wallets.id, wallets.balance, wallets.gift_balance, wallets.total_adjusted - FROM wallets - JOIN users inviter ON inviter.id = wallets.user_id - AND inviter.is_active = 1 AND inviter.is_deleted = 0 - WHERE wallets.user_id = ? AND wallets.status = 'active'", - "SELECT po.amount_usd, - CASE - WHEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) >= - COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - THEN COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'succeeded' - ), 0.0) - ELSE COALESCE(po.refunded_amount_usd, 0.0) - COALESCE(( - SELECT SUM(rr.amount_usd) - FROM refund_requests rr - WHERE rr.payment_order_id = po.id - AND rr.status = 'processing' - ), 0.0) - END AS refunded_amount_usd - FROM payment_orders po - WHERE po.id = ?" -); - #[cfg(test)] mod tests { use super::*; @@ -4926,1160 +2852,4 @@ mod tests { assert_eq!(referral_page_bounds(999, 4), (200, 4)); assert_eq!(referral_page_bounds(20, usize::MAX), (20, i64::MAX)); } - - #[cfg(feature = "sqlite")] - #[tokio::test] - async fn referral_refund_context_combines_legacy_and_settled_refunds() { - let config = crate::DataLayerConfig::from_database(crate::SqlDatabaseConfig { - driver: crate::DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: crate::SqlPoolConfig { - max_connections: 1, - ..crate::SqlPoolConfig::default() - }, - }); - let backends = - crate::DataBackends::from_config(config).expect("sqlite data backends should build"); - let pool = backends - .sqlite() - .expect("sqlite backend should exist") - .pool(); - crate::lifecycle::migrate::run_sqlite_migrations(pool) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, 'user', ?, ?)", - ) - .bind("refund-context-user") - .bind("refund-context@example.test") - .bind("refund-context-user") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("refund context user should insert"); - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, created_at, updated_at) VALUES (?, ?, ?, ?, 'active', ?, ?)", - ) - .bind("refund-context-wallet") - .bind("refund-context-user") - .bind(0.0_f64) - .bind(0.0_f64) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("refund context wallet should insert"); - sqlx::query( - "INSERT INTO payment_orders (id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, refundable_amount_usd, payment_method, status, created_at, credited_at, order_kind) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("refund-context-order") - .bind("refund-context-order-no") - .bind("refund-context-wallet") - .bind("refund-context-user") - .bind(10.0_f64) - .bind(2.0_f64) - .bind(8.0_f64) - .bind("stripe") - .bind("credited") - .bind(1_i64) - .bind(1_i64) - .bind("wallet_recharge") - .execute(pool) - .await - .expect("refund context order should insert"); - - let state = ReferralDataState::new(Some(&backends)); - let context = state - .find_referral_payment_order_refund_context("refund-context-order") - .await - .expect("legacy refund context should query") - .expect("refund context order should exist"); - assert!((context.refunded_amount_usd - 2.0).abs() < f64::EPSILON); - - // The processing request has already increased the legacy order - // counter, but it must not authorize a referral reversal yet. - sqlx::query( - "INSERT INTO refund_requests (id, refund_no, wallet_id, user_id, payment_order_id, source_type, refund_mode, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("refund-context-processing") - .bind("refund-context-processing-no") - .bind("refund-context-wallet") - .bind("refund-context-user") - .bind("refund-context-order") - .bind("wallet") - .bind("offline_payout") - .bind(3.0_f64) - .bind("processing") - .bind(2_i64) - .bind(2_i64) - .execute(pool) - .await - .expect("processing refund should insert"); - sqlx::query( - "UPDATE payment_orders SET refunded_amount_usd = ?, refundable_amount_usd = ? WHERE id = ?", - ) - .bind(5.0_f64) - .bind(5.0_f64) - .bind("refund-context-order") - .execute(pool) - .await - .expect("processing order counter should update"); - let context = state - .find_referral_payment_order_refund_context("refund-context-order") - .await - .expect("processing refund context should query") - .expect("processing refund context order should exist"); - assert!((context.refunded_amount_usd - 2.0).abs() < f64::EPSILON); - - // A settled request is additive to the historical counter. While the - // newer request is still processing, the effective settled amount must - // retain the legacy two dollars rather than dropping to zero. - sqlx::query( - "INSERT INTO refund_requests (id, refund_no, wallet_id, user_id, payment_order_id, source_type, refund_mode, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("refund-context-succeeded") - .bind("refund-context-succeeded-no") - .bind("refund-context-wallet") - .bind("refund-context-user") - .bind("refund-context-order") - .bind("wallet") - .bind("offline_payout") - .bind(2.0_f64) - .bind("succeeded") - .bind(3_i64) - .bind(3_i64) - .execute(pool) - .await - .expect("succeeded refund should insert"); - let context = state - .find_referral_payment_order_refund_context("refund-context-order") - .await - .expect("mixed refund context should query") - .expect("mixed refund context order should exist"); - assert!((context.refunded_amount_usd - 2.0).abs() < f64::EPSILON); - - sqlx::query( - "UPDATE refund_requests SET status = 'succeeded', processed_at = ? WHERE id = ?", - ) - .bind(4_i64) - .bind("refund-context-processing") - .execute(pool) - .await - .expect("processing refund should settle"); - let context = state - .find_referral_payment_order_refund_context("refund-context-order") - .await - .expect("settled refund context should query") - .expect("settled refund context order should exist"); - assert!((context.refunded_amount_usd - 5.0).abs() < f64::EPSILON); - - // Even if an imported order counter is stale, a durable succeeded - // request must not be erased by the aggregate fallback. - sqlx::query( - "UPDATE payment_orders SET refunded_amount_usd = ?, refundable_amount_usd = ? WHERE id = ?", - ) - .bind(0.0_f64) - .bind(10.0_f64) - .bind("refund-context-order") - .execute(pool) - .await - .expect("stale order counter should update"); - let context = state - .find_referral_payment_order_refund_context("refund-context-order") - .await - .expect("stale counter refund context should query") - .expect("stale counter refund context order should exist"); - assert!((context.refunded_amount_usd - 5.0).abs() < f64::EPSILON); - } - - #[cfg(feature = "sqlite")] - #[tokio::test] - async fn admin_referral_lists_are_not_truncated_at_fetch_limit() { - let config = crate::DataLayerConfig::from_database(crate::SqlDatabaseConfig { - driver: crate::DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: crate::SqlPoolConfig { - max_connections: 1, - ..crate::SqlPoolConfig::default() - }, - }); - let backends = - crate::DataBackends::from_config(config).expect("sqlite data backends should build"); - let pool = backends - .sqlite() - .expect("sqlite backend should exist") - .pool(); - crate::lifecycle::migrate::run_sqlite_migrations(pool) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, 'user', ?, ?)", - ) - .bind("large-list-inviter") - .bind("large-list-inviter@example.test") - .bind("large-list-inviter") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("inviter should insert"); - - let mut tx = pool.begin().await.expect("bulk transaction should begin"); - for index in 0..=REFERRAL_FETCH_LIMIT { - let user_id = format!("large-list-invitee-{index}"); - let email = format!("large-list-invitee-{index}@example.test"); - let username = format!("large-list-invitee-{index}"); - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, 'user', ?, ?)", - ) - .bind(&user_id) - .bind(&email) - .bind(&username) - .bind(index as i64 + 2) - .bind(index as i64 + 2) - .execute(&mut *tx) - .await - .expect("invitee should insert"); - sqlx::query( - "INSERT INTO user_referrals (id, inviter_user_id, invitee_user_id, invite_code_snapshot, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind(format!("large-list-referral-{index}")) - .bind("large-list-inviter") - .bind(&user_id) - .bind("AE-LARGE-LIST") - .bind(index as i64 + 2) - .bind(index as i64 + 2) - .execute(&mut *tx) - .await - .expect("referral relationship should insert"); - sqlx::query( - "INSERT INTO referral_rewards (id, referral_id, inviter_user_id, invitee_user_id, reward_type, trigger_point, idempotency_key, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, 'percent', 'paid_order', ?, ?, 'pending', ?, ?)", - ) - .bind(format!("large-list-reward-{index}")) - .bind(format!("large-list-referral-{index}")) - .bind("large-list-inviter") - .bind(&user_id) - .bind(format!("large-list-reward-key-{index}")) - .bind(1.0_f64) - .bind(index as i64 + 2) - .bind(index as i64 + 2) - .execute(&mut *tx) - .await - .expect("referral reward should insert"); - } - tx.commit().await.expect("bulk transaction should commit"); - - let state = ReferralDataState::new(Some(&backends)); - let (items, total, stats) = state - .list_admin_referral_relationships(ReferralRelationshipListQuery { - inviter: Some("large-list-inviter".to_string()), - limit: 1, - offset: REFERRAL_FETCH_LIMIT, - ..ReferralRelationshipListQuery::default() - }) - .await - .expect("large relationship list should succeed") - .expect("sqlite referral backend should be available"); - assert_eq!(total, (REFERRAL_FETCH_LIMIT + 1) as u64); - assert_eq!(items.len(), 1); - assert_eq!(stats.total_invites, (REFERRAL_FETCH_LIMIT + 1) as u64); - - let (reward_items, reward_total, reward_stats) = state - .list_admin_referral_rewards(ReferralRewardListQuery { - order_id: None, - reward_type: Some("percent".to_string()), - status: Some("pending".to_string()), - limit: 1, - offset: REFERRAL_FETCH_LIMIT, - }) - .await - .expect("large reward list should succeed") - .expect("sqlite referral backend should be available"); - assert_eq!(reward_total, (REFERRAL_FETCH_LIMIT + 1) as u64); - assert_eq!(reward_items.len(), 1); - assert_eq!( - reward_stats.total_invites, - (REFERRAL_FETCH_LIMIT + 1) as u64 - ); - assert_eq!( - reward_stats.pending_reward_usd, - (REFERRAL_FETCH_LIMIT + 1) as f64 - ); - } - - #[cfg(feature = "sqlite")] - #[tokio::test] - async fn reconciliation_recovers_applying_rewards_from_wallet_transaction_facts() { - let config = crate::DataLayerConfig::from_database(crate::SqlDatabaseConfig { - driver: crate::DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: crate::SqlPoolConfig { - max_connections: 1, - ..crate::SqlPoolConfig::default() - }, - }); - let backends = - crate::DataBackends::from_config(config).expect("sqlite data backends should build"); - let pool = backends - .sqlite() - .expect("sqlite backend should exist") - .pool(); - crate::lifecycle::migrate::run_sqlite_migrations(pool) - .await - .expect("sqlite migrations should run"); - - for (id, email, username) in [ - ( - "applying-inviter", - "applying-inviter@example.test", - "applying-inviter", - ), - ( - "applying-invitee", - "applying-invitee@example.test", - "applying-invitee", - ), - ] { - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, 'user', ?, ?)", - ) - .bind(id) - .bind(email) - .bind(username) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("referral user should insert"); - } - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, total_adjusted, created_at, updated_at) VALUES (?, ?, ?, ?, 'active', ?, ?, ?)", - ) - .bind("applying-wallet") - .bind("applying-inviter") - .bind(0.0_f64) - // This is the already-committed credit represented by existing-tx. - .bind(2.0_f64) - .bind(2.0_f64) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("inviter wallet should insert"); - sqlx::query( - "INSERT INTO user_referrals (id, inviter_user_id, invitee_user_id, invite_code_snapshot, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("applying-referral") - .bind("applying-inviter") - .bind("applying-invitee") - .bind("AE-APPLYING") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("referral relationship should insert"); - for (id, key, amount) in [ - ("applying-with-tx", "applying-key-with-tx", 2.0_f64), - ("applying-without-tx", "applying-key-without-tx", 3.0_f64), - ( - "applying-with-invalid-tx", - "applying-key-with-invalid-tx", - 4.0_f64, - ), - ] { - sqlx::query( - "INSERT INTO referral_rewards (id, referral_id, inviter_user_id, invitee_user_id, reward_type, trigger_point, idempotency_key, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, 'percent', 'paid_order', ?, ?, 'applying', ?, ?)", - ) - .bind(id) - .bind("applying-referral") - .bind("applying-inviter") - .bind("applying-invitee") - .bind(key) - .bind(amount) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("applying reward should insert"); - } - sqlx::query( - r#" -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, - link_type, link_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward', ?, ?, ?, ?, ?, ?, ?, - 'referral_reward', ?, 'existing referral credit', ?) -"#, - ) - .bind("existing-referral-tx") - .bind("applying-wallet") - .bind(2.0_f64) - .bind(0.0_f64) - .bind(2.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(2.0_f64) - .bind("applying-with-tx") - .bind(1_i64) - .execute(pool) - .await - .expect("existing referral transaction should insert"); - // A row with the right link id and amount is still not proof of a - // credit when its balance snapshot is inconsistent. Recovery must - // validate the complete transaction shape before moving an `applying` - // reward to `applied`. Because this is still evidence of an ambiguous - // historical write, it must not be downgraded to `failed` (which would - // make the next pass credit the wallet a second time). - sqlx::query( - r#" -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, - link_type, link_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward', 4, 2, 6, 0, 0, 2, 2, - 'referral_reward', ?, 'non-credit fact', ?) -"#, - ) - .bind("non-credit-referral-tx") - .bind("applying-wallet") - .bind("applying-with-invalid-tx") - .bind(2_i64) - .execute(pool) - .await - .expect("non-credit transaction should insert"); - let state = ReferralDataState::new(Some(&backends)); - let first = state - .reconcile_referral_rewards_once(None) - .await - .expect("applying recovery should succeed"); - assert_eq!(first.reward_attempted, 3); - assert_eq!(first.reward_applied, 1); - assert_eq!(first.deferred, 2); - let dashboard_after_recovery = state - .referral_dashboard("applying-inviter") - .await - .expect("applying dashboard should aggregate") - .expect("applying inviter dashboard should exist"); - assert!((dashboard_after_recovery.paid_reward_usd - 2.0).abs() < f64::EPSILON); - assert!((dashboard_after_recovery.pending_reward_usd - 7.0).abs() < f64::EPSILON); - - let (gift_after_recovery, transaction_count): (f64, i64) = sqlx::query_as( - "SELECT gift_balance, (SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?) FROM wallets WHERE id = ?", - ) - .bind("applying-wallet") - .bind("applying-wallet") - .fetch_one(pool) - .await - .expect("wallet should remain readable"); - assert!((gift_after_recovery - 2.0).abs() < f64::EPSILON); - assert_eq!(transaction_count, 2); - - let (with_tx_status, recovered_tx_id): (String, Option) = sqlx::query_as( - "SELECT status, wallet_transaction_id FROM referral_rewards WHERE id = ?", - ) - .bind("applying-with-tx") - .fetch_one(pool) - .await - .expect("recovered reward should be readable"); - assert_eq!(with_tx_status, "applied"); - assert_eq!(recovered_tx_id.as_deref(), Some("existing-referral-tx")); - let (without_tx_status, missing_tx_id): (String, Option) = sqlx::query_as( - "SELECT status, wallet_transaction_id FROM referral_rewards WHERE id = ?", - ) - .bind("applying-without-tx") - .fetch_one(pool) - .await - .expect("failed reward should be readable"); - assert_eq!(without_tx_status, "failed"); - assert!(missing_tx_id.is_none()); - let (invalid_tx_status, invalid_tx_id): (String, Option) = sqlx::query_as( - "SELECT status, wallet_transaction_id FROM referral_rewards WHERE id = ?", - ) - .bind("applying-with-invalid-tx") - .fetch_one(pool) - .await - .expect("ambiguous reward should be readable"); - assert_eq!(invalid_tx_status, "applying"); - assert!(invalid_tx_id.is_none()); - - // Only the evidence-free reward may retry through the normal credit - // transaction. The ambiguous reward is inspected again but remains - // applying and never credits a second time. - let second = state - .reconcile_referral_rewards_once(None) - .await - .expect("failed reward retry should succeed"); - assert_eq!(second.reward_attempted, 2); - assert_eq!(second.reward_applied, 1); - assert_eq!(second.deferred, 1); - let (gift_after_retry, transaction_count): (f64, i64) = sqlx::query_as( - "SELECT gift_balance, (SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?) FROM wallets WHERE id = ?", - ) - .bind("applying-wallet") - .bind("applying-wallet") - .fetch_one(pool) - .await - .expect("retried wallet should be readable"); - assert!((gift_after_retry - 5.0).abs() < f64::EPSILON); - assert_eq!(transaction_count, 3); - - let third = state - .reconcile_referral_rewards_once(None) - .await - .expect("settled rewards should be idempotent"); - assert_eq!(third.reward_attempted, 1); - assert_eq!(third.reward_applied, 0); - assert_eq!(third.deferred, 1); - let final_gift: f64 = sqlx::query_scalar("SELECT gift_balance FROM wallets WHERE id = ?") - .bind("applying-wallet") - .fetch_one(pool) - .await - .expect("final wallet balance should be readable"); - assert!((final_gift - 5.0).abs() < f64::EPSILON); - let invalid_final_status: String = - sqlx::query_scalar("SELECT status FROM referral_rewards WHERE id = ?") - .bind("applying-with-invalid-tx") - .fetch_one(pool) - .await - .expect("ambiguous reward should remain readable"); - assert_eq!(invalid_final_status, "applying"); - } - - #[cfg(feature = "sqlite")] - #[tokio::test] - async fn reconciliation_rotates_ambiguous_applying_rows_without_starving_valid_facts() { - let config = crate::DataLayerConfig::from_database(crate::SqlDatabaseConfig { - driver: crate::DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: crate::SqlPoolConfig { - max_connections: 1, - ..crate::SqlPoolConfig::default() - }, - }); - let backends = - crate::DataBackends::from_config(config).expect("sqlite data backends should build"); - let pool = backends - .sqlite() - .expect("sqlite backend should exist") - .pool(); - crate::lifecycle::migrate::run_sqlite_migrations(pool) - .await - .expect("sqlite migrations should run"); - - for (id, email, username) in [ - ( - "rotation-inviter", - "rotation-inviter@example.test", - "rotation-inviter", - ), - ( - "rotation-invitee", - "rotation-invitee@example.test", - "rotation-invitee", - ), - ] { - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, 'user', ?, ?)", - ) - .bind(id) - .bind(email) - .bind(username) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("rotation user should insert"); - } - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, total_adjusted, created_at, updated_at) VALUES (?, ?, ?, ?, 'active', ?, ?, ?)", - ) - .bind("rotation-wallet") - .bind("rotation-inviter") - .bind(0.0_f64) - .bind(1.0_f64) - .bind(1.0_f64) - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("rotation wallet should insert"); - sqlx::query( - "INSERT INTO user_referrals (id, inviter_user_id, invitee_user_id, invite_code_snapshot, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("rotation-referral") - .bind("rotation-inviter") - .bind("rotation-invitee") - .bind("AE-ROTATION") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("rotation referral should insert"); - - // Fill the bounded page with malformed applying rows. Their durable - // amount is invalid, so recovery must leave them applying and rotate - // their updated_at instead of allowing them to monopolise the queue. - // Use a future timestamp to ensure rotation never moves a corrupted - // imported value backwards into the queue's oldest position. - let imported_updated_at = 4_102_444_800_i64; - for index in 0..REFERRAL_RECONCILIATION_LIMIT { - sqlx::query( - "INSERT INTO referral_rewards (id, referral_id, inviter_user_id, invitee_user_id, reward_type, trigger_point, idempotency_key, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, 'percent', 'paid_order', ?, ?, 'applying', ?, ?)", - ) - .bind(format!("rotation-noise-{index:03}")) - .bind("rotation-referral") - .bind("rotation-inviter") - .bind("rotation-invitee") - .bind(format!("rotation-noise-key-{index:03}")) - .bind(0.0_f64) - .bind(1_i64) - .bind(imported_updated_at) - .execute(pool) - .await - .expect("malformed applying row should insert"); - } - sqlx::query( - "INSERT INTO referral_rewards (id, referral_id, inviter_user_id, invitee_user_id, reward_type, trigger_point, idempotency_key, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, 'percent', 'paid_order', ?, ?, 'applying', ?, ?)", - ) - .bind("rotation-valid") - .bind("rotation-referral") - .bind("rotation-inviter") - .bind("rotation-invitee") - .bind("rotation-valid-key") - .bind(1.0_f64) - .bind(2_i64) - .bind(imported_updated_at) - .execute(pool) - .await - .expect("valid applying row should insert"); - sqlx::query( - r#" -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, - link_type, link_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward', ?, ?, ?, ?, ?, ?, ?, - 'referral_reward', ?, 'rotation credit', ?) -"#, - ) - .bind("rotation-valid-tx") - .bind("rotation-wallet") - .bind(1.0_f64) - .bind(0.0_f64) - .bind(1.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(1.0_f64) - .bind("rotation-valid") - .bind(2_i64) - .execute(pool) - .await - .expect("valid wallet fact should insert"); - - let state = ReferralDataState::new(Some(&backends)); - let first = state - .reconcile_referral_rewards_once(None) - .await - .expect("first rotation pass should succeed"); - assert_eq!(first.reward_attempted, REFERRAL_RECONCILIATION_LIMIT as u64); - assert_eq!(first.reward_applied, 0); - assert_eq!(first.deferred, REFERRAL_RECONCILIATION_LIMIT as u64); - let first_valid_status: String = - sqlx::query_scalar("SELECT status FROM referral_rewards WHERE id = ?") - .bind("rotation-valid") - .fetch_one(pool) - .await - .expect("valid row should remain readable"); - assert_eq!(first_valid_status, "applying"); - - // Two independently committed, internally consistent facts for the - // same reward are ambiguous: the wallet may already have been - // credited twice. Recovery must refuse to hide that duplicate. - sqlx::query( - r#" -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, - link_type, link_id, description, created_at -) -VALUES (?, ?, 'adjust', 'referral_reward', ?, ?, ?, ?, ?, ?, ?, - 'referral_reward', ?, 'duplicate rotation credit', ?) -"#, - ) - .bind("rotation-valid-tx-duplicate") - .bind("rotation-wallet") - .bind(1.0_f64) - .bind(0.0_f64) - .bind(1.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(0.0_f64) - .bind(1.0_f64) - .bind("rotation-valid") - .bind(3_i64) - .execute(pool) - .await - .expect("duplicate wallet fact should insert"); - - let second = state - .reconcile_referral_rewards_once(None) - .await - .expect("second rotation pass should succeed"); - assert_eq!(second.reward_applied, 0); - let duplicate_status: String = - sqlx::query_scalar("SELECT status FROM referral_rewards WHERE id = ?") - .bind("rotation-valid") - .fetch_one(pool) - .await - .expect("duplicate reward should remain readable"); - assert_eq!(duplicate_status, "applying"); - - sqlx::query("DELETE FROM wallet_transactions WHERE id = ?") - .bind("rotation-valid-tx-duplicate") - .execute(pool) - .await - .expect("duplicate wallet fact should be removed for recovery test"); - let third = state - .reconcile_referral_rewards_once(None) - .await - .expect("unambiguous rotation pass should succeed"); - assert_eq!(third.reward_applied, 1); - let (valid_status, valid_tx_id, gift_balance): (String, Option, f64) = - sqlx::query_as( - "SELECT (SELECT status FROM referral_rewards WHERE id = ?), (SELECT wallet_transaction_id FROM referral_rewards WHERE id = ?), (SELECT gift_balance FROM wallets WHERE id = ?)", - ) - .bind("rotation-valid") - .bind("rotation-valid") - .bind("rotation-wallet") - .fetch_one(pool) - .await - .expect("rotated valid fact should be readable"); - assert_eq!(valid_status, "applied"); - assert_eq!(valid_tx_id.as_deref(), Some("rotation-valid-tx")); - assert!((gift_balance - 1.0).abs() < f64::EPSILON); - } - - #[cfg(feature = "sqlite")] - #[tokio::test] - async fn reconciliation_does_not_infer_missing_historical_rewards() { - let config = crate::DataLayerConfig::from_database(crate::SqlDatabaseConfig { - driver: crate::DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: crate::SqlPoolConfig { - max_connections: 1, - ..crate::SqlPoolConfig::default() - }, - }); - let backends = - crate::DataBackends::from_config(config).expect("sqlite data backends should build"); - let pool = backends - .sqlite() - .expect("sqlite backend should exist") - .pool(); - crate::lifecycle::migrate::run_sqlite_migrations(pool) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("referral-inviter") - .bind("inviter@example.test") - .bind("inviter") - .bind("user") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("inviter should insert"); - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("referral-invitee") - .bind("invitee@example.test") - .bind("invitee") - .bind("user") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("invitee should insert"); - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .bind("wallet-inviter") - .bind("referral-inviter") - // Recharge balance may be negative when the account has overdraft; - // referral credit must still be able to add to its gift balance. - .bind(-5.0_f64) - .bind(0.0_f64) - .bind("active") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("inviter wallet should insert"); - sqlx::query( - "INSERT INTO wallets (id, user_id, balance, gift_balance, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .bind("wallet-invitee") - .bind("referral-invitee") - .bind(0.0_f64) - .bind(0.0_f64) - .bind("active") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("invitee wallet should insert"); - sqlx::query( - "INSERT INTO user_referrals (id, inviter_user_id, invitee_user_id, invite_code_snapshot, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", - ) - .bind("referral-link") - .bind("referral-inviter") - .bind("referral-invitee") - .bind("AE-TEST") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("referral link should insert"); - sqlx::query( - "INSERT INTO payment_orders (id, order_no, wallet_id, user_id, amount_usd, refundable_amount_usd, payment_method, status, created_at, credited_at, order_kind) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("paid-order-repair") - .bind("paid-order-repair-no") - .bind("wallet-invitee") - .bind("referral-invitee") - .bind(10.0_f64) - .bind(10.0_f64) - .bind("stripe") - .bind("credited") - .bind(1_i64) - .bind(2_i64) - .bind("wallet_recharge") - .execute(pool) - .await - .expect("credited payment order should insert"); - - let state = ReferralDataState::new(Some(&backends)); - let reward_config = ReferralRewardConfig { - percent_enabled: true, - percent_rate: 10.0, - headcount_enabled: false, - headcount_amount_usd: 0.0, - headcount_trigger: "registration".to_string(), - }; - // A credited order with only a referral relationship is not durable - // evidence that the referral feature was enabled for that order. The - // periodic worker must not apply the current configuration to it. - let historical_pass = state - .reconcile_referral_rewards_once(Some(reward_config.clone())) - .await - .expect("historical reconciliation should succeed"); - assert_eq!(historical_pass.order_attempted, 0); - assert_eq!(historical_pass.order_repaired, 0); - let (recharge_balance, gift_balance): (f64, f64) = - sqlx::query_as("SELECT balance, gift_balance FROM wallets WHERE id = ?") - .bind("wallet-inviter") - .fetch_one(pool) - .await - .expect("credited inviter wallet should be readable"); - assert!((recharge_balance + 5.0).abs() < f64::EPSILON); - assert!(gift_balance.abs() < f64::EPSILON); - - // The normal callback path still applies a reward with the - // configuration that was active at payment time. Reconciliation is - // intentionally limited to rows created by that durable path. - let applied = state - .apply_paid_order_referral_rewards("paid-order-repair", reward_config.clone()) - .await - .expect("normal paid-order application should succeed"); - assert_eq!(applied.len(), 1); - let (recharge_balance, gift_balance): (f64, f64) = - sqlx::query_as("SELECT balance, gift_balance FROM wallets WHERE id = ?") - .bind("wallet-inviter") - .fetch_one(pool) - .await - .expect("applied inviter wallet should be readable"); - assert!((recharge_balance + 5.0).abs() < f64::EPSILON); - assert!((gift_balance - 1.0).abs() < f64::EPSILON); - - let dashboard = state - .referral_dashboard("referral-inviter") - .await - .expect("referral dashboard should use the aggregate path") - .expect("inviter dashboard should be available"); - assert_eq!(dashboard.total_invites, 1); - assert_eq!(dashboard.effective_invites, 1); - assert!((dashboard.paid_reward_usd - 1.0).abs() < f64::EPSILON); - - // Headline admin metrics are global and must not become empty or - // filter-scoped just because one of the list queries is narrowed. - let (_, relationship_total, relationship_stats) = state - .list_admin_referral_relationships(ReferralRelationshipListQuery { - inviter: Some("does-not-match".to_string()), - limit: 100, - offset: 0, - ..ReferralRelationshipListQuery::default() - }) - .await - .expect("filtered relationship list should succeed") - .expect("sqlite referral backend should be available"); - assert_eq!(relationship_total, 0); - assert_eq!(relationship_stats.total_invites, 1); - assert_eq!(relationship_stats.effective_invites, 1); - assert!((relationship_stats.paid_reward_usd - 1.0).abs() < f64::EPSILON); - - let (_, reward_total, reward_stats) = state - .list_admin_referral_rewards(ReferralRewardListQuery { - status: Some("voided".to_string()), - limit: 100, - offset: 0, - ..ReferralRewardListQuery::default() - }) - .await - .expect("filtered reward list should succeed") - .expect("sqlite referral backend should be available"); - assert_eq!(reward_total, 0); - assert_eq!(reward_stats.total_invites, 1); - assert_eq!(reward_stats.effective_invites, 1); - assert!((reward_stats.paid_reward_usd - 1.0).abs() < f64::EPSILON); - - let second = state - .reconcile_referral_rewards_once(Some(reward_config.clone())) - .await - .expect("second reconciliation should succeed"); - assert_eq!(second.order_attempted, 0); - assert_eq!(second.order_repaired, 0); - - // A deleted inviter must never receive a delayed reward, even if an - // old pending row and an otherwise active wallet remain in storage. - sqlx::query( - "INSERT INTO referral_rewards (id, referral_id, inviter_user_id, invitee_user_id, reward_type, trigger_point, idempotency_key, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("deleted-inviter-reward") - .bind("referral-link") - .bind("referral-inviter") - .bind("referral-invitee") - .bind("percent") - .bind("paid_order") - .bind("referral:deleted-inviter-reward") - .bind(2.0_f64) - .bind("pending") - .bind(1_i64) - .bind(1_i64) - .execute(pool) - .await - .expect("pending reward should insert"); - sqlx::query("UPDATE users SET is_deleted = 1, is_active = 0 WHERE id = ?") - .bind("referral-inviter") - .execute(pool) - .await - .expect("inviter should be marked deleted"); - let deleted_pass = state - .reconcile_referral_rewards_once(None) - .await - .expect("deleted inviter reconciliation should succeed"); - assert_eq!(deleted_pass.reward_attempted, 0); - assert_eq!(deleted_pass.reward_applied, 0); - let deleted_status: String = - sqlx::query_scalar("SELECT status FROM referral_rewards WHERE id = ?") - .bind("deleted-inviter-reward") - .fetch_one(pool) - .await - .expect("deleted reward should remain readable"); - assert_eq!(deleted_status, "pending"); - sqlx::query("UPDATE referral_rewards SET status = 'voided' WHERE id = ?") - .bind("deleted-inviter-reward") - .execute(pool) - .await - .expect("deleted reward should be voided after the assertion"); - sqlx::query("UPDATE users SET is_deleted = 0, is_active = 1 WHERE id = ?") - .bind("referral-inviter") - .execute(pool) - .await - .expect("inviter should be restored for reversal test"); - - // A refund can be completed after the reward transaction. The reward - // starts with zero pending debt, so this exercises the refund-aware - // candidate query rather than the pending-only retry path. - sqlx::query("UPDATE wallets SET status = 'disabled' WHERE id = ?") - .bind("wallet-inviter") - .execute(pool) - .await - .expect("inviter wallet should be disabled"); - // Processing reserves the user's refund amount before the provider - // settles it. That intermediate state must not authorize a referral - // reversal, even when the legacy payment-order counter is already - // populated. - sqlx::query( - "INSERT INTO refund_requests (id, refund_no, wallet_id, user_id, payment_order_id, source_type, refund_mode, amount_usd, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("refund-processing-reward") - .bind("refund-processing-reward-no") - .bind("wallet-invitee") - .bind("referral-invitee") - .bind("paid-order-repair") - .bind("wallet") - .bind("offline_payout") - .bind(5.0_f64) - .bind("processing") - .bind(3_i64) - .bind(3_i64) - .execute(pool) - .await - .expect("processing refund should insert"); - sqlx::query("UPDATE payment_orders SET refunded_amount_usd = ? WHERE id = ?") - .bind(5.0_f64) - .bind("paid-order-repair") - .execute(pool) - .await - .expect("payment refund should update"); - let processing_reversal = state - .reverse_referral_rewards_for_order("paid-order-repair", 5.0) - .await - .expect("processing refund should not fail referral reconciliation"); - assert!(processing_reversal.is_empty()); - let processing_candidates = state - .list_referral_reversal_candidates() - .await - .expect("processing refund candidates should be queryable"); - assert!(processing_candidates.is_empty()); - - sqlx::query("UPDATE refund_requests SET status = 'succeeded' WHERE id = ?") - .bind("refund-processing-reward") - .execute(pool) - .await - .expect("refund should settle successfully"); - let immediate_reversal = state - .reverse_referral_rewards_for_order("paid-order-repair", 5.0) - .await - .expect("completed refund should persist reversal debt"); - assert_eq!(immediate_reversal.len(), 1); - let disabled_candidates = state - .list_referral_reversal_candidates() - .await - .expect("disabled wallet candidates should be queryable"); - assert!( - disabled_candidates.is_empty(), - "a disabled wallet must not consume the bounded reversal page" - ); - let reversal = state - .reconcile_referral_rewards_once(None) - .await - .expect("refund reconciliation should succeed"); - assert_eq!(reversal.reversal_attempted, 0); - assert_eq!(reversal.reversal_applied, 0); - - let (disabled_gift, disabled_pending): (f64, f64) = sqlx::query_as( - "SELECT (SELECT gift_balance FROM wallets WHERE id = ?), (SELECT pending_reversal_amount_usd FROM referral_rewards WHERE source_order_id = ?)", - ) - .bind("wallet-inviter") - .bind("paid-order-repair") - .fetch_one(pool) - .await - .expect("disabled wallet reversal state should be readable"); - assert!((disabled_gift - 1.0).abs() < f64::EPSILON); - assert!((disabled_pending - 0.5).abs() < f64::EPSILON); - - sqlx::query("UPDATE wallets SET status = 'active' WHERE id = ?") - .bind("wallet-inviter") - .execute(pool) - .await - .expect("inviter wallet should be restored"); - let retry_reversal = state - .reconcile_referral_rewards_once(None) - .await - .expect("restored wallet reversal should succeed"); - assert_eq!(retry_reversal.reversal_attempted, 1); - assert_eq!(retry_reversal.reversal_applied, 1); - let reversal_candidates = state - .list_referral_reversal_candidates() - .await - .expect("fully reconciled reversal should not remain a candidate"); - assert!(reversal_candidates.is_empty()); - - let (gift_balance, transaction_count, oldest_transaction_at): (f64, i64, i64) = - sqlx::query_as( - "SELECT gift_balance, (SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = ?), (SELECT MIN(created_at) FROM wallet_transactions WHERE wallet_id = ?) FROM wallets WHERE id = ?", - ) - .bind("wallet-inviter") - .bind("wallet-inviter") - .bind("wallet-inviter") - .fetch_one(pool) - .await - .expect("wallet state should be readable"); - assert!((gift_balance - 0.5).abs() < f64::EPSILON); - assert_eq!(transaction_count, 2); - // `wallet_transactions.created_at` is stored as Unix seconds by all - // SQL adapters (the public field name retains a historical `_ms` - // suffix). A millisecond value would be roughly three orders larger. - let now_unix_secs = chrono::Utc::now().timestamp(); - assert!(oldest_transaction_at >= now_unix_secs - 60); - assert!(oldest_transaction_at <= now_unix_secs + 60); - - let (reward_count, reversed, pending, status): (i64, f64, f64, String) = - sqlx::query_as( - "SELECT COUNT(*), MAX(reversed_amount_usd), MAX(pending_reversal_amount_usd), MAX(status) FROM referral_rewards WHERE source_order_id = ?", - ) - .bind("paid-order-repair") - .fetch_one(pool) - .await - .expect("reward row should be readable"); - assert_eq!(reward_count, 1); - assert!((reversed - 0.5).abs() < f64::EPSILON); - assert!(pending.abs() < f64::EPSILON); - assert_eq!(status, "applied"); - - // A pending debt must not bypass source-order validation. This can - // happen after an operator/import corrupts a historical order while - // its inviter wallet is active again. - sqlx::query( - "UPDATE referral_rewards SET reversed_amount_usd = 0, pending_reversal_amount_usd = 0.5, status = 'applied' WHERE source_order_id = ?", - ) - .bind("paid-order-repair") - .execute(pool) - .await - .expect("pending reversal fixture should update"); - sqlx::query("UPDATE payment_orders SET amount_usd = 0 WHERE id = ?") - .bind("paid-order-repair") - .execute(pool) - .await - .expect("corrupt order fixture should update"); - let invalid_refund_pass = state - .reconcile_referral_rewards_once(None) - .await - .expect("invalid refund context should be deferred"); - assert_eq!(invalid_refund_pass.reversal_attempted, 0); - assert_eq!(invalid_refund_pass.reversal_applied, 0); - assert_eq!(invalid_refund_pass.deferred, 1); - let (gift_after_invalid, pending_after_invalid): (f64, f64) = sqlx::query_as( - "SELECT (SELECT gift_balance FROM wallets WHERE id = ?), (SELECT pending_reversal_amount_usd FROM referral_rewards WHERE source_order_id = ?)", - ) - .bind("wallet-inviter") - .bind("paid-order-repair") - .fetch_one(pool) - .await - .expect("invalid refund state should be readable"); - assert!((gift_after_invalid - 0.5).abs() < f64::EPSILON); - assert!((pending_after_invalid - 0.5).abs() < f64::EPSILON); - } } diff --git a/crates/aether-data/runtime/src/backend/sqlite.rs b/crates/aether-data/runtime/src/backend/sqlite.rs deleted file mode 100644 index 996399560..000000000 --- a/crates/aether-data/runtime/src/backend/sqlite.rs +++ /dev/null @@ -1,1207 +0,0 @@ -use std::sync::Arc; - -use crate::database::SqlDatabaseConfig; -use crate::driver::sqlite::{SqlitePool, SqlitePoolFactory}; -use crate::repository::announcements::{ - AnnouncementReadRepository, AnnouncementWriteRepository, SqliteAnnouncementRepository, -}; -use crate::repository::audit::{AuditLogReadRepository, SqliteAuditLogReadRepository}; -use crate::repository::auth::{ - AuthApiKeyReadRepository, AuthApiKeyWriteRepository, SqliteAuthApiKeyReadRepository, -}; -use crate::repository::auth_modules::{ - AuthModuleReadRepository, AuthModuleWriteRepository, SqliteAuthModuleReadRepository, - SqliteAuthModuleRepository, -}; -use crate::repository::background_tasks::{ - BackgroundTaskReadRepository, BackgroundTaskWriteRepository, SqliteBackgroundTaskRepository, -}; -use crate::repository::billing::{BillingReadRepository, SqliteBillingReadRepository}; -use crate::repository::candidate_selection::{ - MinimalCandidateSelectionReadRepository, SqliteMinimalCandidateSelectionReadRepository, -}; -use crate::repository::candidates::{ - RequestCandidateReadRepository, RequestCandidateWriteRepository, - SqliteRequestCandidateRepository, -}; -use crate::repository::gemini_file_mappings::{ - GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository, - SqliteGeminiFileMappingRepository, -}; -use crate::repository::global_models::{ - GlobalModelReadRepository, GlobalModelWriteRepository, SqliteGlobalModelReadRepository, -}; -use crate::repository::management_tokens::{ - ManagementTokenReadRepository, ManagementTokenWriteRepository, SqliteManagementTokenRepository, -}; -use crate::repository::oauth_providers::{ - OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository, -}; -use crate::repository::pool_scores::{ - PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository, -}; -use crate::repository::provider_catalog::{ - ProviderCatalogReadRepository, ProviderCatalogWriteRepository, - SqliteProviderCatalogReadRepository, -}; -use crate::repository::proxy_nodes::{ - ProxyNodeReadRepository, ProxyNodeWriteRepository, SqliteProxyNodeReadRepository, -}; -use crate::repository::quota::{ - ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqliteProviderQuotaRepository, -}; -use crate::repository::routing_profiles::{ - RoutingGroupReadRepository, RoutingGroupWriteRepository, SqliteRoutingGroupRepository, -}; -use crate::repository::settlement::{SettlementWriteRepository, SqliteSettlementRepository}; -use crate::repository::usage::{ - SqliteUsageReadRepository, SqliteUsageWriteRepository, UsageReadRepository, - UsageWriteRepository, -}; -use crate::repository::users::{SqliteUserReadRepository, UserReadRepository}; -use crate::repository::video_tasks::{ - SqliteVideoTaskRepository, VideoTaskReadRepository, VideoTaskWriteRepository, -}; -use crate::repository::wallet::{ - SqliteWalletReadRepository, WalletReadRepository, WalletWriteRepository, -}; -use crate::DataLayerError; - -#[derive(Debug, Clone)] -pub struct SqliteBackend { - config: SqlDatabaseConfig, - pool: SqlitePool, -} - -impl SqliteBackend { - pub fn from_config(config: SqlDatabaseConfig) -> Result { - let factory = SqlitePoolFactory::new(config.clone())?; - let pool = factory.connect_lazy()?; - - Ok(Self { config, pool }) - } - - pub fn config(&self) -> &SqlDatabaseConfig { - &self.config - } - - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - - pub fn pool_clone(&self) -> SqlitePool { - self.pool.clone() - } - - pub fn auth_api_key_read_repository(&self) -> Arc { - Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone())) - } - - pub fn announcement_read_repository(&self) -> Arc { - Arc::new(SqliteAnnouncementRepository::new(self.pool_clone())) - } - - pub fn audit_log_read_repository(&self) -> Arc { - Arc::new(SqliteAuditLogReadRepository::new(self.pool_clone())) - } - - pub fn announcement_write_repository(&self) -> Arc { - Arc::new(SqliteAnnouncementRepository::new(self.pool_clone())) - } - - pub fn auth_api_key_write_repository(&self) -> Arc { - Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone())) - } - - pub fn management_token_read_repository(&self) -> Arc { - Arc::new(SqliteManagementTokenRepository::new(self.pool_clone())) - } - - pub fn management_token_write_repository(&self) -> Arc { - Arc::new(SqliteManagementTokenRepository::new(self.pool_clone())) - } - - pub fn auth_module_read_repository(&self) -> Arc { - Arc::new(SqliteAuthModuleReadRepository::new(self.pool_clone())) - } - - pub fn auth_module_write_repository(&self) -> Arc { - Arc::new(SqliteAuthModuleRepository::new(self.pool_clone())) - } - - pub fn billing_read_repository(&self) -> Arc { - Arc::new(SqliteBillingReadRepository::new(self.pool_clone())) - } - - pub fn background_task_read_repository(&self) -> Arc { - Arc::new(SqliteBackgroundTaskRepository::new(self.pool_clone())) - } - - pub fn background_task_write_repository(&self) -> Arc { - Arc::new(SqliteBackgroundTaskRepository::new(self.pool_clone())) - } - - pub fn request_candidate_read_repository(&self) -> Arc { - Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone())) - } - - pub fn request_candidate_write_repository(&self) -> Arc { - Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone())) - } - - pub fn minimal_candidate_selection_read_repository( - &self, - ) -> Arc { - Arc::new(SqliteMinimalCandidateSelectionReadRepository::new( - self.pool_clone(), - )) - } - - pub fn gemini_file_mapping_read_repository(&self) -> Arc { - Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone())) - } - - pub fn gemini_file_mapping_write_repository( - &self, - ) -> Arc { - Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone())) - } - - pub fn global_model_read_repository(&self) -> Arc { - Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone())) - } - - pub fn global_model_write_repository(&self) -> Arc { - Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone())) - } - - pub fn user_read_repository(&self) -> Arc { - Arc::new(SqliteUserReadRepository::new(self.pool_clone())) - } - - pub fn video_task_read_repository(&self) -> Arc { - Arc::new(SqliteVideoTaskRepository::new(self.pool_clone())) - } - - pub fn video_task_write_repository(&self) -> Arc { - Arc::new(SqliteVideoTaskRepository::new(self.pool_clone())) - } - - pub fn oauth_provider_read_repository(&self) -> Arc { - Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone())) - } - - pub fn oauth_provider_write_repository(&self) -> Arc { - Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone())) - } - - pub fn provider_catalog_read_repository(&self) -> Arc { - Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone())) - } - - pub fn provider_catalog_write_repository(&self) -> Arc { - Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone())) - } - - pub fn pool_score_read_repository(&self) -> Arc { - Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone())) - } - - pub fn pool_score_write_repository(&self) -> Arc { - Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone())) - } - - pub fn routing_group_read_repository(&self) -> Arc { - Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone())) - } - - pub fn routing_group_write_repository(&self) -> Arc { - Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone())) - } - - pub fn proxy_node_read_repository(&self) -> Arc { - Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone())) - } - - pub fn proxy_node_write_repository(&self) -> Arc { - Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone())) - } - - pub fn provider_quota_read_repository(&self) -> Arc { - Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone())) - } - - pub fn provider_quota_write_repository(&self) -> Arc { - Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone())) - } - - pub fn settlement_write_repository(&self) -> Arc { - Arc::new(SqliteSettlementRepository::new(self.pool_clone())) - } - - pub fn usage_write_repository(&self) -> Arc { - Arc::new(SqliteUsageWriteRepository::new(self.pool_clone())) - } - - pub fn usage_read_repository(&self) -> Arc { - Arc::new(SqliteUsageReadRepository::new(self.pool_clone())) - } - - pub fn wallet_read_repository(&self) -> Arc { - Arc::new(SqliteWalletReadRepository::new(self.pool_clone())) - } - - pub fn wallet_write_repository(&self) -> Arc { - Arc::new(SqliteWalletReadRepository::new(self.pool_clone())) - } -} - -#[cfg(test)] -mod tests { - use std::collections::BTreeMap; - - use super::SqliteBackend; - use crate::lifecycle::migrate::run_sqlite_migrations; - use crate::repository::system::{ - AdminSystemPurgeTarget, AdminSystemStatsDailyAggregate, - AdminSystemStatsDailyApiKeyAggregate, AdminSystemStatsUserDailyAggregate, - AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateSnapshot, - }; - use crate::{ - DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, StatsDailyAggregationInput, - StatsHourlyAggregationInput, WalletDailyUsageAggregationInput, - }; - - #[tokio::test] - async fn backend_retains_config_and_pool() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite://./data/aether.db".to_string(), - pool: SqlPoolConfig::default(), - }; - - let backend = SqliteBackend::from_config(config.clone()).expect("backend should build"); - - assert_eq!(backend.config(), &config); - let _pool = backend.pool(); - let _pool_clone = backend.pool_clone(); - } - - #[tokio::test] - async fn system_config_round_trips_after_sqlite_migrations() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - let value = serde_json::json!("enabled"); - let stored = backend - .upsert_system_config_entry("feature.local", &value, Some("local flag")) - .await - .expect("system config should upsert"); - assert_eq!(stored.value, value); - assert_eq!( - backend - .find_system_config_value("feature.local") - .await - .expect("system config should read"), - Some(value.clone()) - ); - let replacement = serde_json::json!("disabled"); - assert!(!backend - .compare_and_set_system_config_string_value("feature.local", "stale", "disabled") - .await - .expect("stale system config compare-and-set should complete")); - assert!(backend - .compare_and_set_system_config_string_value("feature.local", "enabled", "disabled") - .await - .expect("matching system config compare-and-set should complete")); - assert_eq!( - backend - .find_system_config_value("feature.local") - .await - .expect("updated system config should read"), - Some(replacement.clone()) - ); - sqlx::query("UPDATE system_configs SET value = ? WHERE key = ?") - .bind(r#""\u5bc6\u94a5""#) - .bind("feature.local") - .execute(backend.pool()) - .await - .expect("legacy escaped JSON string should persist"); - assert!(backend - .compare_and_set_system_config_string_value("feature.local", "密钥", "encrypted-value",) - .await - .expect("escaped JSON string compare-and-set should complete")); - assert_eq!( - backend - .find_system_config_value("feature.local") - .await - .expect("escaped JSON string replacement should read"), - Some(serde_json::json!("encrypted-value")) - ); - assert_eq!( - backend - .list_system_config_entries() - .await - .expect("system config should list") - .len(), - 2 - ); - assert!(backend - .delete_system_config_value("feature.local") - .await - .expect("system config should delete")); - } - - #[tokio::test] - async fn table_maintenance_runs_after_sqlite_migrations() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - let summary = backend - .run_table_maintenance(&["usage", "request_candidates", "audit_logs"]) - .await - .expect("sqlite table maintenance should run"); - - assert_eq!(summary.attempted, 3); - assert_eq!(summary.succeeded, 3); - } - - #[tokio::test] - async fn admin_system_config_purge_deletes_config_scope_and_preserves_users() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - "INSERT INTO users (id, email, username, role, created_at, updated_at) VALUES ('admin-1', 'admin@example.com', 'admin', 'admin', 1, 1)", - ) - .execute(backend.pool()) - .await - .expect("user should insert"); - sqlx::query( - "INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES ('provider-1', 'OpenAI', 'openai', 1, 1)", - ) - .execute(backend.pool()) - .await - .expect("provider should insert"); - sqlx::query( - "INSERT INTO system_configs (id, key, value, created_at, updated_at) VALUES ('config-1', 'site_name', '\"Aether\"', 1, 1)", - ) - .execute(backend.pool()) - .await - .expect("system config should insert"); - - let summary = backend - .purge_admin_system_data(AdminSystemPurgeTarget::Config) - .await - .expect("config purge should run"); - assert!(summary.total() >= 2); - assert_eq!(sqlite_count(backend.pool(), "system_configs").await, 0); - assert_eq!(sqlite_count(backend.pool(), "providers").await, 0); - assert_eq!(sqlite_count(backend.pool(), "users").await, 1); - } - - #[tokio::test] - async fn admin_system_users_purge_deletes_only_non_admin_users_and_keys() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - r#" -INSERT INTO users (id, email, username, role, created_at, updated_at) -VALUES - ('admin-1', 'admin@example.com', 'admin', 'admin', 1, 1), - ('user-1', 'user@example.com', 'alice', 'user', 1, 1) -"#, - ) - .execute(backend.pool()) - .await - .expect("users should insert"); - sqlx::query( - r#" -INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at, total_requests, total_tokens, total_cost_usd) -VALUES - ('admin-key-1', 'admin-1', 'hash-admin', 'admin-key', 1, 1, 5, 50, 0.5), - ('user-key-1', 'user-1', 'hash-user', 'user-key', 1, 1, 7, 70, 0.7) -"#, - ) - .execute(backend.pool()) - .await - .expect("api keys should insert"); - sqlx::query( - r#" -INSERT INTO stats_daily_api_key (id, api_key_id, "date", total_requests, created_at, updated_at) -VALUES - ('admin-key-stats-1', 'admin-key-1', 1, 5, 1, 1), - ('user-key-stats-1', 'user-key-1', 1, 7, 1, 1) -"#, - ) - .execute(backend.pool()) - .await - .expect("api key stats should insert"); - - let summary = backend - .purge_admin_system_data(AdminSystemPurgeTarget::Users) - .await - .expect("users purge should run"); - assert!(summary.total() >= 2); - assert_eq!(sqlite_count(backend.pool(), "users").await, 1); - assert_eq!(sqlite_count(backend.pool(), "api_keys").await, 1); - assert_eq!(sqlite_count(backend.pool(), "stats_daily_api_key").await, 1); - let admin_exists: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM users WHERE id = 'admin-1'") - .fetch_one(backend.pool()) - .await - .expect("admin count should load"); - assert_eq!(admin_exists, 1); - } - - #[tokio::test] - async fn admin_system_usage_aggregates_round_trip_after_sqlite_migrations() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - r#" -INSERT INTO users (id, email, username, role, created_at, updated_at) -VALUES ('target-user-1', 'target@example.com', 'target', 'user', 1, 1) -"#, - ) - .execute(backend.pool()) - .await - .expect("target user should insert"); - sqlx::query( - r#" -INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at) -VALUES ('target-key-1', 'target-user-1', 'hash-target-key', 'target key', 1, 1) -"#, - ) - .execute(backend.pool()) - .await - .expect("target key should insert"); - - let snapshot = AdminSystemUsageAggregateSnapshot { - stats_daily: vec![AdminSystemStatsDailyAggregate { - date_unix_secs: 86_400, - total_requests: 9, - success_requests: 8, - error_requests: 1, - input_tokens: 100, - output_tokens: 200, - cache_creation_tokens: 3, - cache_read_tokens: 4, - total_cost: 1.25, - actual_total_cost: 1.0, - is_complete: true, - aggregated_at_unix_secs: Some(90_000), - }], - stats_user_daily: vec![AdminSystemStatsUserDailyAggregate { - user_id: "source-user-1".to_string(), - username: Some("source".to_string()), - date_unix_secs: 86_400, - total_requests: 5, - success_requests: 5, - error_requests: 0, - input_tokens: 50, - output_tokens: 60, - cache_creation_tokens: 1, - cache_read_tokens: 2, - total_cost: 0.5, - }], - stats_daily_api_key: vec![AdminSystemStatsDailyApiKeyAggregate { - api_key_id: "source-key-1".to_string(), - api_key_name: Some("source key".to_string()), - date_unix_secs: 86_400, - total_requests: 4, - success_requests: 3, - error_requests: 1, - input_tokens: 40, - output_tokens: 30, - cache_creation_tokens: 2, - cache_read_tokens: 1, - total_cost: 0.75, - }], - }; - let user_id_map = - BTreeMap::from([("source-user-1".to_string(), "target-user-1".to_string())]); - let api_key_id_map = - BTreeMap::from([("source-key-1".to_string(), "target-key-1".to_string())]); - - let validation_summary = backend - .import_admin_system_usage_aggregates( - &snapshot, - &user_id_map, - &api_key_id_map, - AdminSystemUsageAggregateImportMode::ValidateError, - ) - .await - .expect("usage aggregates should validate"); - assert_eq!(validation_summary.stats_daily.created, 1); - assert_eq!(validation_summary.stats_user_daily.created, 1); - assert_eq!(validation_summary.stats_daily_api_key.created, 1); - assert_eq!(sqlite_count(backend.pool(), "stats_daily").await, 0); - assert_eq!(sqlite_count(backend.pool(), "stats_user_daily").await, 0); - assert_eq!(sqlite_count(backend.pool(), "stats_daily_api_key").await, 0); - - let summary = backend - .import_admin_system_usage_aggregates( - &snapshot, - &user_id_map, - &api_key_id_map, - AdminSystemUsageAggregateImportMode::Overwrite, - ) - .await - .expect("usage aggregates should import"); - assert_eq!(summary.stats_daily.created, 1); - assert_eq!(summary.stats_user_daily.created, 1); - assert_eq!(summary.stats_daily_api_key.created, 1); - - let exported = backend - .export_admin_system_usage_aggregates() - .await - .expect("usage aggregates should export"); - assert_eq!(exported.stats_daily.len(), 1); - assert_eq!(exported.stats_daily[0].total_requests, 9); - assert_eq!(exported.stats_daily[0].actual_total_cost, 1.0); - assert_eq!(exported.stats_user_daily.len(), 1); - assert_eq!(exported.stats_user_daily[0].user_id, "target-user-1"); - assert_eq!(exported.stats_user_daily[0].total_requests, 5); - assert_eq!(exported.stats_daily_api_key.len(), 1); - assert_eq!(exported.stats_daily_api_key[0].api_key_id, "target-key-1"); - assert_eq!(exported.stats_daily_api_key[0].total_requests, 4); - } - - #[tokio::test] - async fn admin_system_request_bodies_purge_clears_inline_usage_body_fields() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - r#" -INSERT INTO "usage" ( - request_id, - provider_name, - model, - request_body, - response_body, - provider_request_body, - client_response_body, - request_body_compressed, - response_body_compressed, - provider_request_body_compressed, - client_response_body_compressed, - created_at_unix_ms -) -VALUES ( - 'request-1', - 'openai', - 'gpt-4.1', - 'client request', - 'provider response', - 'provider request', - 'client response', - X'01', - X'02', - X'03', - X'04', - 1 -) -"#, - ) - .execute(backend.pool()) - .await - .expect("usage row should insert"); - - let summary = backend - .purge_admin_system_data(AdminSystemPurgeTarget::RequestBodies) - .await - .expect("request body purge should run"); - - assert_eq!(summary.affected.get("usage_body_fields_cleaned"), Some(&1)); - let remaining: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM "usage" -WHERE 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 -"#, - ) - .fetch_one(backend.pool()) - .await - .expect("remaining body count should load"); - assert_eq!(remaining, 0); - } - - async fn sqlite_count(pool: &sqlx::SqlitePool, table: &str) -> i64 { - let sql = format!("SELECT COUNT(*) FROM \"{table}\""); - sqlx::query_scalar::<_, i64>(&sql) - .fetch_one(pool) - .await - .expect("count should load") - } - - #[tokio::test] - async fn wallet_daily_usage_aggregation_uses_settlement_wallets_after_sqlite_migrations() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - r#" -INSERT INTO wallets (id, user_id, balance, gift_balance, limit_mode, created_at, updated_at) -VALUES - ('wallet-1', 'user-1', 10.0, 2.0, 'finite', 1, 1), - ('wallet-stale', 'user-stale', 0.0, 0.0, 'finite', 1, 1) -"#, - ) - .execute(backend.pool()) - .await - .expect("wallets should seed"); - - sqlx::query( - r#" -INSERT INTO "usage" ( - request_id, wallet_id, provider_name, model, status, billing_status, - total_cost_usd, input_tokens, output_tokens, cache_creation_input_tokens, - cache_read_input_tokens, finalized_at, created_at_unix_ms, updated_at_unix_secs -) VALUES - ('request-1', 'wrong-wallet', 'provider', 'model', 'completed', 'pending', - 1.25, 10, 20, 3, 4, 900, 900000, 900), - ('request-2', NULL, 'provider', 'model', 'completed', 'pending', - 2.00, 5, 7, 1, 2, 901, 901000, 901), - ('request-zero', NULL, 'provider', 'model', 'completed', 'pending', - 0.00, 100, 100, 0, 0, 902, 902000, 902), - ('request-outside', NULL, 'provider', 'model', 'completed', 'pending', - 9.00, 50, 50, 0, 0, 903, 903000, 903) -"#, - ) - .execute(backend.pool()) - .await - .expect("usage should seed"); - - sqlx::query( - r#" -INSERT INTO usage_settlement_snapshots ( - request_id, billing_status, wallet_id, finalized_at, created_at, updated_at -) VALUES - ('request-1', 'settled', 'wallet-1', 1000, 1000, 1000), - ('request-2', 'settled', 'wallet-1', 1100, 1100, 1100), - ('request-zero', 'settled', 'wallet-1', 1150, 1150, 1150), - ('request-outside', 'settled', 'wallet-1', 1200, 1200, 1200) -"#, - ) - .execute(backend.pool()) - .await - .expect("settlement snapshots should seed"); - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, wallet_id, billing_date, billing_timezone, total_cost_usd, - total_requests, input_tokens, output_tokens, cache_creation_tokens, - cache_read_tokens, aggregated_at, created_at, updated_at -) VALUES ( - 'stale-ledger', 'wallet-stale', '2026-05-03', 'Asia/Shanghai', - 7.0, 3, 1, 1, 0, 0, 999, 999, 999 -) -"#, - ) - .execute(backend.pool()) - .await - .expect("stale ledger should seed"); - - let summary = backend - .aggregate_wallet_daily_usage(&WalletDailyUsageAggregationInput { - billing_date: "2026-05-03".to_string(), - billing_timezone: "Asia/Shanghai".to_string(), - window_start_unix_secs: 1000, - window_end_unix_secs: 1200, - aggregated_at_unix_secs: 1300, - }) - .await - .expect("wallet daily usage aggregation should run"); - - assert_eq!(summary.aggregated_wallets, 1); - assert_eq!(summary.deleted_stale_ledgers, 1); - - let ledger = sqlx::query_as::< - _, - ( - String, - String, - f64, - i64, - i64, - i64, - i64, - i64, - Option, - Option, - i64, - ), - >( - r#" -SELECT - id, - wallet_id, - total_cost_usd, - total_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - first_finalized_at, - last_finalized_at, - aggregated_at -FROM wallet_daily_usage_ledgers -WHERE billing_date = '2026-05-03' - AND billing_timezone = 'Asia/Shanghai' -"#, - ) - .fetch_one(backend.pool()) - .await - .expect("aggregated ledger should load"); - - assert_eq!(ledger.0.len(), 64); - assert_eq!(ledger.1, "wallet-1"); - assert!((ledger.2 - 3.25).abs() < f64::EPSILON); - assert_eq!(ledger.3, 2); - assert_eq!(ledger.4, 15); - assert_eq!(ledger.5, 27); - assert_eq!(ledger.6, 4); - assert_eq!(ledger.7, 6); - assert_eq!(ledger.8, Some(1000)); - assert_eq!(ledger.9, Some(1100)); - assert_eq!(ledger.10, 1300); - - let stale_count: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM wallet_daily_usage_ledgers WHERE id = 'stale-ledger'", - ) - .fetch_one(backend.pool()) - .await - .expect("stale ledger count should load"); - assert_eq!(stale_count, 0); - } - - #[tokio::test] - async fn stats_aggregation_runs_after_sqlite_migrations() { - let config = SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite::memory:".to_string(), - pool: SqlPoolConfig { - max_connections: 1, - ..SqlPoolConfig::default() - }, - }; - let backend = SqliteBackend::from_config(config).expect("backend should build"); - run_sqlite_migrations(backend.pool()) - .await - .expect("sqlite migrations should run"); - - sqlx::query( - r#" -INSERT INTO "usage" ( - 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, 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', '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) - .expect("aggregation time should be valid"); - let hourly = backend - .aggregate_stats_hourly(&StatsHourlyAggregationInput { - target_hour_utc: target_hour, - aggregated_at, - }) - .await - .expect("hourly stats aggregation should run") - .expect("hourly bucket should aggregate"); - assert_eq!(hourly.hour_utc, target_hour); - assert_eq!(hourly.total_requests, 2); - assert_eq!(hourly.user_rows, 2); - assert_eq!(hourly.user_model_rows, 2); - assert_eq!(hourly.model_rows, 2); - assert_eq!(hourly.provider_rows, 2); - - let hourly_row = sqlx::query_as::<_, (i64, i64, i64, i64, f64)>( - r#" -SELECT total_requests, success_requests, error_requests, input_tokens, total_cost -FROM stats_hourly -WHERE hour_utc = 3600 -"#, - ) - .fetch_one(backend.pool()) - .await - .expect("hourly stats row should load"); - assert_eq!(hourly_row.0, 2); - assert_eq!(hourly_row.1, 1); - assert_eq!(hourly_row.2, 1); - 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, - aggregated_at, - }) - .await - .expect("second hourly aggregation should run"); - assert!(second_hourly.is_none()); - - let target_day = chrono::DateTime::::from_timestamp(0, 0) - .expect("target day should be valid"); - let daily = backend - .aggregate_stats_daily(&StatsDailyAggregationInput { - target_day_utc: target_day, - aggregated_at, - }) - .await - .expect("daily stats aggregation should run") - .expect("daily bucket should aggregate"); - assert_eq!(daily.day_start_utc, target_day); - assert_eq!(daily.total_requests, 2); - assert_eq!(daily.model_rows, 2); - assert_eq!(daily.provider_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, i64)>( - r#" -SELECT total_requests, success_requests, error_requests, unique_models, fallback_count -FROM stats_daily -WHERE "date" = 0 -"#, - ) - .fetch_one(backend.pool()) - .await - .expect("daily stats row should load"); - 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.rs b/crates/aether-data/runtime/src/backend/stats.rs index 9d148408d..69f2dcfeb 100644 --- a/crates/aether-data/runtime/src/backend/stats.rs +++ b/crates/aether-data/runtime/src/backend/stats.rs @@ -1,8 +1,4 @@ -#[cfg(feature = "mysql")] -pub(crate) mod mysql; #[cfg(feature = "postgres")] pub(crate) mod postgres_daily; #[cfg(feature = "postgres")] pub(crate) mod postgres_hourly; -#[cfg(feature = "sqlite")] -pub(crate) mod sqlite; diff --git a/crates/aether-data/runtime/src/backend/stats/mysql.rs b/crates/aether-data/runtime/src/backend/stats/mysql.rs deleted file mode 100644 index 82c1be857..000000000 --- a/crates/aether-data/runtime/src/backend/stats/mysql.rs +++ /dev/null @@ -1,831 +0,0 @@ -use chrono::{DateTime, Utc}; -use sqlx::Row; - -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; -use crate::{ - DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary, - StatsHourlyAggregationInput, StatsHourlyAggregationSummary, -}; - -mod advanced; - -impl MysqlBackend { - pub async fn aggregate_stats_hourly( - &self, - input: &StatsHourlyAggregationInput, - ) -> Result, DataLayerError> { - let Some(hour_utc_unix_secs) = - next_mysql_stats_hourly_bucket(self.pool(), input.target_hour_utc).await? - else { - return Ok(None); - }; - perform_mysql_stats_hourly_aggregation(self.pool(), hour_utc_unix_secs, input.aggregated_at) - .await - .map(Some) - } - - pub async fn aggregate_stats_daily( - &self, - input: &StatsDailyAggregationInput, - ) -> Result, DataLayerError> { - let Some(day_start_unix_secs) = - next_mysql_stats_daily_bucket(self.pool(), input.target_day_utc).await? - else { - return Ok(None); - }; - perform_mysql_stats_daily_aggregation(self.pool(), day_start_unix_secs, input.aggregated_at) - .await - .map(Some) - } -} - -async fn next_mysql_stats_hourly_bucket( - pool: &MysqlPool, - target_hour_utc: DateTime, -) -> Result, DataLayerError> { - let latest_hour: Option = - sqlx::query_scalar("SELECT MAX(hour_utc) FROM stats_hourly WHERE is_complete <> 0") - .fetch_one(pool) - .await - .map_sql_err()?; - let search_from = latest_hour.map(|value| value + 3600).unwrap_or(0); - let search_until = unix_secs(target_hour_utc) + 3600; - if search_from >= search_until { - return Ok(None); - } - let next_bucket: Option = sqlx::query_scalar( - r#" -SELECT CAST(MIN(FLOOR(created_at_unix_ms / 3600) * 3600) AS SIGNED) -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') -"#, - ) - .bind(search_from) - .bind(search_until) - .fetch_one(pool) - .await - .map_sql_err()?; - Ok(next_bucket.filter(|value| *value <= unix_secs(target_hour_utc))) -} - -async fn next_mysql_stats_daily_bucket( - pool: &MysqlPool, - target_day_utc: DateTime, -) -> Result, DataLayerError> { - let latest_day: Option = - sqlx::query_scalar("SELECT MAX(`date`) FROM stats_daily WHERE is_complete <> 0") - .fetch_one(pool) - .await - .map_sql_err()?; - let search_from = latest_day.map(|value| value + 86_400).unwrap_or(0); - let search_until = unix_secs(target_day_utc) + 86_400; - if search_from >= search_until { - return Ok(None); - } - let next_bucket: Option = sqlx::query_scalar( - r#" -SELECT CAST(MIN(FLOOR(created_at_unix_ms / 86400) * 86400) AS SIGNED) -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') -"#, - ) - .bind(search_from) - .bind(search_until) - .fetch_one(pool) - .await - .map_sql_err()?; - Ok(next_bucket.filter(|value| *value <= unix_secs(target_day_utc))) -} - -const MYSQL_STATS_AGGREGATE_SQL: &str = r#" -SELECT - CAST(COUNT(*) AS SIGNED) AS total_requests, - CAST(COALESCE(SUM(CASE - WHEN status = 'failed' - OR status_code >= 400 - 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, - CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS SIGNED) AS cache_creation_tokens, - CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS SIGNED) AS cache_read_tokens, - CAST(COALESCE(SUM(total_cost_usd), 0.0) AS DOUBLE) AS total_cost, - CAST(COALESCE(SUM(actual_total_cost_usd), 0.0) AS DOUBLE) AS actual_total_cost, - CAST(COALESCE(AVG(response_time_ms), 0.0) AS DOUBLE) AS avg_response_time_ms -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') -"#; - -async fn perform_mysql_stats_hourly_aggregation( - pool: &MysqlPool, - hour_utc_unix_secs: i64, - aggregated_at: DateTime, -) -> Result { - 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_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()?; - - sqlx::query( - r#" -INSERT INTO stats_hourly ( - id, hour_utc, total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, avg_response_time_ms, is_complete, - aggregated_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, TRUE, ?, ?, ?) -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), - actual_total_cost = VALUES(actual_total_cost), - avg_response_time_ms = VALUES(avg_response_time_ms), - is_complete = VALUES(is_complete), - aggregated_at = VALUES(aggregated_at), - updated_at = VALUES(updated_at) -"#, - ) - .bind(stats_id(&format!("stats-hourly:{hour_utc_unix_secs}"))) - .bind(hour_utc_unix_secs) - .bind(total_requests) - .bind(total_requests.saturating_sub(error_requests)) - .bind(error_requests) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .bind(row.try_get::("total_cost").map_sql_err()?) - .bind(row.try_get::("actual_total_cost").map_sql_err()?) - .bind( - row.try_get::("avg_response_time_ms") - .map_sql_err()?, - ) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .execute(&mut *tx) - .await - .map_sql_err()?; - - 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 { - hour_utc: utc_from_unix_secs(hour_utc_unix_secs, "stats_hourly.hour_utc")?, - total_requests, - user_rows, - user_model_rows, - model_rows, - provider_rows, - }) -} - -async fn perform_mysql_stats_daily_aggregation( - pool: &MysqlPool, - day_start_unix_secs: i64, - aggregated_at: DateTime, -) -> Result { - 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_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_unix_secs, end_unix_secs).await? as i64; - let unique_providers = - 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#" -INSERT INTO stats_daily ( - id, `date`, total_requests, success_requests, error_requests, - 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, TRUE, ?, ?, ?) -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), - actual_total_cost = VALUES(actual_total_cost), - avg_response_time_ms = VALUES(avg_response_time_ms), - fallback_count = VALUES(fallback_count), - unique_models = VALUES(unique_models), - unique_providers = VALUES(unique_providers), - is_complete = VALUES(is_complete), - aggregated_at = VALUES(aggregated_at), - updated_at = VALUES(updated_at) -"#, - ) - .bind(stats_id(&format!("stats-daily:{day_start_unix_secs}"))) - .bind(day_start_unix_secs) - .bind(total_requests) - .bind(total_requests.saturating_sub(error_requests)) - .bind(error_requests) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .bind(row.try_get::("total_cost").map_sql_err()?) - .bind(row.try_get::("actual_total_cost").map_sql_err()?) - .bind( - row.try_get::("avg_response_time_ms") - .map_sql_err()?, - ) - .bind(fallback_count) - .bind(unique_models) - .bind(unique_providers) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .execute(&mut *tx) - .await - .map_sql_err()?; - - 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 { - day_start_utc: utc_from_unix_secs(day_start_unix_secs, "stats_daily.date")?, - total_requests, - model_rows, - provider_rows, - api_key_rows, - error_rows, - user_rows, - }) -} - -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_unix_secs: i64, - end_unix_secs: i64, -) -> Result { - let not_empty = group_columns - .split(',') - .map(str::trim) - .map(|column| format!("{column} IS NOT NULL AND {column} <> ''")) - .collect::>() - .join(" AND "); - let sql = format!( - 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 {not_empty} - GROUP BY {group_columns} -) AS grouped -"# - ); - let count: i64 = sqlx::query_scalar(&sql) - .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)) -} diff --git a/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs b/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs deleted file mode 100644 index 525f29026..000000000 --- a/crates/aether-data/runtime/src/backend/stats/mysql/advanced.rs +++ /dev/null @@ -1,973 +0,0 @@ -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 deleted file mode 100644 index 571f8c82d..000000000 --- a/crates/aether-data/runtime/src/backend/stats/sqlite.rs +++ /dev/null @@ -1,871 +0,0 @@ -use chrono::{DateTime, Utc}; -use sqlx::Row; - -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; -use crate::{ - DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary, - StatsHourlyAggregationInput, StatsHourlyAggregationSummary, -}; - -mod advanced; - -impl SqliteBackend { - pub async fn aggregate_stats_hourly( - &self, - input: &StatsHourlyAggregationInput, - ) -> Result, DataLayerError> { - let Some(hour_utc_unix_secs) = - next_sqlite_stats_hourly_bucket(self.pool(), input.target_hour_utc).await? - else { - return Ok(None); - }; - perform_sqlite_stats_hourly_aggregation( - self.pool(), - hour_utc_unix_secs, - input.aggregated_at, - ) - .await - .map(Some) - } - - pub async fn aggregate_stats_daily( - &self, - input: &StatsDailyAggregationInput, - ) -> Result, DataLayerError> { - let Some(day_start_unix_secs) = - next_sqlite_stats_daily_bucket(self.pool(), input.target_day_utc).await? - else { - return Ok(None); - }; - perform_sqlite_stats_daily_aggregation( - self.pool(), - day_start_unix_secs, - input.aggregated_at, - ) - .await - .map(Some) - } -} - -async fn next_sqlite_stats_hourly_bucket( - pool: &SqlitePool, - target_hour_utc: DateTime, -) -> Result, DataLayerError> { - let latest_hour: Option = - sqlx::query_scalar("SELECT MAX(hour_utc) FROM stats_hourly WHERE is_complete <> 0") - .fetch_one(pool) - .await - .map_sql_err()?; - let search_from = latest_hour.map(|value| value + 3600).unwrap_or(0); - let search_until = unix_secs(target_hour_utc) + 3600; - if search_from >= search_until { - return Ok(None); - } - let next_bucket: Option = sqlx::query_scalar( - r#" -SELECT MIN(CAST(created_at_unix_ms / 3600 AS INTEGER) * 3600) -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') -"#, - ) - .bind(search_from) - .bind(search_until) - .fetch_one(pool) - .await - .map_sql_err()?; - Ok(next_bucket.filter(|value| *value <= unix_secs(target_hour_utc))) -} - -async fn next_sqlite_stats_daily_bucket( - pool: &SqlitePool, - target_day_utc: DateTime, -) -> Result, DataLayerError> { - let latest_day: Option = - sqlx::query_scalar(r#"SELECT MAX("date") FROM stats_daily WHERE is_complete <> 0"#) - .fetch_one(pool) - .await - .map_sql_err()?; - let search_from = latest_day.map(|value| value + 86_400).unwrap_or(0); - let search_until = unix_secs(target_day_utc) + 86_400; - if search_from >= search_until { - return Ok(None); - } - let next_bucket: Option = sqlx::query_scalar( - r#" -SELECT MIN(CAST(created_at_unix_ms / 86400 AS INTEGER) * 86400) -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') -"#, - ) - .bind(search_from) - .bind(search_until) - .fetch_one(pool) - .await - .map_sql_err()?; - Ok(next_bucket.filter(|value| *value <= unix_secs(target_day_utc))) -} - -const SQLITE_STATS_AGGREGATE_SQL: &str = r#" -SELECT - COUNT(*) AS total_requests, - COALESCE(SUM(CASE - WHEN status = 'failed' - OR status_code >= 400 - 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, - COALESCE(SUM(cache_creation_input_tokens), 0) AS cache_creation_tokens, - COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens, - CAST(COALESCE(SUM(total_cost_usd), 0) AS REAL) AS total_cost, - CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS REAL) AS actual_total_cost, - CAST(COALESCE(AVG(response_time_ms), 0) AS REAL) AS avg_response_time_ms -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') -"#; - -async fn perform_sqlite_stats_hourly_aggregation( - pool: &SqlitePool, - hour_utc_unix_secs: i64, - aggregated_at: DateTime, -) -> Result { - 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_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()?; - - sqlx::query( - r#" -INSERT INTO stats_hourly ( - id, hour_utc, total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, avg_response_time_ms, is_complete, - aggregated_at, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?) -ON CONFLICT (hour_utc) 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, - actual_total_cost = excluded.actual_total_cost, - avg_response_time_ms = excluded.avg_response_time_ms, - is_complete = excluded.is_complete, - aggregated_at = excluded.aggregated_at, - updated_at = excluded.updated_at -"#, - ) - .bind(stats_id(&format!("stats-hourly:{hour_utc_unix_secs}"))) - .bind(hour_utc_unix_secs) - .bind(total_requests) - .bind(total_requests.saturating_sub(error_requests)) - .bind(error_requests) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .bind(sqlite_real(&row, "total_cost")?) - .bind(sqlite_real(&row, "actual_total_cost")?) - .bind(sqlite_real(&row, "avg_response_time_ms")?) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .execute(&mut *tx) - .await - .map_sql_err()?; - - 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 { - hour_utc: utc_from_unix_secs(hour_utc_unix_secs, "stats_hourly.hour_utc")?, - total_requests, - user_rows, - user_model_rows, - model_rows, - provider_rows, - }) -} - -async fn perform_sqlite_stats_daily_aggregation( - pool: &SqlitePool, - day_start_unix_secs: i64, - aggregated_at: DateTime, -) -> Result { - 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_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_unix_secs, end_unix_secs).await? as i64; - let unique_providers = - 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#" -INSERT INTO stats_daily ( - id, "date", total_requests, success_requests, error_requests, - 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?) -ON CONFLICT ("date") 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, - actual_total_cost = excluded.actual_total_cost, - avg_response_time_ms = excluded.avg_response_time_ms, - fallback_count = excluded.fallback_count, - unique_models = excluded.unique_models, - unique_providers = excluded.unique_providers, - is_complete = excluded.is_complete, - aggregated_at = excluded.aggregated_at, - updated_at = excluded.updated_at -"#, - ) - .bind(stats_id(&format!("stats-daily:{day_start_unix_secs}"))) - .bind(day_start_unix_secs) - .bind(total_requests) - .bind(total_requests.saturating_sub(error_requests)) - .bind(error_requests) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .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) - .bind(aggregated_at_unix_secs) - .bind(aggregated_at_unix_secs) - .execute(&mut *tx) - .await - .map_sql_err()?; - - 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 { - day_start_utc: utc_from_unix_secs(day_start_unix_secs, "stats_daily.date")?, - total_requests, - model_rows, - provider_rows, - api_key_rows, - error_rows, - user_rows, - }) -} - -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_unix_secs: i64, - end_unix_secs: i64, -) -> Result { - let not_empty = group_columns - .split(',') - .map(str::trim) - .map(|column| format!("{column} IS NOT NULL AND {column} <> ''")) - .collect::>() - .join(" AND "); - let sql = format!( - 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 {not_empty} - GROUP BY {group_columns} -) -"# - ); - let count: i64 = sqlx::query_scalar(&sql) - .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)) -} diff --git a/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs b/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs deleted file mode 100644 index 75b5839c0..000000000 --- a/crates/aether-data/runtime/src/backend/stats/sqlite/advanced.rs +++ /dev/null @@ -1,986 +0,0 @@ -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 deleted file mode 100644 index 0881d580c..000000000 --- a/crates/aether-data/runtime/src/backend/stats_common.rs +++ /dev/null @@ -1,27 +0,0 @@ -use chrono::{DateTime, Utc}; -use sha2::{Digest, Sha256}; - -use crate::DataLayerError; - -pub(crate) fn unix_secs(value: DateTime) -> i64 { - value.timestamp().max(0) -} - -pub(crate) fn utc_from_unix_secs( - value: i64, - field_name: &str, -) -> Result, DataLayerError> { - DateTime::::from_timestamp(value, 0).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!("{field_name} contains invalid timestamp {value}")) - }) -} - -pub(crate) fn stats_id(value: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(value.as_bytes()); - hasher - .finalize() - .iter() - .map(|byte| format!("{byte:02x}")) - .collect() -} diff --git a/crates/aether-data/runtime/src/backend/system.rs b/crates/aether-data/runtime/src/backend/system.rs index 76f2c02f3..6f88f440e 100644 --- a/crates/aether-data/runtime/src/backend/system.rs +++ b/crates/aether-data/runtime/src/backend/system.rs @@ -1,11 +1,7 @@ use std::collections::BTreeMap; -#[cfg(feature = "mysql")] -use super::MysqlBackend; #[cfg(feature = "postgres")] use super::PostgresBackend; -#[cfg(feature = "sqlite")] -use super::SqliteBackend; use crate::repository::system::{ AdminSystemPurgeSummary, AdminSystemPurgeTarget, AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateImportSummary, AdminSystemUsageAggregateSnapshot, @@ -13,12 +9,8 @@ use crate::repository::system::{ }; use crate::DataLayerError; -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; const ADMIN_CONFIG_PURGE_TABLES: &[&str] = &[ "api_key_provider_mappings", @@ -124,11 +116,6 @@ fn checked_sql_identifier(value: &str) -> Result<&str, DataLayerError> { } } -#[cfg(any(feature = "mysql", feature = "sqlite"))] -fn current_unix_secs() -> u64 { - chrono::Utc::now().timestamp().max(0) as u64 -} - fn i64_from_u64(value: u64, field_name: &str) -> Result { i64::try_from(value) .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64 range"))) @@ -176,17 +163,3 @@ fn should_skip_imported_aggregate( )), } } - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -fn serialize_json_value(value: &serde_json::Value) -> Result { - serde_json::to_string(value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("invalid system config JSON value: {err}")) - }) -} - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -fn parse_json_value(value: String) -> Result { - serde_json::from_str(&value).map_err(|err| { - DataLayerError::UnexpectedValue(format!("invalid system config JSON value: {err}")) - }) -} diff --git a/crates/aether-data/runtime/src/backend/system/mysql.rs b/crates/aether-data/runtime/src/backend/system/mysql.rs deleted file mode 100644 index f62bb847f..000000000 --- a/crates/aether-data/runtime/src/backend/system/mysql.rs +++ /dev/null @@ -1,1317 +0,0 @@ -use sqlx::Row; - -use crate::error::SqlResultExt; -use crate::repository::system::{ - AdminSystemStats, AdminSystemStatsDailyAggregate, AdminSystemStatsDailyApiKeyAggregate, - AdminSystemStatsUserDailyAggregate, -}; -use crate::DataLayerError; - -use super::u64_from_i64; -use super::*; - -const READ_ADMIN_SYSTEM_STATS_SQL: &str = r#" -SELECT - (SELECT COUNT(*) FROM users) AS total_users, - (SELECT COUNT(*) FROM users WHERE is_active = 1) AS active_users, - (SELECT COUNT(*) FROM api_keys) AS total_api_keys, - (SELECT COUNT(*) FROM `usage`) AS total_requests -"#; - -async fn export_mysql_admin_system_usage_aggregates( - pool: &sqlx::MySqlPool, -) -> Result { - let mut snapshot = AdminSystemUsageAggregateSnapshot::default(); - - let daily_rows = sqlx::query( - r#" -SELECT - `date` AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost, - actual_total_cost, - is_complete, - aggregated_at AS aggregated_at_unix_secs -FROM stats_daily -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY `date` ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in daily_rows { - snapshot.stats_daily.push(map_stats_daily_aggregate(&row)?); - } - - let user_daily_rows = sqlx::query( - r#" -SELECT - user_id, - username, - `date` AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost -FROM stats_user_daily -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY user_id ASC, `date` ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in user_daily_rows { - snapshot - .stats_user_daily - .push(map_stats_user_daily_aggregate(&row)?); - } - - let api_key_daily_rows = sqlx::query( - r#" -SELECT - api_key_id, - api_key_name, - `date` AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost -FROM stats_daily_api_key -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY api_key_id ASC, `date` ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in api_key_daily_rows { - snapshot - .stats_daily_api_key - .push(map_stats_daily_api_key_aggregate(&row)?); - } - - Ok(snapshot) -} - -async fn import_mysql_admin_system_usage_aggregates( - pool: &sqlx::MySqlPool, - snapshot: &AdminSystemUsageAggregateSnapshot, - user_id_map: &BTreeMap, - api_key_id_map: &BTreeMap, - mode: AdminSystemUsageAggregateImportMode, -) -> Result { - let mut tx = pool.begin().await.map_sql_err()?; - let mut summary = AdminSystemUsageAggregateImportSummary::default(); - let now = current_unix_secs(); - - for row in &snapshot.stats_daily { - let existing: Option = - sqlx::query_scalar("SELECT id FROM stats_daily WHERE `date` = ? LIMIT 1") - .bind(i64_from_u64(row.date_unix_secs, "stats_daily.date")?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_daily", - row.date_unix_secs, - )? { - summary.stats_daily.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_daily ( - id, `date`, total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, is_complete, aggregated_at, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -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), - actual_total_cost = VALUES(actual_total_cost), - is_complete = VALUES(is_complete), - aggregated_at = VALUES(aggregated_at), - updated_at = VALUES(updated_at) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(i64_from_u64(row.date_unix_secs, "stats_daily.date")?) - .bind(i64_from_u64( - row.total_requests, - "stats_daily.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_daily.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_daily.error_requests", - )?) - .bind(i64_from_u64(row.input_tokens, "stats_daily.input_tokens")?) - .bind(i64_from_u64( - row.output_tokens, - "stats_daily.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_daily.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_daily.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(row.actual_total_cost) - .bind(row.is_complete || row.total_requests > 0) - .bind(optional_i64_from_u64( - row.aggregated_at_unix_secs, - "stats_daily.aggregated_at", - )?) - .bind(i64_from_u64(now, "stats_daily.created_at")?) - .bind(i64_from_u64(now, "stats_daily.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_daily, existing.is_some()); - } - - for row in &snapshot.stats_user_daily { - let Some(target_user_id) = user_id_map.get(&row.user_id) else { - summary.skipped_unmapped_user_daily += 1; - summary.stats_user_daily.skipped += 1; - continue; - }; - let existing: Option = sqlx::query_scalar( - "SELECT id FROM stats_user_daily WHERE user_id = ? AND `date` = ? LIMIT 1", - ) - .bind(target_user_id) - .bind(i64_from_u64(row.date_unix_secs, "stats_user_daily.date")?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_user_daily", - row.date_unix_secs, - )? { - summary.stats_user_daily.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_user_daily ( - id, user_id, username, `date`, total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - username = VALUES(username), - 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), - updated_at = VALUES(updated_at) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(target_user_id) - .bind(row.username.as_deref()) - .bind(i64_from_u64(row.date_unix_secs, "stats_user_daily.date")?) - .bind(i64_from_u64( - row.total_requests, - "stats_user_daily.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_user_daily.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_user_daily.error_requests", - )?) - .bind(i64_from_u64( - row.input_tokens, - "stats_user_daily.input_tokens", - )?) - .bind(i64_from_u64( - row.output_tokens, - "stats_user_daily.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_user_daily.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_user_daily.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(i64_from_u64(now, "stats_user_daily.created_at")?) - .bind(i64_from_u64(now, "stats_user_daily.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_user_daily, existing.is_some()); - } - - for row in &snapshot.stats_daily_api_key { - let Some(target_api_key_id) = api_key_id_map.get(&row.api_key_id) else { - summary.skipped_unmapped_api_key_daily += 1; - summary.stats_daily_api_key.skipped += 1; - continue; - }; - let existing: Option = sqlx::query_scalar( - "SELECT id FROM stats_daily_api_key WHERE api_key_id = ? AND `date` = ? LIMIT 1", - ) - .bind(target_api_key_id) - .bind(i64_from_u64( - row.date_unix_secs, - "stats_daily_api_key.date", - )?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_daily_api_key", - row.date_unix_secs, - )? { - summary.stats_daily_api_key.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_daily_api_key ( - id, api_key_id, api_key_name, `date`, total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - api_key_name = VALUES(api_key_name), - 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), - updated_at = VALUES(updated_at) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(target_api_key_id) - .bind(row.api_key_name.as_deref()) - .bind(i64_from_u64( - row.date_unix_secs, - "stats_daily_api_key.date", - )?) - .bind(i64_from_u64( - row.total_requests, - "stats_daily_api_key.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_daily_api_key.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_daily_api_key.error_requests", - )?) - .bind(i64_from_u64( - row.input_tokens, - "stats_daily_api_key.input_tokens", - )?) - .bind(i64_from_u64( - row.output_tokens, - "stats_daily_api_key.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_daily_api_key.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_daily_api_key.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(i64_from_u64(now, "stats_daily_api_key.created_at")?) - .bind(i64_from_u64(now, "stats_daily_api_key.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_daily_api_key, existing.is_some()); - } - - if mode == AdminSystemUsageAggregateImportMode::ValidateError { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - Ok(summary) -} - -impl MysqlBackend { - pub async fn purge_admin_system_data( - &self, - target: AdminSystemPurgeTarget, - ) -> Result { - let mut tx = self.pool().begin().await.map_sql_err()?; - let mut summary = AdminSystemPurgeSummary::default(); - purge_mysql_admin_system_data(&mut tx, target, &mut summary).await?; - tx.commit().await.map_sql_err()?; - Ok(summary) - } - - pub async fn purge_admin_request_bodies_batch( - &self, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(AdminSystemPurgeSummary::default()); - } - let mut tx = self.pool().begin().await.map_sql_err()?; - let mut summary = AdminSystemPurgeSummary::default(); - purge_mysql_request_bodies_batch(&mut tx, batch_size, &mut summary).await?; - tx.commit().await.map_sql_err()?; - Ok(summary) - } - - pub async fn find_system_config_value( - &self, - key: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT value -FROM system_configs -WHERE `key` = ? -LIMIT 1 -"#, - ) - .bind(key) - .fetch_optional(self.pool()) - .await - .map_sql_err()?; - - row.map(|row| { - row.try_get("value") - .map_sql_err() - .and_then(parse_json_value) - }) - .transpose() - } - - pub async fn compare_and_set_system_config_string_value( - &self, - key: &str, - expected: &str, - replacement: &str, - ) -> Result { - let now = current_unix_secs(); - let replacement = - serialize_json_value(&serde_json::Value::String(replacement.to_string()))?; - let result = sqlx::query( - r#" -UPDATE system_configs -SET value = ?, updated_at = ? -WHERE `key` = ? - AND JSON_TYPE(value) = 'STRING' - AND BINARY JSON_UNQUOTE(value) = BINARY ? -"#, - ) - .bind(replacement) - .bind(now as i64) - .bind(key) - .bind(expected) - .execute(self.pool()) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - pub async fn upsert_system_config_value( - &self, - key: &str, - value: &serde_json::Value, - description: Option<&str>, - ) -> Result { - Ok(self - .upsert_system_config_entry(key, value, description) - .await? - .value) - } - - pub async fn list_system_config_entries( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT `key`, value, description, updated_at -FROM system_configs -ORDER BY `key` ASC -"#, - ) - .fetch_all(self.pool()) - .await - .map_sql_err()?; - - rows.into_iter() - .map(|row| { - Ok(StoredSystemConfigEntry { - key: row.try_get("key").map_sql_err()?, - value: parse_json_value(row.try_get("value").map_sql_err()?)?, - description: row.try_get("description").map_sql_err()?, - updated_at_unix_secs: row - .try_get::, _>("updated_at") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }) - }) - .collect() - } - - pub async fn upsert_system_config_entry( - &self, - key: &str, - value: &serde_json::Value, - description: Option<&str>, - ) -> Result { - let now = current_unix_secs(); - let serialized = serialize_json_value(value)?; - sqlx::query( - r#" -INSERT INTO system_configs (id, `key`, value, description, created_at, updated_at) -VALUES (?, ?, ?, ?, ?, ?) -ON DUPLICATE KEY UPDATE - `key` = VALUES(`key`), - value = VALUES(value), - description = COALESCE(VALUES(description), description), - updated_at = VALUES(updated_at) -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(key) - .bind(serialized) - .bind(description) - .bind(now as i64) - .bind(now as i64) - .execute(self.pool()) - .await - .map_sql_err()?; - - self.list_system_config_entries() - .await? - .into_iter() - .find(|entry| entry.key == key) - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "system config key '{key}' missing after mysql upsert" - )) - }) - } - - pub async fn delete_system_config_value(&self, key: &str) -> Result { - let result = sqlx::query( - r#" -DELETE FROM system_configs -WHERE `key` = ? -"#, - ) - .bind(key) - .execute(self.pool()) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - pub async fn read_admin_system_stats(&self) -> Result { - let row = sqlx::query(READ_ADMIN_SYSTEM_STATS_SQL) - .fetch_one(self.pool()) - .await - .map_sql_err()?; - map_admin_system_stats(row) - } - - pub async fn export_admin_system_usage_aggregates( - &self, - ) -> Result { - export_mysql_admin_system_usage_aggregates(self.pool()).await - } - - pub async fn import_admin_system_usage_aggregates( - &self, - snapshot: &AdminSystemUsageAggregateSnapshot, - user_id_map: &BTreeMap, - api_key_id_map: &BTreeMap, - mode: AdminSystemUsageAggregateImportMode, - ) -> Result { - import_mysql_admin_system_usage_aggregates( - self.pool(), - snapshot, - user_id_map, - api_key_id_map, - mode, - ) - .await - } -} - -pub(super) fn map_stats_daily_aggregate( - row: &sqlx::mysql::MySqlRow, -) -> Result { - Ok(AdminSystemStatsDailyAggregate { - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - actual_total_cost: row.try_get("actual_total_cost").map_sql_err()?, - is_complete: row.try_get("is_complete").map_sql_err()?, - aggregated_at_unix_secs: row - .try_get::, _>("aggregated_at_unix_secs") - .map_sql_err()? - .map(u64_from_i64), - }) -} - -pub(super) fn map_stats_user_daily_aggregate( - row: &sqlx::mysql::MySqlRow, -) -> Result { - Ok(AdminSystemStatsUserDailyAggregate { - user_id: row.try_get("user_id").map_sql_err()?, - username: row.try_get("username").map_sql_err()?, - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - }) -} - -pub(super) fn map_stats_daily_api_key_aggregate( - row: &sqlx::mysql::MySqlRow, -) -> Result { - Ok(AdminSystemStatsDailyApiKeyAggregate { - api_key_id: row.try_get("api_key_id").map_sql_err()?, - api_key_name: row.try_get("api_key_name").map_sql_err()?, - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - }) -} - -pub(super) fn map_admin_system_stats( - row: sqlx::mysql::MySqlRow, -) -> Result { - Ok(AdminSystemStats { - total_users: row.try_get::("total_users").map_sql_err()?.max(0) as u64, - active_users: row.try_get::("active_users").map_sql_err()?.max(0) as u64, - total_api_keys: row - .try_get::("total_api_keys") - .map_sql_err()? - .max(0) as u64, - total_requests: row - .try_get::("total_requests") - .map_sql_err()? - .max(0) as u64, - }) -} - -pub(super) async fn mysql_delete_table( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !mysql_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!("DELETE FROM `{table}`"); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn mysql_delete_non_admin_user_rows( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !mysql_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!( - "DELETE FROM `{table}` WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin')" - ); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn mysql_delete_non_admin_api_key_rows( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !mysql_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!( - "DELETE FROM `{table}` WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'))" - ); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn mysql_execute_if_table( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - if !mysql_table_exists(tx, checked_sql_identifier(table)?).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn mysql_execute_if_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - columns: &[&str], - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - if !mysql_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn mysql_execute_batch_if_table( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, - limit: i64, -) -> Result<(), DataLayerError> { - if !mysql_table_exists(tx, checked_sql_identifier(table)?).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .bind(limit) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn mysql_execute_batch_if_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - columns: &[&str], - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, - limit: i64, -) -> Result<(), DataLayerError> { - if !mysql_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .bind(limit) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn mysql_table_exists( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, -) -> Result { - let table = checked_sql_identifier(table)?; - let total: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = ?", - ) - .bind(table) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - Ok(total > 0) -} - -pub(super) async fn mysql_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table: &str, - columns: &[&str], -) -> Result { - let table = checked_sql_identifier(table)?; - if !mysql_table_exists(tx, table).await? { - return Ok(false); - } - for column in columns { - let column = checked_sql_identifier(column)?; - let total: i64 = sqlx::query_scalar( - "SELECT COUNT(*) - FROM information_schema.columns - WHERE table_schema = DATABASE() - AND table_name = ? - AND column_name = ?", - ) - .bind(table) - .bind(column) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - if total == 0 { - return Ok(false); - } - } - Ok(true) -} - -async fn purge_mysql_admin_system_data( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - target: AdminSystemPurgeTarget, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - match target { - AdminSystemPurgeTarget::Config => { - mysql_execute_if_table( - tx, - "usage", - "usage_provider_refs_cleared", - r#" -UPDATE `usage` -SET provider_id = NULL, - provider_endpoint_id = NULL, - provider_api_key_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, - ) - .await?; - mysql_execute_if_table( - tx, - "request_candidates", - "request_candidate_provider_refs_cleared", - r#" -UPDATE request_candidates -SET provider_id = NULL, - endpoint_id = NULL, - key_id = NULL -WHERE provider_id IS NOT NULL - OR endpoint_id IS NOT NULL - OR key_id IS NOT NULL -"#, - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "video_tasks", - "video_task_provider_refs_cleared", - r#" -UPDATE video_tasks -SET provider_id = NULL, - endpoint_id = NULL, - key_id = NULL -WHERE provider_id IS NOT NULL - OR endpoint_id IS NOT NULL - OR key_id IS NOT NULL -"#, - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "user_preferences", - "user_default_provider_refs_cleared", - "UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id IS NOT NULL", - summary, - ) - .await?; - for table in ADMIN_CONFIG_PURGE_TABLES { - mysql_delete_table(tx, table, summary).await?; - } - } - AdminSystemPurgeTarget::Users => purge_mysql_non_admin_users(tx, summary).await?, - AdminSystemPurgeTarget::Usage => { - mysql_delete_table(tx, "request_candidates", summary).await?; - for table in ADMIN_USAGE_CHILD_TABLES { - mysql_delete_table(tx, table, summary).await?; - } - mysql_delete_table(tx, "usage", summary).await?; - mysql_execute_if_table( - tx, - "api_keys", - "api_key_usage_stats_reset", - r#" -UPDATE api_keys -SET total_requests = 0, - total_tokens = 0, - total_cost_usd = 0, - last_used_at = NULL -WHERE total_requests <> 0 - OR total_tokens <> 0 - OR total_cost_usd <> 0 - OR last_used_at IS NOT NULL -"#, - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "provider_api_keys", - "provider_key_usage_stats_reset", - r#" -UPDATE provider_api_keys -SET request_count = 0, - success_count = 0, - error_count = 0, - total_tokens = 0, - total_cost_usd = 0, - total_response_time_ms = 0, - last_used_at = NULL -WHERE request_count <> 0 - OR success_count <> 0 - OR error_count <> 0 - OR total_tokens <> 0 - OR total_cost_usd <> 0 - OR total_response_time_ms <> 0 - OR last_used_at IS NOT NULL -"#, - 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?; - } - AdminSystemPurgeTarget::RequestBodies => { - mysql_delete_table(tx, "usage_body_blobs", summary).await?; - mysql_execute_if_table_has_columns( - tx, - "usage", - USAGE_BODY_FIELD_COLUMNS, - "usage_body_fields_cleaned", - 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 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 -"#, - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "usage_http_audits", - "usage_http_audit_body_refs_cleaned", - r#" -UPDATE usage_http_audits -SET request_body_ref = NULL, - provider_request_body_ref = NULL, - response_body_ref = NULL, - client_response_body_ref = NULL, - request_body_state = NULL, - provider_request_body_state = NULL, - response_body_state = NULL, - client_response_body_state = NULL, - body_capture_mode = 'none' -WHERE request_body_ref IS NOT NULL - OR provider_request_body_ref IS NOT NULL - OR response_body_ref IS NOT NULL - OR client_response_body_ref IS NOT NULL - OR request_body_state IS NOT NULL - OR provider_request_body_state IS NOT NULL - OR response_body_state IS NOT NULL - OR client_response_body_state IS NOT NULL - OR body_capture_mode <> 'none' -"#, - summary, - ) - .await?; - } - AdminSystemPurgeTarget::Stats => { - for table in ADMIN_STATS_PURGE_TABLES { - mysql_delete_table(tx, table, summary).await?; - } - } - } - Ok(()) -} - -async fn purge_mysql_non_admin_users( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let non_admin_users = "SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'"; - let non_admin_keys = "SELECT id FROM api_keys WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin')"; - - mysql_execute_if_table( - tx, - "api_key_provider_mappings", - "api_key_provider_mappings", - &format!("DELETE FROM api_key_provider_mappings WHERE api_key_id IN ({non_admin_keys})"), - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "usage", - "usage_user_refs_cleared", - &format!("UPDATE `usage` SET user_id = NULL WHERE user_id IN ({non_admin_users})"), - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "usage", - "usage_api_key_refs_cleared", - &format!("UPDATE `usage` SET api_key_id = NULL WHERE api_key_id IN ({non_admin_keys})"), - summary, - ) - .await?; - for (table, column, key) in [ - ( - "request_candidates", - "user_id", - "request_candidate_user_refs_cleared", - ), - ("audit_logs", "user_id", "audit_log_user_refs_cleared"), - ("video_tasks", "user_id", "video_task_user_refs_cleared"), - ("wallets", "user_id", "wallet_user_refs_cleared"), - ( - "payment_orders", - "user_id", - "payment_order_user_refs_cleared", - ), - ( - "wallet_transactions", - "operator_id", - "wallet_transaction_operator_refs_cleared", - ), - ( - "announcements", - "author_id", - "announcement_author_refs_cleared", - ), - ( - "proxy_nodes", - "registered_by", - "proxy_node_registrant_refs_cleared", - ), - ] { - mysql_execute_if_table( - tx, - table, - key, - &format!( - "UPDATE `{table}` SET `{column}` = NULL WHERE `{column}` IN ({non_admin_users})" - ), - summary, - ) - .await?; - } - for (table, key) in [ - ( - "request_candidates", - "request_candidate_api_key_refs_cleared", - ), - ("audit_logs", "audit_log_api_key_refs_cleared"), - ("video_tasks", "video_task_api_key_refs_cleared"), - ("wallets", "wallet_api_key_refs_cleared"), - ] { - mysql_execute_if_table( - tx, - table, - key, - &format!( - "UPDATE `{table}` SET api_key_id = NULL WHERE api_key_id IN ({non_admin_keys})" - ), - summary, - ) - .await?; - } - mysql_delete_non_admin_api_key_rows(tx, "stats_daily_api_key", summary).await?; - mysql_execute_if_table( - tx, - "refund_requests", - "refund_request_user_refs_cleared", - &format!( - r#" -UPDATE refund_requests -SET user_id = CASE WHEN user_id IN ({non_admin_users}) THEN NULL ELSE user_id END, - requested_by = CASE WHEN requested_by IN ({non_admin_users}) THEN NULL ELSE requested_by END, - approved_by = CASE WHEN approved_by IN ({non_admin_users}) THEN NULL ELSE approved_by END, - processed_by = CASE WHEN processed_by IN ({non_admin_users}) THEN NULL ELSE processed_by END -WHERE user_id IN ({non_admin_users}) - OR requested_by IN ({non_admin_users}) - OR approved_by IN ({non_admin_users}) - OR processed_by IN ({non_admin_users}) -"# - ), - summary, - ) - .await?; - for table in ADMIN_USER_SCOPED_TABLES { - mysql_delete_non_admin_user_rows(tx, table, summary).await?; - } - mysql_execute_if_table( - tx, - "api_keys", - "api_keys", - &format!("DELETE FROM api_keys WHERE user_id IN ({non_admin_users})"), - summary, - ) - .await?; - mysql_execute_if_table( - tx, - "users", - "users", - "DELETE FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'", - summary, - ) - .await?; - Ok(()) -} - -async fn purge_mysql_request_bodies_batch( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - batch_size: usize, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - mysql_execute_batch_if_table( - tx, - "usage_body_blobs", - "usage_body_blobs", - r#" -DELETE FROM usage_body_blobs -WHERE body_ref IN ( - SELECT body_ref FROM ( - SELECT body_ref - FROM usage_body_blobs - ORDER BY body_ref ASC - LIMIT ? - ) AS doomed -) -"#, - summary, - limit, - ) - .await?; - mysql_execute_batch_if_table_has_columns( - tx, - "usage", - USAGE_BODY_FIELD_COLUMNS, - "usage_body_fields_cleaned", - 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 request_id IN ( - SELECT request_id FROM ( - SELECT request_id - FROM `usage` - WHERE 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 - ORDER BY created_at_unix_ms ASC, request_id ASC - LIMIT ? - ) AS batch -) -"#, - summary, - limit, - ) - .await?; - mysql_execute_batch_if_table( - tx, - "usage_http_audits", - "usage_http_audit_body_refs_cleaned", - r#" -UPDATE usage_http_audits -SET request_body_ref = NULL, - provider_request_body_ref = NULL, - response_body_ref = NULL, - client_response_body_ref = NULL, - request_body_state = NULL, - provider_request_body_state = NULL, - response_body_state = NULL, - client_response_body_state = NULL, - body_capture_mode = 'none' -WHERE request_id IN ( - SELECT request_id FROM ( - SELECT request_id - FROM usage_http_audits - WHERE request_body_ref IS NOT NULL - OR provider_request_body_ref IS NOT NULL - OR response_body_ref IS NOT NULL - OR client_response_body_ref IS NOT NULL - OR request_body_state IS NOT NULL - OR provider_request_body_state IS NOT NULL - OR response_body_state IS NOT NULL - OR client_response_body_state IS NOT NULL - OR body_capture_mode <> 'none' - ORDER BY request_id ASC - LIMIT ? - ) AS batch -) -"#, - summary, - limit, - ) - .await?; - Ok(()) -} diff --git a/crates/aether-data/runtime/src/backend/system/sqlite.rs b/crates/aether-data/runtime/src/backend/system/sqlite.rs deleted file mode 100644 index db0a41959..000000000 --- a/crates/aether-data/runtime/src/backend/system/sqlite.rs +++ /dev/null @@ -1,1438 +0,0 @@ -use sqlx::Row; - -use crate::error::SqlResultExt; -use crate::repository::system::{ - AdminSystemStats, AdminSystemStatsDailyAggregate, AdminSystemStatsDailyApiKeyAggregate, - AdminSystemStatsUserDailyAggregate, -}; -use crate::DataLayerError; - -use super::u64_from_i64; -use super::*; - -const READ_ADMIN_SYSTEM_STATS_SQL: &str = r#" -SELECT - (SELECT COUNT(*) FROM users) AS total_users, - (SELECT COUNT(*) FROM users WHERE is_active = 1) AS active_users, - (SELECT COUNT(*) FROM api_keys) AS total_api_keys, - (SELECT COUNT(*) FROM "usage") AS total_requests -"#; - -async fn export_sqlite_admin_system_usage_aggregates( - pool: &sqlx::SqlitePool, -) -> Result { - let mut snapshot = AdminSystemUsageAggregateSnapshot::default(); - - let daily_rows = sqlx::query( - r#" -SELECT - "date" AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost, - actual_total_cost, - is_complete, - aggregated_at AS aggregated_at_unix_secs -FROM stats_daily -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY "date" ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in daily_rows { - snapshot.stats_daily.push(map_stats_daily_aggregate(&row)?); - } - - let user_daily_rows = sqlx::query( - r#" -SELECT - user_id, - username, - "date" AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost -FROM stats_user_daily -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY user_id ASC, "date" ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in user_daily_rows { - snapshot - .stats_user_daily - .push(map_stats_user_daily_aggregate(&row)?); - } - - let api_key_daily_rows = sqlx::query( - r#" -SELECT - api_key_id, - api_key_name, - "date" AS date_unix_secs, - total_requests, - success_requests, - error_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - total_cost -FROM stats_daily_api_key -WHERE total_requests <> 0 - OR input_tokens <> 0 - OR output_tokens <> 0 - OR cache_creation_tokens <> 0 - OR cache_read_tokens <> 0 - OR total_cost <> 0 -ORDER BY api_key_id ASC, "date" ASC -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - for row in api_key_daily_rows { - snapshot - .stats_daily_api_key - .push(map_stats_daily_api_key_aggregate(&row)?); - } - - Ok(snapshot) -} - -async fn import_sqlite_admin_system_usage_aggregates( - pool: &sqlx::SqlitePool, - snapshot: &AdminSystemUsageAggregateSnapshot, - user_id_map: &BTreeMap, - api_key_id_map: &BTreeMap, - mode: AdminSystemUsageAggregateImportMode, -) -> Result { - let mut tx = pool.begin().await.map_sql_err()?; - let mut summary = AdminSystemUsageAggregateImportSummary::default(); - let now = current_unix_secs(); - - for row in &snapshot.stats_daily { - let existing: Option = - sqlx::query_scalar(r#"SELECT id FROM stats_daily WHERE "date" = ? LIMIT 1"#) - .bind(i64_from_u64(row.date_unix_secs, "stats_daily.date")?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_daily", - row.date_unix_secs, - )? { - summary.stats_daily.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_daily ( - id, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, actual_total_cost, is_complete, aggregated_at, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT ("date") 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, - actual_total_cost = excluded.actual_total_cost, - is_complete = excluded.is_complete, - aggregated_at = excluded.aggregated_at, - updated_at = excluded.updated_at -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(i64_from_u64(row.date_unix_secs, "stats_daily.date")?) - .bind(i64_from_u64( - row.total_requests, - "stats_daily.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_daily.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_daily.error_requests", - )?) - .bind(i64_from_u64(row.input_tokens, "stats_daily.input_tokens")?) - .bind(i64_from_u64( - row.output_tokens, - "stats_daily.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_daily.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_daily.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(row.actual_total_cost) - .bind(if row.is_complete || row.total_requests > 0 { - 1_i64 - } else { - 0_i64 - }) - .bind(optional_i64_from_u64( - row.aggregated_at_unix_secs, - "stats_daily.aggregated_at", - )?) - .bind(i64_from_u64(now, "stats_daily.created_at")?) - .bind(i64_from_u64(now, "stats_daily.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_daily, existing.is_some()); - } - - for row in &snapshot.stats_user_daily { - let Some(target_user_id) = user_id_map.get(&row.user_id) else { - summary.skipped_unmapped_user_daily += 1; - summary.stats_user_daily.skipped += 1; - continue; - }; - let existing: Option = sqlx::query_scalar( - r#"SELECT id FROM stats_user_daily WHERE user_id = ? AND "date" = ? LIMIT 1"#, - ) - .bind(target_user_id) - .bind(i64_from_u64(row.date_unix_secs, "stats_user_daily.date")?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_user_daily", - row.date_unix_secs, - )? { - summary.stats_user_daily.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_user_daily ( - id, user_id, username, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT ("date", user_id) DO UPDATE -SET username = excluded.username, - 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, - updated_at = excluded.updated_at -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(target_user_id) - .bind(row.username.as_deref()) - .bind(i64_from_u64(row.date_unix_secs, "stats_user_daily.date")?) - .bind(i64_from_u64( - row.total_requests, - "stats_user_daily.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_user_daily.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_user_daily.error_requests", - )?) - .bind(i64_from_u64( - row.input_tokens, - "stats_user_daily.input_tokens", - )?) - .bind(i64_from_u64( - row.output_tokens, - "stats_user_daily.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_user_daily.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_user_daily.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(i64_from_u64(now, "stats_user_daily.created_at")?) - .bind(i64_from_u64(now, "stats_user_daily.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_user_daily, existing.is_some()); - } - - for row in &snapshot.stats_daily_api_key { - let Some(target_api_key_id) = api_key_id_map.get(&row.api_key_id) else { - summary.skipped_unmapped_api_key_daily += 1; - summary.stats_daily_api_key.skipped += 1; - continue; - }; - let existing: Option = sqlx::query_scalar( - r#"SELECT id FROM stats_daily_api_key WHERE api_key_id = ? AND "date" = ? LIMIT 1"#, - ) - .bind(target_api_key_id) - .bind(i64_from_u64( - row.date_unix_secs, - "stats_daily_api_key.date", - )?) - .fetch_optional(&mut *tx) - .await - .map_sql_err()?; - if should_skip_imported_aggregate( - existing.is_some(), - mode, - "stats_daily_api_key", - row.date_unix_secs, - )? { - summary.stats_daily_api_key.skipped += 1; - continue; - } - sqlx::query( - r#" -INSERT INTO stats_daily_api_key ( - id, api_key_id, api_key_name, "date", total_requests, success_requests, error_requests, - input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, - total_cost, created_at, updated_at -) -VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -ON CONFLICT ("date", api_key_id) DO UPDATE -SET api_key_name = excluded.api_key_name, - 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, - updated_at = excluded.updated_at -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(target_api_key_id) - .bind(row.api_key_name.as_deref()) - .bind(i64_from_u64( - row.date_unix_secs, - "stats_daily_api_key.date", - )?) - .bind(i64_from_u64( - row.total_requests, - "stats_daily_api_key.total_requests", - )?) - .bind(i64_from_u64( - row.success_requests, - "stats_daily_api_key.success_requests", - )?) - .bind(i64_from_u64( - row.error_requests, - "stats_daily_api_key.error_requests", - )?) - .bind(i64_from_u64( - row.input_tokens, - "stats_daily_api_key.input_tokens", - )?) - .bind(i64_from_u64( - row.output_tokens, - "stats_daily_api_key.output_tokens", - )?) - .bind(i64_from_u64( - row.cache_creation_tokens, - "stats_daily_api_key.cache_creation_tokens", - )?) - .bind(i64_from_u64( - row.cache_read_tokens, - "stats_daily_api_key.cache_read_tokens", - )?) - .bind(row.total_cost) - .bind(i64_from_u64(now, "stats_daily_api_key.created_at")?) - .bind(i64_from_u64(now, "stats_daily_api_key.updated_at")?) - .execute(&mut *tx) - .await - .map_sql_err()?; - add_aggregate_import_count(&mut summary.stats_daily_api_key, existing.is_some()); - } - - if mode == AdminSystemUsageAggregateImportMode::ValidateError { - tx.rollback().await.map_sql_err()?; - } else { - tx.commit().await.map_sql_err()?; - } - Ok(summary) -} - -pub(super) async fn sqlite_delete_table( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !sqlite_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!("DELETE FROM \"{table}\""); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn sqlite_delete_non_admin_user_rows( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !sqlite_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!( - "DELETE FROM \"{table}\" WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin')" - ); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn sqlite_delete_non_admin_api_key_rows( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let table = checked_sql_identifier(table)?; - if !sqlite_table_exists(tx, table).await? { - return Ok(()); - } - let sql = format!( - "DELETE FROM \"{table}\" WHERE api_key_id IN (SELECT id FROM api_keys WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'))" - ); - let rows = sqlx::query(&sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(table, rows); - Ok(()) -} - -pub(super) async fn sqlite_execute_if_table( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - if !sqlite_table_exists(tx, checked_sql_identifier(table)?).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn sqlite_execute_if_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - columns: &[&str], - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - if !sqlite_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn sqlite_execute_batch_if_table( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, - limit: i64, -) -> Result<(), DataLayerError> { - if !sqlite_table_exists(tx, checked_sql_identifier(table)?).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .bind(limit) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn sqlite_execute_batch_if_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - columns: &[&str], - key: &str, - sql: &str, - summary: &mut AdminSystemPurgeSummary, - limit: i64, -) -> Result<(), DataLayerError> { - if !sqlite_table_has_columns(tx, checked_sql_identifier(table)?, columns).await? { - return Ok(()); - } - let rows = sqlx::query(sql) - .bind(limit) - .execute(&mut **tx) - .await - .map_sql_err()? - .rows_affected(); - summary.add(key, rows); - Ok(()) -} - -pub(super) async fn sqlite_table_exists( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, -) -> Result { - let table = checked_sql_identifier(table)?; - let total: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?") - .bind(table) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - Ok(total > 0) -} - -pub(super) async fn sqlite_table_has_columns( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table: &str, - columns: &[&str], -) -> Result { - let table = checked_sql_identifier(table)?; - if !sqlite_table_exists(tx, table).await? { - return Ok(false); - } - for column in columns { - let column = checked_sql_identifier(column)?; - let total: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name = ?") - .bind(table) - .bind(column) - .fetch_one(&mut **tx) - .await - .map_sql_err()?; - if total == 0 { - return Ok(false); - } - } - Ok(true) -} - -async fn purge_sqlite_admin_system_data( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - target: AdminSystemPurgeTarget, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - match target { - AdminSystemPurgeTarget::Config => { - sqlite_execute_if_table( - tx, - "usage", - "usage_provider_refs_cleared", - r#" -UPDATE "usage" -SET provider_id = NULL, - provider_endpoint_id = NULL, - provider_api_key_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, - ) - .await?; - sqlite_execute_if_table( - tx, - "request_candidates", - "request_candidate_provider_refs_cleared", - r#" -UPDATE request_candidates -SET provider_id = NULL, - endpoint_id = NULL, - key_id = NULL -WHERE provider_id IS NOT NULL - OR endpoint_id IS NOT NULL - OR key_id IS NOT NULL -"#, - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "video_tasks", - "video_task_provider_refs_cleared", - r#" -UPDATE video_tasks -SET provider_id = NULL, - endpoint_id = NULL, - key_id = NULL -WHERE provider_id IS NOT NULL - OR endpoint_id IS NOT NULL - OR key_id IS NOT NULL -"#, - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "user_preferences", - "user_default_provider_refs_cleared", - "UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id IS NOT NULL", - summary, - ) - .await?; - for table in ADMIN_CONFIG_PURGE_TABLES { - sqlite_delete_table(tx, table, summary).await?; - } - } - AdminSystemPurgeTarget::Users => purge_sqlite_non_admin_users(tx, summary).await?, - AdminSystemPurgeTarget::Usage => { - sqlite_delete_table(tx, "request_candidates", summary).await?; - for table in ADMIN_USAGE_CHILD_TABLES { - sqlite_delete_table(tx, table, summary).await?; - } - sqlite_delete_table(tx, "usage", summary).await?; - sqlite_execute_if_table( - tx, - "api_keys", - "api_key_usage_stats_reset", - r#" -UPDATE api_keys -SET total_requests = 0, - total_tokens = 0, - total_cost_usd = 0, - last_used_at = NULL -WHERE total_requests <> 0 - OR total_tokens <> 0 - OR total_cost_usd <> 0 - OR last_used_at IS NOT NULL -"#, - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "provider_api_keys", - "provider_key_usage_stats_reset", - r#" -UPDATE provider_api_keys -SET request_count = 0, - success_count = 0, - error_count = 0, - total_tokens = 0, - total_cost_usd = 0, - total_response_time_ms = 0, - last_used_at = NULL -WHERE request_count <> 0 - OR success_count <> 0 - OR error_count <> 0 - OR total_tokens <> 0 - OR total_cost_usd <> 0 - OR total_response_time_ms <> 0 - OR last_used_at IS NOT NULL -"#, - 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?; - } - AdminSystemPurgeTarget::RequestBodies => { - sqlite_delete_table(tx, "usage_body_blobs", summary).await?; - sqlite_execute_if_table_has_columns( - tx, - "usage", - USAGE_BODY_FIELD_COLUMNS, - "usage_body_fields_cleaned", - 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 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 -"#, - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "usage_http_audits", - "usage_http_audit_body_refs_cleaned", - r#" -UPDATE usage_http_audits -SET request_body_ref = NULL, - provider_request_body_ref = NULL, - response_body_ref = NULL, - client_response_body_ref = NULL, - request_body_state = NULL, - provider_request_body_state = NULL, - response_body_state = NULL, - client_response_body_state = NULL, - body_capture_mode = 'none' -WHERE request_body_ref IS NOT NULL - OR provider_request_body_ref IS NOT NULL - OR response_body_ref IS NOT NULL - OR client_response_body_ref IS NOT NULL - OR request_body_state IS NOT NULL - OR provider_request_body_state IS NOT NULL - OR response_body_state IS NOT NULL - OR client_response_body_state IS NOT NULL - OR body_capture_mode <> 'none' -"#, - summary, - ) - .await?; - } - AdminSystemPurgeTarget::Stats => { - for table in ADMIN_STATS_PURGE_TABLES { - sqlite_delete_table(tx, table, summary).await?; - } - } - } - Ok(()) -} - -async fn purge_sqlite_non_admin_users( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let non_admin_users = "SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'"; - let non_admin_keys = "SELECT id FROM api_keys WHERE user_id IN (SELECT id FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin')"; - - sqlite_execute_if_table( - tx, - "api_key_provider_mappings", - "api_key_provider_mappings", - &format!("DELETE FROM api_key_provider_mappings WHERE api_key_id IN ({non_admin_keys})"), - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "usage", - "usage_user_refs_cleared", - &format!(r#"UPDATE "usage" SET user_id = NULL WHERE user_id IN ({non_admin_users})"#), - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "usage", - "usage_api_key_refs_cleared", - &format!(r#"UPDATE "usage" SET api_key_id = NULL WHERE api_key_id IN ({non_admin_keys})"#), - summary, - ) - .await?; - for (table, column, key) in [ - ( - "request_candidates", - "user_id", - "request_candidate_user_refs_cleared", - ), - ("audit_logs", "user_id", "audit_log_user_refs_cleared"), - ("video_tasks", "user_id", "video_task_user_refs_cleared"), - ("wallets", "user_id", "wallet_user_refs_cleared"), - ( - "payment_orders", - "user_id", - "payment_order_user_refs_cleared", - ), - ( - "wallet_transactions", - "operator_id", - "wallet_transaction_operator_refs_cleared", - ), - ( - "announcements", - "author_id", - "announcement_author_refs_cleared", - ), - ( - "proxy_nodes", - "registered_by", - "proxy_node_registrant_refs_cleared", - ), - ] { - sqlite_execute_if_table( - tx, - table, - key, - &format!( - r#"UPDATE "{table}" SET "{column}" = NULL WHERE "{column}" IN ({non_admin_users})"# - ), - summary, - ) - .await?; - } - for (table, key) in [ - ( - "request_candidates", - "request_candidate_api_key_refs_cleared", - ), - ("audit_logs", "audit_log_api_key_refs_cleared"), - ("video_tasks", "video_task_api_key_refs_cleared"), - ("wallets", "wallet_api_key_refs_cleared"), - ] { - sqlite_execute_if_table( - tx, - table, - key, - &format!( - r#"UPDATE "{table}" SET api_key_id = NULL WHERE api_key_id IN ({non_admin_keys})"# - ), - summary, - ) - .await?; - } - sqlite_delete_non_admin_api_key_rows(tx, "stats_daily_api_key", summary).await?; - sqlite_execute_if_table( - tx, - "refund_requests", - "refund_request_user_refs_cleared", - &format!( - r#" -UPDATE refund_requests -SET user_id = CASE WHEN user_id IN ({non_admin_users}) THEN NULL ELSE user_id END, - requested_by = CASE WHEN requested_by IN ({non_admin_users}) THEN NULL ELSE requested_by END, - approved_by = CASE WHEN approved_by IN ({non_admin_users}) THEN NULL ELSE approved_by END, - processed_by = CASE WHEN processed_by IN ({non_admin_users}) THEN NULL ELSE processed_by END -WHERE user_id IN ({non_admin_users}) - OR requested_by IN ({non_admin_users}) - OR approved_by IN ({non_admin_users}) - OR processed_by IN ({non_admin_users}) -"# - ), - summary, - ) - .await?; - for table in ADMIN_USER_SCOPED_TABLES { - sqlite_delete_non_admin_user_rows(tx, table, summary).await?; - } - sqlite_execute_if_table( - tx, - "api_keys", - "api_keys", - &format!("DELETE FROM api_keys WHERE user_id IN ({non_admin_users})"), - summary, - ) - .await?; - sqlite_execute_if_table( - tx, - "users", - "users", - "DELETE FROM users WHERE LOWER(COALESCE(role, '')) <> 'admin'", - summary, - ) - .await?; - Ok(()) -} - -async fn purge_sqlite_request_bodies_batch( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - batch_size: usize, - summary: &mut AdminSystemPurgeSummary, -) -> Result<(), DataLayerError> { - let limit = i64::try_from(batch_size).unwrap_or(i64::MAX); - sqlite_execute_batch_if_table( - tx, - "usage_body_blobs", - "usage_body_blobs", - r#" -DELETE FROM usage_body_blobs -WHERE body_ref IN ( - SELECT body_ref - FROM usage_body_blobs - ORDER BY body_ref ASC - LIMIT ? -) -"#, - summary, - limit, - ) - .await?; - sqlite_execute_batch_if_table_has_columns( - tx, - "usage", - USAGE_BODY_FIELD_COLUMNS, - "usage_body_fields_cleaned", - 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 request_id IN ( - SELECT request_id - FROM "usage" - WHERE 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 - ORDER BY created_at_unix_ms ASC, request_id ASC - LIMIT ? -) -"#, - summary, - limit, - ) - .await?; - sqlite_execute_batch_if_table( - tx, - "usage_http_audits", - "usage_http_audit_body_refs_cleaned", - r#" -UPDATE usage_http_audits -SET request_body_ref = NULL, - provider_request_body_ref = NULL, - response_body_ref = NULL, - client_response_body_ref = NULL, - request_body_state = NULL, - provider_request_body_state = NULL, - response_body_state = NULL, - client_response_body_state = NULL, - body_capture_mode = 'none' -WHERE request_id IN ( - SELECT request_id - FROM usage_http_audits - WHERE request_body_ref IS NOT NULL - OR provider_request_body_ref IS NOT NULL - OR response_body_ref IS NOT NULL - OR client_response_body_ref IS NOT NULL - OR request_body_state IS NOT NULL - OR provider_request_body_state IS NOT NULL - OR response_body_state IS NOT NULL - OR client_response_body_state IS NOT NULL - OR body_capture_mode <> 'none' - ORDER BY request_id ASC - LIMIT ? -) -"#, - summary, - limit, - ) - .await?; - Ok(()) -} - -impl SqliteBackend { - pub async fn purge_admin_system_data( - &self, - target: AdminSystemPurgeTarget, - ) -> Result { - let mut tx = self.pool().begin().await.map_sql_err()?; - let mut summary = AdminSystemPurgeSummary::default(); - purge_sqlite_admin_system_data(&mut tx, target, &mut summary).await?; - tx.commit().await.map_sql_err()?; - Ok(summary) - } - - pub async fn purge_admin_request_bodies_batch( - &self, - batch_size: usize, - ) -> Result { - if batch_size == 0 { - return Ok(AdminSystemPurgeSummary::default()); - } - let mut tx = self.pool().begin().await.map_sql_err()?; - let mut summary = AdminSystemPurgeSummary::default(); - purge_sqlite_request_bodies_batch(&mut tx, batch_size, &mut summary).await?; - tx.commit().await.map_sql_err()?; - Ok(summary) - } - - pub async fn find_system_config_value( - &self, - key: &str, - ) -> Result, DataLayerError> { - let row = sqlx::query( - r#" -SELECT value -FROM system_configs -WHERE key = ? -LIMIT 1 -"#, - ) - .bind(key) - .fetch_optional(self.pool()) - .await - .map_sql_err()?; - - row.map(|row| { - row.try_get("value") - .map_sql_err() - .and_then(parse_json_value) - }) - .transpose() - } - - pub async fn compare_and_set_system_config_string_value( - &self, - key: &str, - expected: &str, - replacement: &str, - ) -> Result { - let now = current_unix_secs(); - let replacement = - serialize_json_value(&serde_json::Value::String(replacement.to_string()))?; - let result = sqlx::query( - r#" -UPDATE system_configs -SET value = ?, updated_at = ? -WHERE key = ? - AND json_valid(value) - AND json_type(value) = 'text' - AND CAST(json_extract(value, '$') AS TEXT) = ? COLLATE BINARY -"#, - ) - .bind(replacement) - .bind(now as i64) - .bind(key) - .bind(expected) - .execute(self.pool()) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - pub async fn upsert_system_config_value( - &self, - key: &str, - value: &serde_json::Value, - description: Option<&str>, - ) -> Result { - Ok(self - .upsert_system_config_entry(key, value, description) - .await? - .value) - } - - pub async fn list_system_config_entries( - &self, - ) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT key, value, description, updated_at -FROM system_configs -ORDER BY key ASC -"#, - ) - .fetch_all(self.pool()) - .await - .map_sql_err()?; - - rows.into_iter() - .map(|row| { - Ok(StoredSystemConfigEntry { - key: row.try_get("key").map_sql_err()?, - value: parse_json_value(row.try_get("value").map_sql_err()?)?, - description: row.try_get("description").map_sql_err()?, - updated_at_unix_secs: row - .try_get::, _>("updated_at") - .map_sql_err()? - .map(|value| value.max(0) as u64), - }) - }) - .collect() - } - - pub async fn upsert_system_config_entry( - &self, - key: &str, - value: &serde_json::Value, - description: Option<&str>, - ) -> Result { - let now = current_unix_secs(); - let serialized = serialize_json_value(value)?; - sqlx::query( - r#" -INSERT INTO system_configs (id, key, value, description, created_at, updated_at) -VALUES (?, ?, ?, ?, ?, ?) -ON CONFLICT (key) DO UPDATE -SET value = excluded.value, - description = COALESCE(excluded.description, system_configs.description), - updated_at = excluded.updated_at -"#, - ) - .bind(uuid::Uuid::new_v4().to_string()) - .bind(key) - .bind(serialized) - .bind(description) - .bind(now as i64) - .bind(now as i64) - .execute(self.pool()) - .await - .map_sql_err()?; - - self.list_system_config_entries() - .await? - .into_iter() - .find(|entry| entry.key == key) - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "system config key '{key}' missing after sqlite upsert" - )) - }) - } - - pub async fn delete_system_config_value(&self, key: &str) -> Result { - let result = sqlx::query( - r#" -DELETE FROM system_configs -WHERE key = ? -"#, - ) - .bind(key) - .execute(self.pool()) - .await - .map_sql_err()?; - Ok(result.rows_affected() > 0) - } - - pub async fn read_admin_system_stats(&self) -> Result { - let row = sqlx::query(READ_ADMIN_SYSTEM_STATS_SQL) - .fetch_one(self.pool()) - .await - .map_sql_err()?; - map_admin_system_stats(row) - } - - pub async fn export_admin_system_usage_aggregates( - &self, - ) -> Result { - export_sqlite_admin_system_usage_aggregates(self.pool()).await - } - - pub async fn import_admin_system_usage_aggregates( - &self, - snapshot: &AdminSystemUsageAggregateSnapshot, - user_id_map: &BTreeMap, - api_key_id_map: &BTreeMap, - mode: AdminSystemUsageAggregateImportMode, - ) -> Result { - import_sqlite_admin_system_usage_aggregates( - self.pool(), - snapshot, - user_id_map, - api_key_id_map, - mode, - ) - .await - } -} - -pub(super) fn map_stats_daily_aggregate( - row: &sqlx::sqlite::SqliteRow, -) -> Result { - Ok(AdminSystemStatsDailyAggregate { - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - actual_total_cost: row.try_get("actual_total_cost").map_sql_err()?, - is_complete: row.try_get::("is_complete").map_sql_err()? != 0, - aggregated_at_unix_secs: row - .try_get::, _>("aggregated_at_unix_secs") - .map_sql_err()? - .map(u64_from_i64), - }) -} - -pub(super) fn map_stats_user_daily_aggregate( - row: &sqlx::sqlite::SqliteRow, -) -> Result { - Ok(AdminSystemStatsUserDailyAggregate { - user_id: row.try_get("user_id").map_sql_err()?, - username: row.try_get("username").map_sql_err()?, - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - }) -} - -pub(super) fn map_stats_daily_api_key_aggregate( - row: &sqlx::sqlite::SqliteRow, -) -> Result { - Ok(AdminSystemStatsDailyApiKeyAggregate { - api_key_id: row.try_get("api_key_id").map_sql_err()?, - api_key_name: row.try_get("api_key_name").map_sql_err()?, - date_unix_secs: u64_from_i64(row.try_get("date_unix_secs").map_sql_err()?), - total_requests: u64_from_i64(row.try_get("total_requests").map_sql_err()?), - success_requests: u64_from_i64(row.try_get("success_requests").map_sql_err()?), - error_requests: u64_from_i64(row.try_get("error_requests").map_sql_err()?), - input_tokens: u64_from_i64(row.try_get("input_tokens").map_sql_err()?), - output_tokens: u64_from_i64(row.try_get("output_tokens").map_sql_err()?), - cache_creation_tokens: u64_from_i64(row.try_get("cache_creation_tokens").map_sql_err()?), - cache_read_tokens: u64_from_i64(row.try_get("cache_read_tokens").map_sql_err()?), - total_cost: row.try_get("total_cost").map_sql_err()?, - }) -} - -pub(super) fn map_admin_system_stats( - row: sqlx::sqlite::SqliteRow, -) -> Result { - Ok(AdminSystemStats { - total_users: row.try_get::("total_users").map_sql_err()?.max(0) as u64, - active_users: row.try_get::("active_users").map_sql_err()?.max(0) as u64, - total_api_keys: row - .try_get::("total_api_keys") - .map_sql_err()? - .max(0) as u64, - total_requests: row - .try_get::("total_requests") - .map_sql_err()? - .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/backend/wallet.rs b/crates/aether-data/runtime/src/backend/wallet.rs index bf0c3ec54..fa916568d 100644 --- a/crates/aether-data/runtime/src/backend/wallet.rs +++ b/crates/aether-data/runtime/src/backend/wallet.rs @@ -1,14 +1,7 @@ //! Driver-specific wallet usage aggregation adapters. -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -use sha2::{Digest, Sha256}; use crate::DataLayerError; @@ -27,43 +20,9 @@ pub(super) fn unix_secs_to_utc( .ok_or_else(|| DataLayerError::InvalidInput(format!("invalid {field_name}: {value}"))) } -#[cfg(any(feature = "mysql", feature = "sqlite"))] -pub(super) fn wallet_daily_usage_id( - wallet_id: &str, - billing_date: &str, - billing_timezone: &str, -) -> String { - let mut hasher = Sha256::new(); - hasher.update(b"wallet-daily-usage:"); - hasher.update(wallet_id.as_bytes()); - hasher.update(b":"); - hasher.update(billing_date.as_bytes()); - hasher.update(b":"); - hasher.update(billing_timezone.as_bytes()); - hasher - .finalize() - .iter() - .map(|byte| format!("{byte:02x}")) - .collect() -} - #[cfg(test)] mod tests { use super::u64_to_i64; - #[cfg(any(feature = "mysql", feature = "sqlite"))] - use super::wallet_daily_usage_id; - - #[cfg(any(feature = "mysql", feature = "sqlite"))] - #[test] - fn wallet_daily_usage_ids_are_stable_and_partition_specific() { - let first = wallet_daily_usage_id("wallet-1", "2026-07-13", "UTC"); - let same = wallet_daily_usage_id("wallet-1", "2026-07-13", "UTC"); - let other_day = wallet_daily_usage_id("wallet-1", "2026-07-14", "UTC"); - - assert_eq!(first, same); - assert_ne!(first, other_day); - assert_eq!(first.len(), 64); - } #[test] fn rejects_timestamps_outside_i64_range() { diff --git a/crates/aether-data/runtime/src/backend/wallet/mysql.rs b/crates/aether-data/runtime/src/backend/wallet/mysql.rs deleted file mode 100644 index 28f6ce76a..000000000 --- a/crates/aether-data/runtime/src/backend/wallet/mysql.rs +++ /dev/null @@ -1,155 +0,0 @@ -use sqlx::Row; - -use crate::backend::MysqlBackend; -use crate::error::SqlResultExt; -use crate::{DataLayerError, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult}; - -use super::{u64_to_i64, wallet_daily_usage_id}; - -const SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL: &str = r#" -SELECT - usage_settlement_snapshots.wallet_id AS wallet_id, - CAST(COUNT(*) AS SIGNED) AS total_requests, - CAST(COALESCE(SUM(`usage`.total_cost_usd), 0) AS DOUBLE) AS total_cost_usd, - CAST(COALESCE(SUM(`usage`.input_tokens), 0) AS SIGNED) AS input_tokens, - CAST(COALESCE(SUM(`usage`.output_tokens), 0) AS SIGNED) AS output_tokens, - CAST(COALESCE(SUM(`usage`.cache_creation_input_tokens), 0) AS SIGNED) AS cache_creation_tokens, - CAST(COALESCE(SUM(`usage`.cache_read_input_tokens), 0) AS SIGNED) AS cache_read_tokens, - MIN(COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at)) AS first_finalized_at, - MAX(COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at)) AS last_finalized_at -FROM `usage` -JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = `usage`.request_id -WHERE usage_settlement_snapshots.wallet_id IS NOT NULL - AND usage_settlement_snapshots.wallet_id <> '' - AND COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) = 'settled' - AND `usage`.total_cost_usd > 0 - AND COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at) >= ? - AND COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at) < ? -GROUP BY usage_settlement_snapshots.wallet_id -"#; - -impl MysqlBackend { - pub async fn aggregate_wallet_daily_usage( - &self, - input: &WalletDailyUsageAggregationInput, - ) -> Result { - let window_start = u64_to_i64(input.window_start_unix_secs, "window_start")?; - let window_end = u64_to_i64(input.window_end_unix_secs, "window_end")?; - let aggregated_at = u64_to_i64(input.aggregated_at_unix_secs, "aggregated_at")?; - let mut tx = self.pool().begin().await.map_sql_err()?; - - let rows = sqlx::query(SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL) - .bind(window_start) - .bind(window_end) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - - let mut aggregated_wallets = 0usize; - for row in rows { - let wallet_id: String = row.try_get("wallet_id").map_sql_err()?; - sqlx::query( - r#" -DELETE FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_date = ? - AND billing_timezone = ? -"#, - ) - .bind(&wallet_id) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, - wallet_id, - billing_date, - billing_timezone, - total_cost_usd, - total_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - first_finalized_at, - last_finalized_at, - aggregated_at, - created_at, - updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(wallet_daily_usage_id( - &wallet_id, - &input.billing_date, - &input.billing_timezone, - )) - .bind(&wallet_id) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .bind(row.try_get::("total_cost_usd").map_sql_err()?) - .bind(row.try_get::("total_requests").map_sql_err()?) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .bind( - row.try_get::, _>("first_finalized_at") - .map_sql_err()?, - ) - .bind( - row.try_get::, _>("last_finalized_at") - .map_sql_err()?, - ) - .bind(aggregated_at) - .bind(aggregated_at) - .bind(aggregated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - aggregated_wallets += 1; - } - - let deleted_stale_ledgers = sqlx::query( - r#" -DELETE FROM wallet_daily_usage_ledgers -WHERE billing_date = ? - AND billing_timezone = ? - AND NOT EXISTS ( - SELECT 1 - FROM `usage` - JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = `usage`.request_id - WHERE usage_settlement_snapshots.wallet_id = wallet_daily_usage_ledgers.wallet_id - AND COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) = 'settled' - AND `usage`.total_cost_usd > 0 - AND COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at) >= ? - AND COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at) < ? - ) -"#, - ) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .bind(window_start) - .bind(window_end) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - - tx.commit().await.map_sql_err()?; - Ok(WalletDailyUsageAggregationResult { - aggregated_wallets, - deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX), - }) - } -} diff --git a/crates/aether-data/runtime/src/backend/wallet/sqlite.rs b/crates/aether-data/runtime/src/backend/wallet/sqlite.rs deleted file mode 100644 index e8233b167..000000000 --- a/crates/aether-data/runtime/src/backend/wallet/sqlite.rs +++ /dev/null @@ -1,156 +0,0 @@ -use sqlx::Row; - -use crate::backend::SqliteBackend; -use crate::driver::sqlite::sqlite_real; -use crate::error::SqlResultExt; -use crate::{DataLayerError, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult}; - -use super::{u64_to_i64, wallet_daily_usage_id}; - -const SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL: &str = r#" -SELECT - usage_settlement_snapshots.wallet_id AS wallet_id, - COUNT(*) AS total_requests, - CAST(COALESCE(SUM("usage".total_cost_usd), 0) AS REAL) AS total_cost_usd, - COALESCE(SUM("usage".input_tokens), 0) AS input_tokens, - COALESCE(SUM("usage".output_tokens), 0) AS output_tokens, - COALESCE(SUM("usage".cache_creation_input_tokens), 0) AS cache_creation_tokens, - COALESCE(SUM("usage".cache_read_input_tokens), 0) AS cache_read_tokens, - MIN(COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at)) AS first_finalized_at, - MAX(COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at)) AS last_finalized_at -FROM "usage" -JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = "usage".request_id -WHERE usage_settlement_snapshots.wallet_id IS NOT NULL - AND usage_settlement_snapshots.wallet_id <> '' - AND COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) = 'settled' - AND "usage".total_cost_usd > 0 - AND COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at) >= ? - AND COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at) < ? -GROUP BY usage_settlement_snapshots.wallet_id -"#; - -impl SqliteBackend { - pub async fn aggregate_wallet_daily_usage( - &self, - input: &WalletDailyUsageAggregationInput, - ) -> Result { - let window_start = u64_to_i64(input.window_start_unix_secs, "window_start")?; - let window_end = u64_to_i64(input.window_end_unix_secs, "window_end")?; - let aggregated_at = u64_to_i64(input.aggregated_at_unix_secs, "aggregated_at")?; - let mut tx = self.pool().begin().await.map_sql_err()?; - - let rows = sqlx::query(SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL) - .bind(window_start) - .bind(window_end) - .fetch_all(&mut *tx) - .await - .map_sql_err()?; - - let mut aggregated_wallets = 0usize; - for row in rows { - let wallet_id: String = row.try_get("wallet_id").map_sql_err()?; - sqlx::query( - r#" -DELETE FROM wallet_daily_usage_ledgers -WHERE wallet_id = ? - AND billing_date = ? - AND billing_timezone = ? -"#, - ) - .bind(&wallet_id) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .execute(&mut *tx) - .await - .map_sql_err()?; - - sqlx::query( - r#" -INSERT INTO wallet_daily_usage_ledgers ( - id, - wallet_id, - billing_date, - billing_timezone, - total_cost_usd, - total_requests, - input_tokens, - output_tokens, - cache_creation_tokens, - cache_read_tokens, - first_finalized_at, - last_finalized_at, - aggregated_at, - created_at, - updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -"#, - ) - .bind(wallet_daily_usage_id( - &wallet_id, - &input.billing_date, - &input.billing_timezone, - )) - .bind(&wallet_id) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .bind(sqlite_real(&row, "total_cost_usd")?) - .bind(row.try_get::("total_requests").map_sql_err()?) - .bind(row.try_get::("input_tokens").map_sql_err()?) - .bind(row.try_get::("output_tokens").map_sql_err()?) - .bind( - row.try_get::("cache_creation_tokens") - .map_sql_err()?, - ) - .bind(row.try_get::("cache_read_tokens").map_sql_err()?) - .bind( - row.try_get::, _>("first_finalized_at") - .map_sql_err()?, - ) - .bind( - row.try_get::, _>("last_finalized_at") - .map_sql_err()?, - ) - .bind(aggregated_at) - .bind(aggregated_at) - .bind(aggregated_at) - .execute(&mut *tx) - .await - .map_sql_err()?; - aggregated_wallets += 1; - } - - let deleted_stale_ledgers = sqlx::query( - r#" -DELETE FROM wallet_daily_usage_ledgers -WHERE billing_date = ? - AND billing_timezone = ? - AND NOT EXISTS ( - SELECT 1 - FROM "usage" - JOIN usage_settlement_snapshots - ON usage_settlement_snapshots.request_id = "usage".request_id - WHERE usage_settlement_snapshots.wallet_id = wallet_daily_usage_ledgers.wallet_id - AND COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) = 'settled' - AND "usage".total_cost_usd > 0 - AND COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at) >= ? - AND COALESCE(usage_settlement_snapshots.finalized_at, "usage".finalized_at) < ? - ) -"#, - ) - .bind(&input.billing_date) - .bind(&input.billing_timezone) - .bind(window_start) - .bind(window_end) - .execute(&mut *tx) - .await - .map_sql_err()? - .rows_affected(); - - tx.commit().await.map_sql_err()?; - Ok(WalletDailyUsageAggregationResult { - aggregated_wallets, - deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX), - }) - } -} diff --git a/crates/aether-data/runtime/src/backend/write.rs b/crates/aether-data/runtime/src/backend/write.rs index dbbeab29c..6161cc588 100644 --- a/crates/aether-data/runtime/src/backend/write.rs +++ b/crates/aether-data/runtime/src/backend/write.rs @@ -1,12 +1,8 @@ use std::fmt; use std::sync::Arc; -#[cfg(feature = "mysql")] -use super::MysqlBackend; #[cfg(feature = "postgres")] use super::PostgresBackend; -#[cfg(feature = "sqlite")] -use super::SqliteBackend; use crate::repository::announcements::AnnouncementWriteRepository; use crate::repository::auth::AuthApiKeyWriteRepository; use crate::repository::auth_modules::AuthModuleWriteRepository; @@ -79,22 +75,12 @@ impl fmt::Debug for DataWriteRepositories { impl DataWriteRepositories { pub(crate) fn from_backends( #[cfg(feature = "postgres")] postgres: Option<&PostgresBackend>, - #[cfg(feature = "mysql")] mysql: Option<&MysqlBackend>, - #[cfg(feature = "sqlite")] sqlite: Option<&SqliteBackend>, ) -> Self { let mut repositories = Self::default(); #[cfg(feature = "postgres")] if let Some(postgres) = postgres { repositories.install_postgres(postgres); } - #[cfg(feature = "mysql")] - if let Some(mysql) = mysql { - repositories.install_mysql(mysql); - } - #[cfg(feature = "sqlite")] - if let Some(sqlite) = sqlite { - repositories.install_sqlite(sqlite); - } repositories } @@ -162,136 +148,10 @@ impl DataWriteRepositories { } } - #[cfg(feature = "mysql")] - fn install_mysql(&mut self, backend: &MysqlBackend) { - if self.announcements.is_none() { - self.announcements = Some(MysqlBackend::announcement_write_repository(backend)); - } - if self.auth_api_keys.is_none() { - self.auth_api_keys = Some(MysqlBackend::auth_api_key_write_repository(backend)); - } - if self.auth_modules.is_none() { - self.auth_modules = Some(MysqlBackend::auth_module_write_repository(backend)); - } - if self.background_tasks.is_none() { - self.background_tasks = Some(MysqlBackend::background_task_write_repository(backend)); - } - if self.request_candidates.is_none() { - self.request_candidates = - Some(MysqlBackend::request_candidate_write_repository(backend)); - } - if self.gemini_file_mappings.is_none() { - self.gemini_file_mappings = - Some(MysqlBackend::gemini_file_mapping_write_repository(backend)); - } - if self.global_models.is_none() { - self.global_models = Some(MysqlBackend::global_model_write_repository(backend)); - } - if self.management_tokens.is_none() { - self.management_tokens = Some(MysqlBackend::management_token_write_repository(backend)); - } - if self.oauth_providers.is_none() { - self.oauth_providers = Some(MysqlBackend::oauth_provider_write_repository(backend)); - } - if self.pool_scores.is_none() { - self.pool_scores = Some(MysqlBackend::pool_score_write_repository(backend)); - } - if self.proxy_nodes.is_none() { - self.proxy_nodes = Some(MysqlBackend::proxy_node_write_repository(backend)); - } - if self.provider_catalog.is_none() { - self.provider_catalog = Some(MysqlBackend::provider_catalog_write_repository(backend)); - } - if self.provider_quotas.is_none() { - self.provider_quotas = Some(MysqlBackend::provider_quota_write_repository(backend)); - } - if self.routing_groups.is_none() { - self.routing_groups = Some(MysqlBackend::routing_group_write_repository(backend)); - } - if self.settlement.is_none() { - self.settlement = Some(MysqlBackend::settlement_write_repository(backend)); - } - if self.usage.is_none() { - self.usage = Some(MysqlBackend::usage_write_repository(backend)); - } - if self.video_tasks.is_none() { - self.video_tasks = Some(MysqlBackend::video_task_write_repository(backend)); - } - if self.wallets.is_none() { - self.wallets = Some(MysqlBackend::wallet_write_repository(backend)); - } - } - - #[cfg(feature = "sqlite")] - fn install_sqlite(&mut self, backend: &SqliteBackend) { - if self.announcements.is_none() { - self.announcements = Some(SqliteBackend::announcement_write_repository(backend)); - } - if self.auth_api_keys.is_none() { - self.auth_api_keys = Some(SqliteBackend::auth_api_key_write_repository(backend)); - } - if self.auth_modules.is_none() { - self.auth_modules = Some(SqliteBackend::auth_module_write_repository(backend)); - } - if self.background_tasks.is_none() { - self.background_tasks = Some(SqliteBackend::background_task_write_repository(backend)); - } - if self.request_candidates.is_none() { - self.request_candidates = - Some(SqliteBackend::request_candidate_write_repository(backend)); - } - if self.gemini_file_mappings.is_none() { - self.gemini_file_mappings = - Some(SqliteBackend::gemini_file_mapping_write_repository(backend)); - } - if self.global_models.is_none() { - self.global_models = Some(SqliteBackend::global_model_write_repository(backend)); - } - if self.management_tokens.is_none() { - self.management_tokens = - Some(SqliteBackend::management_token_write_repository(backend)); - } - if self.oauth_providers.is_none() { - self.oauth_providers = Some(SqliteBackend::oauth_provider_write_repository(backend)); - } - if self.pool_scores.is_none() { - self.pool_scores = Some(SqliteBackend::pool_score_write_repository(backend)); - } - if self.proxy_nodes.is_none() { - self.proxy_nodes = Some(SqliteBackend::proxy_node_write_repository(backend)); - } - if self.provider_catalog.is_none() { - self.provider_catalog = Some(SqliteBackend::provider_catalog_write_repository(backend)); - } - if self.provider_quotas.is_none() { - self.provider_quotas = Some(SqliteBackend::provider_quota_write_repository(backend)); - } - if self.routing_groups.is_none() { - self.routing_groups = Some(SqliteBackend::routing_group_write_repository(backend)); - } - if self.settlement.is_none() { - self.settlement = Some(SqliteBackend::settlement_write_repository(backend)); - } - if self.usage.is_none() { - self.usage = Some(SqliteBackend::usage_write_repository(backend)); - } - if self.video_tasks.is_none() { - self.video_tasks = Some(SqliteBackend::video_task_write_repository(backend)); - } - if self.wallets.is_none() { - self.wallets = Some(SqliteBackend::wallet_write_repository(backend)); - } - } #[cfg(test)] #[cfg(feature = "postgres")] pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self { - Self::from_backends( - postgres, - #[cfg(feature = "mysql")] - None, - #[cfg(feature = "sqlite")] - None, - ) + Self::from_backends(postgres) } pub fn announcements(&self) -> Option> { diff --git a/crates/aether-data/runtime/src/config.rs b/crates/aether-data/runtime/src/config.rs index d7451acc6..466665ee1 100644 --- a/crates/aether-data/runtime/src/config.rs +++ b/crates/aether-data/runtime/src/config.rs @@ -95,8 +95,8 @@ mod tests { fn new_database_config_takes_priority_over_legacy_postgres_config() { let config = DataLayerConfig { database: Some(SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: "sqlite://./data/aether.db".to_string(), + driver: DatabaseDriver::Postgres, + url: "postgres://localhost/preferred".to_string(), pool: SqlPoolConfig::default(), }), postgres: Some(PostgresPoolConfig { @@ -114,6 +114,7 @@ mod tests { let effective = config .effective_database() .expect("database config should exist"); - assert_eq!(effective.driver, DatabaseDriver::Sqlite); + assert_eq!(effective.driver, DatabaseDriver::Postgres); + assert_eq!(effective.url, "postgres://localhost/preferred"); } } diff --git a/crates/aether-data/runtime/src/driver/mod.rs b/crates/aether-data/runtime/src/driver/mod.rs index 60eaeb878..ccf68b2a9 100644 --- a/crates/aether-data/runtime/src/driver/mod.rs +++ b/crates/aether-data/runtime/src/driver/mod.rs @@ -1,12 +1,7 @@ //! Compatibility paths for database adapter crates. //! -//! New adapter code belongs in `aether-data-postgres`, `aether-data-mysql`, or -//! `aether-data-sqlite`. These modules preserve existing `aether_data::driver` +//! Adapter code belongs in `aether-data-postgres`. These modules preserve existing `aether_data::driver` //! imports while application-facing composition remains in `backend`. -#[cfg(feature = "mysql")] -pub mod mysql; #[cfg(feature = "postgres")] pub mod postgres; -#[cfg(feature = "sqlite")] -pub mod sqlite; diff --git a/crates/aether-data/runtime/src/driver/mysql.rs b/crates/aether-data/runtime/src/driver/mysql.rs deleted file mode 100644 index 664e2097f..000000000 --- a/crates/aether-data/runtime/src/driver/mysql.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! Compatibility facade for the MySQL adapter crate. - -pub use aether_data_mysql::*; diff --git a/crates/aether-data/runtime/src/driver/sqlite.rs b/crates/aether-data/runtime/src/driver/sqlite.rs deleted file mode 100644 index 39259b74a..000000000 --- a/crates/aether-data/runtime/src/driver/sqlite.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! Compatibility facade for the SQLite adapter crate. - -pub use aether_data_sqlite::*; diff --git a/crates/aether-data/runtime/src/lib.rs b/crates/aether-data/runtime/src/lib.rs index 6302656e2..2080ba5e3 100644 --- a/crates/aether-data/runtime/src/lib.rs +++ b/crates/aether-data/runtime/src/lib.rs @@ -17,18 +17,14 @@ pub mod lifecycle; pub mod maintenance; pub mod repository; -#[cfg(feature = "mysql")] -pub use backend::MysqlBackend; #[cfg(feature = "postgres")] pub use backend::PostgresBackend; -#[cfg(feature = "sqlite")] -pub use backend::SqliteBackend; pub use backend::{ DataBackends, DataLeaseBackends, DataReadRepositories, DataTransactionBackends, DataWriteRepositories, }; pub use config::DataLayerConfig; -pub use database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL}; +pub use database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; pub use error::DataLayerError; pub use maintenance::{ DatabaseMaintenanceSummary, DatabasePoolSummary, DatabasePostgresActivityGroup, diff --git a/crates/aether-data/runtime/src/lifecycle/backfill.rs b/crates/aether-data/runtime/src/lifecycle/backfill.rs index a8ac3c0c8..f1555236f 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill.rs @@ -1,25 +1,13 @@ -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; mod types; -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] mod tests; -#[cfg(feature = "mysql")] -pub use mysql::{ - pending_backfills as pending_mysql_backfills, run_backfills as run_mysql_backfills, -}; #[cfg(feature = "postgres")] pub use postgres::{pending_backfills, run_backfills}; -#[cfg(feature = "sqlite")] -pub use sqlite::{ - pending_backfills as pending_sqlite_backfills, run_backfills as run_sqlite_backfills, -}; pub use types::PendingBackfillInfo; -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] use postgres::{pending_backfills_from_applied, AppliedBackfill}; diff --git a/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs b/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs deleted file mode 100644 index 66e37b23f..000000000 --- a/crates/aether-data/runtime/src/lifecycle/backfill/mysql.rs +++ /dev/null @@ -1,263 +0,0 @@ -use std::collections::{HashMap, HashSet}; - -use sqlx::{ - migrate::{Migrate, MigrateError, Migrator}, - query, query_scalar, Connection, MySqlConnection, Row, -}; -use tracing::{error, info, warn}; - -use super::types::PendingBackfillInfo; -use crate::driver::mysql::MysqlPool; - -static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/mysql"); - -const SCHEMA_BACKFILLS_TABLE_EXISTS_SQL: &str = "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = 'schema_backfills'"; -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(()) -} - -async fn pending_backfills_locked( - conn: &mut MySqlConnection, -) -> Result, MigrateError> { - if !schema_backfills_table_exists(conn).await? { - return Ok(pending_backfills_from_applied(&[])); - } - let applied_backfills = list_applied_backfills(conn).await?; - validate_applied_backfills(&applied_backfills)?; - Ok(pending_backfills_from_applied(&applied_backfills)) -} - -async fn schema_backfills_table_exists(conn: &mut MySqlConnection) -> Result { - let total: i64 = query_scalar(SCHEMA_BACKFILLS_TABLE_EXISTS_SQL) - .fetch_one(&mut *conn) - .await?; - Ok(total > 0) -} - -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 deleted file mode 100644 index d34875b12..000000000 --- a/crates/aether-data/runtime/src/lifecycle/backfill/sqlite.rs +++ /dev/null @@ -1,265 +0,0 @@ -use std::collections::{HashMap, HashSet}; - -use sqlx::{ - migrate::{Migrate, MigrateError, Migrator}, - query, query_scalar, Connection, Row, SqliteConnection, -}; -use tracing::{error, info, warn}; - -use super::types::PendingBackfillInfo; -use crate::driver::sqlite::SqlitePool; - -static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/sqlite"); - -const SCHEMA_BACKFILLS_TABLE_EXISTS_SQL: &str = - "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'schema_backfills'"; -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, -) -> Result, MigrateError> { - 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> { - if !schema_backfills_table_exists(conn).await? { - return Ok(pending_backfills_from_applied(&[])); - } - let applied_backfills = list_applied_backfills(conn).await?; - validate_applied_backfills(&applied_backfills)?; - Ok(pending_backfills_from_applied(&applied_backfills)) -} - -async fn schema_backfills_table_exists(conn: &mut SqliteConnection) -> Result { - let total: i64 = query_scalar(SCHEMA_BACKFILLS_TABLE_EXISTS_SQL) - .fetch_one(&mut *conn) - .await?; - Ok(total > 0) -} - -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 02c0d311a..dbecf72dd 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill/tests.rs @@ -4,14 +4,10 @@ use std::{ time::{Duration, Instant}, }; -use sqlx::{query, query_as, query_scalar, Connection, PgConnection, PgPool}; +use sqlx::{query, 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, run_sqlite_migrations}; +use super::{pending_backfills, pending_backfills_from_applied, run_backfills, AppliedBackfill}; +use crate::lifecycle::migrate::prepare_database_for_startup; const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION: i64 = 20260517012000; const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL: &str = @@ -88,586 +84,6 @@ fn corrected_legacy_backfill_is_not_requeued_after_application() { assert!(!pending_versions.contains(&LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION)); } -#[tokio::test] -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_versions, - vec![ - 20260422120000, - 20260505120000, - 20260517012000, - 20260716010000 - ] - ); - - run_mysql_backfills(&pool) - .await - .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 pending_sqlite_backfills_does_not_create_tracking_table() { - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("sqlite backfill status pool should connect"); - run_sqlite_migrations(&pool) - .await - .expect("sqlite schema should migrate"); - - let pending = pending_sqlite_backfills(&pool) - .await - .expect("sqlite pending backfills should load"); - assert!(!pending.is_empty()); - - let tracking_tables: i64 = query_scalar( - "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'schema_backfills'", - ) - .fetch_one(&pool) - .await - .expect("sqlite tracking table state should load"); - assert_eq!(tracking_tables, 0); -} - -#[tokio::test] -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 -) -"#, - ) - .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_versions, - vec![ - 20260422120000, - 20260505120000, - 20260517012000, - 20260716010000 - ] - ); - - run_sqlite_backfills(&pool) - .await - .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)] struct ManagedPostgresServer { child: Option, diff --git a/crates/aether-data/runtime/src/lifecycle/export.rs b/crates/aether-data/runtime/src/lifecycle/export.rs index a97bc1b08..e9523e36b 100644 --- a/crates/aether-data/runtime/src/lifecycle/export.rs +++ b/crates/aether-data/runtime/src/lifecycle/export.rs @@ -1,14 +1,8 @@ use std::collections::{BTreeMap, BTreeSet}; -#[cfg(all(feature = "postgres", feature = "sqlite"))] -use futures_util::TryStreamExt; use serde_json::Value; use sha2::{Digest, Sha256}; -#[cfg(all(feature = "postgres", feature = "sqlite"))] -use sqlx::Acquire; use sqlx::Row; -#[cfg(any(feature = "mysql", feature = "sqlite"))] -use sqlx::{Column, TypeInfo, ValueRef}; use aether_data_contracts::repository::candidates::{ sanitize_request_candidate_error_type, sanitize_request_candidate_extra_data, @@ -18,35 +12,18 @@ use aether_data_contracts::repository::candidates::{ use crate::error::SqlResultExt; use crate::{DataLayerError, DatabaseDriver, SqlDatabaseConfig}; -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] mod tests; -#[cfg(feature = "mysql")] -pub use mysql::{ - export_mysql_core_jsonl, export_mysql_jsonl, import_mysql_jsonl, import_mysql_plan, -}; #[cfg(feature = "postgres")] pub use postgres::{ export_postgres_core_jsonl, export_postgres_jsonl, import_postgres_jsonl, import_postgres_plan, }; -#[cfg(feature = "sqlite")] -pub use sqlite::{ - export_sqlite_core_jsonl, export_sqlite_jsonl, import_sqlite_jsonl, import_sqlite_plan, -}; -#[cfg(all(feature = "postgres", feature = "sqlite"))] -use postgres::{ - is_postgres_boolean_column, is_postgres_timestamp_column, load_postgres_import_columns, -}; - -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] use postgres::normalize_postgres_import_payload; pub const EXPORT_FORMAT_VERSION: u32 = 2; @@ -566,40 +543,6 @@ pub struct DataCopyOptions { pub omit_request_body_details: bool, } -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg(all(feature = "postgres", feature = "sqlite"))] -struct SqliteCopyColumn { - name: String, - declared_type: String, - not_null: bool, - has_default: bool, - primary_key_position: i64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[cfg(all(feature = "postgres", feature = "sqlite"))] -enum SqliteCopyAffinity { - Integer, - Real, - Text, - Blob, - Numeric, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg(all(feature = "postgres", feature = "sqlite"))] -struct SchemaCopyColumn { - sqlite: SqliteCopyColumn, - postgres: PostgresImportColumn, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg(all(feature = "postgres", feature = "sqlite"))] -struct SchemaCopyTable { - table_name: String, - columns: Vec, -} - #[derive(Debug, Clone, PartialEq, Eq)] #[cfg(feature = "postgres")] struct PostgresImportColumn { @@ -611,8 +554,6 @@ struct PostgresImportColumn { #[cfg(feature = "postgres")] type PostgresImportColumns = BTreeMap; -#[cfg(any(feature = "mysql", feature = "sqlite"))] -type ImportColumnNames = BTreeSet; const IMPORTED_CREDENTIAL_REVOKE_REASON: &str = "imported_credentials_revoked"; @@ -770,12 +711,7 @@ const USAGE_HTTP_BODY_DETAIL_COLUMNS: &[&str] = &[ "body_capture_mode", ]; -#[cfg(all(feature = "postgres", feature = "sqlite"))] -const REQUEST_BODY_DETAIL_TABLES: &[&str] = &["usage_body_blobs"]; -#[cfg(all(feature = "postgres", feature = "sqlite"))] -const LIFECYCLE_TABLES: &[&str] = &["_sqlx_migrations", "schema_backfills"]; - -#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn import_column_stores_timestamp(column_name: &str) -> bool { column_name.ends_with("_at") || column_name.ends_with("_unix_secs") @@ -787,7 +723,7 @@ fn import_column_stores_timestamp(column_name: &str) -> bool { ) } -#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool { if !column_name.ends_with("_unix_ms") { return false; @@ -803,7 +739,7 @@ fn import_timestamp_uses_millis(table_name: &str, column_name: &str) -> bool { !(relation_name == "usage" && column_name == "created_at_unix_ms") } -#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn normalize_imported_integer_timestamp( driver_name: &str, table_name: &str, @@ -839,7 +775,7 @@ fn normalize_imported_integer_timestamp( Ok(Some(timestamp)) } -#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn parse_imported_datetime(value: &str) -> Option> { let value = value.trim(); if let Ok(datetime) = chrono::DateTime::parse_from_rfc3339(value) { @@ -859,7 +795,7 @@ fn parse_imported_datetime(value: &str) -> Option> .map(|datetime| datetime.and_utc()) } -#[cfg(any(feature = "mysql", feature = "postgres", feature = "sqlite"))] +#[cfg(feature = "postgres")] fn normalize_imported_binary( driver_name: &str, column_name: &str, @@ -1099,7 +1035,7 @@ pub fn validate_export_records(records: &[DataExportRecord]) -> Result<(), DataL Ok(()) } -pub fn sqlite_core_export_domains() -> Vec { +pub fn postgres_core_export_domains() -> Vec { vec![ ExportDomain::Users, ExportDomain::ApiKeys, @@ -1122,38 +1058,12 @@ pub fn sqlite_core_export_domains() -> Vec { ] } -pub fn mysql_core_export_domains() -> Vec { - sqlite_core_export_domains() -} - -pub fn postgres_core_export_domains() -> Vec { - sqlite_core_export_domains() -} - pub async fn export_database_jsonl( database: SqlDatabaseConfig, domains: Vec, created_at_unix_secs: u64, ) -> Result { match database.driver { - #[cfg(feature = "sqlite")] - DatabaseDriver::Sqlite => { - let pool = crate::driver::sqlite::SqlitePoolFactory::new(database)?.connect_lazy()?; - if domains.is_empty() { - export_sqlite_core_jsonl(&pool, created_at_unix_secs).await - } else { - export_sqlite_jsonl(&pool, domains, created_at_unix_secs).await - } - } - #[cfg(feature = "mysql")] - DatabaseDriver::Mysql => { - let pool = crate::driver::mysql::MysqlPoolFactory::new(database)?.connect_lazy()?; - if domains.is_empty() { - export_mysql_core_jsonl(&pool, created_at_unix_secs).await - } else { - export_mysql_jsonl(&pool, domains, created_at_unix_secs).await - } - } #[cfg(feature = "postgres")] DatabaseDriver::Postgres => { let pool = @@ -1165,14 +1075,6 @@ pub async fn export_database_jsonl( export_postgres_jsonl(&pool, domains, created_at_unix_secs).await } } - #[cfg(not(feature = "sqlite"))] - DatabaseDriver::Sqlite => Err(DataLayerError::InvalidInput( - "SQLite driver is not enabled for aether-data".to_string(), - )), - #[cfg(not(feature = "mysql"))] - DatabaseDriver::Mysql => Err(DataLayerError::InvalidInput( - "MySQL driver is not enabled for aether-data".to_string(), - )), #[cfg(not(feature = "postgres"))] DatabaseDriver::Postgres => Err(DataLayerError::InvalidInput( "PostgreSQL driver is not enabled for aether-data".to_string(), @@ -1185,16 +1087,6 @@ pub async fn import_database_jsonl( input: &str, ) -> Result { match database.driver { - #[cfg(feature = "sqlite")] - DatabaseDriver::Sqlite => { - let pool = crate::driver::sqlite::SqlitePoolFactory::new(database)?.connect_lazy()?; - import_sqlite_jsonl(&pool, input).await - } - #[cfg(feature = "mysql")] - DatabaseDriver::Mysql => { - let pool = crate::driver::mysql::MysqlPoolFactory::new(database)?.connect_lazy()?; - import_mysql_jsonl(&pool, input).await - } #[cfg(feature = "postgres")] DatabaseDriver::Postgres => { let pool = @@ -1202,14 +1094,6 @@ pub async fn import_database_jsonl( .connect_lazy()?; import_postgres_jsonl(&pool, input).await } - #[cfg(not(feature = "sqlite"))] - DatabaseDriver::Sqlite => Err(DataLayerError::InvalidInput( - "SQLite driver is not enabled for aether-data".to_string(), - )), - #[cfg(not(feature = "mysql"))] - DatabaseDriver::Mysql => Err(DataLayerError::InvalidInput( - "MySQL driver is not enabled for aether-data".to_string(), - )), #[cfg(not(feature = "postgres"))] DatabaseDriver::Postgres => Err(DataLayerError::InvalidInput( "PostgreSQL driver is not enabled for aether-data".to_string(), @@ -1224,14 +1108,6 @@ pub async fn copy_database_records( created_at_unix_secs: u64, options: DataCopyOptions, ) -> Result { - #[cfg(all(feature = "postgres", feature = "sqlite"))] - if domains.is_empty() - && source.driver == DatabaseDriver::Postgres - && target.driver == DatabaseDriver::Sqlite - { - return copy_postgres_to_sqlite_from_target_schema(source, target, options).await; - } - let mut records = decode_jsonl(&export_database_jsonl(source, domains, created_at_unix_secs).await?)?; if options.omit_request_body_details { @@ -1275,651 +1151,11 @@ fn omit_request_body_details_from_records(records: &mut Vec) { }); } -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn copy_postgres_to_sqlite_from_target_schema( - source: SqlDatabaseConfig, - mut target: SqlDatabaseConfig, - options: DataCopyOptions, -) -> Result { - target.pool.min_connections = 1; - target.pool.max_connections = 1; - - let postgres_pool = - 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(&mut postgres_tx).await?; - let target_tables = load_sqlite_copy_table_names(&sqlite_pool).await?; - - ensure_no_nonempty_source_tables_outside_target_schema( - &mut postgres_tx, - &source_tables, - &target_tables, - options, - ) - .await?; - - 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) - || !source_tables.contains(&table_name) - || (options.omit_request_body_details && copy_table_is_request_body_detail(&table_name)) - { - continue; - } - - let table_plan = build_postgres_sqlite_copy_table_plan( - &mut postgres_tx, - &sqlite_pool, - &table_name, - options, - ) - .await?; - if table_plan.columns.is_empty() { - continue; - } - table_plans.push(table_plan); - } - - let mut connection = sqlite_pool.acquire().await.map_sql_err()?; - sqlx::raw_sql("PRAGMA foreign_keys = OFF") - .execute(&mut *connection) - .await - .map_sql_err()?; - 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_tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - source_tables: &BTreeSet, - target_tables: &BTreeSet, - options: DataCopyOptions, -) -> Result<(), DataLayerError> { - let mut missing = Vec::new(); - for table_name in source_tables { - if copy_table_is_lifecycle(table_name) - || (options.omit_request_body_details && copy_table_is_request_body_detail(table_name)) - || target_tables.contains(table_name) - { - continue; - } - if postgres_public_table_has_rows(postgres_tx, table_name).await? { - missing.push(table_name.clone()); - } - } - - if !missing.is_empty() { - return Err(DataLayerError::InvalidInput(format!( - "source Postgres has non-empty public tables that do not exist in the target SQLite schema: {}", - missing.join(", ") - ))); - } - Ok(()) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn build_postgres_sqlite_copy_table_plan( - 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(&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 { - if options.omit_request_body_details - && table_name == "usage" - && USAGE_REQUEST_BODY_DETAIL_COLUMNS.contains(&sqlite_column.name.as_str()) - { - 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 { - sqlite: sqlite_column, - postgres: postgres_column.clone(), - }); - continue; - } - - if source_has_rows && sqlite_copy_column_is_required(&sqlite_column) { - return Err(DataLayerError::InvalidInput(format!( - "target SQLite table '{table_name}' has required column '{}' that does not exist in source Postgres", - sqlite_column.name - ))); - } - } - - if source_has_rows && columns.is_empty() { - return Err(DataLayerError::InvalidInput(format!( - "source Postgres table '{table_name}' has rows, but none of its columns exist in target SQLite" - ))); - } - - Ok(SchemaCopyTable { - table_name: table_name.to_string(), - columns, - }) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn copy_postgres_sqlite_table( - 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(&mut **postgres_tx); - let mut imported = 0usize; - - while let Some(row) = rows.try_next().await.map_sql_err()? { - let mut payload = row.try_get::("payload").map_sql_err()?; - let object = payload.as_object_mut().ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "postgres copy row for table '{}' did not produce a JSON object", - table.table_name - )) - })?; - prepare_postgres_sqlite_copy_payload(table, object); - let mut query = sqlx::query(&target_sql); - for column in &table.columns { - let value = object.get(&column.sqlite.name).ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "postgres copy row for table '{}' is missing column '{}'", - table.table_name, column.sqlite.name - )) - })?; - query = bind_sqlite_copy_value(query, value, &column.sqlite)?; - } - query.execute(&mut **sqlite_tx).await.map_sql_err()?; - imported = imported.saturating_add(1); - } - - Ok(imported) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn prepare_postgres_sqlite_copy_payload( - table: &SchemaCopyTable, - object: &mut serde_json::Map, -) { - deactivate_imported_credentials(&table.table_name, object, |column_name| { - table - .columns - .iter() - .any(|column| column.sqlite.name == column_name) - }); - sanitize_request_candidate_auxiliary_payload(&table.table_name, object); - sanitize_payment_security_payload(&table.table_name, object); -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn postgres_schema_copy_select_sql(table: &SchemaCopyTable) -> Result { - let table_sql = format!( - "public.{}", - postgres_quote_identifier(table.table_name.as_str())? - ); - let mut payload_parts = Vec::new(); - for column in &table.columns { - if let Some(expr) = postgres_schema_copy_override_expr(&table.table_name, column)? { - payload_parts.push(sql_string_literal(&column.sqlite.name)); - payload_parts.push(expr); - } - } - let payload_sql = if payload_parts.is_empty() { - "to_jsonb(t)".to_string() - } else { - format!( - "to_jsonb(t) || jsonb_build_object({})", - payload_parts.join(", ") - ) - }; - - let order_by = table - .columns - .iter() - .filter(|column| column.sqlite.primary_key_position > 0) - .map(|column| { - postgres_quote_identifier(&column.sqlite.name).map(|quoted| format!("t.{quoted} ASC")) - }) - .collect::, _>>()?; - let order_sql = if order_by.is_empty() { - String::new() - } else { - format!(" ORDER BY {}", order_by.join(", ")) - }; - - Ok(format!( - "SELECT {payload_sql} AS payload FROM {table_sql} AS t{order_sql}" - )) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn postgres_schema_copy_override_expr( - table_name: &str, - column: &SchemaCopyColumn, -) -> Result, DataLayerError> { - let column_sql = format!("t.{}", postgres_quote_identifier(&column.sqlite.name)?); - let affinity = sqlite_copy_affinity(&column.sqlite); - - if affinity == SqliteCopyAffinity::Blob && is_postgres_bytea_column(&column.postgres) { - return Ok(Some(format!( - "CASE WHEN {column_sql} IS NULL THEN NULL ELSE encode({column_sql}, 'hex') END" - ))); - } - - if affinity == SqliteCopyAffinity::Integer && is_postgres_boolean_column(&column.postgres) { - return Ok(Some(format!( - "CASE WHEN {column_sql} IS NULL THEN NULL WHEN {column_sql} THEN 1 ELSE 0 END" - ))); - } - - if affinity == SqliteCopyAffinity::Integer - && (is_postgres_timestamp_column(&column.postgres) - || is_postgres_date_column(&column.postgres)) - { - let timestamp_sql = if is_postgres_date_column(&column.postgres) { - format!("{column_sql}::timestamp") - } else { - column_sql.clone() - }; - let multiplier = if import_timestamp_uses_millis(table_name, &column.sqlite.name) { - " * 1000" - } else { - "" - }; - return Ok(Some(format!( - "CASE WHEN {column_sql} IS NULL THEN NULL ELSE FLOOR(EXTRACT(EPOCH FROM {timestamp_sql}){multiplier})::bigint END" - ))); - } - - Ok(None) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn sqlite_schema_copy_insert_sql(table: &SchemaCopyTable) -> Result { - let table_sql = sqlite_quote_identifier(&table.table_name)?; - let column_sql = table - .columns - .iter() - .map(|column| sqlite_quote_identifier(&column.sqlite.name)) - .collect::, _>>()? - .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 INTO {table_sql} ({column_sql}) VALUES ({placeholder_sql}) {conflict_sql}" - )) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn load_postgres_public_table_names( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, -) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT table_name -FROM information_schema.tables -WHERE table_schema = 'public' - AND table_type = 'BASE TABLE' -ORDER BY table_name -"#, - ) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - - let mut tables = BTreeSet::new(); - for row in rows { - tables.insert(row.try_get::("table_name").map_sql_err()?); - } - Ok(tables) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn load_sqlite_copy_table_names( - pool: &crate::driver::sqlite::SqlitePool, -) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT name -FROM sqlite_schema -WHERE type = 'table' - AND name NOT LIKE 'sqlite_%' -ORDER BY name -"#, - ) - .fetch_all(pool) - .await - .map_sql_err()?; - - let mut tables = BTreeSet::new(); - for row in rows { - let table_name = row.try_get::("name").map_sql_err()?; - if !copy_table_is_lifecycle(&table_name) && !copy_table_is_sqlite_internal(&table_name) { - tables.insert(table_name); - } - } - Ok(tables) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn load_sqlite_copy_columns( - pool: &crate::driver::sqlite::SqlitePool, - table_name: &str, -) -> Result, DataLayerError> { - let table_sql = sqlite_quote_identifier(table_name)?; - let rows = sqlx::query(&format!("PRAGMA table_info({table_sql})")) - .fetch_all(pool) - .await - .map_sql_err()?; - - let mut columns = Vec::new(); - for row in rows { - columns.push(SqliteCopyColumn { - name: row.try_get::("name").map_sql_err()?, - declared_type: row - .try_get::, _>("type") - .map_sql_err()? - .unwrap_or_default(), - not_null: row.try_get::("notnull").map_sql_err()? != 0, - has_default: row - .try_get::, _>("dflt_value") - .map_sql_err()? - .is_some(), - primary_key_position: row.try_get::("pk").map_sql_err()?, - }); - } - - if columns.is_empty() { - return Err(DataLayerError::UnexpectedValue(format!( - "target SQLite table '{table_name}' has no visible columns" - ))); - } - Ok(columns) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn postgres_public_table_has_rows( - 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(&mut **tx) - .await - .map_sql_err() -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -async fn ensure_sqlite_foreign_key_check_passes( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, -) -> Result<(), DataLayerError> { - let rows = sqlx::query("PRAGMA foreign_key_check") - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - if rows.is_empty() { - return Ok(()); - } - - let mut violations = Vec::new(); - for row in rows.iter().take(10) { - let table = row - .try_get::, _>("table") - .map_sql_err()? - .unwrap_or_else(|| "".to_string()); - let rowid = row.try_get::, _>("rowid").map_sql_err()?; - let parent = row - .try_get::, _>("parent") - .map_sql_err()? - .unwrap_or_else(|| "".to_string()); - violations.push(format!("{table} rowid={rowid:?} parent={parent}")); - } - Err(DataLayerError::InvalidInput(format!( - "target SQLite foreign key check failed after copy: {}", - violations.join("; ") - ))) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn copy_table_is_lifecycle(table_name: &str) -> bool { - LIFECYCLE_TABLES.contains(&table_name) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn copy_table_is_sqlite_internal(table_name: &str) -> bool { - table_name.starts_with("sqlite_") -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn copy_table_is_request_body_detail(table_name: &str) -> bool { - REQUEST_BODY_DETAIL_TABLES.contains(&table_name) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -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_affinity(column: &SqliteCopyColumn) -> SqliteCopyAffinity { - let declared_type = column.declared_type.to_ascii_uppercase(); - if declared_type.contains("INT") { - SqliteCopyAffinity::Integer - } else if declared_type.contains("CHAR") - || declared_type.contains("CLOB") - || declared_type.contains("TEXT") - { - SqliteCopyAffinity::Text - } else if declared_type.contains("BLOB") || declared_type.trim().is_empty() { - SqliteCopyAffinity::Blob - } else if declared_type.contains("REAL") - || declared_type.contains("FLOA") - || declared_type.contains("DOUB") - { - SqliteCopyAffinity::Real - } else { - SqliteCopyAffinity::Numeric - } -} - #[cfg(feature = "postgres")] fn is_postgres_bytea_column(column: &PostgresImportColumn) -> bool { column.data_type == "bytea" || column.udt_name == "bytea" } -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn is_postgres_date_column(column: &PostgresImportColumn) -> bool { - column.data_type == "date" || column.udt_name == "date" -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn bind_sqlite_copy_value<'q>( - query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>, - value: &'q Value, - column: &SqliteCopyColumn, -) -> Result>, DataLayerError> -{ - Ok(match sqlite_copy_affinity(column) { - SqliteCopyAffinity::Integer => match value { - Value::Null => query.bind(Option::::None), - Value::Bool(value) => query.bind(i64::from(*value)), - Value::Number(number) => { - let value = number - .as_i64() - .or_else(|| number.as_u64().and_then(|value| i64::try_from(value).ok())) - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected integer, got {number}", - column.name - )) - })?; - query.bind(value) - } - Value::String(value) => query.bind(value.parse::().map_err(|err| { - DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected integer string: {err}", - column.name - )) - })?), - Value::Array(_) | Value::Object(_) => { - return Err(DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected integer-compatible value", - column.name - ))); - } - }, - SqliteCopyAffinity::Real => match value { - Value::Null => query.bind(Option::::None), - Value::Number(number) => query.bind(number.as_f64().ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected finite real value", - column.name - )) - })?), - Value::String(value) => query.bind(value.parse::().map_err(|err| { - DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected real string: {err}", - column.name - )) - })?), - Value::Bool(value) => query.bind(if *value { 1.0 } else { 0.0 }), - Value::Array(_) | Value::Object(_) => { - return Err(DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected real-compatible value", - column.name - ))); - } - }, - SqliteCopyAffinity::Blob => match value { - Value::Null => query.bind(Option::>::None), - Value::String(value) => query.bind(hex_decode(value, &column.name)?), - Value::Array(values) => { - let mut bytes = Vec::with_capacity(values.len()); - for value in values { - let Some(byte) = value.as_u64().and_then(|value| u8::try_from(value).ok()) - else { - return Err(DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' contains non-byte array value", - column.name - ))); - }; - bytes.push(byte); - } - query.bind(bytes) - } - Value::Bool(_) | Value::Number(_) | Value::Object(_) => { - return Err(DataLayerError::InvalidInput(format!( - "sqlite copy column '{}' expected blob-compatible value", - column.name - ))); - } - }, - SqliteCopyAffinity::Text | SqliteCopyAffinity::Numeric => { - bind_sqlite_json_value(query, value)? - } - }) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn sql_string_literal(value: &str) -> String { - format!("'{}'", value.replace('\'', "''")) -} - -#[cfg(all(feature = "postgres", feature = "sqlite"))] -fn hex_decode(value: &str, column_name: &str) -> Result, DataLayerError> { - let value = value.trim(); - if !value.len().is_multiple_of(2) { - return Err(DataLayerError::InvalidInput(format!( - "sqlite copy column '{column_name}' has odd-length hex data" - ))); - } - - let mut bytes = Vec::with_capacity(value.len() / 2); - for index in (0..value.len()).step_by(2) { - let byte = u8::from_str_radix(&value[index..index + 2], 16).map_err(|err| { - DataLayerError::InvalidInput(format!( - "sqlite copy column '{column_name}' has invalid hex data at byte {}: {err}", - index / 2 - )) - })?; - bytes.push(byte); - } - Ok(bytes) -} - fn export_order_by(domain: ExportDomain, id_column: &str) -> String { if domain == ExportDomain::UserGroupMembers { "group_id ASC, user_id ASC".to_string() @@ -1928,60 +1164,6 @@ fn export_order_by(domain: ExportDomain, id_column: &str) -> String { } } -#[cfg(feature = "sqlite")] -fn sqlite_quote_identifier(identifier: &str) -> Result { - if identifier.trim().is_empty() { - return Err(DataLayerError::InvalidInput( - "sqlite import column name cannot be empty".to_string(), - )); - } - if !identifier - .chars() - .all(|ch| ch.is_ascii_alphanumeric() || ch == '_') - { - return Err(DataLayerError::InvalidInput(format!( - "sqlite import column name '{identifier}' contains unsupported characters" - ))); - } - Ok(format!(r#""{identifier}""#)) -} - -#[cfg(feature = "sqlite")] -fn bind_sqlite_json_value<'q>( - query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>, - value: &'q Value, -) -> Result>, DataLayerError> -{ - Ok(match value { - Value::Null => query.bind(Option::::None), - Value::Bool(value) => query.bind(i64::from(*value)), - Value::Number(value) => { - if let Some(value) = value.as_i64() { - query.bind(value) - } else if let Some(value) = value.as_u64() { - let value = i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!( - "sqlite import integer value {value} exceeds i64" - )) - })?; - query.bind(value) - } else if let Some(value) = value.as_f64() { - query.bind(value) - } else { - return Err(DataLayerError::InvalidInput( - "sqlite import number is not representable".to_string(), - )); - } - } - Value::String(value) => query.bind(value), - Value::Array(_) | Value::Object(_) => { - let value = serde_json::to_string(value) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; - query.bind(value) - } - }) -} - #[cfg(feature = "postgres")] pub(super) fn postgres_quote_identifier(identifier: &str) -> Result { if identifier.trim().is_empty() { @@ -2000,57 +1182,6 @@ pub(super) fn postgres_quote_identifier(identifier: &str) -> Result Result, DataLayerError> { - let object = row.payload.as_object().ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "{} export row '{}' payload must be a JSON object", - domain.as_str(), - row.id - )) - })?; - if object.is_empty() { - return Err(DataLayerError::InvalidInput(format!( - "{} export row '{}' payload cannot be empty", - domain.as_str(), - row.id - ))); - } - - let mut filtered = serde_json::Map::new(); - 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() { - return Err(DataLayerError::InvalidInput(format!( - "{} export row '{}' has no columns supported by {driver_name} table '{table_name}'", - domain.as_str(), - row.id - ))); - } - - Ok(filtered) -} - 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()) @@ -2250,17 +1381,8 @@ mod payment_export_security_tests { mod request_candidate_export_security_tests { use serde_json::json; - #[cfg(all(feature = "postgres", feature = "sqlite"))] - use serde_json::Value; - use super::{domain_payload_table, payload_with_table, ExportRow}; - #[cfg(all(feature = "postgres", feature = "sqlite"))] - use super::{ - prepare_postgres_sqlite_copy_payload, PostgresImportColumn, SchemaCopyColumn, - SchemaCopyTable, SqliteCopyColumn, - }; - #[test] fn request_candidate_auxiliary_export_and_import_drop_sensitive_diagnostics() { let raw = json!({ @@ -2312,73 +1434,4 @@ mod request_candidate_export_security_tests { assert_eq!(imported["required_capabilities"], json!({"vision": true})); assert!(!imported.to_string().contains("import-secret")); } - - #[cfg(all(feature = "postgres", feature = "sqlite"))] - #[test] - fn postgres_to_sqlite_fast_copy_sanitizes_request_candidate_diagnostics() { - let table = SchemaCopyTable { - table_name: "request_candidates".to_string(), - columns: [ - "error_message", - "skip_reason", - "error_type", - "extra_data", - "required_capabilities", - ] - .into_iter() - .map(|name| SchemaCopyColumn { - sqlite: SqliteCopyColumn { - name: name.to_string(), - declared_type: "TEXT".to_string(), - not_null: false, - has_default: false, - primary_key_position: 0, - }, - postgres: PostgresImportColumn { - data_type: "text".to_string(), - udt_name: "text".to_string(), - is_nullable: true, - has_default: false, - }, - }) - .collect(), - }; - let mut payload = json!({ - "error_message": "Bearer fast-copy-secret", - "skip_reason": "fast-copy-secret", - "error_type": "fast-copy-secret", - "extra_data": { - "upstream_url": "https://user:pass@example.com/private?token=fast-copy-secret", - "image_progress": { - "phase": "upstream_streaming", - "message": "fast-copy-secret" - } - }, - "required_capabilities": { - "vision": 1, - "tenant_secret": "fast-copy-secret" - } - }) - .as_object() - .cloned() - .expect("copy payload should be an object"); - - prepare_postgres_sqlite_copy_payload(&table, &mut payload); - - assert!(payload["error_message"].is_null()); - assert_eq!(payload["skip_reason"], "unclassified_skip"); - assert_eq!(payload["error_type"], "unclassified_error"); - assert_eq!( - payload["extra_data"]["upstream_url"], - "https://example.com/" - ); - assert_eq!( - payload["extra_data"]["image_progress"], - json!({"phase": "upstream_streaming"}) - ); - assert_eq!(payload["required_capabilities"], json!({"vision": true})); - assert!(!Value::Object(payload) - .to_string() - .contains("fast-copy-secret")); - } } diff --git a/crates/aether-data/runtime/src/lifecycle/export/mysql.rs b/crates/aether-data/runtime/src/lifecycle/export/mysql.rs deleted file mode 100644 index b241a0464..000000000 --- a/crates/aether-data/runtime/src/lifecycle/export/mysql.rs +++ /dev/null @@ -1,914 +0,0 @@ -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, - created_at_unix_secs: u64, -) -> Result { - export_mysql_jsonl(pool, mysql_core_export_domains(), created_at_unix_secs).await -} - -pub async fn export_mysql_jsonl( - pool: &crate::driver::mysql::MysqlPool, - 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), - domains.clone(), - ); - 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(&mut tx, &mut records).await?; - continue; - } - if domain == ExportDomain::Wallets { - 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(&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) -} - -pub async fn import_mysql_jsonl( - pool: &crate::driver::mysql::MysqlPool, - input: &str, -) -> Result { - let plan = build_import_plan(input)?; - import_mysql_plan(pool, &plan).await -} - -pub async fn import_mysql_plan( - pool: &crate::driver::mysql::MysqlPool, - plan: &DataImportPlan, -) -> Result { - let identity_scope = IdentityImportScope::from_plan(plan)?; - let mut tx = pool.begin().await.map_sql_err()?; - let identity_state = capture_mysql_identity_import_state(&mut tx, &identity_scope).await?; - 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_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(&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(&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(&mut tx, &mut column_cache, table_name).await?; - for row in plan.rows(*domain) { - import_mysql_row(&mut tx, table_name, *domain, row, &target_columns).await?; - imported = imported.saturating_add(1); - } - } - enforce_mysql_identity_import_invariants(&mut tx, &identity_scope, identity_state).await?; - tx.commit().await.map_sql_err()?; - Ok(imported) -} - -async fn capture_mysql_identity_import_state( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - scope: &IdentityImportScope, -) -> Result { - let mut affected_user_ids = if scope.finalizes_oauth_links { - scope.user_ids.iter().cloned().collect::>() - } else { - BTreeSet::new() - }; - for link_id in &scope.oauth_link_ids { - if let Some(user_id) = - sqlx::query_scalar::<_, String>("SELECT user_id FROM user_oauth_links WHERE id = ?") - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - affected_user_ids.insert(user_id); - } - } - for provider_type in &scope.oauth_provider_types { - let user_ids = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ?", - ) - .bind(provider_type) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - affected_user_ids.extend(user_ids); - } - Ok(IdentityImportState { affected_user_ids }) -} - -async fn enforce_mysql_identity_import_invariants( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - scope: &IdentityImportScope, - mut state: IdentityImportState, -) -> Result<(), DataLayerError> { - for user_id in &scope.user_ids { - let auth_source = - sqlx::query_scalar::<_, String>("SELECT auth_source FROM users WHERE id = ? LIMIT 1") - .bind(user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "imported users row '{user_id}' did not produce a user record" - )) - })?; - if !matches!(auth_source.as_str(), "local" | "ldap" | "oauth") { - return Err(DataLayerError::InvalidInput(format!( - "imported user '{user_id}' has unsupported auth_source '{auth_source}'" - ))); - } - if auth_source == "oauth" { - sqlx::query("UPDATE users SET email_verified = 0 WHERE id = ?") - .bind(user_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - } - - for link_id in &scope.oauth_link_ids { - let user_id = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE id = ? LIMIT 1", - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "imported OAuth link row '{link_id}' did not produce a link record" - )) - })?; - state.affected_user_ids.insert(user_id); - } - - for link_id in &scope.oauth_link_ids { - if let Some((provider_type, provider_user_id)) = sqlx::query_as::<_, (String, String)>( - r#" -SELECT imported.provider_type, imported.provider_user_id -FROM user_oauth_links imported -JOIN user_oauth_links duplicate - ON duplicate.provider_type = imported.provider_type - AND duplicate.provider_user_id = imported.provider_user_id - AND duplicate.id <> imported.id -WHERE imported.id = ? -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import assigns provider identity '{provider_type}:{provider_user_id}' more than once" - ))); - } - } - - for link_id in &scope.oauth_link_ids { - if let Some((user_id, provider_type)) = sqlx::query_as::<_, (String, String)>( - r#" -SELECT imported.user_id, imported.provider_type -FROM user_oauth_links imported -JOIN user_oauth_links duplicate - ON duplicate.user_id = imported.user_id - AND duplicate.provider_type = imported.provider_type - AND duplicate.id <> imported.id -WHERE imported.id = ? -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import links user '{user_id}' to provider '{provider_type}' more than once" - ))); - } - } - - for link_id in &scope.oauth_link_ids { - if let Some(invalid_id) = sqlx::query_scalar::<_, String>( - r#" -SELECT links.id -FROM user_oauth_links links -LEFT JOIN users ON users.id = links.user_id -LEFT JOIN oauth_providers providers ON providers.provider_type = links.provider_type -WHERE links.id = ? - AND ( - users.id IS NULL - OR providers.provider_type IS NULL - OR BINARY links.provider_type <> BINARY LOWER(TRIM(links.provider_type)) - OR links.provider_type = '' - OR BINARY links.provider_user_id <> BINARY TRIM(links.provider_user_id) - OR links.provider_user_id = '' - OR BINARY providers.provider_type <> BINARY LOWER(TRIM(providers.provider_type)) - ) -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import produced invalid or orphaned link '{invalid_id}'" - ))); - } - } - - if !scope.validates_oauth_login_methods { - return Ok(()); - } - for user_id in state.affected_user_ids { - if sqlx::query_scalar::<_, String>( - r#" -SELECT users.id -FROM users -WHERE users.id = ? - AND users.auth_source = 'oauth' - AND users.is_active = 1 - AND users.is_deleted = 0 - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links links - JOIN oauth_providers providers ON providers.provider_type = links.provider_type - WHERE links.user_id = users.id - AND providers.is_enabled = 1 - AND BINARY links.provider_type = BINARY LOWER(TRIM(links.provider_type)) - AND BINARY links.provider_user_id = BINARY TRIM(links.provider_user_id) - AND links.provider_user_id <> '' - ) -LIMIT 1 -"#, - ) - .bind(&user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .is_some() - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import would leave active user '{user_id}' without an enabled identity binding" - ))); - } - } - - Ok(()) -} - -fn mysql_domain_table( - domain: ExportDomain, -) -> Result<(&'static str, &'static str), DataLayerError> { - match domain { - ExportDomain::Users => Ok(("users", "id")), - ExportDomain::ApiKeys => Ok(("api_keys", "id")), - ExportDomain::Providers => Ok(("providers", "id")), - ExportDomain::ProviderKeys => Ok(("provider_api_keys", "id")), - ExportDomain::Endpoints => Ok(("provider_endpoints", "id")), - ExportDomain::Models => Ok(("models", "id")), - ExportDomain::GlobalModels => Ok(("global_models", "id")), - ExportDomain::AuthModules => Ok(("auth_modules", "id")), - ExportDomain::OAuthProviders => Ok(("oauth_providers", "provider_type")), - ExportDomain::UserOAuthLinks => Ok(("user_oauth_links", "id")), - ExportDomain::UserGroups => Ok(("user_groups", "id")), - ExportDomain::UserGroupMembers => Ok(("user_group_members", "group_id")), - ExportDomain::ProxyNodes => Ok(("proxy_nodes", "id")), - ExportDomain::SystemConfigs => Ok(("system_configs", "id")), - ExportDomain::Wallets => Err(DataLayerError::InvalidInput( - "mysql wallet export uses multiple tables and must be handled as a domain".to_string(), - )), - ExportDomain::Usage => Ok(("`usage`", "request_id")), - 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, - id_column: &str, -) -> Result { - if domain == ExportDomain::UserGroupMembers { - let group_id = mysql_required_export_text(row, "group_id", domain)?; - let user_id = mysql_required_export_text(row, "user_id", domain)?; - return Ok(format!("{group_id}:{user_id}")); - } - mysql_required_export_text(row, id_column, domain) -} - -fn mysql_required_export_text( - row: &sqlx::mysql::MySqlRow, - column: &str, - domain: ExportDomain, -) -> Result { - row.try_get::, _>(column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{} export row has null id column '{}'", - domain.as_str(), - column - )) - }) -} - -async fn export_mysql_billing_records( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - records: &mut Vec, -) -> Result<(), DataLayerError> { - for (table_name, id_column) in [ - ("billing_rules", "id"), - ("dimension_collectors", "id"), - ("usage_settlement_snapshots", "request_id"), - ] { - let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; - for row in rows { - let id = row - .try_get::, _>(id_column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "billing export row in table '{table_name}' has null id" - )) - })?; - records.push(DataExportRecord::row( - ExportDomain::Billing, - format!("{table_name}:{id}"), - payload_with_table(mysql_row_payload(&row)?, table_name)?, - )); - } - } - Ok(()) -} - -async fn export_mysql_wallet_records( - 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(&mut **tx).await.map_sql_err()?; - for row in rows { - let id = row - .try_get::, _>(id_column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "wallet export row in table '{table_name}' has null id" - )) - })?; - records.push(DataExportRecord::row( - ExportDomain::Wallets, - format!("{table_name}:{id}"), - payload_with_table(mysql_row_payload(&row)?, table_name)?, - )); - } - } - Ok(()) -} - -async fn import_mysql_row( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table_name: &str, - domain: ExportDomain, - row: &ExportRow, - target_columns: &MysqlImportColumns, -) -> Result<(), DataLayerError> { - let mut object = - filter_import_payload("mysql", table_name, domain, row, &target_columns.names)?; - deactivate_imported_credentials(table_name, &mut object, |column_name| { - target_columns.names.contains(column_name) - }); - - 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 sql = format!("INSERT INTO {table_name} ({column_sql}) VALUES ({placeholder_sql})"); - let mut query = sqlx::query(&sql); - for column in columns { - query = bind_mysql_import_column(query, &object, target_columns, table_name, column)?; - } - 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( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - row: &ExportRow, - 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(tx, column_cache, table_name).await?; - import_mysql_row( - tx, - table_name, - ExportDomain::Billing, - &ExportRow { - id: row.id.clone(), - payload, - }, - &target_columns, - ) - .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"), - "dimension_collectors" => Ok("dimension_collectors"), - "usage_settlement_snapshots" => Ok("usage_settlement_snapshots"), - other => Err(DataLayerError::InvalidInput(format!( - "unsupported mysql billing export table '{other}'" - ))), - } -} - -async fn import_mysql_wallet_row( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - row: &ExportRow, - 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(tx, column_cache, table_name).await?; - import_mysql_row( - tx, - table_name, - ExportDomain::Wallets, - &ExportRow { - id: row.id.clone(), - payload, - }, - &target_columns, - ) - .await -} - -fn mysql_wallet_tables() -> &'static [(&'static str, &'static str)] { - &[ - ("wallets", "id"), - ("wallet_transactions", "id"), - ("wallet_daily_usage_ledgers", "id"), - ("payment_orders", "id"), - ("payment_callbacks", "id"), - ("refund_requests", "id"), - ("redeem_code_batches", "id"), - ("redeem_codes", "id"), - ] -} - -fn mysql_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { - mysql_wallet_tables() - .iter() - .find(|(candidate, _)| *candidate == table_name) - .map(|(table, _)| *table) - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "unsupported mysql wallet export table '{table_name}'" - )) - }) -} - -async fn mysql_import_columns_cached( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - cache: &mut BTreeMap, - table_name: &str, -) -> Result { - if let Some(columns) = cache.get(table_name) { - return Ok(columns.clone()); - } - - 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( - tx: &mut sqlx::Transaction<'_, sqlx::MySql>, - table_name: &str, -) -> Result { - let relation_name = table_name.trim_matches('`'); - let rows = sqlx::query( - r#" -SELECT - CAST(COLUMN_NAME AS CHAR) AS column_name, - CAST(DATA_TYPE AS CHAR) AS data_type, - CAST(COLUMN_KEY AS CHAR) AS column_key, - ORDINAL_POSITION AS ordinal_position -FROM information_schema.columns -WHERE table_schema = DATABASE() - AND table_name = ? -"#, - ) - .bind(relation_name) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - - let mut columns = MysqlImportColumns::default(); - let mut primary_key = BTreeMap::new(); - for row in rows { - 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.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( - "mysql import column name cannot be empty".to_string(), - )); - } - if !identifier - .chars() - .all(|ch| ch.is_ascii_alphanumeric() || ch == '_') - { - return Err(DataLayerError::InvalidInput(format!( - "mysql import column name '{identifier}' contains unsupported characters" - ))); - } - Ok(format!("`{identifier}`")) -} - -fn bind_mysql_json_value<'q>( - query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, - value: &'q Value, -) -> Result, DataLayerError> { - Ok(match value { - Value::Null => query.bind(Option::::None), - Value::Bool(value) => query.bind(i64::from(*value)), - Value::Number(value) => { - if let Some(value) = value.as_i64() { - query.bind(value) - } else if let Some(value) = value.as_u64() { - let value = i64::try_from(value).map_err(|_| { - DataLayerError::InvalidInput(format!( - "mysql import integer value {value} exceeds i64" - )) - })?; - query.bind(value) - } else if let Some(value) = value.as_f64() { - query.bind(value) - } else { - return Err(DataLayerError::InvalidInput( - "mysql import number is not representable".to_string(), - )); - } - } - Value::String(value) => query.bind(value), - Value::Array(_) | Value::Object(_) => { - let value = serde_json::to_string(value) - .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?; - query.bind(value) - } - }) -} - -fn mysql_row_payload(row: &sqlx::mysql::MySqlRow) -> Result { - let mut object = serde_json::Map::new(); - for (index, column) in row.columns().iter().enumerate() { - object.insert(column.name().to_string(), mysql_value_to_json(row, index)?); - } - Ok(Value::Object(object)) -} - -fn mysql_value_to_json(row: &sqlx::mysql::MySqlRow, index: usize) -> Result { - let raw = row.try_get_raw(index).map_sql_err()?; - if raw.is_null() { - return Ok(Value::Null); - } - - match raw.type_info().name().to_ascii_uppercase().as_str() { - "BOOL" | "BOOLEAN" => Ok(Value::Bool(row.try_get::(index).map_sql_err()?)), - "TINYINT" | "TINY" | "SMALLINT" | "SHORT" | "MEDIUMINT" | "INT24" | "INT" | "INTEGER" - | "LONG" | "BIGINT" | "LONGLONG" | "YEAR" => { - Ok(Value::from(row.try_get::(index).map_sql_err()?)) - } - "FLOAT" | "DOUBLE" => { - let value = row.try_get::(index).map_sql_err()?; - serde_json::Number::from_f64(value) - .map(Value::Number) - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "mysql export column {} contains non-finite float", - index - )) - }) - } - "DECIMAL" | "NEWDECIMAL" => Ok(Value::String( - row.try_get::(index) - .map_sql_err()? - .to_string(), - )), - "VARCHAR" | "VAR_STRING" | "STRING" | "TEXT" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" - | "JSON" | "ENUM" | "SET" | "DATE" | "DATETIME" | "TIMESTAMP" | "TIME" => Ok( - Value::String(row.try_get::(index).map_sql_err()?), - ), - "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BIT" | "GEOMETRY" => { - let bytes = row.try_get::, _>(index).map_sql_err()?; - Ok(Value::Array(bytes.into_iter().map(Value::from).collect())) - } - other => Err(DataLayerError::UnexpectedValue(format!( - "unsupported mysql export column type '{other}' at index {index}" - ))), - } -} - -#[cfg(test)] -mod tests { - use super::normalize_mysql_decimal_value; - use serde_json::json; - - #[test] - fn decimal_import_binds_numbers_and_strings_as_decimal_text() { - let value = json!(12345.12345678); - assert_eq!( - normalize_mysql_decimal_value("billing_total_cost_usd", &value) - .expect("decimal value should normalize") - .as_deref(), - Some("12345.12345678") - ); - assert_eq!( - normalize_mysql_decimal_value( - "billing_total_cost_usd", - &json!("123456789012.12345678") - ) - .expect("decimal string should normalize") - .as_deref(), - Some("123456789012.12345678") - ); - assert!(normalize_mysql_decimal_value("billing_total_cost_usd", &json!(true)).is_err()); - } -} diff --git a/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs b/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs deleted file mode 100644 index a15020646..000000000 --- a/crates/aether-data/runtime/src/lifecycle/export/sqlite.rs +++ /dev/null @@ -1,772 +0,0 @@ -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, -) -> Result { - export_sqlite_jsonl(pool, sqlite_core_export_domains(), created_at_unix_secs).await -} - -pub async fn export_sqlite_jsonl( - pool: &crate::driver::sqlite::SqlitePool, - 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), - domains.clone(), - ); - 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(&mut tx, &mut records).await?; - continue; - } - if domain == ExportDomain::Wallets { - 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(&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) -} - -pub async fn import_sqlite_jsonl( - pool: &crate::driver::sqlite::SqlitePool, - input: &str, -) -> Result { - let plan = build_import_plan(input)?; - import_sqlite_plan(pool, &plan).await -} - -pub async fn import_sqlite_plan( - pool: &crate::driver::sqlite::SqlitePool, - plan: &DataImportPlan, -) -> Result { - let identity_scope = IdentityImportScope::from_plan(plan)?; - let mut tx = pool.begin().await.map_sql_err()?; - let identity_state = capture_sqlite_identity_import_state(&mut tx, &identity_scope).await?; - 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_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(&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(&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(&mut tx, &mut column_cache, table_name).await?; - for row in plan.rows(*domain) { - import_sqlite_row(&mut tx, table_name, *domain, row, &target_columns).await?; - imported = imported.saturating_add(1); - } - } - enforce_sqlite_identity_import_invariants(&mut tx, &identity_scope, identity_state).await?; - tx.commit().await.map_sql_err()?; - Ok(imported) -} - -async fn capture_sqlite_identity_import_state( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - scope: &IdentityImportScope, -) -> Result { - let mut affected_user_ids = if scope.finalizes_oauth_links { - scope.user_ids.iter().cloned().collect::>() - } else { - BTreeSet::new() - }; - for link_id in &scope.oauth_link_ids { - if let Some(user_id) = - sqlx::query_scalar::<_, String>("SELECT user_id FROM user_oauth_links WHERE id = ?") - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - affected_user_ids.insert(user_id); - } - } - for provider_type in &scope.oauth_provider_types { - let user_ids = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE provider_type = ?", - ) - .bind(provider_type) - .fetch_all(&mut **tx) - .await - .map_sql_err()?; - affected_user_ids.extend(user_ids); - } - Ok(IdentityImportState { affected_user_ids }) -} - -async fn enforce_sqlite_identity_import_invariants( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - scope: &IdentityImportScope, - mut state: IdentityImportState, -) -> Result<(), DataLayerError> { - for user_id in &scope.user_ids { - let auth_source = - sqlx::query_scalar::<_, String>("SELECT auth_source FROM users WHERE id = ? LIMIT 1") - .bind(user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "imported users row '{user_id}' did not produce a user record" - )) - })?; - if !matches!(auth_source.as_str(), "local" | "ldap" | "oauth") { - return Err(DataLayerError::InvalidInput(format!( - "imported user '{user_id}' has unsupported auth_source '{auth_source}'" - ))); - } - if auth_source == "oauth" { - sqlx::query("UPDATE users SET email_verified = 0 WHERE id = ?") - .bind(user_id) - .execute(&mut **tx) - .await - .map_sql_err()?; - } - } - - for link_id in &scope.oauth_link_ids { - let user_id = sqlx::query_scalar::<_, String>( - "SELECT user_id FROM user_oauth_links WHERE id = ? LIMIT 1", - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "imported OAuth link row '{link_id}' did not produce a link record" - )) - })?; - state.affected_user_ids.insert(user_id); - } - - for link_id in &scope.oauth_link_ids { - if let Some((provider_type, provider_user_id)) = sqlx::query_as::<_, (String, String)>( - r#" -SELECT imported.provider_type, imported.provider_user_id -FROM user_oauth_links imported -JOIN user_oauth_links duplicate - ON duplicate.provider_type = imported.provider_type - AND duplicate.provider_user_id = imported.provider_user_id - AND duplicate.id <> imported.id -WHERE imported.id = ? -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import assigns provider identity '{provider_type}:{provider_user_id}' more than once" - ))); - } - } - - for link_id in &scope.oauth_link_ids { - if let Some((user_id, provider_type)) = sqlx::query_as::<_, (String, String)>( - r#" -SELECT imported.user_id, imported.provider_type -FROM user_oauth_links imported -JOIN user_oauth_links duplicate - ON duplicate.user_id = imported.user_id - AND duplicate.provider_type = imported.provider_type - AND duplicate.id <> imported.id -WHERE imported.id = ? -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import links user '{user_id}' to provider '{provider_type}' more than once" - ))); - } - } - - for link_id in &scope.oauth_link_ids { - if let Some(invalid_id) = sqlx::query_scalar::<_, String>( - r#" -SELECT links.id -FROM user_oauth_links links -LEFT JOIN users ON users.id = links.user_id -LEFT JOIN oauth_providers providers ON providers.provider_type = links.provider_type -WHERE links.id = ? - AND ( - users.id IS NULL - OR providers.provider_type IS NULL - OR links.provider_type <> LOWER(TRIM(links.provider_type)) - OR links.provider_type = '' - OR links.provider_user_id <> TRIM(links.provider_user_id) - OR links.provider_user_id = '' - OR providers.provider_type <> LOWER(TRIM(providers.provider_type)) - ) -LIMIT 1 -"#, - ) - .bind(link_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import produced invalid or orphaned link '{invalid_id}'" - ))); - } - } - - if !scope.validates_oauth_login_methods { - return Ok(()); - } - for user_id in state.affected_user_ids { - if sqlx::query_scalar::<_, String>( - r#" -SELECT users.id -FROM users -WHERE users.id = ? - AND users.auth_source = 'oauth' - AND users.is_active = 1 - AND users.is_deleted = 0 - AND NOT EXISTS ( - SELECT 1 - FROM user_oauth_links links - JOIN oauth_providers providers ON providers.provider_type = links.provider_type - WHERE links.user_id = users.id - AND providers.is_enabled = 1 - AND links.provider_type = LOWER(TRIM(links.provider_type)) - AND links.provider_user_id = TRIM(links.provider_user_id) - AND links.provider_user_id <> '' - ) -LIMIT 1 -"#, - ) - .bind(&user_id) - .fetch_optional(&mut **tx) - .await - .map_sql_err()? - .is_some() - { - return Err(DataLayerError::InvalidInput(format!( - "OAuth import would leave active user '{user_id}' without an enabled identity binding" - ))); - } - } - - Ok(()) -} - -fn sqlite_domain_table( - domain: ExportDomain, -) -> Result<(&'static str, &'static str), DataLayerError> { - match domain { - ExportDomain::Users => Ok(("users", "id")), - ExportDomain::ApiKeys => Ok(("api_keys", "id")), - ExportDomain::Providers => Ok(("providers", "id")), - ExportDomain::ProviderKeys => Ok(("provider_api_keys", "id")), - ExportDomain::Endpoints => Ok(("provider_endpoints", "id")), - ExportDomain::Models => Ok(("models", "id")), - ExportDomain::GlobalModels => Ok(("global_models", "id")), - ExportDomain::AuthModules => Ok(("auth_modules", "id")), - ExportDomain::OAuthProviders => Ok(("oauth_providers", "provider_type")), - ExportDomain::UserOAuthLinks => Ok(("user_oauth_links", "id")), - ExportDomain::UserGroups => Ok(("user_groups", "id")), - ExportDomain::UserGroupMembers => Ok(("user_group_members", "group_id")), - ExportDomain::ProxyNodes => Ok(("proxy_nodes", "id")), - ExportDomain::SystemConfigs => Ok(("system_configs", "id")), - ExportDomain::Wallets => Err(DataLayerError::InvalidInput( - "sqlite wallet export uses multiple tables and must be handled as a domain".to_string(), - )), - ExportDomain::Usage => Ok((r#""usage""#, "request_id")), - ExportDomain::Billing => Err(DataLayerError::InvalidInput( - "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, - id_column: &str, -) -> Result { - if domain == ExportDomain::UserGroupMembers { - let group_id = sqlite_required_export_text(row, "group_id", domain)?; - let user_id = sqlite_required_export_text(row, "user_id", domain)?; - return Ok(format!("{group_id}:{user_id}")); - } - sqlite_required_export_text(row, id_column, domain) -} - -fn sqlite_required_export_text( - row: &sqlx::sqlite::SqliteRow, - column: &str, - domain: ExportDomain, -) -> Result { - row.try_get::, _>(column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "{} export row has null id column '{}'", - domain.as_str(), - column - )) - }) -} - -async fn export_sqlite_billing_records( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - records: &mut Vec, -) -> Result<(), DataLayerError> { - for table_name in [ - "billing_rules", - "dimension_collectors", - "usage_settlement_snapshots", - ] { - let id_column = if table_name == "usage_settlement_snapshots" { - "request_id" - } else { - "id" - }; - let sql = format!("SELECT * FROM {table_name} ORDER BY {id_column} ASC"); - let rows = sqlx::query(&sql).fetch_all(&mut **tx).await.map_sql_err()?; - for row in rows { - let id = row - .try_get::, _>(id_column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "billing export row in table '{table_name}' has null id" - )) - })?; - records.push(DataExportRecord::row( - ExportDomain::Billing, - format!("{table_name}:{id}"), - payload_with_table(sqlite_row_payload(&row)?, table_name)?, - )); - } - } - Ok(()) -} - -async fn export_sqlite_wallet_records( - 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(&mut **tx).await.map_sql_err()?; - for row in rows { - let id = row - .try_get::, _>(id_column) - .map_sql_err()? - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "wallet export row in table '{table_name}' has null id" - )) - })?; - records.push(DataExportRecord::row( - ExportDomain::Wallets, - format!("{table_name}:{id}"), - payload_with_table(sqlite_row_payload(&row)?, table_name)?, - )); - } - } - Ok(()) -} - -async fn import_sqlite_row( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table_name: &str, - domain: ExportDomain, - row: &ExportRow, - target_columns: &SqliteImportColumns, -) -> Result<(), DataLayerError> { - let mut object = - filter_import_payload("sqlite", table_name, domain, row, &target_columns.names)?; - deactivate_imported_credentials(table_name, &mut object, |column_name| { - target_columns.names.contains(column_name) - }); - - let columns = object.keys().map(String::as_str).collect::>(); - let column_sql = columns - .iter() - .map(|column| sqlite_quote_identifier(column)) - .collect::, _>>()? - .join(", "); - let placeholder_sql = vec!["?"; columns.len()].join(", "); - 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"); - 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(&mut **tx).await.map_sql_err()?; - Ok(()) -} - -async fn import_sqlite_billing_row( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - row: &ExportRow, - 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(tx, column_cache, table_name).await?; - import_sqlite_row( - tx, - table_name, - ExportDomain::Billing, - &ExportRow { - id: row.id.clone(), - payload, - }, - &target_columns, - ) - .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"), - "dimension_collectors" => Ok("dimension_collectors"), - "usage_settlement_snapshots" => Ok("usage_settlement_snapshots"), - other => Err(DataLayerError::InvalidInput(format!( - "unsupported sqlite billing export table '{other}'" - ))), - } -} - -async fn import_sqlite_wallet_row( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - row: &ExportRow, - 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(tx, column_cache, table_name).await?; - import_sqlite_row( - tx, - table_name, - ExportDomain::Wallets, - &ExportRow { - id: row.id.clone(), - payload, - }, - &target_columns, - ) - .await -} - -fn sqlite_wallet_tables() -> &'static [(&'static str, &'static str)] { - &[ - ("wallets", "id"), - ("wallet_transactions", "id"), - ("wallet_daily_usage_ledgers", "id"), - ("payment_orders", "id"), - ("payment_callbacks", "id"), - ("refund_requests", "id"), - ("redeem_code_batches", "id"), - ("redeem_codes", "id"), - ] -} - -fn sqlite_wallet_table_name(table_name: &str) -> Result<&'static str, DataLayerError> { - sqlite_wallet_tables() - .iter() - .find(|(candidate, _)| *candidate == table_name) - .map(|(table, _)| *table) - .ok_or_else(|| { - DataLayerError::InvalidInput(format!( - "unsupported sqlite wallet export table '{table_name}'" - )) - }) -} - -async fn sqlite_import_columns_cached( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - cache: &mut BTreeMap, - table_name: &str, -) -> Result { - if let Some(columns) = cache.get(table_name) { - return Ok(columns.clone()); - } - - 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( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table_name: &str, -) -> Result { - let sql = format!("PRAGMA table_info({table_name})"); - 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 { - 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.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() { - object.insert( - column.name().to_string(), - sqlite_value_to_json(row, index, column.name())?, - ); - } - Ok(Value::Object(object)) -} - -fn sqlite_value_to_json( - row: &sqlx::sqlite::SqliteRow, - index: usize, - column_name: &str, -) -> Result { - let raw = row.try_get_raw(index).map_sql_err()?; - if raw.is_null() { - return Ok(Value::Null); - } - - match raw.type_info().name().to_ascii_uppercase().as_str() { - "INTEGER" => { - let value = row.try_get::(index).map_sql_err()?; - if sqlite_integer_column_is_boolean(column_name) { - match value { - 0 => return Ok(Value::Bool(false)), - 1 => return Ok(Value::Bool(true)), - _ => {} - } - } - Ok(Value::from(value)) - } - "REAL" | "FLOAT" | "DOUBLE" => { - let value = row.try_get::(index).map_sql_err()?; - serde_json::Number::from_f64(value) - .map(Value::Number) - .ok_or_else(|| { - DataLayerError::UnexpectedValue(format!( - "sqlite export column {} contains non-finite float", - index - )) - }) - } - "TEXT" => Ok(Value::String( - row.try_get::(index).map_sql_err()?, - )), - "BLOB" => { - let bytes = row.try_get::, _>(index).map_sql_err()?; - Ok(Value::Array(bytes.into_iter().map(Value::from).collect())) - } - other => Err(DataLayerError::UnexpectedValue(format!( - "unsupported sqlite export column type '{other}' at index {index}" - ))), - } -} - -fn sqlite_integer_column_is_boolean(column_name: &str) -> bool { - column_name.starts_with("is_") - || column_name.starts_with("has_") - || column_name.starts_with("supports_") - || column_name.starts_with("enable_") - || column_name.starts_with("use_") - || matches!( - column_name, - "announcement_notifications" - | "auto_delete_on_expiry" - | "auto_fetch_models" - | "email_notifications" - | "email_verified" - | "format_converted" - | "keep_priority_on_conversion" - | "signature_valid" - | "tunnel_connected" - | "tunnel_mode" - | "usage_alerts" - | "webhook_sent" - ) -} diff --git a/crates/aether-data/runtime/src/lifecycle/export/tests.rs b/crates/aether-data/runtime/src/lifecycle/export/tests.rs index 2ba3f6b65..48686a98a 100644 --- a/crates/aether-data/runtime/src/lifecycle/export/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/export/tests.rs @@ -4,19 +4,13 @@ use serde_json::{json, Value}; use super::{ build_import_plan, deactivate_imported_credentials, decode_jsonl, decode_jsonl_with_limits, - encode_jsonl, export_mysql_core_jsonl, export_mysql_jsonl, 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, + encode_jsonl, export_postgres_core_jsonl, 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, + postgres_bytea_json_value, postgres_core_export_domains, DataExportManifest, DataExportRecord, + ExportDomain, ExportRow, PostgresImportColumn, }; use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory}; -use crate::lifecycle::migrate::{ - run_migrations as run_postgres_migrations, run_mysql_migrations, run_sqlite_migrations, -}; +use crate::lifecycle::migrate::run_migrations as run_postgres_migrations; use crate::DatabaseDriver; #[test] @@ -64,73 +58,8 @@ fn jsonl_round_trips_manifest_and_domain_rows() { } #[test] -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); +fn core_export_domains_include_auxiliary_tables() { + assert!(postgres_core_export_domains().contains(&ExportDomain::Auxiliary)); } #[test] @@ -157,7 +86,7 @@ fn jsonl_rejects_rows_outside_manifest_domains() { let records = vec![ DataExportRecord::manifest(DataExportManifest::new( 1_700_000_000, - Some(DatabaseDriver::Sqlite), + Some(DatabaseDriver::Postgres), vec![ExportDomain::Users], )), DataExportRecord::row( @@ -218,7 +147,7 @@ fn jsonl_rejects_duplicate_domain_ids() { } #[test] -fn postgres_import_payload_normalizes_sqlite_values_for_target_columns() { +fn postgres_import_payload_normalizes_imported_values_for_target_columns() { let target_columns = BTreeMap::from([ ( "id".to_string(), @@ -272,7 +201,7 @@ fn postgres_import_payload_normalizes_sqlite_values_for_target_columns() { fn cross_driver_timestamp_normalization_preserves_usage_second_contract() { assert_eq!( normalize_imported_integer_timestamp( - "sqlite", + "postgres", r#""usage""#, "created_at_unix_ms", &json!("1970-01-01T00:00:01.234900Z"), @@ -282,7 +211,7 @@ fn cross_driver_timestamp_normalization_preserves_usage_second_contract() { ); assert_eq!( normalize_imported_integer_timestamp( - "mysql", + "postgres", "request_candidates", "created_at_unix_ms", &json!("1970-01-01T00:00:01.234900Z"), @@ -332,16 +261,16 @@ fn cross_driver_timestamp_normalization_preserves_usage_second_contract() { #[test] fn cross_driver_binary_normalization_preserves_raw_bytes() { assert_eq!( - normalize_imported_binary("sqlite", "payload_gzip", &json!([0, 1, 127, 255])) + normalize_imported_binary("postgres", "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")) + normalize_imported_binary("postgres", "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!(normalize_imported_binary("postgres", "payload_gzip", &json!([256])).is_err()); assert_eq!( postgres_bytea_json_value("payload_gzip", &json!([0, 1, 127, 255])) .expect("postgres bytea should normalize"), @@ -375,60 +304,6 @@ 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 imported_identity_credentials_are_replaced_with_disabled_tombstones() { let columns = BTreeSet::from([ @@ -541,154 +416,6 @@ fn imported_proxy_nodes_receive_a_new_offline_tunnel_generation() { ); } -#[tokio::test] -async fn sqlite_import_rotates_proxy_node_generations_and_clears_online_state() { - 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( - r#" -INSERT INTO proxy_nodes ( - id, tunnel_generation, name, ip, port, status, active_connections, - tunnel_mode, tunnel_connected, proxy_metadata, created_at, updated_at -) VALUES ( - 'import-existing-node', 'target-live-generation', 'existing node', '127.0.0.1', - 8080, 'online', 9, 1, 1, '{"target":"metadata"}', 1, 1 -) -"#, - ) - .execute(&pool) - .await - .expect("existing proxy node should insert"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::ProxyNodes], - )), - DataExportRecord::row( - ExportDomain::ProxyNodes, - "import-existing-node", - json!({ - "id": "import-existing-node", - "tunnel_generation": "backup-stale-generation", - "name": "restored existing node", - "ip": "127.0.0.1", - "port": 8080, - "status": "online", - "active_connections": 42, - "tunnel_mode": true, - "tunnel_connected": true, - "proxy_metadata": { - "tunnel_security": {"encryption_key": "preserved-psk"} - }, - "created_at": 1, - "updated_at": 2 - }), - ), - DataExportRecord::row( - ExportDomain::ProxyNodes, - "import-legacy-node", - json!({ - "id": "import-legacy-node", - "name": "legacy backup node", - "ip": "127.0.0.2", - "port": 8081, - "status": "online", - "active_connections": 7, - "tunnel_mode": true, - "tunnel_connected": true, - "created_at": 1, - "updated_at": 2 - }), - ), - ]) - .expect("proxy node import fixture should encode"); - - assert_eq!( - import_sqlite_jsonl(&pool, &encoded) - .await - .expect("proxy nodes should import"), - 2 - ); - - let restored = sqlx::query_as::<_, (String, String, bool, i32, Option)>( - r#" -SELECT tunnel_generation, status, tunnel_connected, active_connections, proxy_metadata -FROM proxy_nodes -WHERE id = 'import-existing-node' -"#, - ) - .fetch_one(&pool) - .await - .expect("restored proxy node should load"); - assert_ne!(restored.0, "target-live-generation"); - assert_ne!(restored.0, "backup-stale-generation"); - assert!(uuid::Uuid::parse_str(&restored.0).is_ok()); - assert_eq!(restored.1, "offline"); - assert!(!restored.2); - assert_eq!(restored.3, 0); - assert_eq!( - restored - .4 - .as_deref() - .and_then(|value| serde_json::from_str::(value).ok()) - .and_then(|value| value["tunnel_security"]["encryption_key"] - .as_str() - .map(str::to_string)), - Some("preserved-psk".to_string()) - ); - - let legacy_generation: String = sqlx::query_scalar( - "SELECT tunnel_generation FROM proxy_nodes WHERE id = 'import-legacy-node'", - ) - .fetch_one(&pool) - .await - .expect("legacy imported proxy node should load"); - assert!(uuid::Uuid::parse_str(&legacy_generation).is_ok()); -} - -#[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(), @@ -707,1037 +434,6 @@ 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_users_import_fails_closed_for_oauth_email_verification() { - 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::Users], - )), - DataExportRecord::row( - ExportDomain::Users, - "oauth-user", - json!({ - "id": "oauth-user", - "email": "oauth@example.test", - "email_verified": true, - "username": "oauth-user", - "role": "user", - "auth_source": "oauth", - "created_at": 1, - "updated_at": 1 - }), - ), - DataExportRecord::row( - ExportDomain::Users, - "local-user", - json!({ - "id": "local-user", - "email": "local@example.test", - "email_verified": true, - "username": "local-user", - "role": "user", - "auth_source": "local", - "created_at": 1, - "updated_at": 1 - }), - ), - ]) - .expect("users fixture should encode"); - - assert_eq!( - import_sqlite_jsonl(&pool, &encoded) - .await - .expect("users-only staged restore should succeed without OAuth links"), - 2 - ); - let verification = - sqlx::query_as::<_, (String, i64)>("SELECT id, email_verified FROM users ORDER BY id ASC") - .fetch_all(&pool) - .await - .expect("verification state should load"); - assert_eq!( - verification, - vec![("local-user".to_string(), 1), ("oauth-user".to_string(), 0)] - ); -} - -#[tokio::test] -async fn sqlite_users_and_providers_can_restore_before_oauth_links() { - 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::Users, ExportDomain::OAuthProviders], - )), - DataExportRecord::row( - ExportDomain::Users, - "oauth-user", - json!({ - "id": "oauth-user", - "email": "oauth@example.test", - "email_verified": true, - "username": "oauth-user", - "role": "user", - "auth_source": "oauth", - "created_at": 1, - "updated_at": 1 - }), - ), - DataExportRecord::row( - ExportDomain::OAuthProviders, - "linuxdo", - json!({ - "provider_type": "linuxdo", - "display_name": "Linux.do", - "client_id": "client", - "redirect_uri": "https://gateway.example.test/oauth/callback", - "frontend_callback_url": "https://app.example.test/auth/callback", - "is_enabled": true, - "created_at": 1, - "updated_at": 1 - }), - ), - ]) - .expect("staged identity fixture should encode"); - - assert_eq!( - import_sqlite_jsonl(&pool, &encoded) - .await - .expect("users and Providers should restore before links"), - 2 - ); - let email_verified: i64 = - sqlx::query_scalar("SELECT email_verified FROM users WHERE id = 'oauth-user'") - .fetch_one(&pool) - .await - .expect("staged OAuth user should load"); - assert_eq!(email_verified, 0); -} - -#[tokio::test] -async fn sqlite_oauth_link_import_rolls_back_without_enabled_login_binding() { - 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::raw_sql( - r#" -INSERT INTO users ( - id, email, email_verified, username, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES ( - 'oauth-user', 'oauth@example.test', 0, 'oauth-user', 'user', 'oauth', - 1, 0, 1, 1 -); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 0, 1, 1 -); -"#, - ) - .execute(&pool) - .await - .expect("OAuth fixtures should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::UserOAuthLinks], - )), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "link-disabled", - json!({ - "id": "link-disabled", - "user_id": "oauth-user", - "provider_type": "linuxdo", - "provider_user_id": "subject-1", - "linked_at": 1 - }), - ), - ]) - .expect("OAuth link fixture should encode"); - - let err = import_sqlite_jsonl(&pool, &encoded) - .await - .expect_err("disabled-only OAuth binding should fail"); - assert!(err - .to_string() - .contains("without an enabled identity binding")); - let link_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links") - .fetch_one(&pool) - .await - .expect("rolled-back OAuth link count should load"); - assert_eq!(link_count, 0); -} - -#[tokio::test] -async fn sqlite_oauth_provider_import_rolls_back_if_it_removes_last_enabled_binding() { - 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::raw_sql( - r#" -INSERT INTO users ( - id, email, email_verified, username, role, auth_source, - is_active, is_deleted, created_at, updated_at -) VALUES ( - 'oauth-user', 'oauth@example.test', 0, 'oauth-user', 'user', 'oauth', - 1, 0, 1, 1 -); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 1, 1, 1 -); -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, linked_at -) VALUES ( - 'existing-link', 'oauth-user', 'linuxdo', 'subject-1', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("OAuth fixtures should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::OAuthProviders], - )), - DataExportRecord::row( - ExportDomain::OAuthProviders, - "linuxdo", - json!({ - "provider_type": "linuxdo", - "display_name": "Linux.do disabled", - "client_id": "client", - "redirect_uri": "https://gateway.example.test/oauth/callback", - "frontend_callback_url": "https://app.example.test/auth/callback", - "is_enabled": false, - "created_at": 1, - "updated_at": 2 - }), - ), - ]) - .expect("disabled Provider fixture should encode"); - - let err = import_sqlite_jsonl(&pool, &encoded) - .await - .expect_err("disabling the last OAuth login method should fail"); - assert!(err - .to_string() - .contains("without an enabled identity binding")); - let provider = sqlx::query_as::<_, (String, i64)>( - "SELECT display_name, is_enabled FROM oauth_providers WHERE provider_type = 'linuxdo'", - ) - .fetch_one(&pool) - .await - .expect("rolled-back Provider should load"); - assert_eq!(provider, ("Linux.do".to_string(), 1)); -} - -#[tokio::test] -async fn sqlite_oauth_link_reassignment_rolls_back_if_old_owner_loses_last_binding() { - 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::raw_sql( - r#" -INSERT INTO users (id, username, role, auth_source, created_at, updated_at) VALUES - ('oauth-owner', 'oauth-owner', 'user', 'oauth', 1, 1), - ('local-target', 'local-target', 'user', 'local', 1, 1); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 1, 1, 1 -); -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, linked_at -) VALUES ( - 'reassigned-link', 'oauth-owner', 'linuxdo', 'subject-1', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("OAuth reassignment fixtures should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::UserOAuthLinks], - )), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "reassigned-link", - json!({ - "id": "reassigned-link", - "user_id": "local-target", - "provider_type": "linuxdo", - "provider_user_id": "subject-1", - "linked_at": 2 - }), - ), - ]) - .expect("OAuth reassignment fixture should encode"); - - let err = import_sqlite_jsonl(&pool, &encoded) - .await - .expect_err("taking the old owner's last OAuth binding should fail"); - assert!(err - .to_string() - .contains("without an enabled identity binding")); - let owner: String = - sqlx::query_scalar("SELECT user_id FROM user_oauth_links WHERE id = 'reassigned-link'") - .fetch_one(&pool) - .await - .expect("rolled-back OAuth link should load"); - assert_eq!(owner, "oauth-owner"); -} - -#[tokio::test] -async fn sqlite_oauth_link_import_ignores_unrelated_legacy_identity_damage() { - 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::raw_sql( - r#" -INSERT INTO users (id, username, role, auth_source, created_at, updated_at) VALUES - ('broken-oauth', 'broken-oauth', 'user', 'oauth', 1, 1), - ('local-user', 'local-user', 'user', 'local', 1, 1); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 1, 1, 1 -); -INSERT INTO user_oauth_links ( - id, user_id, provider_type, provider_user_id, linked_at -) VALUES ( - 'legacy-orphan', 'missing-user', 'missing-provider', 'legacy-subject', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("legacy damaged identity fixtures should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::UserOAuthLinks], - )), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "valid-link", - json!({ - "id": "valid-link", - "user_id": "local-user", - "provider_type": "linuxdo", - "provider_user_id": "valid-subject", - "linked_at": 1 - }), - ), - ]) - .expect("valid OAuth link fixture should encode"); - - assert_eq!( - import_sqlite_jsonl(&pool, &encoded) - .await - .expect("unrelated legacy damage must not block a valid scoped import"), - 1 - ); - let valid_link_count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links WHERE id = 'valid-link'") - .fetch_one(&pool) - .await - .expect("valid OAuth link count should load"); - assert_eq!(valid_link_count, 1); -} - -#[tokio::test] -async fn sqlite_oauth_link_import_rejects_duplicate_identity_in_legacy_schema() { - 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::raw_sql( - r#" -DROP INDEX uq_user_oauth_links_provider_user; -INSERT INTO users (id, username, role, auth_source, created_at, updated_at) VALUES - ('local-a', 'local-a', 'user', 'local', 1, 1), - ('local-b', 'local-b', 'user', 'local', 1, 1); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 1, 1, 1 -); -"#, - ) - .execute(&pool) - .await - .expect("legacy schema fixture should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::UserOAuthLinks], - )), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "link-a", - json!({ - "id": "link-a", - "user_id": "local-a", - "provider_type": "linuxdo", - "provider_user_id": "same-subject", - "linked_at": 1 - }), - ), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "link-b", - json!({ - "id": "link-b", - "user_id": "local-b", - "provider_type": "linuxdo", - "provider_user_id": "same-subject", - "linked_at": 1 - }), - ), - ]) - .expect("duplicate identity fixture should encode"); - - let err = import_sqlite_jsonl(&pool, &encoded) - .await - .expect_err("duplicate provider identity should fail"); - assert!(err.to_string().contains("more than once")); - let link_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links") - .fetch_one(&pool) - .await - .expect("rolled-back OAuth link count should load"); - assert_eq!(link_count, 0); -} - -#[tokio::test] -async fn sqlite_oauth_link_import_rejects_duplicate_user_provider_in_legacy_schema() { - 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::raw_sql( - r#" -DROP INDEX uq_user_oauth_links_user_provider; -INSERT INTO users (id, username, role, auth_source, created_at, updated_at) -VALUES ('local-user', 'local-user', 'user', 'local', 1, 1); -INSERT INTO oauth_providers ( - provider_type, display_name, client_id, redirect_uri, frontend_callback_url, - is_enabled, created_at, updated_at -) VALUES ( - 'linuxdo', 'Linux.do', 'client', 'https://gateway.example.test/oauth/callback', - 'https://app.example.test/auth/callback', 1, 1, 1 -); -"#, - ) - .execute(&pool) - .await - .expect("legacy schema fixture should seed"); - - let encoded = encode_jsonl(&[ - DataExportRecord::manifest(DataExportManifest::new( - 1_700_000_000, - Some(DatabaseDriver::Postgres), - vec![ExportDomain::UserOAuthLinks], - )), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "link-a", - json!({ - "id": "link-a", - "user_id": "local-user", - "provider_type": "linuxdo", - "provider_user_id": "subject-a", - "linked_at": 1 - }), - ), - DataExportRecord::row( - ExportDomain::UserOAuthLinks, - "link-b", - json!({ - "id": "link-b", - "user_id": "local-user", - "provider_type": "linuxdo", - "provider_user_id": "subject-b", - "linked_at": 1 - }), - ), - ]) - .expect("duplicate user-provider fixture should encode"); - - let err = import_sqlite_jsonl(&pool, &encoded) - .await - .expect_err("duplicate user provider should fail"); - assert!(err.to_string().contains("more than once")); - let link_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM user_oauth_links") - .fetch_one(&pool) - .await - .expect("rolled-back OAuth link count should load"); - assert_eq!(link_count, 0); -} - -#[tokio::test] -async fn sqlite_core_export_reads_migrated_database_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( - r#" -INSERT INTO users (id, email, username, auth_source, created_at, updated_at) -VALUES ('user-1', 'owner@example.com', 'owner', 'local', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO user_groups (id, name, normalized_name, description, priority, allowed_models, allowed_models_mode, created_at, updated_at) -VALUES ('group-1', 'Export Group', 'export group', 'Exported group', 10, '["gpt-test"]', 'specific', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO user_group_members (group_id, user_id, created_at) -VALUES ('group-1', 'user-1', '1970-01-01T00:00:01Z'); -INSERT INTO api_keys (id, user_id, key_hash, key_encrypted, name, created_at, updated_at) -VALUES ('api-key-1', 'user-1', 'hash-1', 'ciphertext-1', 'Default', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO providers (id, name, provider_type, created_at, updated_at) -VALUES ('provider-1', 'Provider One', 'openai', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO provider_api_keys (id, provider_id, name, encrypted_key, created_at, updated_at) -VALUES ('provider-key-1', 'provider-1', 'Provider Key', 'ciphertext-provider', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -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-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) -VALUES ('billing-rule-1', 'global-model-1', 'Rule One', 'chat', 'input_tokens * 0.01', '{}', '{"input":"input_tokens"}', 1, '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO dimension_collectors (id, api_format, task_type, dimension_name, source_type, value_type, transform_expression, priority, is_enabled, created_at, updated_at) -VALUES ('collector-1', 'openai', 'chat', 'input_tokens', 'computed', 'float', 'usage.input_tokens', 10, 1, '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -INSERT INTO system_configs (id, key, value, created_at, updated_at) -VALUES ('config-1', 'billing.enabled', 'true', '1970-01-01T00:00:01Z', '1970-01-01T00:00:02Z'); -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', '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' -); -INSERT INTO usage_cost_reservations ( - request_id, subject_id, reservation_token, admitted_at, - reserved_cost_units, state, reservation_expires_at, retain_until, - created_at, updated_at -) -VALUES ( - 'request-1', 'user-1', 'reservation-1', 1, - 500, 'reserved', 2, 3, - 1, 1 -); -INSERT INTO usage_request_admissions ( - request_id, subject_id, event_token, admitted_at, - retain_until, state, created_at -) -VALUES ( - 'request-1', 'user-1', 'admission-1', 1, - 3, 'active', 1 -); -"#, - ) - .execute(&pool) - .await - .expect("sqlite export rows should seed"); - - let encoded = export_sqlite_core_jsonl(&pool, 1_700_000_000) - .await - .expect("sqlite export should encode"); - let import_plan = build_import_plan(&encoded).expect("sqlite export should decode"); - - assert_eq!( - import_plan.manifest.source_driver, - Some(DatabaseDriver::Sqlite) - ); - assert_eq!(import_plan.manifest.domains, sqlite_core_export_domains()); - assert_eq!( - import_plan.rows(ExportDomain::Users)[0].payload["email"], - "owner@example.com" - ); - assert!(import_plan - .rows(ExportDomain::UserGroups) - .iter() - .any(|row| row.id == "group-1" && row.payload["name"] == "Export Group")); - assert!(import_plan - .rows(ExportDomain::UserGroupMembers) - .iter() - .any(|row| row.id == "group-1:user-1" - && row.payload["group_id"] == "group-1" - && row.payload["user_id"] == "user-1")); - assert_eq!( - import_plan.rows(ExportDomain::ApiKeys)[0].payload["key_encrypted"], - "ciphertext-1" - ); - assert_eq!( - import_plan.rows(ExportDomain::ProviderKeys)[0].payload["encrypted_key"], - "ciphertext-provider" - ); - assert_eq!(import_plan.rows(ExportDomain::Usage)[0].id, "request-1"); - assert_eq!(import_plan.rows(ExportDomain::Billing).len(), 2); - assert_eq!( - import_plan.rows(ExportDomain::Billing)[0].payload["__table"], - "billing_rules" - ); - assert_eq!( - 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")); - assert!(import_plan - .rows(ExportDomain::Auxiliary) - .iter() - .any(|row| row.payload["__table"] == "usage_cost_reservations" - && row.payload["reservation_token"] == "reservation-1" - && row.payload["reserved_cost_units"] == 500 - && row.payload["state"] == "reserved")); - assert!(import_plan - .rows(ExportDomain::Auxiliary) - .iter() - .any(|row| row.payload["__table"] == "usage_request_admissions" - && row.payload["event_token"] == "admission-1" - && row.payload["state"] == "active")); - - let target_pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("target sqlite pool should connect"); - run_sqlite_migrations(&target_pool) - .await - .expect("target sqlite migrations should run"); - let imported = import_sqlite_jsonl(&target_pool, &encoded) - .await - .expect("sqlite import should load exported rows"); - assert_eq!(imported, 22); - - let imported_api_key = sqlx::query_as::<_, (String, Option, bool, bool, String)>( - "SELECT key_hash, key_encrypted, is_active, is_locked, status FROM api_keys WHERE id = 'api-key-1'", - ) - .fetch_one(&target_pool) - .await - .expect("imported api key should load"); - assert_ne!(imported_api_key.0, "hash-1"); - assert_eq!(imported_api_key.1, None); - assert!(!imported_api_key.2); - assert!(imported_api_key.3); - assert_eq!(imported_api_key.4, "disabled"); - - 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, - ("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'", - ) - .fetch_one(&target_pool) - .await - .expect("imported user group member should load"); - assert_eq!(imported_group_member.0, "group-1"); - assert_eq!(imported_group_member.1, "user-1"); - - let imported_billing_rule = sqlx::query_as::<_, (String,)>( - "SELECT expression FROM billing_rules WHERE id = 'billing-rule-1'", - ) - .fetch_one(&target_pool) - .await - .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()) - ); - - let imported_reservation = sqlx::query_as::<_, (String, i64, String)>( - r#" -SELECT subject_id, reserved_cost_units, state -FROM usage_cost_reservations -WHERE reservation_token = 'reservation-1' -"#, - ) - .fetch_one(&target_pool) - .await - .expect("imported usage cost reservation should load"); - assert_eq!( - imported_reservation, - ("user-1".to_string(), 500, "reserved".to_string()) - ); - - let imported_admission = sqlx::query_as::<_, (String, String, Option)>( - r#" -SELECT subject_id, state, released_at -FROM usage_request_admissions -WHERE event_token = 'admission-1' -"#, - ) - .fetch_one(&target_pool) - .await - .expect("imported usage request admission should load"); - assert_eq!( - imported_admission, - ("user-1".to_string(), "active".to_string(), None) - ); - - if let Some(database_url) = std::env::var("AETHER_TEST_POSTGRES_URL") - .ok() - .filter(|value| !value.trim().is_empty()) - { - let config = PostgresPoolConfig { - database_url, - min_connections: 1, - max_connections: 1, - acquire_timeout_ms: 1_000, - idle_timeout_ms: 5_000, - max_lifetime_ms: 30_000, - statement_cache_capacity: 64, - require_ssl: false, - }; - let postgres_pool = PostgresPoolFactory::new(config) - .expect("postgres factory should build") - .connect_lazy() - .expect("postgres pool should build"); - run_postgres_migrations(&postgres_pool) - .await - .expect("postgres migrations should run"); - - let imported = import_postgres_jsonl(&postgres_pool, &encoded) - .await - .expect("postgres import should load exported rows"); - assert_eq!(imported, 22); - - let imported_api_key = sqlx::query_as::<_, (String, Option, bool, bool, String)>( - "SELECT key_hash, key_encrypted, is_active, is_locked, status FROM api_keys WHERE id = 'api-key-1'", - ) - .fetch_one(&postgres_pool) - .await - .expect("imported postgres api key should load"); - assert_ne!(imported_api_key.0, "hash-1"); - assert_eq!(imported_api_key.1, None); - assert!(!imported_api_key.2); - assert!(imported_api_key.3); - assert_eq!(imported_api_key.4, "disabled"); - } -} - #[tokio::test] async fn postgres_core_export_reads_migrated_database_rows_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_POSTGRES_URL") @@ -1946,249 +642,6 @@ async fn postgres_core_export_reads_migrated_database_rows_when_url_is_set() { .rows(ExportDomain::Usage) .iter() .any(|row| row.id == request_id)); - - let target_pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("target sqlite pool should connect"); - run_sqlite_migrations(&target_pool) - .await - .expect("target sqlite migrations should run"); - let imported = import_sqlite_jsonl(&target_pool, &encoded) - .await - .expect("sqlite import should load postgres exported rows"); - assert_eq!(imported, import_plan_row_count(&import_plan)); - - let imported_api_key = - sqlx::query_as::<_, (Option,)>("SELECT key_encrypted FROM api_keys WHERE id = $1") - .bind(&api_key_id) - .fetch_one(&target_pool) - .await - .expect("imported sqlite api key should load"); - assert_eq!( - imported_api_key.0, None, - "cross-driver imports must revoke recoverable API-key ciphertext" - ); - 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 = ?", - ) - .bind(&group_id) - .bind(&user_id) - .fetch_one(&target_pool) - .await - .expect("imported sqlite user group member should load"); - assert_eq!(imported_group_member.0, group_id); - assert_eq!(imported_group_member.1, user_id); -} - -#[tokio::test] -async fn mysql_core_export_reads_migrated_database_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 core export smoke 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_mysql_migrations(&pool) - .await - .expect("mysql migrations should run"); - - let suffix = unique_suffix(); - let user_id = format!("export-user-{suffix}"); - let api_key_id = format!("export-api-key-{suffix}"); - let provider_id = format!("export-provider-{suffix}"); - let provider_key_id = format!("export-provider-key-{suffix}"); - let endpoint_id = format!("export-endpoint-{suffix}"); - let global_model_id = format!("export-global-model-{suffix}"); - let model_id = format!("export-model-{suffix}"); - let config_id = format!("export-config-{suffix}"); - let wallet_id = format!("export-wallet-{suffix}"); - let request_id = format!("export-request-{suffix}"); - let group_id = format!("export-group-{suffix}"); - - sqlx::query( - "INSERT INTO users (id, email, username, auth_source, created_at, updated_at) VALUES (?, ?, ?, 'local', 1, 2)", - ) - .bind(&user_id) - .bind(format!("{user_id}@example.com")) - .bind(format!("owner-{suffix}")) - .execute(&pool) - .await - .expect("user should seed"); - sqlx::query( - "INSERT INTO user_groups (id, name, normalized_name, priority, allowed_models, allowed_models_mode, created_at, updated_at) VALUES (?, ?, ?, 10, '[\"provider-model\"]', 'specific', 1, 2)", - ) - .bind(&group_id) - .bind(format!("Export Group {suffix}")) - .bind(format!("export group {suffix}")) - .execute(&pool) - .await - .expect("user group should seed"); - sqlx::query("INSERT INTO user_group_members (group_id, user_id, created_at) VALUES (?, ?, 1)") - .bind(&group_id) - .bind(&user_id) - .execute(&pool) - .await - .expect("user group member should seed"); - sqlx::query( - "INSERT INTO api_keys (id, user_id, key_hash, key_encrypted, name, created_at, updated_at) VALUES (?, ?, ?, 'ciphertext-1', 'Default', 1, 2)", - ) - .bind(&api_key_id) - .bind(&user_id) - .bind(format!("hash-{api_key_id}")) - .execute(&pool) - .await - .expect("api key should seed"); - sqlx::query( - "INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES (?, ?, 'openai', 1, 2)", - ) - .bind(&provider_id) - .bind(format!("Provider {suffix}")) - .execute(&pool) - .await - .expect("provider should seed"); - sqlx::query( - "INSERT INTO provider_api_keys (id, provider_id, name, encrypted_key, created_at, updated_at) VALUES (?, ?, 'Provider Key', 'ciphertext-provider', 1, 2)", - ) - .bind(&provider_key_id) - .bind(&provider_id) - .execute(&pool) - .await - .expect("provider key should seed"); - sqlx::query( - "INSERT INTO provider_endpoints (id, provider_id, name, base_url, created_at, updated_at) VALUES (?, ?, 'Primary', 'https://example.test', 1, 2)", - ) - .bind(&endpoint_id) - .bind(&provider_id) - .execute(&pool) - .await - .expect("endpoint should seed"); - sqlx::query("INSERT INTO global_models (id, name, created_at, updated_at) VALUES (?, ?, 1, 2)") - .bind(&global_model_id) - .bind(format!("global-model-{suffix}")) - .execute(&pool) - .await - .expect("global model should seed"); - sqlx::query( - "INSERT INTO models (id, provider_id, global_model_id, provider_model_name, created_at, updated_at) VALUES (?, ?, ?, 'provider-model', 1, 2)", - ) - .bind(&model_id) - .bind(&provider_id) - .bind(&global_model_id) - .execute(&pool) - .await - .expect("model should seed"); - sqlx::query( - "INSERT INTO system_configs (id, `key`, value, created_at, updated_at) VALUES (?, ?, 'true', 1, 2)", - ) - .bind(&config_id) - .bind(format!("export.config.{suffix}")) - .execute(&pool) - .await - .expect("system config should seed"); - sqlx::query("INSERT INTO wallets (id, user_id, created_at, updated_at) VALUES (?, ?, 1, 2)") - .bind(&wallet_id) - .bind(&user_id) - .execute(&pool) - .await - .expect("wallet should seed"); - sqlx::query( - "INSERT INTO `usage` (request_id, id, user_id, provider_name, model, status, billing_status, created_at_unix_ms, updated_at_unix_secs) VALUES (?, ?, ?, 'Provider One', 'provider-model', 'completed', 'settled', 1, 2)", - ) - .bind(&request_id) - .bind(&request_id) - .bind(&user_id) - .execute(&pool) - .await - .expect("usage should seed"); - - let encoded = export_mysql_core_jsonl(&pool, 1_700_000_000) - .await - .expect("mysql export should encode"); - let import_plan = build_import_plan(&encoded).expect("mysql export should decode"); - - assert_eq!( - import_plan.manifest.source_driver, - Some(DatabaseDriver::Mysql) - ); - assert_eq!(import_plan.manifest.domains, mysql_core_export_domains()); - assert!(import_plan - .rows(ExportDomain::Users) - .iter() - .any(|row| row.id == user_id)); - assert!(import_plan - .rows(ExportDomain::UserGroups) - .iter() - .any(|row| row.id == group_id)); - assert!(import_plan - .rows(ExportDomain::UserGroupMembers) - .iter() - .any(|row| row.id == format!("{group_id}:{user_id}"))); - assert!(import_plan - .rows(ExportDomain::ApiKeys) - .iter() - .any(|row| row.id == api_key_id && row.payload["key_encrypted"] == "ciphertext-1")); - assert!(import_plan - .rows(ExportDomain::ProviderKeys) - .iter() - .any(|row| { - row.id == provider_key_id && row.payload["encrypted_key"] == "ciphertext-provider" - })); - assert!(import_plan - .rows(ExportDomain::Usage) - .iter() - .any(|row| row.id == request_id)); - - let selected_export = export_mysql_jsonl( - &pool, - vec![ - ExportDomain::Users, - ExportDomain::UserGroups, - ExportDomain::UserGroupMembers, - ExportDomain::ApiKeys, - ExportDomain::ProviderKeys, - ExportDomain::Usage, - ], - 1_700_000_001, - ) - .await - .expect("selected mysql export should encode"); - let imported = import_mysql_jsonl(&pool, &selected_export) - .await - .expect("mysql import should be idempotent"); - assert!(imported >= 6); - - let imported_api_key = sqlx::query_as::<_, (String, Option, bool, bool, String)>( - "SELECT key_hash, key_encrypted, is_active, is_locked, status FROM api_keys WHERE id = ?", - ) - .bind(&api_key_id) - .fetch_one(&pool) - .await - .expect("imported mysql api key should load"); - assert_ne!(imported_api_key.0, "hash-1"); - assert_eq!(imported_api_key.1, None); - assert!(!imported_api_key.2); - assert!(imported_api_key.3); - assert_eq!(imported_api_key.4, "disabled"); } fn unique_suffix() -> String { @@ -2201,11 +654,3 @@ fn unique_suffix() -> String { let counter = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); format!("{:016x}", nanos ^ counter.rotate_left(17)) } - -fn import_plan_row_count(plan: &DataImportPlan) -> usize { - plan.manifest - .domains - .iter() - .map(|domain| plan.rows(*domain).len()) - .sum() -} diff --git a/crates/aether-data/runtime/src/lifecycle/migrate.rs b/crates/aether-data/runtime/src/lifecycle/migrate.rs index 566259075..07139bdae 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate.rs @@ -3,58 +3,13 @@ //! Each driver owns its migrator and startup preparation. The facade keeps //! the established public entry points used by gateway bootstrap code. -#[cfg(feature = "mysql")] -mod mysql; #[cfg(feature = "postgres")] mod postgres; -#[cfg(feature = "sqlite")] -mod sqlite; mod types; -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] mod tests; #[cfg(feature = "postgres")] pub use postgres::{pending_migrations, prepare_database_for_startup, run_migrations}; pub use types::PendingMigrationInfo; - -#[cfg(any(feature = "mysql", feature = "sqlite"))] -use sqlx::migrate::MigrateError; - -#[cfg(feature = "mysql")] -pub async fn run_mysql_migrations(pool: &sqlx::MySqlPool) -> Result<(), MigrateError> { - mysql::run_migrations(pool).await -} - -#[cfg(feature = "mysql")] -pub async fn pending_mysql_migrations( - pool: &sqlx::MySqlPool, -) -> Result, MigrateError> { - mysql::pending_migrations(pool).await -} - -#[cfg(feature = "mysql")] -pub async fn prepare_mysql_database_for_startup( - pool: &sqlx::MySqlPool, -) -> Result, MigrateError> { - mysql::prepare_database_for_startup(pool).await -} - -#[cfg(feature = "sqlite")] -pub async fn run_sqlite_migrations(pool: &sqlx::SqlitePool) -> Result<(), MigrateError> { - sqlite::run_migrations(pool).await -} - -#[cfg(feature = "sqlite")] -pub async fn pending_sqlite_migrations( - pool: &sqlx::SqlitePool, -) -> Result, MigrateError> { - sqlite::pending_migrations(pool).await -} - -#[cfg(feature = "sqlite")] -pub async fn prepare_sqlite_database_for_startup( - pool: &sqlx::SqlitePool, -) -> Result, MigrateError> { - sqlite::prepare_database_for_startup(pool).await -} diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/mysql.rs b/crates/aether-data/runtime/src/lifecycle/migrate/mysql.rs deleted file mode 100644 index 7bc42a06f..000000000 --- a/crates/aether-data/runtime/src/lifecycle/migrate/mysql.rs +++ /dev/null @@ -1,5 +0,0 @@ -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] -pub(super) use aether_data_mysql::MIGRATOR; -pub(super) use aether_data_mysql::{ - pending_migrations, prepare_database_for_startup, run_migrations, -}; diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/postgres.rs b/crates/aether-data/runtime/src/lifecycle/migrate/postgres.rs index bf12dc76d..8569949ff 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/postgres.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/postgres.rs @@ -6,7 +6,7 @@ use sqlx::{ use super::types::PendingMigrationInfo; pub use aether_data_postgres::pending_migrations; -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] +#[cfg(all(test, feature = "postgres"))] pub(super) use aether_data_postgres::{ all_up_migrations, pending_migrations_from_applied, POSTGRES_MIGRATOR, }; diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/sqlite.rs b/crates/aether-data/runtime/src/lifecycle/migrate/sqlite.rs deleted file mode 100644 index 50ee78ff0..000000000 --- a/crates/aether-data/runtime/src/lifecycle/migrate/sqlite.rs +++ /dev/null @@ -1,5 +0,0 @@ -#[cfg(all(test, feature = "postgres", feature = "mysql", feature = "sqlite"))] -pub(super) use aether_data_sqlite::MIGRATOR; -pub(super) use aether_data_sqlite::{ - pending_migrations, prepare_database_for_startup, run_migrations, -}; diff --git a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs index 42180e917..90a64c6fb 100644 --- a/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs +++ b/crates/aether-data/runtime/src/lifecycle/migrate/tests.rs @@ -7,7 +7,7 @@ use std::time::{Duration, Instant}; use sqlx::{ migrate::{AppliedMigration, Migrate}, - query, query_scalar, Connection, PgConnection, PgPool, SqlitePool, + query, query_scalar, Connection, PgConnection, PgPool, }; use aether_data_contracts::repository::{ @@ -501,161 +501,6 @@ fn create_table_names(sql: &str) -> BTreeSet { .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())) - // SQLite rebuilds tables to add foreign keys. These staging tables are - // renamed to the canonical table names before the migration finishes. - .filter(|table| !table.ends_with("_with_user_fk")) - .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'")); @@ -989,70 +834,38 @@ fn split_baseline_sources_match_executable_migrations() { EMPTY_DATABASE_SNAPSHOT_SQL, compose_manifest("bootstrap/postgres/manifest.txt") ); - assert_eq!( - include_str!("../../../../adapters/mysql/migrations/20260403000000_baseline.sql"), - compose_manifest("drivers/mysql/baseline/manifest.txt") - ); - assert_eq!( - include_str!("../../../../adapters/sqlite/migrations/20260403000000_baseline.sql"), - compose_manifest("drivers/sqlite/baseline/manifest.txt") - ); } #[test] -fn mysql_and_sqlite_migrations_do_not_use_postgres_jsonb() { - let mysql_sources = super::mysql::MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .map(|migration| migration.sql.as_ref()); - let sqlite_sources = super::sqlite::MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .map(|migration| migration.sql.as_ref()); - - for source in mysql_sources.chain(sqlite_sources) { - assert!( - !source.to_ascii_lowercase().contains("jsonb"), - "Postgres jsonb must stay out of MySQL/SQLite migrations" - ); - } -} - -#[test] -fn worker_boot_cleanup_migration_is_enabled_for_every_driver() { +fn worker_boot_cleanup_migration_is_enabled_for_postgres() { const VERSION: i64 = 20260731000000; - for (driver, migrator) in [ - ("postgres", &POSTGRES_MIGRATOR), - ("mysql", &super::mysql::MIGRATOR), - ("sqlite", &super::sqlite::MIGRATOR), + let (driver, migrator) = ("postgres", &POSTGRES_MIGRATOR); + let migration = migrator + .iter() + .find(|migration| migration.version == VERSION) + .unwrap_or_else(|| panic!("{driver} worker boot cleanup migration should be embedded")); + let sql = migration.sql.as_ref(); + + for required in [ + "DELETE FROM background_task_events", + "DELETE FROM background_task_runs", + "id LIKE 'boot:%'", + "owner_instance IS NOT NULL", + "created_by = 'system'", + "progress_message = 'worker booted'", ] { - let migration = migrator - .iter() - .find(|migration| migration.version == VERSION) - .unwrap_or_else(|| panic!("{driver} worker boot cleanup migration should be embedded")); - let sql = migration.sql.as_ref(); - - for required in [ - "DELETE FROM background_task_events", - "DELETE FROM background_task_runs", - "id LIKE 'boot:%'", - "owner_instance IS NOT NULL", - "created_by = 'system'", - "progress_message = 'worker booted'", - ] { - assert!( - sql.contains(required), - "{driver} worker boot cleanup migration is missing {required}" - ); - } - assert!( - sql.find("DELETE FROM background_task_events") - < sql.find("DELETE FROM background_task_runs"), - "{driver} must delete child events before worker boot runs" + sql.contains(required), + "{driver} worker boot cleanup migration is missing {required}" ); } + + assert!( + sql.find("DELETE FROM background_task_events") + < sql.find("DELETE FROM background_task_runs"), + "{driver} must delete child events before worker boot runs" + ); } const UNPUBLISHED_LEGACY_DATA_REWRITE_MIGRATION_VERSIONS: &[i64] = &[ @@ -1068,20 +881,15 @@ const UNPUBLISHED_LEGACY_DATA_REWRITE_MIGRATION_VERSIONS: &[i64] = &[ ]; #[test] -fn unpublished_legacy_data_rewrite_migrations_are_absent_for_every_driver() { - for (driver, migrator) in [ - ("postgres", &POSTGRES_MIGRATOR), - ("mysql", &super::mysql::MIGRATOR), - ("sqlite", &super::sqlite::MIGRATOR), - ] { - for version in UNPUBLISHED_LEGACY_DATA_REWRITE_MIGRATION_VERSIONS { - assert!( - migrator - .iter() - .all(|migration| migration.version != *version), - "{driver} must not embed unpublished legacy rewrite migration {version}" - ); - } +fn unpublished_legacy_data_rewrite_migrations_are_absent_for_postgres() { + let (driver, migrator) = ("postgres", &POSTGRES_MIGRATOR); + for version in UNPUBLISHED_LEGACY_DATA_REWRITE_MIGRATION_VERSIONS { + assert!( + migrator + .iter() + .all(|migration| migration.version != *version), + "{driver} must not embed unpublished legacy rewrite migration {version}" + ); } } @@ -1089,31 +897,25 @@ fn unpublished_legacy_data_rewrite_migrations_are_absent_for_every_driver() { fn deleted_user_history_schema_decoupling_does_not_rewrite_history() { const VERSION: i64 = 20260827050000; - for (driver, migrator) in [ - ("postgres", &POSTGRES_MIGRATOR), - ("mysql", &super::mysql::MIGRATOR), - ("sqlite", &super::sqlite::MIGRATOR), - ] { - let migration = migrator - .iter() - .find(|migration| migration.version == VERSION) - .unwrap_or_else(|| panic!("{driver} user-history schema migration should be embedded")); - let sql = migration - .sql - .lines() - .map(str::trim) - .filter(|line| !line.is_empty() && !line.starts_with("--")) - .collect::>() - .join("\n") - .to_ascii_uppercase(); - for history_rewrite in ["UPDATE ", "DELETE ", "TRUNCATE ", "REPLACE ", "MERGE "] { - assert!( - !sql - .split(';') - .any(|statement| statement.trim_start().starts_with(history_rewrite)), - "{driver} user-history schema migration rewrites legacy rows with {history_rewrite}" - ); - } + let (driver, migrator) = ("postgres", &POSTGRES_MIGRATOR); + let migration = migrator + .iter() + .find(|migration| migration.version == VERSION) + .unwrap_or_else(|| panic!("{driver} user-history schema migration should be embedded")); + let sql = migration + .sql + .lines() + .map(str::trim) + .filter(|line| !line.is_empty() && !line.starts_with("--")) + .collect::>() + .join("\n") + .to_ascii_uppercase(); + for history_rewrite in ["UPDATE ", "DELETE ", "TRUNCATE ", "REPLACE ", "MERGE "] { + assert!( + !sql.split(';') + .any(|statement| statement.trim_start().starts_with(history_rewrite)), + "{driver} user-history schema migration rewrites legacy rows with {history_rewrite}" + ); } let postgres_migration = POSTGRES_MIGRATOR @@ -1160,590 +962,19 @@ fn deleted_user_history_schema_decoupling_does_not_rewrite_history() { "postgres migration must decouple {constraint}" ); } - - let mysql_migration = super::mysql::MIGRATOR - .iter() - .find(|migration| migration.version == VERSION) - .expect("mysql user-history schema migration should be embedded"); - for constraint in [ - "user_plan_entitlements_user_id_fkey", - "entitlement_usage_ledgers_user_id_fkey", - "user_referrals_inviter_user_id_fkey", - "user_referrals_invitee_user_id_fkey", - "referral_rewards_inviter_user_id_fkey", - "referral_rewards_invitee_user_id_fkey", - ] { - assert!( - mysql_migration.sql.contains(constraint), - "mysql migration must decouple {constraint}" - ); - } - - let sqlite_migration = super::sqlite::MIGRATOR - .iter() - .find(|migration| migration.version == VERSION) - .expect("sqlite user-history schema migration should be embedded"); - for table in [ - "user_plan_entitlements", - "entitlement_usage_ledgers", - "user_referrals", - "referral_rewards", - ] { - assert!( - sqlite_migration - .sql - .contains(&format!("ALTER TABLE {table} RENAME TO")), - "sqlite migration must rebuild {table} without the legacy user foreign key" - ); - } -} - -#[test] -fn mysql_and_sqlite_migrations_include_enabled_incrementals() { - let mysql_versions = super::mysql::MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .map(|migration| migration.version) - .collect::>(); - let sqlite_versions = super::sqlite::MIGRATOR - .iter() - .filter(|migration| migration.migration_type.is_up_migration()) - .map(|migration| migration.version) - .collect::>(); - - assert_eq!( - mysql_versions, - vec![ - 20260403000000, - 20260507120000, - 20260508000000, - 20260509000000, - 20260509120000, - 20260510120000, - 20260511120000, - 20260511130000, - 20260512000000, - 20260512090000, - 20260512110000, - 20260516000000, - 20260518000000, - 20260519000000, - 20260519120000, - 20260519130000, - 20260520000000, - 20260520010000, - 20260524000000, - 20260527000000, - 20260528000000, - 20260528020000, - 20260725010000, - 20260725020000, - 20260725030000, - 20260727000000, - 20260731000000, - 20260814000000, - 20260815000000, - 20260816000000, - 20260817000000, - 20260821000000, - 20260821120000, - 20260821130000, - 20260827040000, - 20260827050000, - 20260831000000, - 20260831010000, - 20260831020000, - 20260831030000, - 20260903000000, - ] - ); - assert_eq!( - sqlite_versions, - vec![ - 20260403000000, - 20260507120000, - 20260508000000, - 20260509000000, - 20260509120000, - 20260510120000, - 20260511120000, - 20260511130000, - 20260512000000, - 20260512090000, - 20260512110000, - 20260516000000, - 20260518000000, - 20260519000000, - 20260519120000, - 20260519130000, - 20260520000000, - 20260520010000, - 20260524000000, - 20260527000000, - 20260528000000, - 20260528020000, - 20260725000000, - 20260725010000, - 20260725020000, - 20260725030000, - 20260725040000, - 20260727000000, - 20260731000000, - 20260814000000, - 20260815000000, - 20260816000000, - 20260821000000, - 20260821120000, - 20260821130000, - 20260827050000, - 20260831000000, - 20260831010000, - 20260831020000, - 20260831030000, - 20260903000000, - ] - ); -} - -#[tokio::test] -async fn sqlite_gateway_order_uniqueness_migration_rejects_historical_duplicates() { - const VERSION: i64 = 20260821120000; - - let pool = SqlitePool::connect("sqlite::memory:") - .await - .expect("sqlite pool should connect"); - let mut connection = pool.acquire().await.expect("sqlite connection should open"); - connection - .ensure_migrations_table() - .await - .expect("migration table should be created"); - for migration in super::sqlite::MIGRATOR - .iter() - .filter(|migration| migration.version < VERSION) - { - connection - .apply(migration) - .await - .expect("pre-uniqueness migration should apply"); - } - drop(connection); - - query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) VALUES - ('duplicate-wallet-a', 'duplicate-user-a', 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, 1, 1), - ('duplicate-wallet-b', 'duplicate-user-b', 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, 1, 1); - -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, gateway_order_id, status, created_at -) VALUES - ('duplicate-order-a', 'duplicate-no-a', 'duplicate-wallet-a', 'duplicate-user-a', 1, 0, 0, ' EPAY ', 'duplicate-gateway-id', 'pending', 1), - ('duplicate-order-b', 'duplicate-no-b', 'duplicate-wallet-b', 'duplicate-user-b', 1, 0, 0, 'epay', 'duplicate-gateway-id', 'pending', 1); -"#, - ) - .execute(&pool) - .await - .expect("historical duplicate fixtures should insert"); - - let migration = super::sqlite::MIGRATOR - .iter() - .find(|migration| migration.version == VERSION) - .expect("gateway-order uniqueness migration should be embedded"); - let error = sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect_err("historical financial duplicates must block migration"); - assert!(error.to_string().to_ascii_lowercase().contains("unique")); - - let order_count: i64 = query_scalar( - "SELECT COUNT(*) FROM payment_orders WHERE gateway_order_id = 'duplicate-gateway-id'", - ) - .fetch_one(&pool) - .await - .expect("duplicate financial records should remain intact"); - assert_eq!(order_count, 2); -} - -#[tokio::test] -async fn sqlite_gateway_order_uniqueness_migration_normalizes_legacy_payment_methods() { - const VERSION: i64 = 20260821120000; - - let pool = SqlitePool::connect("sqlite::memory:") - .await - .expect("sqlite pool should connect"); - let mut connection = pool.acquire().await.expect("sqlite connection should open"); - connection - .ensure_migrations_table() - .await - .expect("migration table should be created"); - for migration in super::sqlite::MIGRATOR - .iter() - .filter(|migration| migration.version < VERSION) - { - connection - .apply(migration) - .await - .expect("pre-uniqueness migration should apply"); - } - drop(connection); - - query( - r#" -INSERT INTO wallets ( - id, user_id, balance, gift_balance, limit_mode, currency, status, - total_recharged, total_consumed, total_refunded, total_adjusted, - created_at, updated_at -) VALUES ('legacy-method-wallet', 'legacy-method-user', 0, 0, 'finite', 'USD', 'active', 0, 0, 0, 0, 1, 1); - -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, gateway_order_id, status, created_at -) VALUES ('legacy-method-order', 'legacy-method-no', 'legacy-method-wallet', 'legacy-method-user', 1, 0, 0, ' EPAY ', 'CaseSensitiveTxn', 'pending', 1); - -INSERT INTO payment_callbacks ( - id, payment_method, callback_key, signature_valid, status, created_at -) VALUES ('legacy-method-callback', ' EPAY ', 'legacy-method-key', 0, 'received', 1); -"#, - ) - .execute(&pool) - .await - .expect("legacy payment methods should insert"); - - let migration = super::sqlite::MIGRATOR - .iter() - .find(|migration| migration.version == VERSION) - .expect("gateway-order uniqueness migration should be embedded"); - sqlx::raw_sql(migration.sql.as_ref()) - .execute(&pool) - .await - .expect("non-conflicting legacy payment methods should normalize"); - - let order_method: String = - query_scalar("SELECT payment_method FROM payment_orders WHERE id = 'legacy-method-order'") - .fetch_one(&pool) - .await - .expect("normalized order should load"); - let callback_method: String = query_scalar( - "SELECT payment_method FROM payment_callbacks WHERE id = 'legacy-method-callback'", - ) - .fetch_one(&pool) - .await - .expect("normalized callback should load"); - assert_eq!(order_method, "epay"); - assert_eq!(callback_method, "epay"); - - query( - r#" -INSERT INTO payment_orders ( - id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd, - refundable_amount_usd, payment_method, gateway_order_id, status, created_at -) VALUES ('case-sensitive-order', 'case-sensitive-no', 'legacy-method-wallet', 'legacy-method-user', 1, 0, 0, 'epay', 'casesensitivetxn', 'pending', 1); -"#, - ) - .execute(&pool) - .await - .expect("case-distinct opaque gateway identifiers should remain distinct"); -} - -#[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:") + let Some(server) = ManagedPostgresServer::try_start() .await - .expect("sqlite pool should connect"); + .expect("postgres fixture should start") + else { + return; + }; + let pool = PgPool::connect(server.database_url()) + .await + .expect("postgres pool should connect"); query( r#" CREATE TABLE providers ( @@ -1821,7 +1052,7 @@ INSERT INTO provider_endpoints (id, provider_id, api_format, base_url, custom_pa .await .expect("endpoint fixture should insert"); - let migration = super::sqlite::MIGRATOR + let migration = POSTGRES_MIGRATOR .iter() .find(|migration| migration.version == 20260528000000) .expect("endpoint API root migration should be embedded"); @@ -2017,14 +1248,7 @@ INSERT INTO provider_endpoints (id, provider_id, api_format, base_url, custom_pa #[test] fn fresh_usage_schema_projects_upstream_stream_mode_for_all_drivers() { - let mysql_baseline = - include_str!("../../../../adapters/mysql/migrations/20260403000000_baseline.sql"); - let sqlite_baseline = - include_str!("../../../../adapters/sqlite/migrations/20260403000000_baseline.sql"); - assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("upstream_is_stream boolean")); - assert!(mysql_baseline.contains("upstream_is_stream TINYINT(1)")); - assert!(sqlite_baseline.contains("upstream_is_stream INTEGER")); } #[tokio::test] @@ -2592,93 +1816,6 @@ fn pending_migrations_from_applied_only_returns_post_snapshot_migrations() { .all(|migration| migration.version > EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION)); } -#[tokio::test] -async fn sqlite_migrations_create_core_config_tables() { - let pool = sqlx::sqlite::SqlitePoolOptions::new() - .max_connections(1) - .connect("sqlite::memory:") - .await - .expect("sqlite in-memory pool should connect"); - - let pending = super::prepare_sqlite_database_for_startup(&pool) - .await - .expect("sqlite startup preparation should inspect pending migrations"); - assert!( - !pending.is_empty(), - "fresh sqlite databases should report pending migrations before migration" - ); - - super::run_sqlite_migrations(&pool) - .await - .expect("sqlite migrations should run"); - - let pending = super::prepare_sqlite_database_for_startup(&pool) - .await - .expect("sqlite startup preparation should inspect applied migrations"); - assert!( - pending.is_empty(), - "sqlite startup preparation should report no pending migrations after migration" - ); - - for table_name in [ - "users", - "user_preferences", - "user_sessions", - "api_keys", - "management_tokens", - "billing_rules", - "dimension_collectors", - "providers", - "provider_api_keys", - "provider_endpoints", - "models", - "global_models", - "system_configs", - "auth_modules", - "oauth_providers", - "proxy_nodes", - "wallets", - "wallet_transactions", - "wallet_daily_usage_ledgers", - "payment_orders", - "payment_callbacks", - "refund_requests", - "redeem_code_batches", - "redeem_codes", - ] { - let exists: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", - ) - .bind(table_name) - .fetch_one(&pool) - .await - .expect("sqlite_master query should succeed"); - assert_eq!(exists, 1, "missing sqlite table {table_name}"); - } - - let total_adjusted_exists: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM pragma_table_info('wallets') WHERE name = ?") - .bind("total_adjusted") - .fetch_one(&pool) - .await - .expect("sqlite wallet column query should succeed"); - assert_eq!( - total_adjusted_exists, 1, - "missing sqlite wallets.total_adjusted" - ); - - let upstream_is_stream_exists: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM pragma_table_info('usage') WHERE name = ?") - .bind("upstream_is_stream") - .fetch_one(&pool) - .await - .expect("sqlite usage column query should succeed"); - assert_eq!( - upstream_is_stream_exists, 1, - "missing sqlite usage.upstream_is_stream" - ); -} - #[tokio::test] async fn postgres_migrations_create_core_config_tables_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_POSTGRES_URL") @@ -2952,112 +2089,6 @@ WHERE id = 'metadata-migration-key' .expect("provider migration fixture should clean up"); } -#[tokio::test] -async fn mysql_migrations_create_core_config_tables_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 migration smoke 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"); - - super::run_mysql_migrations(&pool) - .await - .expect("mysql migrations should run"); - - let pending = super::prepare_mysql_database_for_startup(&pool) - .await - .expect("mysql startup preparation should inspect applied migrations"); - assert!( - pending.is_empty(), - "mysql startup preparation should report no pending migrations after migration" - ); - - for table_name in [ - "users", - "user_preferences", - "user_sessions", - "api_keys", - "management_tokens", - "billing_rules", - "dimension_collectors", - "providers", - "provider_api_keys", - "provider_endpoints", - "models", - "global_models", - "system_configs", - "auth_modules", - "oauth_providers", - "proxy_nodes", - "usage", - "usage_settlement_snapshots", - "wallets", - "wallet_transactions", - "wallet_daily_usage_ledgers", - "payment_orders", - "payment_callbacks", - "refund_requests", - "redeem_code_batches", - "redeem_codes", - ] { - let exists: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM information_schema.tables -WHERE table_schema = DATABASE() - AND table_name = ? -"#, - ) - .bind(table_name) - .fetch_one(&pool) - .await - .expect("mysql information_schema query should succeed"); - assert_eq!(exists, 1, "missing mysql table {table_name}"); - } - - let total_adjusted_exists: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM information_schema.columns -WHERE table_schema = DATABASE() - AND table_name = 'wallets' - AND column_name = 'total_adjusted' -"#, - ) - .fetch_one(&pool) - .await - .expect("mysql information_schema column query should succeed"); - assert_eq!( - total_adjusted_exists, 1, - "missing mysql wallets.total_adjusted" - ); - - let upstream_is_stream_exists: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM information_schema.columns -WHERE table_schema = DATABASE() - AND table_name = 'usage' - AND column_name = 'upstream_is_stream' -"#, - ) - .fetch_one(&pool) - .await - .expect("mysql usage column query should succeed"); - assert_eq!( - upstream_is_stream_exists, 1, - "missing mysql usage.upstream_is_stream" - ); -} - #[tokio::test] async fn prepare_database_for_startup_bootstraps_clean_database() { let Some(server) = ManagedPostgresServer::try_start() diff --git a/crates/aether-data/runtime/src/repository/announcements/mod.rs b/crates/aether-data/runtime/src/repository/announcements/mod.rs index ff54aa7d9..3f9c18689 100644 --- a/crates/aether-data/runtime/src/repository/announcements/mod.rs +++ b/crates/aether-data/runtime/src/repository/announcements/mod.rs @@ -4,10 +4,6 @@ pub use aether_data_contracts::repository::announcements::{ AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository, CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlAnnouncementRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxAnnouncementReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteAnnouncementRepository; pub use memory::InMemoryAnnouncementReadRepository; diff --git a/crates/aether-data/runtime/src/repository/audit.rs b/crates/aether-data/runtime/src/repository/audit.rs index 4bfd491ab..f0a87c409 100644 --- a/crates/aether-data/runtime/src/repository/audit.rs +++ b/crates/aether-data/runtime/src/repository/audit.rs @@ -8,10 +8,6 @@ pub use aether_data_contracts::repository::audit::{ StoredAdminAuditLogPage, StoredSuspiciousActivity, StoredUserAuditLog, StoredUserAuditLogPage, SUSPICIOUS_EVENT_TYPES, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlAuditLogReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::PostgresAuditLogReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteAuditLogReadRepository; pub use types::{read_request_audit_bundle, RequestAuditBundle, RequestAuditReader}; diff --git a/crates/aether-data/runtime/src/repository/auth/mod.rs b/crates/aether-data/runtime/src/repository/auth/mod.rs index 46c9081c2..3bb4e469a 100644 --- a/crates/aether-data/runtime/src/repository/auth/mod.rs +++ b/crates/aether-data/runtime/src/repository/auth/mod.rs @@ -9,10 +9,6 @@ pub use aether_data_contracts::repository::auth::{ StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlAuthApiKeyReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxAuthApiKeySnapshotReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteAuthApiKeyReadRepository; pub use memory::InMemoryAuthApiKeySnapshotRepository; diff --git a/crates/aether-data/runtime/src/repository/auth_modules/mod.rs b/crates/aether-data/runtime/src/repository/auth_modules/mod.rs index 22d053007..3a0d37a2b 100644 --- a/crates/aether-data/runtime/src/repository/auth_modules/mod.rs +++ b/crates/aether-data/runtime/src/repository/auth_modules/mod.rs @@ -4,10 +4,6 @@ pub use aether_data_contracts::repository::auth_modules::{ AuthModuleReadRepository, AuthModuleWriteRepository, CompareAndSwapLdapConfigResult, LdapBindPasswordUpdate, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::{MysqlAuthModuleReadRepository, MysqlAuthModuleRepository}; #[cfg(feature = "postgres")] pub use aether_data_postgres::{SqlxAuthModuleReadRepository, SqlxAuthModuleRepository}; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::{SqliteAuthModuleReadRepository, SqliteAuthModuleRepository}; pub use memory::InMemoryAuthModuleReadRepository; diff --git a/crates/aether-data/runtime/src/repository/background_tasks/mod.rs b/crates/aether-data/runtime/src/repository/background_tasks/mod.rs index fcce2ae41..b70200a53 100644 --- a/crates/aether-data/runtime/src/repository/background_tasks/mod.rs +++ b/crates/aether-data/runtime/src/repository/background_tasks/mod.rs @@ -8,10 +8,6 @@ pub(crate) use aether_data_contracts::repository::background_tasks::{ StoredBackgroundTaskRunPage, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlBackgroundTaskRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxBackgroundTaskRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteBackgroundTaskRepository; pub use memory::InMemoryBackgroundTaskRepository; diff --git a/crates/aether-data/runtime/src/repository/billing/mod.rs b/crates/aether-data/runtime/src/repository/billing/mod.rs index 223ed8dc2..5618a9530 100644 --- a/crates/aether-data/runtime/src/repository/billing/mod.rs +++ b/crates/aether-data/runtime/src/repository/billing/mod.rs @@ -1,9 +1,5 @@ mod memory; pub use aether_data_contracts::repository::billing::*; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlBillingReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxBillingReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteBillingReadRepository; pub use memory::InMemoryBillingReadRepository; diff --git a/crates/aether-data/runtime/src/repository/candidate_selection/mod.rs b/crates/aether-data/runtime/src/repository/candidate_selection/mod.rs index 68428a5c9..19f4a2daf 100644 --- a/crates/aether-data/runtime/src/repository/candidate_selection/mod.rs +++ b/crates/aether-data/runtime/src/repository/candidate_selection/mod.rs @@ -8,10 +8,6 @@ pub(crate) use aether_data_contracts::repository::candidate_selection::{ StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlMinimalCandidateSelectionReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxMinimalCandidateSelectionReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteMinimalCandidateSelectionReadRepository; pub use memory::InMemoryMinimalCandidateSelectionReadRepository; diff --git a/crates/aether-data/runtime/src/repository/candidates/mod.rs b/crates/aether-data/runtime/src/repository/candidates/mod.rs index 9768e55ff..6d5dcb01a 100644 --- a/crates/aether-data/runtime/src/repository/candidates/mod.rs +++ b/crates/aether-data/runtime/src/repository/candidates/mod.rs @@ -9,10 +9,6 @@ pub(crate) use aether_data_contracts::repository::candidates::{ RequestCandidateTrace, RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlRequestCandidateRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxRequestCandidateReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteRequestCandidateRepository; pub use memory::InMemoryRequestCandidateRepository; diff --git a/crates/aether-data/runtime/src/repository/gemini_file_mappings/mod.rs b/crates/aether-data/runtime/src/repository/gemini_file_mappings/mod.rs index e6b27013c..eb3645f0a 100644 --- a/crates/aether-data/runtime/src/repository/gemini_file_mappings/mod.rs +++ b/crates/aether-data/runtime/src/repository/gemini_file_mappings/mod.rs @@ -1,16 +1,8 @@ pub mod memory; -#[cfg(feature = "mysql")] -pub mod mysql { - pub use aether_data_mysql::MysqlGeminiFileMappingRepository; -} #[cfg(feature = "postgres")] pub mod postgres { pub use aether_data_postgres::SqlxGeminiFileMappingRepository; } -#[cfg(feature = "sqlite")] -pub mod sqlite { - pub use aether_data_sqlite::SqliteGeminiFileMappingRepository; -} pub mod types { pub use aether_data_contracts::repository::gemini_file_mappings::*; } @@ -20,10 +12,6 @@ pub use aether_data_contracts::repository::gemini_file_mappings::{ GeminiFileMappingRepository, GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping, StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlGeminiFileMappingRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxGeminiFileMappingRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteGeminiFileMappingRepository; pub use memory::InMemoryGeminiFileMappingRepository; diff --git a/crates/aether-data/runtime/src/repository/global_models/mod.rs b/crates/aether-data/runtime/src/repository/global_models/mod.rs index abbf7d0a9..c79555d83 100644 --- a/crates/aether-data/runtime/src/repository/global_models/mod.rs +++ b/crates/aether-data/runtime/src/repository/global_models/mod.rs @@ -10,10 +10,6 @@ pub(crate) use aether_data_contracts::repository::global_models::{ StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlGlobalModelReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxGlobalModelReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteGlobalModelReadRepository; pub use memory::InMemoryGlobalModelReadRepository; diff --git a/crates/aether-data/runtime/src/repository/management_tokens/mod.rs b/crates/aether-data/runtime/src/repository/management_tokens/mod.rs index 3b0442b55..a22a014ef 100644 --- a/crates/aether-data/runtime/src/repository/management_tokens/mod.rs +++ b/crates/aether-data/runtime/src/repository/management_tokens/mod.rs @@ -6,10 +6,6 @@ pub use aether_data_contracts::repository::management_tokens::{ StoredManagementToken, StoredManagementTokenListPage, StoredManagementTokenUserSummary, StoredManagementTokenWithUser, UpdateManagementTokenRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlManagementTokenRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxManagementTokenRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteManagementTokenRepository; pub use memory::InMemoryManagementTokenRepository; diff --git a/crates/aether-data/runtime/src/repository/oauth_providers/mod.rs b/crates/aether-data/runtime/src/repository/oauth_providers/mod.rs index f6258518b..2f4f042ee 100644 --- a/crates/aether-data/runtime/src/repository/oauth_providers/mod.rs +++ b/crates/aether-data/runtime/src/repository/oauth_providers/mod.rs @@ -6,10 +6,6 @@ pub use aether_data_contracts::repository::oauth_providers::{ OAuthProviderRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig, UpsertOAuthProviderConfigOutcome, UpsertOAuthProviderConfigRecord, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlOAuthProviderRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxOAuthProviderRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteOAuthProviderRepository; pub use memory::InMemoryOAuthProviderRepository; diff --git a/crates/aether-data/runtime/src/repository/pool_scores/mod.rs b/crates/aether-data/runtime/src/repository/pool_scores/mod.rs index 3bb3ee0fd..4811b1396 100644 --- a/crates/aether-data/runtime/src/repository/pool_scores/mod.rs +++ b/crates/aether-data/runtime/src/repository/pool_scores/mod.rs @@ -2,10 +2,6 @@ pub use aether_data_contracts::repository::pool_scores::*; mod memory; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlPoolMemberScoreRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::PostgresPoolMemberScoreRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqlitePoolMemberScoreRepository; pub use memory::InMemoryPoolMemberScoreRepository; diff --git a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs index 5de56b4e8..0eed80c25 100644 --- a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs @@ -15,10 +15,6 @@ pub(crate) use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlProviderCatalogReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxProviderCatalogReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteProviderCatalogReadRepository; pub use memory::InMemoryProviderCatalogReadRepository; diff --git a/crates/aether-data/runtime/src/repository/proxy_nodes/mod.rs b/crates/aether-data/runtime/src/repository/proxy_nodes/mod.rs index 44ada043d..d998c515b 100644 --- a/crates/aether-data/runtime/src/repository/proxy_nodes/mod.rs +++ b/crates/aether-data/runtime/src/repository/proxy_nodes/mod.rs @@ -1,12 +1,8 @@ mod memory; pub use aether_data_contracts::repository::proxy_nodes::*; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlProxyNodeReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxProxyNodeRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteProxyNodeReadRepository; pub use memory::InMemoryProxyNodeRepository; pub fn log_reported_tunnel_error_event( diff --git a/crates/aether-data/runtime/src/repository/quota/mod.rs b/crates/aether-data/runtime/src/repository/quota/mod.rs index eb46f31df..9ad47a602 100644 --- a/crates/aether-data/runtime/src/repository/quota/mod.rs +++ b/crates/aether-data/runtime/src/repository/quota/mod.rs @@ -5,10 +5,6 @@ pub(crate) use aether_data_contracts::repository::quota::{ ProviderQuotaReadRepository, ProviderQuotaRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlProviderQuotaRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxProviderQuotaRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteProviderQuotaRepository; pub use memory::InMemoryProviderQuotaRepository; diff --git a/crates/aether-data/runtime/src/repository/routing_profiles/mod.rs b/crates/aether-data/runtime/src/repository/routing_profiles/mod.rs index ada2f523a..c8f5f035a 100644 --- a/crates/aether-data/runtime/src/repository/routing_profiles/mod.rs +++ b/crates/aether-data/runtime/src/repository/routing_profiles/mod.rs @@ -8,9 +8,5 @@ pub(crate) use aether_data_contracts::repository::routing_profiles::{ }; pub use memory::InMemoryRoutingGroupRepository; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlRoutingGroupRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::PostgresRoutingGroupRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteRoutingGroupRepository; diff --git a/crates/aether-data/runtime/src/repository/settlement/mod.rs b/crates/aether-data/runtime/src/repository/settlement/mod.rs index ad3c377fa..72a41c092 100644 --- a/crates/aether-data/runtime/src/repository/settlement/mod.rs +++ b/crates/aether-data/runtime/src/repository/settlement/mod.rs @@ -1,11 +1,7 @@ mod memory; pub use aether_data_contracts::repository::settlement::*; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlSettlementRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxSettlementRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteSettlementRepository; pub use memory::InMemorySettlementRepository; #[cfg(test)] diff --git a/crates/aether-data/runtime/src/repository/usage/mod.rs b/crates/aether-data/runtime/src/repository/usage/mod.rs index 1a0a441bb..abd7d9128 100644 --- a/crates/aether-data/runtime/src/repository/usage/mod.rs +++ b/crates/aether-data/runtime/src/repository/usage/mod.rs @@ -1,6 +1,4 @@ mod memory; -#[cfg(feature = "mysql")] -mod mysql; #[allow(unused_imports)] pub(crate) use aether_data_contracts::repository::usage::{ @@ -43,11 +41,7 @@ pub mod cleanup { } #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxUsageReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::{SqliteUsageReadRepository, SqliteUsageWriteRepository}; pub use memory::InMemoryUsageReadRepository; -#[cfg(feature = "mysql")] -pub use mysql::{MysqlUsageReadRepository, MysqlUsageWriteRepository}; #[cfg(test)] mod tests { diff --git a/crates/aether-data/runtime/src/repository/usage/mysql.rs b/crates/aether-data/runtime/src/repository/usage/mysql.rs deleted file mode 100644 index 74bd34c29..000000000 --- a/crates/aether-data/runtime/src/repository/usage/mysql.rs +++ /dev/null @@ -1,457 +0,0 @@ -use std::collections::BTreeMap; - -use aether_data_contracts::repository::usage::{ - 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; -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: MysqlUsageStorage, -} - -impl MysqlUsageReadRepository { - pub fn new(pool: MysqlPool) -> Self { - Self { - storage: MysqlUsageStorage::new(pool), - } - } - - async fn materialize_read_model( - &self, - filter: MysqlUsageReadFilter, - ) -> Result { - Ok(InMemoryUsageReadRepository::seed( - self.storage.load_usage_records_in_range(&filter).await?, - )) - } - - 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, - 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( - &self, - user_ids: &[String], - ) -> Result, DataLayerError> { - self.storage - .summarize_usage_totals_by_user_ids(user_ids) - .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) - .await? - { - return Ok(summary); - } - 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 - } - - async fn list_dashboard_daily_breakdown( - &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) - .await?; - if !rows.is_empty() { - return Ok(rows); - } - 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 - } -} - -#[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")); - } -} diff --git a/crates/aether-data/runtime/src/repository/users/mod.rs b/crates/aether-data/runtime/src/repository/users/mod.rs index 11dcb77b3..07007b3f5 100644 --- a/crates/aether-data/runtime/src/repository/users/mod.rs +++ b/crates/aether-data/runtime/src/repository/users/mod.rs @@ -11,10 +11,6 @@ pub use aether_data_contracts::repository::users::{ UserExportSortOrder, UserExportSummary, UserReadRepository, LAST_ACTIVE_ADMIN_DELETE_DENIED, LAST_ACTIVE_ADMIN_UPDATE_DENIED, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlUserReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxUserReadRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteUserReadRepository; pub use memory::InMemoryUserReadRepository; diff --git a/crates/aether-data/runtime/src/repository/video_tasks/mod.rs b/crates/aether-data/runtime/src/repository/video_tasks/mod.rs index 58116096e..1283886fc 100644 --- a/crates/aether-data/runtime/src/repository/video_tasks/mod.rs +++ b/crates/aether-data/runtime/src/repository/video_tasks/mod.rs @@ -6,10 +6,6 @@ pub(crate) use aether_data_contracts::repository::video_tasks::{ VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskRepository, VideoTaskStatus, VideoTaskStatusCount, VideoTaskWriteRepository, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlVideoTaskRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::{SqlxVideoTaskReadRepository, SqlxVideoTaskRepository}; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteVideoTaskRepository; pub use memory::InMemoryVideoTaskRepository; diff --git a/crates/aether-data/runtime/src/repository/wallet/mod.rs b/crates/aether-data/runtime/src/repository/wallet/mod.rs index 1706deca4..8029a4e77 100644 --- a/crates/aether-data/runtime/src/repository/wallet/mod.rs +++ b/crates/aether-data/runtime/src/repository/wallet/mod.rs @@ -42,10 +42,6 @@ pub use aether_data_contracts::repository::wallet::{ WalletReadRepository, WalletReadSeed, WalletReadSnapshot, WalletRepository, WalletWriteRepository, WALLET_RECHARGE_CHECKOUT_CLAIM_LEASE_SECS, }; -#[cfg(feature = "mysql")] -pub use aether_data_mysql::MysqlWalletReadRepository; #[cfg(feature = "postgres")] pub use aether_data_postgres::SqlxWalletRepository; -#[cfg(feature = "sqlite")] -pub use aether_data_sqlite::SqliteWalletReadRepository; pub use memory::InMemoryWalletRepository; diff --git a/crates/aether-data/runtime/tests/public_entrypoints.rs b/crates/aether-data/runtime/tests/public_entrypoints.rs index e1d362bda..cbc880a26 100644 --- a/crates/aether-data/runtime/tests/public_entrypoints.rs +++ b/crates/aether-data/runtime/tests/public_entrypoints.rs @@ -4,7 +4,7 @@ use std::{ }; const OLD_ENTRYPOINTS: &[&str] = &[ - "backends", "backfill", "export", "migrate", "mysql", "postgres", "redis", "sqlite", + "backends", "backfill", "export", "migrate", "postgres", "redis", ]; const SELF_TEST_PATH: &str = "crates/aether-data/runtime/tests/public_entrypoints.rs"; @@ -212,7 +212,7 @@ fn relative_path(path: &Path, workspace_root: &Path) -> PathBuf { fn grouped_import_scanner_allows_nested_new_paths() { let source = r#" use aether_data::{ - driver::{postgres::PostgresPool, mysql::MySqlPool}, + driver::{postgres::PostgresPool}, lifecycle::{backfill::PendingBackfillInfo, migrate::PendingMigrationInfo}, }; "#; diff --git a/crates/aether-data/schema/src/bin/aether-schema.rs b/crates/aether-data/schema/src/bin/aether-schema.rs index 80b751a7c..afb1dbbea 100644 --- a/crates/aether-data/schema/src/bin/aether-schema.rs +++ b/crates/aether-data/schema/src/bin/aether-schema.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use aether_data_schema::dialect::{mysql, postgres, sqlite}; +use aether_data_schema::dialect::postgres; use aether_data_schema::{check_generated_dir, generate_loaded_to_dir, load_schema_sources}; use clap::{Parser, Subcommand, ValueEnum}; @@ -39,8 +39,6 @@ enum Command { #[derive(Debug, Clone, Copy, ValueEnum)] enum Driver { Postgres, - Mysql, - Sqlite, } fn main() -> Result<(), Box> { @@ -66,8 +64,6 @@ fn main() -> Result<(), Box> { let schema = load_schema_sources(schema_dir)?.schema; let output = match driver { Driver::Postgres => postgres::emit_schema(&schema), - Driver::Mysql => mysql::emit_schema(&schema), - Driver::Sqlite => sqlite::emit_schema(&schema), }; print!("{output}"); } diff --git a/crates/aether-data/schema/src/dialect/mod.rs b/crates/aether-data/schema/src/dialect/mod.rs index 388a85b93..e91dcca8d 100644 --- a/crates/aether-data/schema/src/dialect/mod.rs +++ b/crates/aether-data/schema/src/dialect/mod.rs @@ -1,22 +1,16 @@ use crate::{Column, DefaultValue, DriverColumnOverride, LogicalType, ReferentialAction}; -pub mod mysql; pub mod postgres; -pub mod sqlite; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Dialect { Postgres, - Mysql, - Sqlite, } impl Dialect { pub const fn as_str(self) -> &'static str { match self { Self::Postgres => "postgres", - Self::Mysql => "mysql", - Self::Sqlite => "sqlite", } } } @@ -29,25 +23,11 @@ fn default_sql(default: &DefaultValue) -> String { match default { DefaultValue::String(value) => quote_string(value), DefaultValue::Integer(value) => value.to_string(), - DefaultValue::Bool(value) => { - if *value { - "1".to_string() - } else { - "0".to_string() - } - } + DefaultValue::Bool(value) => value.to_string(), DefaultValue::Raw { raw } => raw.clone(), } } -fn default_sql_bool_keywords(default: &DefaultValue) -> String { - match default { - DefaultValue::Bool(true) => "true".to_string(), - DefaultValue::Bool(false) => "false".to_string(), - _ => default_sql(default), - } -} - fn column_default<'a>( column: &'a Column, override_: Option<&'a DriverColumnOverride>, @@ -94,50 +74,6 @@ fn postgres_type(column: &Column) -> String { } } -fn mysql_type(column: &Column) -> String { - let override_ = column.driver.mysql.as_ref(); - if let Some(sql_type) = override_type(override_) { - return sql_type.to_string(); - } - match column.logical_type { - LogicalType::TextId | LogicalType::Text => match column.length { - Some(length) => format!("VARCHAR({length})"), - None => "TEXT".to_string(), - }, - LogicalType::LongText => "LONGTEXT".to_string(), - LogicalType::Bool => "TINYINT(1)".to_string(), - LogicalType::Int32 => "INT".to_string(), - LogicalType::Int64 | LogicalType::UnixSeconds | LogicalType::UnixMillis => { - "BIGINT".to_string() - } - LogicalType::Float64 => "DOUBLE".to_string(), - LogicalType::DecimalMoney => "DECIMAL(20,8)".to_string(), - LogicalType::Timestamp => "BIGINT".to_string(), - LogicalType::Json => "JSON".to_string(), - LogicalType::Bytes => "LONGBLOB".to_string(), - } -} - -fn sqlite_type(column: &Column) -> String { - let override_ = column.driver.sqlite.as_ref(); - if let Some(sql_type) = override_type(override_) { - return sql_type.to_string(); - } - match column.logical_type { - LogicalType::TextId | LogicalType::Text | LogicalType::LongText | LogicalType::Json => { - "TEXT".to_string() - } - LogicalType::Bool - | LogicalType::Int32 - | LogicalType::Int64 - | LogicalType::UnixSeconds - | LogicalType::UnixMillis - | LogicalType::Timestamp => "INTEGER".to_string(), - LogicalType::Float64 | LogicalType::DecimalMoney => "REAL".to_string(), - LogicalType::Bytes => "BLOB".to_string(), - } -} - fn referential_action_sql(action: &ReferentialAction) -> &'static str { match action { ReferentialAction::Cascade => "CASCADE", diff --git a/crates/aether-data/schema/src/dialect/mysql.rs b/crates/aether-data/schema/src/dialect/mysql.rs deleted file mode 100644 index 82bcd0507..000000000 --- a/crates/aether-data/schema/src/dialect/mysql.rs +++ /dev/null @@ -1,117 +0,0 @@ -use crate::dialect::{ - column_default, column_nullable, default_sql, mysql_type, referential_action_sql, -}; -use crate::LogicalSchema; - -pub fn emit_schema(schema: &LogicalSchema) -> String { - emit_named_schema(schema, &schema.ordered_table_names()) -} - -pub fn emit_named_schema(schema: &LogicalSchema, table_names: &[String]) -> String { - let mut out = String::new(); - for table_name in table_names { - let table = schema - .tables - .get(table_name) - .expect("named schema table should exist"); - let mut definitions = Vec::new(); - for column in &table.columns { - let mut definition = format!(" `{}` {}", column.name, mysql_type(column)); - if !column_nullable(column, column.driver.mysql.as_ref()) { - definition.push_str(" NOT NULL"); - } - if column.auto_increment { - definition.push_str(" AUTO_INCREMENT"); - } - if let Some(default) = column_default(column, column.driver.mysql.as_ref()) { - definition.push_str(" DEFAULT "); - definition.push_str(&default_sql(default)); - } - definitions.push(definition); - } - if !table.primary_key.is_empty() { - definitions.push(format!( - " PRIMARY KEY ({})", - table - .primary_key - .iter() - .map(|column| format!("`{column}`")) - .collect::>() - .join(", ") - )); - } - for unique in &table.uniques { - definitions.push(format!( - " UNIQUE KEY {} ({})", - unique.name, - unique - .columns - .iter() - .map(|column| format!("`{column}`")) - .collect::>() - .join(", ") - )); - } - for index in &table.indexes { - let unique = if index.unique { "UNIQUE " } else { "" }; - definitions.push(format!( - " {unique}KEY {} ({})", - index.name, - index - .columns - .iter() - .map(|column| format!("`{column}`")) - .collect::>() - .join(", ") - )); - } - for foreign_key in &table.foreign_keys { - let mut definition = format!( - " CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({})", - foreign_key.name, - foreign_key - .columns - .iter() - .map(|column| format!("`{column}`")) - .collect::>() - .join(", "), - foreign_key.references_table, - foreign_key - .references_columns - .iter() - .map(|column| format!("`{column}`")) - .collect::>() - .join(", ") - ); - if let Some(action) = &foreign_key.on_delete { - definition.push_str(" ON DELETE "); - definition.push_str(referential_action_sql(action)); - } - definitions.push(definition); - } - - out.push_str(&format!( - "CREATE TABLE IF NOT EXISTS {} (\n", - quote_identifier_if_needed(table_name) - )); - out.push_str(&definitions.join(",\n")); - out.push_str("\n);\n\n"); - } - out -} - -fn quote_identifier_if_needed(identifier: &str) -> String { - if needs_quoting(identifier) { - quote_identifier(identifier) - } else { - identifier.to_string() - } -} - -fn needs_quoting(identifier: &str) -> bool { - matches!(identifier, "date" | "usage") -} - -fn quote_identifier(identifier: &str) -> String { - format!("`{}`", identifier.replace('`', "``")) -} diff --git a/crates/aether-data/schema/src/dialect/postgres.rs b/crates/aether-data/schema/src/dialect/postgres.rs index 6eed85604..5d490c2f6 100644 --- a/crates/aether-data/schema/src/dialect/postgres.rs +++ b/crates/aether-data/schema/src/dialect/postgres.rs @@ -1,6 +1,5 @@ use crate::dialect::{ - column_default, column_nullable, default_sql_bool_keywords, postgres_type, - referential_action_sql, + column_default, column_nullable, default_sql, postgres_type, referential_action_sql, }; use crate::LogicalSchema; @@ -30,7 +29,7 @@ pub fn emit_named_schema(schema: &LogicalSchema, table_names: &[String]) -> Stri out.push_str(&postgres_type(column)); if let Some(default) = column_default(column, column.driver.postgres.as_ref()) { out.push_str(" DEFAULT "); - out.push_str(&default_sql_bool_keywords(default)); + out.push_str(&default_sql(default)); } if !column_nullable(column, column.driver.postgres.as_ref()) { out.push_str(" NOT NULL"); diff --git a/crates/aether-data/schema/src/dialect/sqlite.rs b/crates/aether-data/schema/src/dialect/sqlite.rs deleted file mode 100644 index 8e64d8694..000000000 --- a/crates/aether-data/schema/src/dialect/sqlite.rs +++ /dev/null @@ -1,92 +0,0 @@ -use crate::dialect::{ - column_default, column_nullable, default_sql, referential_action_sql, sqlite_type, -}; -use crate::LogicalSchema; - -pub fn emit_schema(schema: &LogicalSchema) -> String { - emit_named_schema(schema, &schema.ordered_table_names()) -} - -pub fn emit_named_schema(schema: &LogicalSchema, table_names: &[String]) -> String { - let mut out = String::new(); - for table_name in table_names { - let table = schema - .tables - .get(table_name) - .expect("named schema table should exist"); - let mut definitions = Vec::new(); - for column in &table.columns { - let mut definition = format!(" {} {}", column.name, sqlite_type(column)); - if table.primary_key.len() == 1 && table.primary_key[0] == column.name { - definition.push_str(" PRIMARY KEY"); - if column.auto_increment { - definition.push_str(" AUTOINCREMENT"); - } - } - if !column.auto_increment && !column_nullable(column, column.driver.sqlite.as_ref()) { - definition.push_str(" NOT NULL"); - } - if let Some(default) = column_default(column, column.driver.sqlite.as_ref()) { - definition.push_str(" DEFAULT "); - definition.push_str(&default_sql(default)); - } - definitions.push(definition); - } - if table.primary_key.len() > 1 { - definitions.push(format!( - " PRIMARY KEY ({})", - table.primary_key.join(", ") - )); - } - for unique in &table.uniques { - definitions.push(format!(" UNIQUE ({})", unique.columns.join(", "))); - } - for foreign_key in &table.foreign_keys { - let mut definition = format!( - " CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({})", - foreign_key.name, - foreign_key.columns.join(", "), - foreign_key.references_table, - foreign_key.references_columns.join(", ") - ); - if let Some(action) = &foreign_key.on_delete { - definition.push_str(" ON DELETE "); - definition.push_str(referential_action_sql(action)); - } - definitions.push(definition); - } - - let quoted_table_name = quote_identifier_if_needed(table_name); - out.push_str(&format!( - "CREATE TABLE IF NOT EXISTS {quoted_table_name} (\n" - )); - out.push_str(&definitions.join(",\n")); - out.push_str("\n);\n"); - for index in &table.indexes { - let unique = if index.unique { "UNIQUE " } else { "" }; - out.push_str(&format!( - "CREATE {unique}INDEX IF NOT EXISTS {} ON {quoted_table_name} ({});\n", - index.name, - index.columns.join(", ") - )); - } - out.push('\n'); - } - out -} - -fn quote_identifier_if_needed(identifier: &str) -> String { - if needs_quoting(identifier) { - quote_identifier(identifier) - } else { - identifier.to_string() - } -} - -fn needs_quoting(identifier: &str) -> bool { - matches!(identifier, "date" | "usage") -} - -fn quote_identifier(identifier: &str) -> String { - format!("\"{}\"", identifier.replace('"', "\"\"")) -} diff --git a/crates/aether-data/schema/src/lib.rs b/crates/aether-data/schema/src/lib.rs index 639c65e5b..e95360842 100644 --- a/crates/aether-data/schema/src/lib.rs +++ b/crates/aether-data/schema/src/lib.rs @@ -139,10 +139,6 @@ pub enum DefaultValue { pub struct DriverColumnOverrides { #[serde(default)] pub postgres: Option, - #[serde(default)] - pub mysql: Option, - #[serde(default)] - pub sqlite: Option, } #[derive(Debug, Clone, serde::Deserialize, PartialEq, Eq)] @@ -430,12 +426,6 @@ pub fn generate_sources_to_dir( write_driver_sources(output_root, "postgres", sources, |tables| { dialect::postgres::emit_named_schema(schema, tables) })?; - write_driver_sources(output_root, "mysql", sources, |tables| { - dialect::mysql::emit_named_schema(schema, tables) - })?; - write_driver_sources(output_root, "sqlite", sources, |tables| { - dialect::sqlite::emit_named_schema(schema, tables) - })?; Ok(()) } @@ -449,12 +439,6 @@ pub fn check_generated_dir( check_driver_sources(output_root, "postgres", &loaded.sources, |tables| { dialect::postgres::emit_named_schema(&loaded.schema, tables) })?; - check_driver_sources(output_root, "mysql", &loaded.sources, |tables| { - dialect::mysql::emit_named_schema(&loaded.schema, tables) - })?; - check_driver_sources(output_root, "sqlite", &loaded.sources, |tables| { - dialect::sqlite::emit_named_schema(&loaded.schema, tables) - })?; Ok(()) } @@ -688,12 +672,7 @@ fn write_generated_readme(output_root: &Path) -> Result<(), SchemaError> { } fn assert_generated_root_files(output_root: &Path) -> Result<(), SchemaError> { - let expected = BTreeSet::from([ - "README.md".to_string(), - "mysql".to_string(), - "postgres".to_string(), - "sqlite".to_string(), - ]); + let expected = BTreeSet::from(["README.md".to_string(), "postgres".to_string()]); let actual = fs::read_dir(output_root) .map_err(|source| SchemaError::Read { path: output_root.to_path_buf(), @@ -774,7 +753,7 @@ fn generated_readme() -> String { bash crates/aether-data/runtime/schema/compose_schema.sh generate\n\ ```\n\n\ Runtime migrations are not loaded from this directory. The executable SQL lives under \ - `crates/aether-data/adapters/{postgres,mysql,sqlite}/migrations`, and the Postgres bootstrap snapshot \ + `crates/aether-data/adapters/postgres/migrations`, and the Postgres bootstrap snapshot \ is generated at build time from `crates/aether-data/runtime/schema/bootstrap/postgres` into the crate \ build output until a generated fragment is deliberately promoted into the driver-specific \ schema manifests.\n" @@ -808,7 +787,7 @@ fn write_generated(path: PathBuf, contents: &str) -> Result<(), SchemaError> { #[cfg(test)] mod tests { use super::*; - use crate::dialect::{mysql, postgres, sqlite}; + use crate::dialect::postgres; fn announcements_schema() -> LogicalSchema { LogicalSchema { @@ -879,42 +858,6 @@ mod tests { assert!(postgres_sql.contains( "CREATE INDEX IF NOT EXISTS announcements_is_active_idx ON public.announcements USING btree (is_active);" )); - - let mysql_sql = mysql::emit_schema(&schema); - assert!(mysql_sql.contains("`id` VARCHAR(64) NOT NULL")); - assert!(mysql_sql.contains("`is_active` TINYINT(1) NOT NULL DEFAULT 1")); - assert!(mysql_sql.contains("KEY announcements_is_active_idx (`is_active`)")); - - let sqlite_sql = sqlite::emit_schema(&schema); - assert!(sqlite_sql.contains("id TEXT PRIMARY KEY NOT NULL")); - assert!(sqlite_sql.contains("is_active INTEGER NOT NULL DEFAULT 1")); - assert!(sqlite_sql.contains( - "CREATE INDEX IF NOT EXISTS announcements_is_active_idx ON announcements (is_active);" - )); - } - - #[test] - fn mysql_column_type_override_can_select_binary_collation() { - let mut schema = announcements_schema(); - schema - .tables - .get_mut("announcements") - .expect("fixture table exists") - .columns - .first_mut() - .expect("fixture id column exists") - .driver - .mysql = Some(DriverColumnOverride { - sql_type: Some( - "VARCHAR(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin".to_string(), - ), - default: None, - nullable: None, - }); - - let mysql_sql = mysql::emit_schema(&schema); - assert!(mysql_sql - .contains("`id` VARCHAR(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin NOT NULL")); } #[test] @@ -1028,19 +971,8 @@ ALTER TABLE users ADD COLUMN ldap_dn VARCHAR(1024); .and_then(Path::parent) .expect("crate should live under workspace/crates/aether-data"); let schema_dir = workspace.join("crates/aether-data/runtime/schema/logical"); - let mut required_sql_paths = vec![workspace + let required_sql_paths = vec![workspace .join("crates/aether-data/adapters/postgres/migrations/20260403000000_baseline.sql")]; - for driver in ["mysql", "sqlite"] { - let driver_dir = - workspace.join(format!("crates/aether-data/adapters/{driver}/migrations")); - let mut paths = std::fs::read_dir(&driver_dir) - .unwrap_or_else(|err| panic!("failed to read {}: {err}", driver_dir.display())) - .map(|entry| entry.expect("migration entry should be readable").path()) - .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("sql")) - .collect::>(); - paths.sort(); - required_sql_paths.extend(paths); - } let loaded = load_schema_sources(schema_dir).expect("workspace logical schema should load"); check_required_tables(&loaded.schema, &required_sql_paths) diff --git a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs index 6e6551940..afc501151 100644 --- a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs +++ b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs @@ -1,6 +1,6 @@ //! Responses WebSocket end-to-end coverage. //! -//! Every test starts a protocol-aware mock upstream, seeds a throwaway SQLite +//! Every test starts a protocol-aware mock upstream, seeds a throwaway PostgreSQL //! store, mounts the real gateway router, and drives the public //! `/v1/responses` WebSocket the way a client would. //! @@ -9,7 +9,6 @@ //! `response.completed` is not evidence that the turn was ever accounted for — //! only the row is. -use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -28,7 +27,7 @@ use aether_data_contracts::repository::provider_catalog::{ }; use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageAuditListQuery}; use aether_gateway::{build_router_with_state, AppState, GatewayDataConfig, UsageRuntimeConfig}; -use aether_testkit::SpawnedServer; +use aether_testkit::{ManagedPostgresServer, SpawnedServer}; use axum::extract::ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade}; use axum::extract::State; use axum::http::{HeaderMap, Uri}; @@ -749,9 +748,9 @@ fn is_billed(audit: &StoredRequestUsageAudit) -> bool { // --------------------------------------------------------------------------- /// A live gateway wired to a mock Responses WebSocket upstream over a throwaway -/// SQLite store. +/// PostgreSQL store. struct Harness { - database: TemporarySqlite, + database: TemporaryPostgres, upstream: Arc, websocket_url: String, _upstream_server: SpawnedServer, @@ -856,7 +855,7 @@ impl Harness { let upstream_server = SpawnedServer::start(mock_upstream_router(Arc::clone(&upstream))).await?; - let database = TemporarySqlite::new(); + let database = TemporaryPostgres::new().await?; prepare_and_seed_database( &database.config, upstream_server.base_url(), @@ -961,7 +960,7 @@ impl Harness { /// Reads the persisted audit rows, oldest first. /// /// Opens its own handle per call rather than holding one for the lifetime of - /// the harness: the gateway keeps its own pool on the same SQLite file for + /// the harness: the gateway keeps its own pool on the same database for /// the whole test, and an idle second pool only adds contention. async fn usage_audits(&self) -> Result, BoxError> { let backends = DataBackends::from_config(DataLayerConfig::from_database( @@ -1359,24 +1358,20 @@ async fn send_mock_event(socket: &mut WebSocket, event: Value) -> Result<(), axu // Seeded data store // --------------------------------------------------------------------------- -struct TemporarySqlite { - directory: PathBuf, +struct TemporaryPostgres { + _server: ManagedPostgresServer, config: SqlDatabaseConfig, } -impl TemporarySqlite { - fn new() -> Self { - let directory = std::env::temp_dir().join(format!( - "aether-responses-ws-e2e-{}-{}", - std::process::id(), - uuid::Uuid::new_v4() - )); - let database_path = directory.join("aether.db"); - Self { - directory, +impl TemporaryPostgres { + async fn new() -> Result { + let server = ManagedPostgresServer::start().await?; + let database_url = server.database_url().to_string(); + Ok(Self { + _server: server, config: SqlDatabaseConfig { - driver: DatabaseDriver::Sqlite, - url: format!("sqlite://{}", database_path.display()), + driver: DatabaseDriver::Postgres, + url: database_url, pool: SqlPoolConfig { min_connections: 1, max_connections: 4, @@ -1387,13 +1382,7 @@ impl TemporarySqlite { require_ssl: false, }, }, - } - } -} - -impl Drop for TemporarySqlite { - fn drop(&mut self) { - let _ = std::fs::remove_dir_all(&self.directory); + }) } } @@ -1727,8 +1716,8 @@ async fn seed_weekly_request_limit( .id; let pool = backends - .sqlite() - .ok_or("SQLite backend unavailable")? + .postgres() + .ok_or("PostgreSQL backend unavailable")? .pool(); let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH)? @@ -1740,8 +1729,9 @@ INSERT INTO billing_plans ( id, title, description, price_amount, price_currency, duration_unit, duration_value, enabled, sort_order, max_active_per_user, purchase_limit_scope, entitlements_json, created_at, updated_at -) VALUES (?, 'WS weekly policy', NULL, 0, 'USD', 'month', 1, 1, 0, 1, - 'active_period', ?, ?, ?) +) VALUES ($1, 'WS weekly policy', NULL, 0, 'USD', 'month', 1, true, 0, 1, + 'active_period', $2::text::jsonb, + TO_TIMESTAMP($3::bigint::double precision), TO_TIMESTAMP($4::bigint::double precision)) "#, ) .bind("plan-responses-ws-weekly") @@ -1765,7 +1755,9 @@ INSERT INTO billing_plans ( INSERT INTO payment_orders ( id, order_no, wallet_id, user_id, amount_usd, pay_currency, status, payment_method, created_at, paid_at, credited_at, expires_at -) VALUES (?, ?, ?, ?, 0, 'USD', 'paid', 'test', ?, ?, ?, ?) +) VALUES ($1, $2, $3, $4, 0, 'USD', 'paid', 'test', + TO_TIMESTAMP($5::bigint::double precision), TO_TIMESTAMP($6::bigint::double precision), + TO_TIMESTAMP($7::bigint::double precision), TO_TIMESTAMP($8::bigint::double precision)) "#, ) .bind("order-responses-ws-weekly") @@ -1783,7 +1775,9 @@ INSERT INTO payment_orders ( INSERT INTO user_plan_entitlements ( id, user_id, plan_id, payment_order_id, status, starts_at, expires_at, entitlements_snapshot, created_at, updated_at -) VALUES (?, ?, ?, ?, 'active', ?, ?, ?, ?, ?) +) VALUES ($1, $2, $3, $4, 'active', + TO_TIMESTAMP($5::bigint::double precision), TO_TIMESTAMP($6::bigint::double precision), + $7::text::jsonb, TO_TIMESTAMP($8::bigint::double precision), TO_TIMESTAMP($9::bigint::double precision)) "#, ) .bind("entitlement-responses-ws-weekly") diff --git a/docker-compose.release-local.yml b/docker-compose.release-local.yml index e39eeaa4c..feed89459 100644 --- a/docker-compose.release-local.yml +++ b/docker-compose.release-local.yml @@ -4,12 +4,49 @@ # 启动: # docker compose -f docker-compose.release-local.yml up -d --build # -# 重置(包含 /opt/aether/current 与历史版本目录): +# 重置(删除测试 PostgreSQL 数据、/opt/aether/current 与历史版本目录): # docker compose -f docker-compose.release-local.yml down -v name: aether-release-local services: + postgres: + # OCI index digest for the official multi-architecture postgres:15.19 image. + image: postgres:15.19@sha256:5f72c7b5bd616308ccfd2e74d6be16fb06364e5eecbb815fe9dc6ab9761d2111 + container_name: aether-release-local-postgres + shm_size: ${POSTGRES_SHM_SIZE:-512mb} + environment: + POSTGRES_DB: aether + POSTGRES_USER: postgres + POSTGRES_PASSWORD: ${DB_PASSWORD:?set DB_PASSWORD in .env} + TZ: Asia/Shanghai + volumes: + - postgres_data:/var/lib/postgresql/data + command: > + postgres + -c idle_in_transaction_session_timeout=30000 + -c tcp_keepalives_idle=30 + -c tcp_keepalives_interval=10 + -c shared_preload_libraries=pg_stat_statements + -c pg_stat_statements.track=all + -c pg_stat_statements.max=10000 + -c shared_buffers=${POSTGRES_SHARED_BUFFERS:-1GB} + -c effective_cache_size=${POSTGRES_EFFECTIVE_CACHE_SIZE:-3GB} + -c work_mem=${POSTGRES_WORK_MEM:-16MB} + -c maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-256MB} + healthcheck: + test: [ "CMD-SHELL", "pg_isready -U postgres" ] + interval: 5s + timeout: 5s + retries: 5 + logging: + driver: local + options: + max-size: "100m" + max-file: "3" + compress: "true" + restart: unless-stopped + release-local-app: build: context: . @@ -23,8 +60,8 @@ services: - ${AETHER_ENV_FILE:-.env} environment: TZ: Asia/Shanghai - AETHER_DATABASE_DRIVER: sqlite - AETHER_DATABASE_URL: sqlite://./data/aether.db + AETHER_DATABASE_DRIVER: postgres + AETHER_DATABASE_URL: postgresql://postgres:${DB_PASSWORD:?set DB_PASSWORD in .env}@postgres:5432/aether AETHER_RUNTIME_BACKEND: memory AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY: single-node AETHER_GATEWAY_NODE_ROLE: all @@ -40,12 +77,15 @@ services: AETHER_BASE_DIR: /opt/aether AETHER_UPDATE_STRATEGY: self AETHER_GATEWAY_AUTO_PREPARE_DATABASE: ${AETHER_GATEWAY_AUTO_PREPARE_DATABASE:-true} + depends_on: + postgres: + condition: service_healthy ports: - "${AETHER_RELEASE_LOCAL_PORT:-18085}:${AETHER_RELEASE_LOCAL_PORT:-18085}" volumes: - - ./data-release-local:/app/data - aether_release_local_root:/opt/aether restart: unless-stopped volumes: + postgres_data: aether_release_local_root: diff --git a/docker-compose.single-node.yml b/docker-compose.single-node.yml index dab9b5e28..15c361eb9 100644 --- a/docker-compose.single-node.yml +++ b/docker-compose.single-node.yml @@ -1,4 +1,60 @@ +# Aether 单节点部署配置 - PostgreSQL + Redis +# 使用方法: docker compose up -d + services: + postgres: + # OCI index digest for the official multi-architecture postgres:15.19 image. + image: postgres:15.19@sha256:5f72c7b5bd616308ccfd2e74d6be16fb06364e5eecbb815fe9dc6ab9761d2111 + container_name: aether-postgres + shm_size: ${POSTGRES_SHM_SIZE:-512mb} + environment: + POSTGRES_DB: aether + POSTGRES_USER: postgres + POSTGRES_PASSWORD: ${DB_PASSWORD:?set DB_PASSWORD in .env} + TZ: Asia/Shanghai + volumes: + - postgres_data:/var/lib/postgresql/data + ports: + - "127.0.0.1:${DB_PORT:-5432}:5432" + command: > + postgres + -c idle_in_transaction_session_timeout=30000 + -c tcp_keepalives_idle=30 + -c tcp_keepalives_interval=10 + -c shared_preload_libraries=pg_stat_statements + -c pg_stat_statements.track=all + -c pg_stat_statements.max=10000 + -c shared_buffers=${POSTGRES_SHARED_BUFFERS:-1GB} + -c effective_cache_size=${POSTGRES_EFFECTIVE_CACHE_SIZE:-3GB} + -c work_mem=${POSTGRES_WORK_MEM:-16MB} + -c maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-256MB} + healthcheck: + test: [ "CMD-SHELL", "pg_isready -U postgres" ] + interval: 5s + timeout: 5s + retries: 5 + logging: + driver: local + options: + max-size: "100m" + max-file: "3" + compress: "true" + restart: unless-stopped + + redis: + # OCI index digest for the official multi-architecture redis:7.4.11-alpine image. + image: redis:7.4.11-alpine@sha256:ff02b58f971e7d7d156a1267e283fcbbeee91773b6aa36c49dac28ecfe28eadf + container_name: aether-redis + command: redis-server --dir /tmp --appendonly no --save "" --requirepass ${REDIS_PASSWORD:?set REDIS_PASSWORD in .env} --maxclients ${REDIS_MAXCLIENTS:-10000} + ports: + - "127.0.0.1:${REDIS_PORT:-6379}:6379" + healthcheck: + test: [ "CMD-SHELL", "redis-cli -a \"${REDIS_PASSWORD:?set REDIS_PASSWORD in .env}\" ping | grep -q PONG" ] + interval: 5s + timeout: 3s + retries: 5 + restart: unless-stopped + app: image: ${APP_IMAGE:-ghcr.io/fawney19/aether:latest} container_name: aether-app @@ -11,28 +67,33 @@ services: tmpfs: - /tmp:rw,nosuid,nodev,noexec,mode=1777 env_file: - - ${AETHER_ENV_FILE:-.env} + - .env environment: + AETHER_DATABASE_DRIVER: postgres + DATABASE_URL: postgresql://postgres:${DB_PASSWORD:?set DB_PASSWORD in .env}@postgres:5432/aether + REDIS_URL: redis://:${REDIS_PASSWORD:?set REDIS_PASSWORD in .env}@redis:6379/0 TZ: Asia/Shanghai AETHER_BASE_DIR: /opt/aether + AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY: single-node AETHER_UPDATE_STRATEGY: docker AETHER_DOCKER_UPDATE_COMMAND: ${AETHER_DOCKER_UPDATE_COMMAND:-./update.sh} - AETHER_DATABASE_DRIVER: sqlite - AETHER_DATABASE_URL: sqlite:///opt/aether/data/aether.db - AETHER_RUNTIME_BACKEND: memory - AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY: single-node - AETHER_GATEWAY_NODE_ROLE: all AETHER_LOG_DESTINATION: stdout AETHER_LOG_FORMAT: ${AETHER_LOG_FORMAT:-pretty} APP_PORT: ${APP_PORT:-8084} AETHER_GATEWAY_AUTO_PREPARE_DATABASE: ${AETHER_GATEWAY_AUTO_PREPARE_DATABASE:-true} + depends_on: + postgres: + condition: service_healthy + redis: + condition: service_healthy ports: - "${APP_PORT:-8084}:${APP_PORT:-8084}" - volumes: - - ./data:/opt/aether/data logging: driver: json-file options: max-size: "100m" max-file: "10" restart: unless-stopped + +volumes: + postgres_data: diff --git a/docker-compose.yml b/docker-compose.yml index 678975106..ff6a5037e 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -55,31 +55,6 @@ services: retries: 5 restart: unless-stopped - mysql: - # OCI index digest for the official multi-architecture mysql:8.0.46 image. - image: mysql:8.0.46@sha256:7dcddc01f13bab2f15cde676d44d01f61fc9f99fe7785e86196dfc07d358ae2b - container_name: aether-mysql - profiles: - - mysql - environment: - MYSQL_DATABASE: ${MYSQL_DATABASE:-aether} - MYSQL_USER: ${MYSQL_USER:-aether} - MYSQL_PASSWORD: ${MYSQL_PASSWORD:?set MYSQL_PASSWORD in .env} - MYSQL_ROOT_PASSWORD: ${MYSQL_ROOT_PASSWORD:?set MYSQL_ROOT_PASSWORD in .env} - TZ: Asia/Shanghai - volumes: - - mysql_data:/var/lib/mysql - healthcheck: - test: - [ - "CMD-SHELL", - "mysqladmin ping -h 127.0.0.1 -u$${MYSQL_USER} -p$${MYSQL_PASSWORD} --silent" - ] - interval: 5s - timeout: 5s - retries: 10 - restart: unless-stopped - app: image: ${APP_IMAGE:-ghcr.io/fawney19/aether:latest} container_name: aether-app @@ -94,6 +69,7 @@ services: env_file: - .env environment: + AETHER_DATABASE_DRIVER: postgres DATABASE_URL: postgresql://postgres:${DB_PASSWORD:?set DB_PASSWORD in .env}@postgres:5432/aether REDIS_URL: redis://:${REDIS_PASSWORD:?set REDIS_PASSWORD in .env}@redis:6379/0 TZ: Asia/Shanghai @@ -120,4 +96,3 @@ services: volumes: postgres_data: - mysql_data: diff --git a/generate_keys.sh b/generate_keys.sh index d8275ff81..a91d9feae 100755 --- a/generate_keys.sh +++ b/generate_keys.sh @@ -9,8 +9,6 @@ jwt_key=$(urlsafe_rand 32) encryption_key=$(urlsafe_rand 32) db_password=$(urlsafe_rand 32) redis_password=$(urlsafe_rand 32) -mysql_password=$(urlsafe_rand 32) -mysql_root_password=$(urlsafe_rand 32) cat </dev/null 2>&1 || return 1 - [[ "$(docker inspect --format '{{.State.Running}}' aether-app 2>/dev/null || true)" == "true" ]] -} - -prepare_compose_single_node_data_directory() { - local data_dir="${COMPOSE_DIR}/data" - local current_uid - local unsafe_path - - resolve_compose_env_container_identity "${COMPOSE_DIR}/.env" - ensure_directory "${data_dir}" 0700 - if compose_app_container_is_running; then - die "stop the existing aether-app container before migrating the Docker Compose SQLite data directory" - fi - - unsafe_path="$(find "${data_dir}" -type l -print -quit)" - [[ -z "${unsafe_path}" ]] \ - || die "Docker Compose data directory may not contain symbolic links: ${unsafe_path}" - unsafe_path="$(find "${data_dir}" ! -type d ! -type f -print -quit)" - [[ -z "${unsafe_path}" ]] \ - || die "Docker Compose data directory may contain only directories and regular files: ${unsafe_path}" - unsafe_path="$(find "${data_dir}" -type f -links +1 -print -quit)" - [[ -z "${unsafe_path}" ]] \ - || die "Docker Compose data directory may not contain multiply-linked files: ${unsafe_path}" - - current_uid="$(id -u)" - if [[ "${EUID}" -eq 0 ]]; then - chown -R -P "${COMPOSE_CONTAINER_UID}:${COMPOSE_CONTAINER_GID}" "${data_dir}" \ - || die "could not assign ${data_dir} to the non-root container identity" - elif [[ "${current_uid}" != "${COMPOSE_CONTAINER_UID}" ]]; then - die "${data_dir} must be owned by container uid ${COMPOSE_CONTAINER_UID}; rerun this installer as root to migrate existing SQLite data" - fi - - find "${data_dir}" -type d -exec chmod 0700 {} + \ - || die "could not secure Docker Compose data directories: ${data_dir}" - find "${data_dir}" -type f -exec chmod 0600 {} + \ - || die "could not secure Docker Compose data files: ${data_dir}" - [[ "$(stat_file_uid "${data_dir}")" == "${COMPOSE_CONTAINER_UID}" ]] \ - || die "Docker Compose data directory owner does not match AETHER_CONTAINER_UID=${COMPOSE_CONTAINER_UID}: ${data_dir}" - info "prepared ${data_dir} for container uid ${COMPOSE_CONTAINER_UID} gid ${COMPOSE_CONTAINER_GID}" -} - append_compose_log_env_defaults() { local output="$1" replace_or_append_env "${output}" "AETHER_LOG_DESTINATION" "${COMPOSE_LOG_DESTINATION_DEFAULT}" @@ -1862,27 +1809,14 @@ append_compose_log_env_defaults() { replace_or_append_env "${output}" "AETHER_LOG_MAX_FILES" "${COMPOSE_LOG_MAX_FILES_DEFAULT}" } -compose_log_env_block() { - cat < "${output}" </dev/null || true; rm -rf "${TEST_ROOT}"' EXIT COMPOSE_DIR="${TEST_ROOT}/compose" -mkdir -p "${COMPOSE_DIR}/data" +mkdir -p "${COMPOSE_DIR}" fake_bin="${TEST_ROOT}/fake-bin" mkdir -p "${fake_bin}" printf '#!/usr/bin/env bash\nprintf "false\\n"\n' >"${fake_bin}/docker" @@ -107,52 +88,6 @@ fixture_uid="$(id -u)" fixture_gid="$(id -g)" [[ "${fixture_uid}" != "0" ]] || fixture_uid="65532" [[ "${fixture_gid}" != "0" ]] || fixture_gid="65532" -printf 'AETHER_CONTAINER_UID=%s\nAETHER_CONTAINER_GID=%s\n' \ - "${fixture_uid}" "${fixture_gid}" >"${COMPOSE_DIR}/.env" -printf 'sqlite fixture\n' >"${COMPOSE_DIR}/data/aether.db" -chmod 0755 "${COMPOSE_DIR}/data/aether.db" -chmod 4755 "${COMPOSE_DIR}/data/aether.db" 2>/dev/null || true - -prepare_compose_single_node_data_directory -[[ "$(stat_file_mode "${COMPOSE_DIR}/data")" == "700" ]] \ - || fail_test "single-node data directory was not restricted to mode 0700" -[[ "$(stat_file_mode "${COMPOSE_DIR}/data/aether.db")" == "600" ]] \ - || fail_test "single-node SQLite file was not restricted to mode 0600" - -printf 'AETHER_CONTAINER_UID=0\nAETHER_CONTAINER_GID=%s\n' \ - "${fixture_gid}" >"${COMPOSE_DIR}/.env" -if (prepare_compose_single_node_data_directory) >/dev/null 2>&1; then - fail_test "root container uid was accepted" -fi - -printf 'AETHER_CONTAINER_UID=%s\nAETHER_CONTAINER_GID=%s\n' \ - "${fixture_uid}" "${fixture_gid}" >"${COMPOSE_DIR}/.env" -ln -s "${TEST_ROOT}" "${COMPOSE_DIR}/data/unsafe-link" -if (prepare_compose_single_node_data_directory) >/dev/null 2>&1; then - fail_test "symbolic link inside the managed SQLite directory was accepted" -fi - -rm -f "${COMPOSE_DIR}/data/unsafe-link" -mkfifo "${COMPOSE_DIR}/data/unsafe-fifo" -if (prepare_compose_single_node_data_directory) >/dev/null 2>&1; then - fail_test "FIFO inside the managed SQLite directory was accepted" -fi -rm -f "${COMPOSE_DIR}/data/unsafe-fifo" - -printf '#!/usr/bin/env bash\nprintf "true\\n"\n' >"${fake_bin}/docker" -chmod 0755 "${fake_bin}/docker" -if (prepare_compose_single_node_data_directory) >/dev/null 2>&1; then - fail_test "running app container did not block SQLite permission migration" -fi -printf '#!/usr/bin/env bash\nprintf "false\\n"\n' >"${fake_bin}/docker" -chmod 0755 "${fake_bin}/docker" - -ln "${COMPOSE_DIR}/data/aether.db" "${COMPOSE_DIR}/data/unsafe-hardlink" -if (prepare_compose_single_node_data_directory) >/dev/null 2>&1; then - fail_test "hard link inside the managed SQLite directory was accepted" -fi -rm -f "${COMPOSE_DIR}/data/unsafe-hardlink" - cp "${REPO_ROOT}/.env.example" "${COMPOSE_DIR}/.env.example" ADMIN_PASSWORD="test-admin-password" APP_IMAGE="example.invalid/aether:test" @@ -165,16 +100,49 @@ generate_compose_env "${generated_env}" generated_secrets=() for key in \ - JWT_SECRET_KEY ENCRYPTION_KEY DB_PASSWORD REDIS_PASSWORD \ - MYSQL_PASSWORD MYSQL_ROOT_PASSWORD; do + JWT_SECRET_KEY ENCRYPTION_KEY DB_PASSWORD REDIS_PASSWORD; do value="$(env_file_value "${generated_env}" "${key}")" [[ "${value}" =~ ^[A-Za-z0-9_-]{40,}$ ]] \ || fail_test "installer did not generate a strong URL-safe ${key}" generated_secrets+=("${value}") done -[[ "$(printf '%s\n' "${generated_secrets[@]}" | sort -u | wc -l | tr -d '[:space:]')" == "6" ]] \ +[[ "$(printf '%s\n' "${generated_secrets[@]}" | sort -u | wc -l | tr -d '[:space:]')" == "4" ]] \ || fail_test "installer reused a generated secret" +generated_single_node_env="${TEST_ROOT}/generated-single-node.env" +generate_compose_single_node_env "${generated_single_node_env}" +assert_line "${generated_single_node_env}" "AETHER_DATABASE_DRIVER=postgres" +assert_line "${generated_single_node_env}" "AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY=single-node" +assert_generated_key_output "single-node installer" "$(cat "${generated_single_node_env}")" + +( + unset AETHER_DATABASE_URL DATABASE_URL AETHER_GATEWAY_DATA_POSTGRES_URL + native_env="${TEST_ROOT}/native.env" + if (generate_first_install_env "${native_env}") >/dev/null 2>&1; then + fail_test "native installer accepted a missing PostgreSQL URL" + fi + for database_url in 'mysql://localhost/aether' 'sqlite::memory:' $'postgres://localhost/aether\nINJECTED=true'; do + if (DATABASE_URL="${database_url}" generate_first_install_env "${native_env}") >/dev/null 2>&1; then + fail_test "native installer accepted an unsupported or unsafe database URL" + fi + done + [[ ! -e "${native_env}" ]] || fail_test "invalid native database config wrote an env file" + DATABASE_URL="postgresql://localhost/aether" generate_first_install_env "${native_env}" + assert_line "${native_env}" "AETHER_DATABASE_DRIVER=postgres" + assert_line "${native_env}" "AETHER_DATABASE_URL=postgresql://localhost/aether" + validate_env_file "${native_env}" + replace_or_append_env "${native_env}" "AETHER_DATABASE_DRIVER" "sqlite" + if (validate_env_file "${native_env}") >/dev/null 2>&1; then + fail_test "native env validation accepted a removed driver" + fi + replace_or_append_env "${native_env}" "AETHER_DATABASE_DRIVER" "postgres" + replace_or_append_env "${native_env}" "AETHER_DATABASE_URL" "sqlite::memory:" + replace_or_append_env "${native_env}" "DATABASE_URL" "sqlite::memory:" + if (validate_env_file "${native_env}") >/dev/null 2>&1; then + fail_test "native env validation accepted a removed database URL" + fi +) + assert_generated_key_output \ "repository generate_keys.sh" "$("${REPO_ROOT}/generate_keys.sh")" generated_key_script="${TEST_ROOT}/generated-keys.sh" diff --git a/tests/release_supply_chain_test.sh b/tests/release_supply_chain_test.sh index ff413daee..a7c3b5d6e 100755 --- a/tests/release_supply_chain_test.sh +++ b/tests/release_supply_chain_test.sh @@ -23,10 +23,8 @@ assert_line "${COMPOSE_FILE}" \ " image: postgres:15.19@sha256:5f72c7b5bd616308ccfd2e74d6be16fb06364e5eecbb815fe9dc6ab9761d2111" assert_line "${COMPOSE_FILE}" \ " image: redis:7.4.11-alpine@sha256:ff02b58f971e7d7d156a1267e283fcbbeee91773b6aa36c49dac28ecfe28eadf" -assert_line "${COMPOSE_FILE}" \ - " image: mysql:8.0.46@sha256:7dcddc01f13bab2f15cde676d44d01f61fc9f99fe7785e86196dfc07d358ae2b" -if grep -Eq '^[[:space:]]+image:[[:space:]]+(postgres|redis|mysql):[^@[:space:]]+[[:space:]]*$' "${COMPOSE_FILE}"; then +if grep -Eq '^[[:space:]]+image:[[:space:]]+(postgres|redis):[^@[:space:]]+[[:space:]]*$' "${COMPOSE_FILE}"; then fail_test "compose contains a mutable third-party image tag" fi