mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 02:17:46 +08:00
Add multi-database data layer
Introduce aether-data-schema and driver-specific schema generation for Postgres, MySQL, and SQLite. Split data backends, lifecycle, repositories, and gateway runtime integration across database drivers. Verified with cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace.
This commit is contained in:
fawney19
@@ -0,0 +1,579 @@
|
||||
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::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::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
SqliteProviderCatalogReadRepository,
|
||||
};
|
||||
use crate::repository::proxy_nodes::{
|
||||
ProxyNodeReadRepository, ProxyNodeWriteRepository, SqliteProxyNodeReadRepository,
|
||||
};
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqliteProviderQuotaRepository,
|
||||
};
|
||||
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<Self, DataLayerError> {
|
||||
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<dyn AuthApiKeyReadRepository> {
|
||||
Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_read_repository(&self) -> Arc<dyn AnnouncementReadRepository> {
|
||||
Arc::new(SqliteAnnouncementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn audit_log_read_repository(&self) -> Arc<dyn AuditLogReadRepository> {
|
||||
Arc::new(SqliteAuditLogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_write_repository(&self) -> Arc<dyn AnnouncementWriteRepository> {
|
||||
Arc::new(SqliteAnnouncementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_api_key_write_repository(&self) -> Arc<dyn AuthApiKeyWriteRepository> {
|
||||
Arc::new(SqliteAuthApiKeyReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_read_repository(&self) -> Arc<dyn ManagementTokenReadRepository> {
|
||||
Arc::new(SqliteManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_write_repository(&self) -> Arc<dyn ManagementTokenWriteRepository> {
|
||||
Arc::new(SqliteManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_read_repository(&self) -> Arc<dyn AuthModuleReadRepository> {
|
||||
Arc::new(SqliteAuthModuleReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_write_repository(&self) -> Arc<dyn AuthModuleWriteRepository> {
|
||||
Arc::new(SqliteAuthModuleRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn billing_read_repository(&self) -> Arc<dyn BillingReadRepository> {
|
||||
Arc::new(SqliteBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_write_repository(&self) -> Arc<dyn RequestCandidateWriteRepository> {
|
||||
Arc::new(SqliteRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection_read_repository(
|
||||
&self,
|
||||
) -> Arc<dyn MinimalCandidateSelectionReadRepository> {
|
||||
Arc::new(SqliteMinimalCandidateSelectionReadRepository::new(
|
||||
self.pool_clone(),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_read_repository(&self) -> Arc<dyn GeminiFileMappingReadRepository> {
|
||||
Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_write_repository(
|
||||
&self,
|
||||
) -> Arc<dyn GeminiFileMappingWriteRepository> {
|
||||
Arc::new(SqliteGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_read_repository(&self) -> Arc<dyn GlobalModelReadRepository> {
|
||||
Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_write_repository(&self) -> Arc<dyn GlobalModelWriteRepository> {
|
||||
Arc::new(SqliteGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn user_read_repository(&self) -> Arc<dyn UserReadRepository> {
|
||||
Arc::new(SqliteUserReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
|
||||
Arc::new(SqliteVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_write_repository(&self) -> Arc<dyn VideoTaskWriteRepository> {
|
||||
Arc::new(SqliteVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_read_repository(&self) -> Arc<dyn OAuthProviderReadRepository> {
|
||||
Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_write_repository(&self) -> Arc<dyn OAuthProviderWriteRepository> {
|
||||
Arc::new(SqliteOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
|
||||
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_write_repository(&self) -> Arc<dyn ProviderCatalogWriteRepository> {
|
||||
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_write_repository(&self) -> Arc<dyn ProxyNodeWriteRepository> {
|
||||
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||
Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_write_repository(&self) -> Arc<dyn ProviderQuotaWriteRepository> {
|
||||
Arc::new(SqliteProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn settlement_write_repository(&self) -> Arc<dyn SettlementWriteRepository> {
|
||||
Arc::new(SqliteSettlementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_write_repository(&self) -> Arc<dyn UsageWriteRepository> {
|
||||
Arc::new(SqliteUsageWriteRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
|
||||
Arc::new(SqliteUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_read_repository(&self) -> Arc<dyn WalletReadRepository> {
|
||||
Arc::new(SqliteWalletReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_write_repository(&self) -> Arc<dyn WalletWriteRepository> {
|
||||
Arc::new(SqliteWalletReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqliteBackend;
|
||||
use crate::lifecycle::migrate::run_sqlite_migrations;
|
||||
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": true});
|
||||
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)
|
||||
);
|
||||
assert_eq!(
|
||||
backend
|
||||
.list_system_config_entries()
|
||||
.await
|
||||
.expect("system config should list")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
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 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<i64>,
|
||||
Option<i64>,
|
||||
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, status, billing_status,
|
||||
status_code, error_category, input_tokens, output_tokens,
|
||||
cache_creation_input_tokens, cache_read_input_tokens, total_cost_usd,
|
||||
actual_total_cost_usd, response_time_ms, created_at_unix_ms, updated_at_unix_secs
|
||||
) VALUES
|
||||
('stats-1', 'user-1', 'key-1', 'provider-a', 'model-a', 'completed', 'settled',
|
||||
200, NULL, 10, 20, 1, 2, 0.30, 0.25, 100, 3600000, 3600),
|
||||
('stats-2', 'user-2', 'key-2', 'provider-b', 'model-b', 'failed', 'void',
|
||||
500, 'upstream_error', 5, 7, 0, 1, 0.20, 0.20, 300, 3610000, 3610),
|
||||
('stats-pending', 'user-3', 'key-3', 'provider-a', 'model-a', 'pending', 'pending',
|
||||
NULL, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3620000, 3620),
|
||||
('stats-unknown-provider', 'user-4', 'key-4', 'unknown', 'model-a', 'completed', 'settled',
|
||||
200, NULL, 100, 100, 0, 0, 9.99, 9.99, 50, 3630000, 3630)
|
||||
"#,
|
||||
)
|
||||
.execute(backend.pool())
|
||||
.await
|
||||
.expect("usage stats rows should seed");
|
||||
|
||||
let target_hour = chrono::DateTime::<chrono::Utc>::from_timestamp(3600, 0)
|
||||
.expect("target hour should be valid");
|
||||
let aggregated_at = chrono::DateTime::<chrono::Utc>::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::<chrono::Utc>::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, 2);
|
||||
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));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user