mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +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
66
crates/aether-data/src/backend/leases.rs
Normal file
66
crates/aether-data/src/backend/leases.rs
Normal file
@@ -0,0 +1,66 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::driver::postgres::{PostgresLeaseRunner, PostgresLeaseRunnerConfig};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataLeaseBackends {
|
||||
postgres: Option<PostgresLeaseRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataLeaseBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataLeaseBackends")
|
||||
.field("has_postgres", &self.postgres.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataLeaseBackends {
|
||||
pub(crate) fn from_postgres(
|
||||
postgres: Option<&PostgresBackend>,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
Ok(Self {
|
||||
postgres: postgres
|
||||
.map(|backend| backend.lease_runner(PostgresLeaseRunnerConfig::default()))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<PostgresLeaseRunner> {
|
||||
self.postgres.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataLeaseBackends;
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::driver::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_lease_runner_from_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let leases =
|
||||
DataLeaseBackends::from_postgres(Some(&backend)).expect("lease backends should build");
|
||||
|
||||
assert!(leases.has_any());
|
||||
assert!(leases.postgres().is_some());
|
||||
}
|
||||
}
|
||||
58
crates/aether-data/src/backend/locks.rs
Normal file
58
crates/aether-data/src/backend/locks.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::RedisBackend;
|
||||
use crate::driver::redis::{RedisLockRunner, RedisLockRunnerConfig};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataLockBackends {
|
||||
redis: Option<RedisLockRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataLockBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataLockBackends")
|
||||
.field("has_redis", &self.redis.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataLockBackends {
|
||||
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
|
||||
Ok(Self {
|
||||
redis: redis
|
||||
.map(|backend| backend.lock_runner(RedisLockRunnerConfig::default()))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn redis(&self) -> Option<RedisLockRunner> {
|
||||
self.redis.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.redis.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataLockBackends;
|
||||
use crate::backend::RedisBackend;
|
||||
use crate::driver::redis::RedisClientConfig;
|
||||
|
||||
#[test]
|
||||
fn builds_redis_lock_runner_from_backend() {
|
||||
let backend = RedisBackend::from_config(RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
})
|
||||
.expect("redis backend should build");
|
||||
|
||||
let locks =
|
||||
DataLockBackends::from_redis(Some(&backend)).expect("lock backends should build");
|
||||
|
||||
assert!(locks.has_any());
|
||||
assert!(locks.redis().is_some());
|
||||
}
|
||||
}
|
||||
474
crates/aether-data/src/backend/maintenance.rs
Normal file
474
crates/aether-data/src/backend/maintenance.rs
Normal file
@@ -0,0 +1,474 @@
|
||||
use super::{
|
||||
summarize_pool, DataBackends, MysqlBackend, PostgresBackend, SqlBackendRef, SqliteBackend,
|
||||
};
|
||||
use crate::error::{SqlResultExt, SqlxResultExt};
|
||||
use crate::maintenance::{
|
||||
DatabaseMaintenanceSummary, DatabasePoolSummary, StatsDailyAggregationInput,
|
||||
StatsDailyAggregationSummary, StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
|
||||
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
|
||||
};
|
||||
use crate::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
|
||||
use crate::DataLayerError;
|
||||
use sqlx::migrate::MigrateError;
|
||||
|
||||
fn maintenance_identifier(value: &str) -> Result<&str, DataLayerError> {
|
||||
let valid = !value.is_empty()
|
||||
&& value
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_');
|
||||
if valid {
|
||||
Ok(value)
|
||||
} else {
|
||||
Err(DataLayerError::InvalidInput(format!(
|
||||
"invalid maintenance table name: {value}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
impl DataBackends {
|
||||
pub fn has_database_maintenance_backend(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub fn has_database_pool_summary(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub fn has_system_config_backend(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub fn has_wallet_daily_usage_aggregation_backend(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub fn has_stats_hourly_aggregation_backend(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub fn has_stats_daily_aggregation_backend(&self) -> bool {
|
||||
self.sql_backend().is_some()
|
||||
}
|
||||
|
||||
pub async fn run_database_maintenance(
|
||||
&self,
|
||||
table_names: &[&str],
|
||||
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.run_database_maintenance(table_names).await,
|
||||
None => Ok(DatabaseMaintenanceSummary::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_database_migrations(&self) -> Result<bool, MigrateError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.run_database_migrations().await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_database_backfills(&self) -> Result<bool, MigrateError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.run_database_backfills().await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn pending_database_migrations(
|
||||
&self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::migrate::PendingMigrationInfo>>, MigrateError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.pending_database_migrations().await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn prepare_database_for_startup(
|
||||
&self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::migrate::PendingMigrationInfo>>, MigrateError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.prepare_database_for_startup().await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn pending_database_backfills(
|
||||
&self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::backfill::PendingBackfillInfo>>, MigrateError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.pending_database_backfills().await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn database_pool_summary(&self) -> Option<DatabasePoolSummary> {
|
||||
self.sql_backend().map(SqlBackendRef::database_pool_summary)
|
||||
}
|
||||
|
||||
pub async fn aggregate_wallet_daily_usage(
|
||||
&self,
|
||||
input: &WalletDailyUsageAggregationInput,
|
||||
) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.aggregate_wallet_daily_usage(input).await,
|
||||
None => Ok(WalletDailyUsageAggregationResult::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<Option<StatsHourlyAggregationSummary>, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.aggregate_stats_hourly(input).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn aggregate_stats_daily(
|
||||
&self,
|
||||
input: &StatsDailyAggregationInput,
|
||||
) -> Result<Option<StatsDailyAggregationSummary>, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.aggregate_stats_daily(input).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.find_system_config_value(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.list_system_config_entries().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config_entry(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<Option<StoredSystemConfigEntry>, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_system_config_value(&self, key: &str) -> Result<bool, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.delete_system_config_value(key).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn read_admin_system_stats(&self) -> Result<AdminSystemStats, DataLayerError> {
|
||||
match self.sql_backend() {
|
||||
Some(backend) => backend.read_admin_system_stats().await,
|
||||
None => Ok(AdminSystemStats::default()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn run_table_maintenance(
|
||||
&self,
|
||||
table_names: &[&str],
|
||||
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
|
||||
let mut summary = DatabaseMaintenanceSummary::default();
|
||||
for table_name in table_names {
|
||||
let table_name = maintenance_identifier(table_name)?;
|
||||
summary.attempted += 1;
|
||||
let statement = format!("VACUUM ANALYZE \"{table_name}\"");
|
||||
if sqlx::raw_sql(&statement)
|
||||
.execute(self.pool())
|
||||
.await
|
||||
.map_postgres_err()
|
||||
.is_ok()
|
||||
{
|
||||
summary.succeeded += 1;
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
impl MysqlBackend {
|
||||
pub async fn run_table_maintenance(
|
||||
&self,
|
||||
table_names: &[&str],
|
||||
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
|
||||
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::raw_sql(&statement)
|
||||
.execute(self.pool())
|
||||
.await
|
||||
.map_sql_err()
|
||||
.is_ok()
|
||||
{
|
||||
summary.succeeded += 1;
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
impl SqliteBackend {
|
||||
pub async fn run_table_maintenance(
|
||||
&self,
|
||||
table_names: &[&str],
|
||||
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> SqlBackendRef<'a> {
|
||||
async fn run_database_maintenance(
|
||||
self,
|
||||
table_names: &[&str],
|
||||
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
|
||||
match self {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_database_migrations(self) -> Result<bool, MigrateError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => {
|
||||
crate::lifecycle::migrate::run_migrations(postgres.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
Self::Mysql(mysql) => {
|
||||
crate::lifecycle::migrate::run_mysql_migrations(mysql.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
Self::Sqlite(sqlite) => {
|
||||
crate::lifecycle::migrate::run_sqlite_migrations(sqlite.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_database_backfills(self) -> Result<bool, MigrateError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => {
|
||||
crate::lifecycle::backfill::run_backfills(postgres.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
Self::Mysql(mysql) => {
|
||||
crate::lifecycle::backfill::run_mysql_backfills(mysql.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
Self::Sqlite(sqlite) => {
|
||||
crate::lifecycle::backfill::run_sqlite_backfills(sqlite.pool()).await?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn pending_database_migrations(
|
||||
self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::migrate::PendingMigrationInfo>>, MigrateError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => Ok(Some(
|
||||
crate::lifecycle::migrate::pending_migrations(postgres.pool()).await?,
|
||||
)),
|
||||
Self::Mysql(mysql) => Ok(Some(
|
||||
crate::lifecycle::migrate::pending_mysql_migrations(mysql.pool()).await?,
|
||||
)),
|
||||
Self::Sqlite(sqlite) => Ok(Some(
|
||||
crate::lifecycle::migrate::pending_sqlite_migrations(sqlite.pool()).await?,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn prepare_database_for_startup(
|
||||
self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::migrate::PendingMigrationInfo>>, MigrateError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => Ok(Some(
|
||||
crate::lifecycle::migrate::prepare_database_for_startup(postgres.pool()).await?,
|
||||
)),
|
||||
Self::Mysql(mysql) => Ok(Some(
|
||||
crate::lifecycle::migrate::prepare_mysql_database_for_startup(mysql.pool()).await?,
|
||||
)),
|
||||
Self::Sqlite(sqlite) => Ok(Some(
|
||||
crate::lifecycle::migrate::prepare_sqlite_database_for_startup(sqlite.pool())
|
||||
.await?,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn pending_database_backfills(
|
||||
self,
|
||||
) -> Result<Option<Vec<crate::lifecycle::backfill::PendingBackfillInfo>>, MigrateError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => Ok(Some(
|
||||
crate::lifecycle::backfill::pending_backfills(postgres.pool()).await?,
|
||||
)),
|
||||
Self::Mysql(mysql) => Ok(Some(
|
||||
crate::lifecycle::backfill::pending_mysql_backfills(mysql.pool()).await?,
|
||||
)),
|
||||
Self::Sqlite(sqlite) => Ok(Some(
|
||||
crate::lifecycle::backfill::pending_sqlite_backfills(sqlite.pool()).await?,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn database_pool_summary(self) -> DatabasePoolSummary {
|
||||
match self {
|
||||
Self::Postgres(postgres) => summarize_pool(
|
||||
crate::database::DatabaseDriver::Postgres,
|
||||
usize::try_from(postgres.pool().size()).unwrap_or(usize::MAX),
|
||||
postgres.pool().num_idle(),
|
||||
postgres.config().max_connections,
|
||||
),
|
||||
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,
|
||||
),
|
||||
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,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
async fn aggregate_wallet_daily_usage(
|
||||
self,
|
||||
input: &WalletDailyUsageAggregationInput,
|
||||
) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
|
||||
match self {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
async fn aggregate_stats_hourly(
|
||||
self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<Option<StatsHourlyAggregationSummary>, DataLayerError> {
|
||||
match self {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
async fn aggregate_stats_daily(
|
||||
self,
|
||||
input: &StatsDailyAggregationInput,
|
||||
) -> Result<Option<StatsDailyAggregationSummary>, DataLayerError> {
|
||||
match self {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
async fn find_system_config_value(
|
||||
self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => postgres.find_system_config_value(key).await,
|
||||
Self::Mysql(mysql) => mysql.find_system_config_value(key).await,
|
||||
Self::Sqlite(sqlite) => sqlite.find_system_config_value(key).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_system_config_entries(
|
||||
self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => postgres.list_system_config_entries().await,
|
||||
Self::Mysql(mysql) => mysql.list_system_config_entries().await,
|
||||
Self::Sqlite(sqlite) => sqlite.list_system_config_entries().await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn upsert_system_config_entry(
|
||||
self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<StoredSystemConfigEntry, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => {
|
||||
postgres
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await
|
||||
}
|
||||
Self::Mysql(mysql) => {
|
||||
mysql
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await
|
||||
}
|
||||
Self::Sqlite(sqlite) => {
|
||||
sqlite
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_system_config_value(self, key: &str) -> Result<bool, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => postgres.delete_system_config_value(key).await,
|
||||
Self::Mysql(mysql) => mysql.delete_system_config_value(key).await,
|
||||
Self::Sqlite(sqlite) => sqlite.delete_system_config_value(key).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_admin_system_stats(self) -> Result<AdminSystemStats, DataLayerError> {
|
||||
match self {
|
||||
Self::Postgres(postgres) => postgres.read_admin_system_stats().await,
|
||||
Self::Mysql(mysql) => mysql.read_admin_system_stats().await,
|
||||
Self::Sqlite(sqlite) => sqlite.read_admin_system_stats().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
434
crates/aether-data/src/backend/mod.rs
Normal file
434
crates/aether-data/src/backend/mod.rs
Normal file
@@ -0,0 +1,434 @@
|
||||
//! Backend composition layer.
|
||||
//!
|
||||
//! `DataBackends` chooses the configured SQL driver, builds low-level pools,
|
||||
//! instantiates concrete repositories, and exposes app-facing read/write,
|
||||
//! lease, lock, worker, and maintenance handles. Request-path repository SQL
|
||||
//! belongs in `repository/*`; backend-owned maintenance SQL lives in focused
|
||||
//! modules such as `stats`, `wallet`, and `system`. Pool/client primitives
|
||||
//! belong in `driver/*`.
|
||||
|
||||
mod leases;
|
||||
mod locks;
|
||||
mod maintenance;
|
||||
mod mysql;
|
||||
mod postgres;
|
||||
mod read;
|
||||
mod redis;
|
||||
mod sqlite;
|
||||
mod stats;
|
||||
mod stats_common;
|
||||
mod system;
|
||||
mod transactions;
|
||||
mod wallet;
|
||||
mod workers;
|
||||
mod write;
|
||||
|
||||
use crate::maintenance::DatabasePoolSummary;
|
||||
pub use leases::DataLeaseBackends;
|
||||
pub use locks::DataLockBackends;
|
||||
pub use mysql::MysqlBackend;
|
||||
pub use postgres::PostgresBackend;
|
||||
pub use read::DataReadRepositories;
|
||||
pub use redis::RedisBackend;
|
||||
pub use sqlite::SqliteBackend;
|
||||
pub use transactions::DataTransactionBackends;
|
||||
pub use workers::DataWorkerBackends;
|
||||
pub use write::DataWriteRepositories;
|
||||
|
||||
use crate::database::DatabaseDriver;
|
||||
use crate::{DataLayerConfig, DataLayerError};
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum SqlBackendRef<'a> {
|
||||
Postgres(&'a PostgresBackend),
|
||||
Mysql(&'a MysqlBackend),
|
||||
Sqlite(&'a SqliteBackend),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct DataBackends {
|
||||
config: DataLayerConfig,
|
||||
postgres: Option<PostgresBackend>,
|
||||
mysql: Option<MysqlBackend>,
|
||||
sqlite: Option<SqliteBackend>,
|
||||
redis: Option<RedisBackend>,
|
||||
leases: DataLeaseBackends,
|
||||
locks: DataLockBackends,
|
||||
read: DataReadRepositories,
|
||||
transactions: DataTransactionBackends,
|
||||
workers: DataWorkerBackends,
|
||||
write: DataWriteRepositories,
|
||||
}
|
||||
|
||||
fn summarize_pool(
|
||||
driver: DatabaseDriver,
|
||||
pool_size: usize,
|
||||
idle: usize,
|
||||
max_connections: u32,
|
||||
) -> DatabasePoolSummary {
|
||||
let max_connections = max_connections.max(1);
|
||||
let checked_out = pool_size.saturating_sub(idle);
|
||||
let usage_rate = checked_out as f64 / f64::from(max_connections) * 100.0;
|
||||
|
||||
DatabasePoolSummary {
|
||||
driver,
|
||||
checked_out,
|
||||
pool_size,
|
||||
idle,
|
||||
max_connections,
|
||||
usage_rate,
|
||||
}
|
||||
}
|
||||
|
||||
impl DataBackends {
|
||||
fn sql_backend(&self) -> Option<SqlBackendRef<'_>> {
|
||||
self.postgres
|
||||
.as_ref()
|
||||
.map(SqlBackendRef::Postgres)
|
||||
.or_else(|| self.mysql.as_ref().map(SqlBackendRef::Mysql))
|
||||
.or_else(|| self.sqlite.as_ref().map(SqlBackendRef::Sqlite))
|
||||
}
|
||||
|
||||
pub fn from_config(config: DataLayerConfig) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
|
||||
let database = config.effective_database();
|
||||
let postgres = match database.clone() {
|
||||
Some(database) if database.driver == DatabaseDriver::Postgres => Some(
|
||||
PostgresBackend::from_config(database.to_postgres_config()?)?,
|
||||
),
|
||||
_ => None,
|
||||
};
|
||||
let mysql = match database.clone() {
|
||||
Some(database) if database.driver == DatabaseDriver::Mysql => {
|
||||
Some(MysqlBackend::from_config(database)?)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let sqlite = match database {
|
||||
Some(database) if database.driver == DatabaseDriver::Sqlite => {
|
||||
Some(SqliteBackend::from_config(database)?)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let redis = config
|
||||
.redis
|
||||
.clone()
|
||||
.map(RedisBackend::from_config)
|
||||
.transpose()?;
|
||||
let leases = DataLeaseBackends::from_postgres(postgres.as_ref())?;
|
||||
let locks = DataLockBackends::from_redis(redis.as_ref())?;
|
||||
let read =
|
||||
DataReadRepositories::from_backends(postgres.as_ref(), mysql.as_ref(), sqlite.as_ref());
|
||||
let transactions = DataTransactionBackends::from_postgres(postgres.as_ref());
|
||||
let workers = DataWorkerBackends::from_redis(redis.as_ref())?;
|
||||
let write = DataWriteRepositories::from_backends(
|
||||
postgres.as_ref(),
|
||||
mysql.as_ref(),
|
||||
sqlite.as_ref(),
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
postgres,
|
||||
mysql,
|
||||
sqlite,
|
||||
redis,
|
||||
leases,
|
||||
locks,
|
||||
read,
|
||||
transactions,
|
||||
workers,
|
||||
write,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &DataLayerConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<&PostgresBackend> {
|
||||
self.postgres.as_ref()
|
||||
}
|
||||
|
||||
pub fn database_driver(&self) -> Option<DatabaseDriver> {
|
||||
self.config
|
||||
.effective_database()
|
||||
.map(|database| database.driver)
|
||||
}
|
||||
|
||||
pub fn mysql(&self) -> Option<&MysqlBackend> {
|
||||
self.mysql.as_ref()
|
||||
}
|
||||
|
||||
pub fn sqlite(&self) -> Option<&SqliteBackend> {
|
||||
self.sqlite.as_ref()
|
||||
}
|
||||
|
||||
pub fn redis(&self) -> Option<&RedisBackend> {
|
||||
self.redis.as_ref()
|
||||
}
|
||||
|
||||
pub fn read(&self) -> &DataReadRepositories {
|
||||
&self.read
|
||||
}
|
||||
|
||||
pub fn leases(&self) -> &DataLeaseBackends {
|
||||
&self.leases
|
||||
}
|
||||
|
||||
pub fn locks(&self) -> &DataLockBackends {
|
||||
&self.locks
|
||||
}
|
||||
|
||||
pub fn transactions(&self) -> &DataTransactionBackends {
|
||||
&self.transactions
|
||||
}
|
||||
|
||||
pub fn workers(&self) -> &DataWorkerBackends {
|
||||
&self.workers
|
||||
}
|
||||
|
||||
pub fn write(&self) -> &DataWriteRepositories {
|
||||
&self.write
|
||||
}
|
||||
|
||||
pub fn has_runtime_backends(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
|| self.mysql.is_some()
|
||||
|| self.sqlite.is_some()
|
||||
|| self.redis.is_some()
|
||||
|| self.leases.has_any()
|
||||
|| self.locks.has_any()
|
||||
|| self.read.has_any()
|
||||
|| self.transactions.has_any()
|
||||
|| self.workers.has_any()
|
||||
|| self.write.has_any()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataBackends;
|
||||
use crate::{
|
||||
driver::postgres::PostgresPoolConfig, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig,
|
||||
SqlPoolConfig,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn builds_empty_backends_from_default_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig::default())
|
||||
.expect("empty config should be accepted");
|
||||
|
||||
assert!(!backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_none());
|
||||
assert!(backends.mysql().is_none());
|
||||
assert!(backends.sqlite().is_none());
|
||||
assert!(backends.redis().is_none());
|
||||
assert!(backends.leases().postgres().is_none());
|
||||
assert!(backends.locks().redis().is_none());
|
||||
assert!(backends.read().auth_api_keys().is_none());
|
||||
assert!(backends.read().auth_modules().is_none());
|
||||
assert!(backends.read().billing().is_none());
|
||||
assert!(backends.read().gemini_file_mappings().is_none());
|
||||
assert!(backends.read().global_models().is_none());
|
||||
assert!(backends.read().management_tokens().is_none());
|
||||
assert!(backends.read().oauth_providers().is_none());
|
||||
assert!(backends.read().proxy_nodes().is_none());
|
||||
assert!(backends.read().minimal_candidate_selection().is_none());
|
||||
assert!(backends.read().request_candidates().is_none());
|
||||
assert!(backends.read().provider_catalog().is_none());
|
||||
assert!(backends.read().usage().is_none());
|
||||
assert!(backends.read().video_tasks().is_none());
|
||||
assert!(backends.transactions().postgres().is_none());
|
||||
assert!(backends.workers().redis().is_none());
|
||||
assert!(backends.write().settlement().is_none());
|
||||
assert!(backends.write().usage().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_backend_from_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig {
|
||||
database: None,
|
||||
postgres: Some(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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,
|
||||
}),
|
||||
redis: None,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
assert!(backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_some());
|
||||
assert!(backends.mysql().is_none());
|
||||
assert!(backends.sqlite().is_none());
|
||||
assert!(backends.leases().postgres().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().proxy_nodes().is_some());
|
||||
assert!(backends.read().minimal_candidate_selection().is_some());
|
||||
assert!(backends.read().request_candidates().is_some());
|
||||
assert!(backends.read().provider_catalog().is_some());
|
||||
assert!(backends.read().provider_quotas().is_some());
|
||||
assert!(backends.read().usage().is_some());
|
||||
assert!(backends.read().video_tasks().is_some());
|
||||
assert!(backends.read().wallets().is_some());
|
||||
assert!(backends.transactions().postgres().is_some());
|
||||
assert!(backends.write().auth_modules().is_some());
|
||||
assert!(backends.write().gemini_file_mappings().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().settlement().is_some());
|
||||
assert!(backends.write().usage().is_some());
|
||||
assert!(backends.write().wallets().is_some());
|
||||
assert!(backends.config().effective_database().is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
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,
|
||||
redis: None,
|
||||
})
|
||||
.expect("mysql backend should build");
|
||||
|
||||
assert!(backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_none());
|
||||
assert!(backends.mysql().is_some());
|
||||
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]
|
||||
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,
|
||||
redis: None,
|
||||
})
|
||||
.expect("sqlite backend should build");
|
||||
|
||||
assert!(backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_none());
|
||||
assert!(backends.mysql().is_none());
|
||||
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());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_redis_backend_from_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig {
|
||||
database: None,
|
||||
postgres: None,
|
||||
redis: Some(crate::driver::redis::RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
}),
|
||||
})
|
||||
.expect("redis backend should build");
|
||||
|
||||
assert!(backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_none());
|
||||
assert!(backends.mysql().is_none());
|
||||
assert!(backends.sqlite().is_none());
|
||||
assert!(backends.redis().is_some());
|
||||
assert!(backends.leases().postgres().is_none());
|
||||
assert!(backends.locks().redis().is_some());
|
||||
assert!(backends.workers().redis().is_some());
|
||||
assert!(backends.read().auth_api_keys().is_none());
|
||||
assert!(backends.read().auth_modules().is_none());
|
||||
assert!(backends.read().global_models().is_none());
|
||||
assert!(backends.read().oauth_providers().is_none());
|
||||
assert!(backends.transactions().postgres().is_none());
|
||||
assert!(backends.write().settlement().is_none());
|
||||
assert!(backends.write().usage().is_none());
|
||||
assert!(backends.config().redis.is_some());
|
||||
}
|
||||
}
|
||||
562
crates/aether-data/src/backend/mysql.rs
Normal file
562
crates/aether-data/src/backend/mysql.rs
Normal file
@@ -0,0 +1,562 @@
|
||||
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::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::provider_catalog::{
|
||||
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
ProviderCatalogWriteRepository,
|
||||
};
|
||||
use crate::repository::proxy_nodes::{
|
||||
MysqlProxyNodeReadRepository, ProxyNodeReadRepository, ProxyNodeWriteRepository,
|
||||
};
|
||||
use crate::repository::quota::{
|
||||
MysqlProviderQuotaRepository, ProviderQuotaReadRepository, ProviderQuotaWriteRepository,
|
||||
};
|
||||
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<Self, DataLayerError> {
|
||||
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<dyn AuthApiKeyReadRepository> {
|
||||
Arc::new(MysqlAuthApiKeyReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_read_repository(&self) -> Arc<dyn AnnouncementReadRepository> {
|
||||
Arc::new(MysqlAnnouncementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn audit_log_read_repository(&self) -> Arc<dyn AuditLogReadRepository> {
|
||||
Arc::new(MysqlAuditLogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_write_repository(&self) -> Arc<dyn AnnouncementWriteRepository> {
|
||||
Arc::new(MysqlAnnouncementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_api_key_write_repository(&self) -> Arc<dyn AuthApiKeyWriteRepository> {
|
||||
Arc::new(MysqlAuthApiKeyReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_read_repository(&self) -> Arc<dyn ManagementTokenReadRepository> {
|
||||
Arc::new(MysqlManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_write_repository(&self) -> Arc<dyn ManagementTokenWriteRepository> {
|
||||
Arc::new(MysqlManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_read_repository(&self) -> Arc<dyn AuthModuleReadRepository> {
|
||||
Arc::new(MysqlAuthModuleReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_write_repository(&self) -> Arc<dyn AuthModuleWriteRepository> {
|
||||
Arc::new(MysqlAuthModuleRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn billing_read_repository(&self) -> Arc<dyn BillingReadRepository> {
|
||||
Arc::new(MysqlBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(MysqlRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_write_repository(&self) -> Arc<dyn RequestCandidateWriteRepository> {
|
||||
Arc::new(MysqlRequestCandidateRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection_read_repository(
|
||||
&self,
|
||||
) -> Arc<dyn MinimalCandidateSelectionReadRepository> {
|
||||
Arc::new(MysqlMinimalCandidateSelectionReadRepository::new(
|
||||
self.pool_clone(),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_read_repository(&self) -> Arc<dyn GeminiFileMappingReadRepository> {
|
||||
Arc::new(MysqlGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_write_repository(
|
||||
&self,
|
||||
) -> Arc<dyn GeminiFileMappingWriteRepository> {
|
||||
Arc::new(MysqlGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_read_repository(&self) -> Arc<dyn GlobalModelReadRepository> {
|
||||
Arc::new(MysqlGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_write_repository(&self) -> Arc<dyn GlobalModelWriteRepository> {
|
||||
Arc::new(MysqlGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_read_repository(&self) -> Arc<dyn OAuthProviderReadRepository> {
|
||||
Arc::new(MysqlOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_write_repository(&self) -> Arc<dyn OAuthProviderWriteRepository> {
|
||||
Arc::new(MysqlOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
|
||||
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_write_repository(&self) -> Arc<dyn ProviderCatalogWriteRepository> {
|
||||
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_write_repository(&self) -> Arc<dyn ProxyNodeWriteRepository> {
|
||||
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||
Arc::new(MysqlProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_write_repository(&self) -> Arc<dyn ProviderQuotaWriteRepository> {
|
||||
Arc::new(MysqlProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn settlement_write_repository(&self) -> Arc<dyn SettlementWriteRepository> {
|
||||
Arc::new(MysqlSettlementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_write_repository(&self) -> Arc<dyn UsageWriteRepository> {
|
||||
Arc::new(MysqlUsageWriteRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
|
||||
Arc::new(MysqlUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn user_read_repository(&self) -> Arc<dyn UserReadRepository> {
|
||||
Arc::new(MysqlUserReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
|
||||
Arc::new(MysqlVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_write_repository(&self) -> Arc<dyn VideoTaskWriteRepository> {
|
||||
Arc::new(MysqlVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_read_repository(&self) -> Arc<dyn WalletReadRepository> {
|
||||
Arc::new(MysqlWalletReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_write_repository(&self) -> Arc<dyn WalletWriteRepository> {
|
||||
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}");
|
||||
|
||||
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, 4099999900, 4099999900000, 4099999900),
|
||||
(?, NULL, 'provider', 'model', 'completed', 'pending',
|
||||
2.00, 5, 7, 1, 2, 4099999901, 4099999901000, 4099999901),
|
||||
(?, NULL, 'provider', 'model', 'completed', 'pending',
|
||||
0.00, 100, 100, 0, 0, 4099999902, 4099999902000, 4099999902),
|
||||
(?, NULL, 'provider', 'model', 'completed', 'pending',
|
||||
9.00, 50, 50, 0, 0, 4099999903, 4099999903000, 4099999903)
|
||||
"#,
|
||||
)
|
||||
.bind(&request_one)
|
||||
.bind(&request_two)
|
||||
.bind(&request_zero)
|
||||
.bind(&request_outside)
|
||||
.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', ?, 4100000000, 4100000000, 4100000000),
|
||||
(?, 'settled', ?, 4100000100, 4100000100, 4100000100),
|
||||
(?, 'settled', ?, 4100000150, 4100000150, 4100000150),
|
||||
(?, 'settled', ?, 4100000200, 4100000200, 4100000200)
|
||||
"#,
|
||||
)
|
||||
.bind(&request_one)
|
||||
.bind(&wallet_id)
|
||||
.bind(&request_two)
|
||||
.bind(&wallet_id)
|
||||
.bind(&request_zero)
|
||||
.bind(&wallet_id)
|
||||
.bind(&request_outside)
|
||||
.bind(&wallet_id)
|
||||
.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, 4099999999, 4099999999, 4099999999)
|
||||
"#,
|
||||
)
|
||||
.bind(&stale_ledger_id)
|
||||
.bind(&stale_wallet_id)
|
||||
.bind(&timezone)
|
||||
.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: 4_100_000_000,
|
||||
window_end_unix_secs: 4_100_000_200,
|
||||
aggregated_at_unix_secs: 4_100_000_300,
|
||||
})
|
||||
.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<i64>,
|
||||
Option<i64>,
|
||||
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(4_100_000_000));
|
||||
assert_eq!(ledger.8, Some(4_100_000_100));
|
||||
assert_eq!(ledger.9, 4_100_000_300);
|
||||
|
||||
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");
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
297
crates/aether-data/src/backend/postgres.rs
Normal file
297
crates/aether-data/src/backend/postgres.rs
Normal file
@@ -0,0 +1,297 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::driver::postgres::{
|
||||
PostgresLeaseRunner, PostgresLeaseRunnerConfig, PostgresPool, PostgresPoolConfig,
|
||||
PostgresPoolFactory, PostgresTransactionRunner,
|
||||
};
|
||||
use crate::repository::announcements::{
|
||||
AnnouncementReadRepository, AnnouncementWriteRepository, SqlxAnnouncementReadRepository,
|
||||
};
|
||||
use crate::repository::audit::{AuditLogReadRepository, PostgresAuditLogReadRepository};
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyReadRepository, AuthApiKeyWriteRepository, SqlxAuthApiKeySnapshotReadRepository,
|
||||
};
|
||||
use crate::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, SqlxAuthModuleReadRepository,
|
||||
SqlxAuthModuleRepository,
|
||||
};
|
||||
use crate::repository::billing::{BillingReadRepository, SqlxBillingReadRepository};
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, SqlxMinimalCandidateSelectionReadRepository,
|
||||
};
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateWriteRepository,
|
||||
SqlxRequestCandidateReadRepository,
|
||||
};
|
||||
use crate::repository::gemini_file_mappings::{
|
||||
GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository,
|
||||
SqlxGeminiFileMappingRepository,
|
||||
};
|
||||
use crate::repository::global_models::{
|
||||
GlobalModelReadRepository, GlobalModelWriteRepository, SqlxGlobalModelReadRepository,
|
||||
};
|
||||
use crate::repository::management_tokens::{
|
||||
ManagementTokenReadRepository, ManagementTokenWriteRepository, SqlxManagementTokenRepository,
|
||||
};
|
||||
use crate::repository::oauth_providers::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
SqlxProviderCatalogReadRepository,
|
||||
};
|
||||
use crate::repository::proxy_nodes::{
|
||||
ProxyNodeReadRepository, ProxyNodeWriteRepository, SqlxProxyNodeRepository,
|
||||
};
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
|
||||
};
|
||||
use crate::repository::settlement::{SettlementWriteRepository, SqlxSettlementRepository};
|
||||
use crate::repository::usage::{
|
||||
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::repository::users::{SqlxUserReadRepository, UserReadRepository};
|
||||
use crate::repository::video_tasks::{
|
||||
SqlxVideoTaskReadRepository, SqlxVideoTaskRepository, VideoTaskReadRepository,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
use crate::repository::wallet::{
|
||||
SqlxWalletRepository, WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresBackend {
|
||||
config: PostgresPoolConfig,
|
||||
pool: PostgresPool,
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
pub fn from_config(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
|
||||
let factory = PostgresPoolFactory::new(config.clone())?;
|
||||
let pool = factory.connect_lazy()?;
|
||||
|
||||
Ok(Self { config, pool })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &PostgresPoolConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PostgresPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn pool_clone(&self) -> PostgresPool {
|
||||
self.pool.clone()
|
||||
}
|
||||
|
||||
pub fn auth_api_key_read_repository(&self) -> Arc<dyn AuthApiKeyReadRepository> {
|
||||
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_read_repository(&self) -> Arc<dyn AnnouncementReadRepository> {
|
||||
Arc::new(SqlxAnnouncementReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn audit_log_read_repository(&self) -> Arc<dyn AuditLogReadRepository> {
|
||||
Arc::new(PostgresAuditLogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn announcement_write_repository(&self) -> Arc<dyn AnnouncementWriteRepository> {
|
||||
Arc::new(SqlxAnnouncementReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_api_key_write_repository(&self) -> Arc<dyn AuthApiKeyWriteRepository> {
|
||||
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_read_repository(&self) -> Arc<dyn AuthModuleReadRepository> {
|
||||
Arc::new(SqlxAuthModuleReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn auth_module_write_repository(&self) -> Arc<dyn AuthModuleWriteRepository> {
|
||||
Arc::new(SqlxAuthModuleRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn billing_read_repository(&self) -> Arc<dyn BillingReadRepository> {
|
||||
Arc::new(SqlxBillingReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection_read_repository(
|
||||
&self,
|
||||
) -> Arc<dyn MinimalCandidateSelectionReadRepository> {
|
||||
Arc::new(SqlxMinimalCandidateSelectionReadRepository::new(
|
||||
self.pool_clone(),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_write_repository(&self) -> Arc<dyn RequestCandidateWriteRepository> {
|
||||
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_read_repository(&self) -> Arc<dyn GeminiFileMappingReadRepository> {
|
||||
Arc::new(SqlxGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn gemini_file_mapping_write_repository(
|
||||
&self,
|
||||
) -> Arc<dyn GeminiFileMappingWriteRepository> {
|
||||
Arc::new(SqlxGeminiFileMappingRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_read_repository(&self) -> Arc<dyn GlobalModelReadRepository> {
|
||||
Arc::new(SqlxGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn global_model_write_repository(&self) -> Arc<dyn GlobalModelWriteRepository> {
|
||||
Arc::new(SqlxGlobalModelReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_read_repository(&self) -> Arc<dyn ManagementTokenReadRepository> {
|
||||
Arc::new(SqlxManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn management_token_write_repository(&self) -> Arc<dyn ManagementTokenWriteRepository> {
|
||||
Arc::new(SqlxManagementTokenRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_read_repository(&self) -> Arc<dyn OAuthProviderReadRepository> {
|
||||
Arc::new(SqlxOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn oauth_provider_write_repository(&self) -> Arc<dyn OAuthProviderWriteRepository> {
|
||||
Arc::new(SqlxOAuthProviderRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(SqlxProxyNodeRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_write_repository(&self) -> Arc<dyn ProxyNodeWriteRepository> {
|
||||
Arc::new(SqlxProxyNodeRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
|
||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_write_repository(&self) -> Arc<dyn ProviderCatalogWriteRepository> {
|
||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
|
||||
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn user_read_repository(&self) -> Arc<dyn UserReadRepository> {
|
||||
Arc::new(SqlxUserReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_write_repository(&self) -> Arc<dyn UsageWriteRepository> {
|
||||
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_read_repository(&self) -> Arc<dyn WalletReadRepository> {
|
||||
Arc::new(SqlxWalletRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn wallet_write_repository(&self) -> Arc<dyn WalletWriteRepository> {
|
||||
Arc::new(SqlxWalletRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn settlement_write_repository(&self) -> Arc<dyn SettlementWriteRepository> {
|
||||
Arc::new(SqlxSettlementRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
|
||||
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_write_repository(&self) -> Arc<dyn VideoTaskWriteRepository> {
|
||||
Arc::new(SqlxVideoTaskRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> PostgresTransactionRunner {
|
||||
PostgresTransactionRunner::new(self.pool_clone())
|
||||
}
|
||||
|
||||
pub fn lease_runner(
|
||||
&self,
|
||||
config: PostgresLeaseRunnerConfig,
|
||||
) -> Result<PostgresLeaseRunner, DataLayerError> {
|
||||
PostgresLeaseRunner::new(self.transaction_runner(), config)
|
||||
}
|
||||
|
||||
pub fn provider_quota_write_repository(&self) -> Arc<dyn ProviderQuotaWriteRepository> {
|
||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::PostgresBackend;
|
||||
use crate::driver::postgres::{PostgresLeaseRunnerConfig, PostgresPoolConfig};
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_retains_config_and_pool() {
|
||||
let config = PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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 backend =
|
||||
PostgresBackend::from_config(config.clone()).expect("backend should build lazily");
|
||||
|
||||
assert_eq!(backend.config(), &config);
|
||||
let _pool = backend.pool();
|
||||
let _pool_clone = backend.pool_clone();
|
||||
let _auth_api_key_reader = backend.auth_api_key_read_repository();
|
||||
let _auth_api_key_writer = backend.auth_api_key_write_repository();
|
||||
let _auth_module_reader = backend.auth_module_read_repository();
|
||||
let _billing_reader = backend.billing_read_repository();
|
||||
let _gemini_file_mapping_reader = backend.gemini_file_mapping_read_repository();
|
||||
let _global_model_reader = backend.global_model_read_repository();
|
||||
let _global_model_writer = backend.global_model_write_repository();
|
||||
let _management_token_reader = backend.management_token_read_repository();
|
||||
let _management_token_writer = backend.management_token_write_repository();
|
||||
let _oauth_provider_reader = backend.oauth_provider_read_repository();
|
||||
let _oauth_provider_writer = backend.oauth_provider_write_repository();
|
||||
let _proxy_node_reader = backend.proxy_node_read_repository();
|
||||
let _proxy_node_writer = backend.proxy_node_write_repository();
|
||||
let _minimal_candidate_selection_reader =
|
||||
backend.minimal_candidate_selection_read_repository();
|
||||
let _request_candidate_reader = backend.request_candidate_read_repository();
|
||||
let _request_candidate_writer = backend.request_candidate_write_repository();
|
||||
let _gemini_file_mapping_writer = backend.gemini_file_mapping_write_repository();
|
||||
let _provider_catalog_reader = backend.provider_catalog_read_repository();
|
||||
let _provider_catalog_writer = backend.provider_catalog_write_repository();
|
||||
let _provider_quota_reader = backend.provider_quota_read_repository();
|
||||
let _usage_reader = backend.usage_read_repository();
|
||||
let _usage_writer = backend.usage_write_repository();
|
||||
let _wallet_reader = backend.wallet_read_repository();
|
||||
let _wallet_writer = backend.wallet_write_repository();
|
||||
let _settlement_writer = backend.settlement_write_repository();
|
||||
let _video_task_reader = backend.video_task_read_repository();
|
||||
let _video_task_writer = backend.video_task_write_repository();
|
||||
let _transaction_runner = backend.transaction_runner();
|
||||
let _lease_runner = backend
|
||||
.lease_runner(PostgresLeaseRunnerConfig::default())
|
||||
.expect("lease runner should build");
|
||||
let _provider_quota_writer = backend.provider_quota_write_repository();
|
||||
}
|
||||
}
|
||||
301
crates/aether-data/src/backend/read.rs
Normal file
301
crates/aether-data/src/backend/read.rs
Normal file
@@ -0,0 +1,301 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::repository::announcements::AnnouncementReadRepository;
|
||||
use crate::repository::audit::AuditLogReadRepository;
|
||||
use crate::repository::auth::AuthApiKeyReadRepository;
|
||||
use crate::repository::auth_modules::AuthModuleReadRepository;
|
||||
use crate::repository::billing::BillingReadRepository;
|
||||
use crate::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
|
||||
use crate::repository::candidates::RequestCandidateReadRepository;
|
||||
use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
|
||||
use crate::repository::global_models::GlobalModelReadRepository;
|
||||
use crate::repository::management_tokens::ManagementTokenReadRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderReadRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||
use crate::repository::quota::ProviderQuotaReadRepository;
|
||||
use crate::repository::usage::UsageReadRepository;
|
||||
use crate::repository::users::UserReadRepository;
|
||||
use crate::repository::video_tasks::VideoTaskReadRepository;
|
||||
use crate::repository::wallet::WalletReadRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataReadRepositories {
|
||||
announcements: Option<Arc<dyn AnnouncementReadRepository>>,
|
||||
audit_logs: Option<Arc<dyn AuditLogReadRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleReadRepository>>,
|
||||
billing: Option<Arc<dyn BillingReadRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingReadRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
|
||||
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||
usage: Option<Arc<dyn UsageReadRepository>>,
|
||||
users: Option<Arc<dyn UserReadRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
wallets: Option<Arc<dyn WalletReadRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataReadRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataReadRepositories")
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_audit_logs", &self.audit_logs.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_billing", &self.billing.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
&self.gemini_file_mappings.is_some(),
|
||||
)
|
||||
.field("has_global_models", &self.global_models.is_some())
|
||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||
.field(
|
||||
"has_minimal_candidate_selection",
|
||||
&self.minimal_candidate_selection.is_some(),
|
||||
)
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_users", &self.users.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
.field("has_wallets", &self.wallets.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataReadRepositories {
|
||||
pub(crate) fn from_backends(
|
||||
postgres: Option<&PostgresBackend>,
|
||||
mysql: Option<&MysqlBackend>,
|
||||
sqlite: Option<&SqliteBackend>,
|
||||
) -> Self {
|
||||
Self {
|
||||
announcements: postgres
|
||||
.map(PostgresBackend::announcement_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::announcement_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::announcement_read_repository)),
|
||||
audit_logs: postgres
|
||||
.map(PostgresBackend::audit_log_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::audit_log_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::audit_log_read_repository)),
|
||||
auth_api_keys: postgres
|
||||
.map(PostgresBackend::auth_api_key_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_api_key_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_api_key_read_repository)),
|
||||
auth_modules: postgres
|
||||
.map(PostgresBackend::auth_module_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_module_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_module_read_repository)),
|
||||
billing: postgres
|
||||
.map(PostgresBackend::billing_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::billing_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::billing_read_repository)),
|
||||
gemini_file_mappings: postgres
|
||||
.map(PostgresBackend::gemini_file_mapping_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::gemini_file_mapping_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::gemini_file_mapping_read_repository)),
|
||||
global_models: postgres
|
||||
.map(PostgresBackend::global_model_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::global_model_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::global_model_read_repository)),
|
||||
management_tokens: postgres
|
||||
.map(PostgresBackend::management_token_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::management_token_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::management_token_read_repository)),
|
||||
oauth_providers: postgres
|
||||
.map(PostgresBackend::oauth_provider_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
|
||||
proxy_nodes: postgres
|
||||
.map(PostgresBackend::proxy_node_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::proxy_node_read_repository)),
|
||||
minimal_candidate_selection: postgres
|
||||
.map(PostgresBackend::minimal_candidate_selection_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::minimal_candidate_selection_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::minimal_candidate_selection_read_repository)),
|
||||
request_candidates: postgres
|
||||
.map(PostgresBackend::request_candidate_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::request_candidate_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::request_candidate_read_repository)),
|
||||
provider_catalog: postgres
|
||||
.map(PostgresBackend::provider_catalog_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_catalog_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_catalog_read_repository)),
|
||||
provider_quotas: postgres
|
||||
.map(PostgresBackend::provider_quota_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_quota_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_quota_read_repository)),
|
||||
usage: postgres
|
||||
.map(PostgresBackend::usage_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::usage_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::usage_read_repository)),
|
||||
users: postgres
|
||||
.map(PostgresBackend::user_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::user_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::user_read_repository)),
|
||||
video_tasks: postgres
|
||||
.map(PostgresBackend::video_task_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::video_task_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::video_task_read_repository)),
|
||||
wallets: postgres
|
||||
.map(PostgresBackend::wallet_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::wallet_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::wallet_read_repository)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self::from_backends(postgres, None, None)
|
||||
}
|
||||
|
||||
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyReadRepository>> {
|
||||
self.auth_api_keys.clone()
|
||||
}
|
||||
|
||||
pub fn announcements(&self) -> Option<Arc<dyn AnnouncementReadRepository>> {
|
||||
self.announcements.clone()
|
||||
}
|
||||
|
||||
pub fn audit_logs(&self) -> Option<Arc<dyn AuditLogReadRepository>> {
|
||||
self.audit_logs.clone()
|
||||
}
|
||||
|
||||
pub fn auth_modules(&self) -> Option<Arc<dyn AuthModuleReadRepository>> {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn billing(&self) -> Option<Arc<dyn BillingReadRepository>> {
|
||||
self.billing.clone()
|
||||
}
|
||||
|
||||
pub fn gemini_file_mappings(&self) -> Option<Arc<dyn GeminiFileMappingReadRepository>> {
|
||||
self.gemini_file_mappings.clone()
|
||||
}
|
||||
|
||||
pub fn global_models(&self) -> Option<Arc<dyn GlobalModelReadRepository>> {
|
||||
self.global_models.clone()
|
||||
}
|
||||
|
||||
pub fn management_tokens(&self) -> Option<Arc<dyn ManagementTokenReadRepository>> {
|
||||
self.management_tokens.clone()
|
||||
}
|
||||
|
||||
pub fn oauth_providers(&self) -> Option<Arc<dyn OAuthProviderReadRepository>> {
|
||||
self.oauth_providers.clone()
|
||||
}
|
||||
|
||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
|
||||
pub fn minimal_candidate_selection(
|
||||
&self,
|
||||
) -> Option<Arc<dyn MinimalCandidateSelectionReadRepository>> {
|
||||
self.minimal_candidate_selection.clone()
|
||||
}
|
||||
|
||||
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateReadRepository>> {
|
||||
self.request_candidates.clone()
|
||||
}
|
||||
|
||||
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogReadRepository>> {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
|
||||
pub fn provider_quotas(&self) -> Option<Arc<dyn ProviderQuotaReadRepository>> {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
|
||||
pub fn users(&self) -> Option<Arc<dyn UserReadRepository>> {
|
||||
self.users.clone()
|
||||
}
|
||||
|
||||
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskReadRepository>> {
|
||||
self.video_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn wallets(&self) -> Option<Arc<dyn WalletReadRepository>> {
|
||||
self.wallets.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.auth_api_keys.is_some()
|
||||
|| self.announcements.is_some()
|
||||
|| self.audit_logs.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.billing.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.minimal_candidate_selection.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.users.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|| self.wallets.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataReadRepositories;
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::driver::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_read_repositories_from_postgres_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let read = DataReadRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(read.has_any());
|
||||
assert!(read.announcements().is_some());
|
||||
assert!(read.audit_logs().is_some());
|
||||
assert!(read.auth_api_keys().is_some());
|
||||
assert!(read.auth_modules().is_some());
|
||||
assert!(read.billing().is_some());
|
||||
assert!(read.gemini_file_mappings().is_some());
|
||||
assert!(read.global_models().is_some());
|
||||
assert!(read.management_tokens().is_some());
|
||||
assert!(read.oauth_providers().is_some());
|
||||
assert!(read.proxy_nodes().is_some());
|
||||
assert!(read.minimal_candidate_selection().is_some());
|
||||
assert!(read.request_candidates().is_some());
|
||||
assert!(read.provider_catalog().is_some());
|
||||
assert!(read.provider_quotas().is_some());
|
||||
assert!(read.usage().is_some());
|
||||
assert!(read.video_tasks().is_some());
|
||||
assert!(read.wallets().is_some());
|
||||
}
|
||||
}
|
||||
86
crates/aether-data/src/backend/redis.rs
Normal file
86
crates/aether-data/src/backend/redis.rs
Normal file
@@ -0,0 +1,86 @@
|
||||
use crate::driver::redis::{
|
||||
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisKvRunner,
|
||||
RedisKvRunnerConfig, RedisLockRunner, RedisLockRunnerConfig, RedisStreamRunner,
|
||||
RedisStreamRunnerConfig,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisBackend {
|
||||
config: RedisClientConfig,
|
||||
client: RedisClient,
|
||||
}
|
||||
|
||||
impl RedisBackend {
|
||||
pub fn from_config(config: RedisClientConfig) -> Result<Self, DataLayerError> {
|
||||
let factory = RedisClientFactory::new(config.clone())?;
|
||||
let client = factory.connect_lazy()?;
|
||||
Ok(Self { config, client })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &RedisClientConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn client_clone(&self) -> RedisClient {
|
||||
self.client.clone()
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> RedisKeyspace {
|
||||
self.config.keyspace()
|
||||
}
|
||||
|
||||
pub fn lock_runner(
|
||||
&self,
|
||||
config: RedisLockRunnerConfig,
|
||||
) -> Result<RedisLockRunner, DataLayerError> {
|
||||
RedisLockRunner::new(self.client_clone(), self.keyspace(), config)
|
||||
}
|
||||
|
||||
pub fn stream_runner(
|
||||
&self,
|
||||
config: RedisStreamRunnerConfig,
|
||||
) -> Result<RedisStreamRunner, DataLayerError> {
|
||||
RedisStreamRunner::new(self.client_clone(), self.keyspace(), config)
|
||||
}
|
||||
|
||||
pub fn kv_runner(&self, config: RedisKvRunnerConfig) -> Result<RedisKvRunner, DataLayerError> {
|
||||
RedisKvRunner::new(self.client_clone(), self.keyspace(), config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisBackend;
|
||||
use crate::driver::redis::{
|
||||
RedisClientConfig, RedisKvRunnerConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn backend_retains_config_client_and_shared_runners() {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
};
|
||||
|
||||
let backend = RedisBackend::from_config(config.clone()).expect("backend should build");
|
||||
|
||||
assert_eq!(backend.config(), &config);
|
||||
assert_eq!(backend.keyspace().key("audit"), "aether:audit");
|
||||
let _client_ref = backend.client();
|
||||
let _client_clone = backend.client_clone();
|
||||
let _lock_runner = backend
|
||||
.lock_runner(RedisLockRunnerConfig::default())
|
||||
.expect("lock runner should build");
|
||||
let _stream_runner = backend
|
||||
.stream_runner(RedisStreamRunnerConfig::default())
|
||||
.expect("stream runner should build");
|
||||
let _kv_runner = backend
|
||||
.kv_runner(RedisKvRunnerConfig::default())
|
||||
.expect("kv runner should build");
|
||||
}
|
||||
}
|
||||
579
crates/aether-data/src/backend/sqlite.rs
Normal file
579
crates/aether-data/src/backend/sqlite.rs
Normal file
@@ -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));
|
||||
}
|
||||
}
|
||||
4
crates/aether-data/src/backend/stats.rs
Normal file
4
crates/aether-data/src/backend/stats.rs
Normal file
@@ -0,0 +1,4 @@
|
||||
pub(crate) mod mysql;
|
||||
pub(crate) mod postgres_daily;
|
||||
pub(crate) mod postgres_hourly;
|
||||
pub(crate) mod sqlite;
|
||||
371
crates/aether-data/src/backend/stats/mysql.rs
Normal file
371
crates/aether-data/src/backend/stats/mysql.rs
Normal file
@@ -0,0 +1,371 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
|
||||
use crate::backend::stats_common::{stats_id, unix_ms, unix_secs, utc_from_unix_secs};
|
||||
use crate::backend::MysqlBackend;
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::{
|
||||
DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary,
|
||||
StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
|
||||
};
|
||||
|
||||
impl MysqlBackend {
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<Option<StatsHourlyAggregationSummary>, 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<Option<StatsDailyAggregationSummary>, 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<Utc>,
|
||||
) -> Result<Option<i64>, DataLayerError> {
|
||||
let latest_hour: Option<i64> =
|
||||
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<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT MIN(FLOOR(created_at_unix_ms / 3600000) * 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(unix_ms(search_from)?)
|
||||
.bind(unix_ms(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<Utc>,
|
||||
) -> Result<Option<i64>, DataLayerError> {
|
||||
let latest_day: Option<i64> =
|
||||
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<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT MIN(FLOOR(created_at_unix_ms / 86400000) * 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(unix_ms(search_from)?)
|
||||
.bind(unix_ms(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
|
||||
COUNT(*) AS total_requests,
|
||||
COALESCE(SUM(CASE
|
||||
WHEN status = 'failed'
|
||||
OR status_code >= 400
|
||||
OR (error_category IS NOT NULL AND error_category <> '')
|
||||
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,
|
||||
COALESCE(SUM(total_cost_usd), 0.0) AS total_cost,
|
||||
COALESCE(SUM(actual_total_cost_usd), 0.0) AS actual_total_cost,
|
||||
COALESCE(AVG(response_time_ms), 0.0) 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<Utc>,
|
||||
) -> Result<StatsHourlyAggregationSummary, DataLayerError> {
|
||||
let start_ms = unix_ms(hour_utc_unix_secs)?;
|
||||
let end_ms = unix_ms(hour_utc_unix_secs + 3600)?;
|
||||
let aggregated_at_unix_secs = unix_secs(aggregated_at);
|
||||
let mut tx = pool.begin().await.map_sql_err()?;
|
||||
let row = sqlx::query(MYSQL_STATS_AGGREGATE_SQL)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.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::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("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 = mysql_group_count(&mut tx, "user_id", start_ms, end_ms).await?;
|
||||
let user_model_rows = mysql_group_count(&mut tx, "user_id, model", start_ms, end_ms).await?;
|
||||
let model_rows = mysql_group_count(&mut tx, "model", start_ms, end_ms).await?;
|
||||
let provider_rows = mysql_group_count(&mut tx, "provider_name", start_ms, end_ms).await?;
|
||||
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<Utc>,
|
||||
) -> Result<StatsDailyAggregationSummary, DataLayerError> {
|
||||
let start_ms = unix_ms(day_start_unix_secs)?;
|
||||
let end_ms = unix_ms(day_start_unix_secs + 86_400)?;
|
||||
let aggregated_at_unix_secs = unix_secs(aggregated_at);
|
||||
let mut tx = pool.begin().await.map_sql_err()?;
|
||||
let row = sqlx::query(MYSQL_STATS_AGGREGATE_SQL)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let total_requests: i64 = row.try_get("total_requests").map_sql_err()?;
|
||||
let error_requests: i64 = row.try_get("error_requests").map_sql_err()?;
|
||||
let unique_models = mysql_group_count(&mut tx, "model", start_ms, end_ms).await? as i64;
|
||||
let unique_providers =
|
||||
mysql_group_count(&mut tx, "provider_name", start_ms, end_ms).await? as i64;
|
||||
|
||||
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, 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::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("avg_response_time_ms")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.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 = usize::try_from(unique_models).unwrap_or(usize::MAX);
|
||||
let provider_rows = usize::try_from(unique_providers).unwrap_or(usize::MAX);
|
||||
let api_key_rows = mysql_group_count(&mut tx, "api_key_id", start_ms, end_ms).await?;
|
||||
let error_rows = mysql_error_group_count(&mut tx, start_ms, end_ms).await?;
|
||||
let user_rows = mysql_group_count(&mut tx, "user_id", start_ms, end_ms).await?;
|
||||
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 mysql_group_count(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
group_columns: &str,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let not_empty = group_columns
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.map(|column| format!("{column} IS NOT NULL AND {column} <> ''"))
|
||||
.collect::<Vec<_>>()
|
||||
.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_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn mysql_error_group_count(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let count: i64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COUNT(*)
|
||||
FROM (
|
||||
SELECT 1
|
||||
FROM `usage`
|
||||
WHERE created_at_unix_ms >= ?
|
||||
AND created_at_unix_ms < ?
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
AND (
|
||||
status = 'failed'
|
||||
OR status_code >= 400
|
||||
OR (error_category IS NOT NULL AND error_category <> '')
|
||||
)
|
||||
GROUP BY COALESCE(NULLIF(error_category, ''), 'unknown_error'), provider_name, model
|
||||
) AS grouped
|
||||
"#,
|
||||
)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX))
|
||||
}
|
||||
638
crates/aether-data/src/backend/stats/postgres_daily/mod.rs
Normal file
638
crates/aether-data/src/backend/stats/postgres_daily/mod.rs
Normal file
@@ -0,0 +1,638 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::{
|
||||
error::postgres_error, DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary,
|
||||
};
|
||||
|
||||
mod percentiles;
|
||||
mod sql;
|
||||
|
||||
use self::percentiles::{percentile_ms_to_i64, PercentileSummary};
|
||||
use self::sql::*;
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn aggregate_stats_daily(
|
||||
&self,
|
||||
input: &StatsDailyAggregationInput,
|
||||
) -> Result<Option<StatsDailyAggregationSummary>, DataLayerError> {
|
||||
let Some(day_start_utc) = next_stats_aggregation_day(self.pool(), input.target_day_utc)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
perform_stats_aggregation_for_day(self.pool(), day_start_utc, input.aggregated_at)
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(postgres_error)
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_stats_aggregation_day(
|
||||
pool: &crate::driver::postgres::PostgresPool,
|
||||
target_day_utc: DateTime<Utc>,
|
||||
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
|
||||
let latest_row = sqlx::query(SELECT_LATEST_STATS_DAILY_DATE_SQL)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
let latest_day = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_date")?;
|
||||
let search_from = latest_day
|
||||
.map(|value| value + chrono::Duration::days(1))
|
||||
.unwrap_or_else(|| {
|
||||
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
|
||||
});
|
||||
let search_until = target_day_utc + chrono::Duration::days(1);
|
||||
if search_from >= search_until {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let next_row = sqlx::query(SELECT_NEXT_STATS_DAILY_BUCKET_SQL)
|
||||
.bind(search_from)
|
||||
.bind(search_until)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
|
||||
Ok(next_bucket.filter(|value| *value <= target_day_utc))
|
||||
}
|
||||
|
||||
async fn perform_stats_aggregation_for_day(
|
||||
pool: &crate::driver::postgres::PostgresPool,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<StatsDailyAggregationSummary, sqlx::Error> {
|
||||
let day_end_utc = day_start_utc + chrono::Duration::days(1);
|
||||
let mut tx = pool.begin().await?;
|
||||
let aggregate_row = sqlx::query(SELECT_STATS_DAILY_AGGREGATE_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_one(&mut *tx)
|
||||
.await?;
|
||||
let total_requests = aggregate_row.try_get::<i64, _>("total_requests")?;
|
||||
let error_requests = aggregate_row.try_get::<i64, _>("error_requests")?;
|
||||
let success_requests = total_requests.saturating_sub(error_requests);
|
||||
let fallback_count = sqlx::query(SELECT_STATS_DAILY_FALLBACK_COUNT_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(vec!["success", "failed"])
|
||||
.fetch_one(&mut *tx)
|
||||
.await?
|
||||
.try_get::<i64, _>("fallback_count")?;
|
||||
let response_percentiles = fetch_stats_daily_percentiles(
|
||||
&mut tx,
|
||||
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
|
||||
day_start_utc,
|
||||
day_end_utc,
|
||||
)
|
||||
.await?;
|
||||
let first_byte_percentiles = fetch_stats_daily_percentiles(
|
||||
&mut tx,
|
||||
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
|
||||
day_start_utc,
|
||||
day_end_utc,
|
||||
)
|
||||
.await?;
|
||||
|
||||
sqlx::query(UPSERT_STATS_DAILY_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(day_start_utc)
|
||||
.bind(total_requests)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_hit_total_requests")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_hit_requests")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_total_requests")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_cache_hit_requests")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_input_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_cache_creation_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_cache_read_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("completed_total_input_context")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("completed_cache_creation_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("completed_cache_read_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("settled_total_cost")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("settled_total_requests")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("settled_input_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("settled_output_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("settled_cache_creation_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("settled_cache_read_tokens")?)
|
||||
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
|
||||
.bind(aggregate_row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
|
||||
.bind(success_requests)
|
||||
.bind(error_requests)
|
||||
.bind(aggregate_row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("effective_input_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_creation_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_5m_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_creation_ephemeral_1h_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("cache_read_tokens")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("total_input_context")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("total_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("actual_total_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("input_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("output_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("cache_creation_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("cache_read_cost")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("response_time_sum_ms")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("response_time_samples")?)
|
||||
.bind(aggregate_row.try_get::<f64, _>("avg_response_time_ms")?)
|
||||
.bind(response_percentiles.p50)
|
||||
.bind(response_percentiles.p90)
|
||||
.bind(response_percentiles.p99)
|
||||
.bind(first_byte_percentiles.p50)
|
||||
.bind(first_byte_percentiles.p90)
|
||||
.bind(first_byte_percentiles.p99)
|
||||
.bind(fallback_count)
|
||||
.bind(aggregate_row.try_get::<i64, _>("unique_models")?)
|
||||
.bind(aggregate_row.try_get::<i64, _>("unique_providers")?)
|
||||
.bind(true)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
let model_rows =
|
||||
upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
let provider_rows =
|
||||
upsert_stats_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_daily_cost_savings_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
|
||||
.await?;
|
||||
upsert_stats_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
|
||||
.await?;
|
||||
upsert_stats_daily_cost_savings_model_provider_rows(
|
||||
&mut tx,
|
||||
day_start_utc,
|
||||
day_end_utc,
|
||||
now_utc,
|
||||
)
|
||||
.await?;
|
||||
let api_key_rows =
|
||||
upsert_stats_daily_api_key_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
let error_rows =
|
||||
refresh_stats_daily_error_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
let user_rows =
|
||||
upsert_stats_user_daily_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_user_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_user_daily_model_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
|
||||
.await?;
|
||||
upsert_stats_user_daily_provider_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_user_daily_cost_savings_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
upsert_stats_user_daily_cost_savings_provider_rows(
|
||||
&mut tx,
|
||||
day_start_utc,
|
||||
day_end_utc,
|
||||
now_utc,
|
||||
)
|
||||
.await?;
|
||||
upsert_stats_user_daily_cost_savings_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc)
|
||||
.await?;
|
||||
upsert_stats_user_daily_cost_savings_model_provider_rows(
|
||||
&mut tx,
|
||||
day_start_utc,
|
||||
day_end_utc,
|
||||
now_utc,
|
||||
)
|
||||
.await?;
|
||||
upsert_stats_user_daily_api_format_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
|
||||
refresh_stats_summary_row(&mut tx, day_end_utc, now_utc).await?;
|
||||
refresh_stats_user_summary_rows(&mut tx, day_end_utc, now_utc).await?;
|
||||
tx.commit().await?;
|
||||
|
||||
Ok(StatsDailyAggregationSummary {
|
||||
day_start_utc,
|
||||
total_requests,
|
||||
model_rows,
|
||||
provider_rows,
|
||||
api_key_rows,
|
||||
error_rows,
|
||||
user_rows,
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_stats_daily_percentiles(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
sql: &str,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
) -> Result<PercentileSummary, sqlx::Error> {
|
||||
let row = sqlx::query(sql)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
let sample_count = row.try_get::<i64, _>("sample_count")?;
|
||||
if sample_count < 10 {
|
||||
return Ok(PercentileSummary::default());
|
||||
}
|
||||
|
||||
Ok(PercentileSummary {
|
||||
p50: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p50")?),
|
||||
p90: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p90")?),
|
||||
p99: percentile_ms_to_i64(row.try_get::<Option<f64>, _>("p99")?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_model_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_cost_savings_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_cost_savings_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_cost_savings_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_cost_savings_model_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_api_key_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn refresh_stats_daily_error_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
sqlx::query(DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL)
|
||||
.bind(day_start_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
let rows_affected = sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_model_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_MODEL_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_cost_savings_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_cost_savings_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_cost_savings_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_cost_savings_model_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_COST_SAVINGS_MODEL_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_api_format_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
day_start_utc: DateTime<Utc>,
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_API_FORMAT_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn refresh_stats_summary_row(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
cutoff_date: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<(), sqlx::Error> {
|
||||
let totals_row = sqlx::query(SELECT_STATS_SUMMARY_TOTALS_SQL)
|
||||
.bind(cutoff_date)
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
let entity_counts_row = sqlx::query(SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL)
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
let existing_summary_id = sqlx::query_scalar::<_, String>(SELECT_EXISTING_STATS_SUMMARY_ID_SQL)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await?;
|
||||
|
||||
let all_time_requests = totals_row.try_get::<i64, _>("all_time_requests")?;
|
||||
let all_time_success_requests = totals_row.try_get::<i64, _>("all_time_success_requests")?;
|
||||
let all_time_error_requests = totals_row.try_get::<i64, _>("all_time_error_requests")?;
|
||||
let all_time_input_tokens = totals_row.try_get::<i64, _>("all_time_input_tokens")?;
|
||||
let all_time_output_tokens = totals_row.try_get::<i64, _>("all_time_output_tokens")?;
|
||||
let all_time_cache_creation_tokens =
|
||||
totals_row.try_get::<i64, _>("all_time_cache_creation_tokens")?;
|
||||
let all_time_cache_read_tokens = totals_row.try_get::<i64, _>("all_time_cache_read_tokens")?;
|
||||
let all_time_cost = totals_row.try_get::<f64, _>("all_time_cost")?;
|
||||
let all_time_actual_cost = totals_row.try_get::<f64, _>("all_time_actual_cost")?;
|
||||
let total_users = entity_counts_row.try_get::<i64, _>("total_users")?;
|
||||
let active_users = entity_counts_row.try_get::<i64, _>("active_users")?;
|
||||
let total_api_keys = entity_counts_row.try_get::<i64, _>("total_api_keys")?;
|
||||
let active_api_keys = entity_counts_row.try_get::<i64, _>("active_api_keys")?;
|
||||
|
||||
if let Some(summary_id) = existing_summary_id {
|
||||
sqlx::query(UPDATE_STATS_SUMMARY_SQL)
|
||||
.bind(summary_id)
|
||||
.bind(cutoff_date)
|
||||
.bind(all_time_requests)
|
||||
.bind(all_time_success_requests)
|
||||
.bind(all_time_error_requests)
|
||||
.bind(all_time_input_tokens)
|
||||
.bind(all_time_output_tokens)
|
||||
.bind(all_time_cache_creation_tokens)
|
||||
.bind(all_time_cache_read_tokens)
|
||||
.bind(all_time_cost)
|
||||
.bind(all_time_actual_cost)
|
||||
.bind(total_users)
|
||||
.bind(active_users)
|
||||
.bind(total_api_keys)
|
||||
.bind(active_api_keys)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
} else {
|
||||
sqlx::query(INSERT_STATS_SUMMARY_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(cutoff_date)
|
||||
.bind(all_time_requests)
|
||||
.bind(all_time_success_requests)
|
||||
.bind(all_time_error_requests)
|
||||
.bind(all_time_input_tokens)
|
||||
.bind(all_time_output_tokens)
|
||||
.bind(all_time_cache_creation_tokens)
|
||||
.bind(all_time_cache_read_tokens)
|
||||
.bind(all_time_cost)
|
||||
.bind(all_time_actual_cost)
|
||||
.bind(total_users)
|
||||
.bind(active_users)
|
||||
.bind(total_api_keys)
|
||||
.bind(active_api_keys)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn refresh_stats_user_summary_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
cutoff_date: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<(), sqlx::Error> {
|
||||
sqlx::query(UPSERT_STATS_USER_SUMMARY_SQL)
|
||||
.bind(cutoff_date)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub(super) struct PercentileSummary {
|
||||
pub(super) p50: Option<i64>,
|
||||
pub(super) p90: Option<i64>,
|
||||
pub(super) p99: Option<i64>,
|
||||
}
|
||||
|
||||
pub(super) fn percentile_ms_to_i64(value: Option<f64>) -> Option<i64> {
|
||||
value.and_then(|raw| raw.is_finite().then_some(raw.floor() as i64))
|
||||
}
|
||||
3110
crates/aether-data/src/backend/stats/postgres_daily/sql.rs
Normal file
3110
crates/aether-data/src/backend/stats/postgres_daily/sql.rs
Normal file
File diff suppressed because it is too large
Load Diff
203
crates/aether-data/src/backend/stats/postgres_hourly/mod.rs
Normal file
203
crates/aether-data/src/backend/stats/postgres_hourly/mod.rs
Normal file
@@ -0,0 +1,203 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::{
|
||||
error::postgres_error, DataLayerError, StatsHourlyAggregationInput,
|
||||
StatsHourlyAggregationSummary,
|
||||
};
|
||||
|
||||
mod sql;
|
||||
|
||||
use self::sql::*;
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<Option<StatsHourlyAggregationSummary>, DataLayerError> {
|
||||
let Some(hour_utc) = next_stats_hourly_bucket(self.pool(), input.target_hour_utc)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
perform_stats_hourly_aggregation_for_hour(self.pool(), hour_utc, input.aggregated_at)
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(postgres_error)
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_stats_hourly_bucket(
|
||||
pool: &crate::driver::postgres::PostgresPool,
|
||||
target_hour_utc: DateTime<Utc>,
|
||||
) -> Result<Option<DateTime<Utc>>, sqlx::Error> {
|
||||
let latest_row = sqlx::query(SELECT_LATEST_STATS_HOURLY_HOUR_SQL)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
let latest_hour = latest_row.try_get::<Option<DateTime<Utc>>, _>("latest_hour")?;
|
||||
let search_from = latest_hour
|
||||
.map(|value| value + chrono::Duration::hours(1))
|
||||
.unwrap_or_else(|| {
|
||||
DateTime::<Utc>::from_timestamp(0, 0).expect("unix epoch should be valid")
|
||||
});
|
||||
let search_until = target_hour_utc + chrono::Duration::hours(1);
|
||||
if search_from >= search_until {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let next_row = sqlx::query(SELECT_NEXT_STATS_HOURLY_BUCKET_SQL)
|
||||
.bind(search_from)
|
||||
.bind(search_until)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
let next_bucket = next_row.try_get::<Option<DateTime<Utc>>, _>("next_bucket")?;
|
||||
Ok(next_bucket.filter(|value| *value <= target_hour_utc))
|
||||
}
|
||||
|
||||
async fn perform_stats_hourly_aggregation_for_hour(
|
||||
pool: &crate::driver::postgres::PostgresPool,
|
||||
hour_utc: DateTime<Utc>,
|
||||
aggregated_at: DateTime<Utc>,
|
||||
) -> Result<StatsHourlyAggregationSummary, sqlx::Error> {
|
||||
let hour_end = hour_utc + chrono::Duration::hours(1);
|
||||
let mut tx = pool.begin().await?;
|
||||
|
||||
let row = sqlx::query(SELECT_STATS_HOURLY_AGGREGATE_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.fetch_one(&mut *tx)
|
||||
.await?;
|
||||
let total_requests = row.try_get::<i64, _>("total_requests")?;
|
||||
let error_requests = row.try_get::<i64, _>("error_requests")?;
|
||||
let success_requests = total_requests.saturating_sub(error_requests);
|
||||
sqlx::query(UPSERT_STATS_HOURLY_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(hour_utc)
|
||||
.bind(total_requests)
|
||||
.bind(row.try_get::<i64, _>("cache_hit_total_requests")?)
|
||||
.bind(row.try_get::<i64, _>("cache_hit_requests")?)
|
||||
.bind(row.try_get::<i64, _>("completed_total_requests")?)
|
||||
.bind(row.try_get::<i64, _>("completed_cache_hit_requests")?)
|
||||
.bind(row.try_get::<i64, _>("completed_input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("completed_cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("completed_cache_read_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("completed_total_input_context")?)
|
||||
.bind(row.try_get::<f64, _>("completed_cache_creation_cost")?)
|
||||
.bind(row.try_get::<f64, _>("completed_cache_read_cost")?)
|
||||
.bind(row.try_get::<f64, _>("settled_total_cost")?)
|
||||
.bind(row.try_get::<i64, _>("settled_total_requests")?)
|
||||
.bind(row.try_get::<i64, _>("settled_input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("settled_output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("settled_cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("settled_cache_read_tokens")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("settled_first_finalized_at_unix_secs")?)
|
||||
.bind(row.try_get::<Option<i64>, _>("settled_last_finalized_at_unix_secs")?)
|
||||
.bind(success_requests)
|
||||
.bind(error_requests)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
|
||||
.bind(row.try_get::<f64, _>("total_cost")?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost")?)
|
||||
.bind(row.try_get::<f64, _>("response_time_sum_ms")?)
|
||||
.bind(row.try_get::<i64, _>("response_time_samples")?)
|
||||
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
|
||||
.bind(true)
|
||||
.bind(aggregated_at)
|
||||
.bind(aggregated_at)
|
||||
.bind(aggregated_at)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
let user_rows =
|
||||
upsert_stats_hourly_user_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
|
||||
let user_model_rows =
|
||||
upsert_stats_hourly_user_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
|
||||
let model_rows =
|
||||
upsert_stats_hourly_model_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
|
||||
let provider_rows =
|
||||
upsert_stats_hourly_provider_rows(&mut tx, hour_utc, hour_end, aggregated_at).await?;
|
||||
tx.commit().await?;
|
||||
|
||||
Ok(StatsHourlyAggregationSummary {
|
||||
hour_utc,
|
||||
total_requests,
|
||||
user_rows,
|
||||
user_model_rows,
|
||||
model_rows,
|
||||
provider_rows,
|
||||
})
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_user_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
hour_utc: DateTime<Utc>,
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
hour_utc: DateTime<Utc>,
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_user_model_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
hour_utc: DateTime<Utc>,
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_MODEL_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_provider_rows(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
hour_utc: DateTime<Utc>,
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
906
crates/aether-data/src/backend/stats/postgres_hourly/sql.rs
Normal file
906
crates/aether-data/src/backend/stats/postgres_hourly/sql.rs
Normal file
@@ -0,0 +1,906 @@
|
||||
pub(super) const SELECT_LATEST_STATS_HOURLY_HOUR_SQL: &str = r#"
|
||||
SELECT MAX(hour_utc) AS latest_hour
|
||||
FROM stats_hourly
|
||||
WHERE is_complete IS TRUE
|
||||
"#;
|
||||
pub(super) const SELECT_NEXT_STATS_HOURLY_BUCKET_SQL: &str = r#"
|
||||
SELECT date_trunc('hour', MIN(created_at)) AS next_bucket
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
"#;
|
||||
pub(super) const SELECT_STATS_HOURLY_AGGREGATE_SQL: &str = r#"
|
||||
SELECT
|
||||
(
|
||||
SELECT CAST(COUNT(cache_hit_usage.id) AS BIGINT)
|
||||
FROM usage_billing_facts AS cache_hit_usage
|
||||
WHERE cache_hit_usage.created_at >= $1
|
||||
AND cache_hit_usage.created_at < $2
|
||||
) AS cache_hit_total_requests,
|
||||
(
|
||||
SELECT CAST(
|
||||
COUNT(cache_hit_usage.id) FILTER (
|
||||
WHERE GREATEST(COALESCE(cache_hit_usage.cache_read_input_tokens, 0), 0) > 0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS cache_hit_usage
|
||||
WHERE cache_hit_usage.created_at >= $1
|
||||
AND cache_hit_usage.created_at < $2
|
||||
) AS cache_hit_requests,
|
||||
(
|
||||
SELECT CAST(COUNT(completed_usage.id) AS BIGINT)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_total_requests,
|
||||
(
|
||||
SELECT CAST(
|
||||
COUNT(completed_usage.id) FILTER (
|
||||
WHERE GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0) > 0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_cache_hit_requests,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(SUM(GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)), 0) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_input_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN COALESCE(completed_usage.cache_creation_input_tokens, 0) = 0
|
||||
AND (
|
||||
COALESCE(completed_usage.cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(completed_usage.cache_creation_input_tokens_1h, 0)
|
||||
) > 0
|
||||
THEN COALESCE(completed_usage.cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(completed_usage.cache_creation_input_tokens_1h, 0)
|
||||
ELSE COALESCE(completed_usage.cache_creation_input_tokens, 0)
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_cache_creation_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0)),
|
||||
0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_cache_read_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN split_part(
|
||||
lower(
|
||||
COALESCE(
|
||||
COALESCE(
|
||||
completed_usage.endpoint_api_format,
|
||||
completed_usage.api_format
|
||||
),
|
||||
''
|
||||
)
|
||||
),
|
||||
':',
|
||||
1
|
||||
) IN ('claude', 'anthropic')
|
||||
THEN GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)
|
||||
+ CASE
|
||||
WHEN COALESCE(completed_usage.cache_creation_input_tokens, 0) = 0
|
||||
AND (
|
||||
COALESCE(completed_usage.cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(completed_usage.cache_creation_input_tokens_1h, 0)
|
||||
) > 0
|
||||
THEN COALESCE(completed_usage.cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(completed_usage.cache_creation_input_tokens_1h, 0)
|
||||
ELSE COALESCE(completed_usage.cache_creation_input_tokens, 0)
|
||||
END
|
||||
+ GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0)
|
||||
WHEN split_part(
|
||||
lower(
|
||||
COALESCE(
|
||||
COALESCE(
|
||||
completed_usage.endpoint_api_format,
|
||||
completed_usage.api_format
|
||||
),
|
||||
''
|
||||
)
|
||||
),
|
||||
':',
|
||||
1
|
||||
) IN ('openai', 'gemini', 'google')
|
||||
THEN (
|
||||
CASE
|
||||
WHEN GREATEST(COALESCE(completed_usage.input_tokens, 0), 0) <= 0
|
||||
THEN 0
|
||||
WHEN GREATEST(
|
||||
COALESCE(completed_usage.cache_read_input_tokens, 0),
|
||||
0
|
||||
) <= 0
|
||||
THEN GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)
|
||||
ELSE GREATEST(
|
||||
GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)
|
||||
- GREATEST(
|
||||
COALESCE(completed_usage.cache_read_input_tokens, 0),
|
||||
0
|
||||
),
|
||||
0
|
||||
)
|
||||
END
|
||||
) + GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0)
|
||||
ELSE CASE
|
||||
WHEN (
|
||||
CASE
|
||||
WHEN COALESCE(
|
||||
completed_usage.cache_creation_input_tokens,
|
||||
0
|
||||
) = 0
|
||||
AND (
|
||||
COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_5m,
|
||||
0
|
||||
)
|
||||
+ COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_1h,
|
||||
0
|
||||
)
|
||||
) > 0
|
||||
THEN COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_5m,
|
||||
0
|
||||
)
|
||||
+ COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_1h,
|
||||
0
|
||||
)
|
||||
ELSE COALESCE(
|
||||
completed_usage.cache_creation_input_tokens,
|
||||
0
|
||||
)
|
||||
END
|
||||
) > 0
|
||||
THEN GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)
|
||||
+ (
|
||||
CASE
|
||||
WHEN COALESCE(
|
||||
completed_usage.cache_creation_input_tokens,
|
||||
0
|
||||
) = 0
|
||||
AND (
|
||||
COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_5m,
|
||||
0
|
||||
)
|
||||
+ COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_1h,
|
||||
0
|
||||
)
|
||||
) > 0
|
||||
THEN COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_5m,
|
||||
0
|
||||
)
|
||||
+ COALESCE(
|
||||
completed_usage.cache_creation_input_tokens_1h,
|
||||
0
|
||||
)
|
||||
ELSE COALESCE(
|
||||
completed_usage.cache_creation_input_tokens,
|
||||
0
|
||||
)
|
||||
END
|
||||
)
|
||||
+ GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0)
|
||||
ELSE GREATEST(COALESCE(completed_usage.input_tokens, 0), 0)
|
||||
+ GREATEST(COALESCE(completed_usage.cache_read_input_tokens, 0), 0)
|
||||
END
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_total_input_context,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
COALESCE(
|
||||
CAST(completed_usage.cache_creation_cost_usd AS DOUBLE PRECISION),
|
||||
0
|
||||
)
|
||||
),
|
||||
0
|
||||
) AS DOUBLE PRECISION
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_cache_creation_cost,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
COALESCE(CAST(completed_usage.cache_read_cost_usd AS DOUBLE PRECISION), 0)
|
||||
),
|
||||
0
|
||||
) AS DOUBLE PRECISION
|
||||
)
|
||||
FROM usage_billing_facts AS completed_usage
|
||||
WHERE completed_usage.created_at >= $1
|
||||
AND completed_usage.created_at < $2
|
||||
AND completed_usage.status = 'completed'
|
||||
) AS completed_cache_read_cost,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(SUM(COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0)), 0)
|
||||
AS DOUBLE PRECISION
|
||||
)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_total_cost,
|
||||
(
|
||||
SELECT CAST(COUNT(settled_usage.id) AS BIGINT)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_total_requests,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(SUM(GREATEST(COALESCE(settled_usage.input_tokens, 0), 0)), 0) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_input_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(SUM(GREATEST(COALESCE(settled_usage.output_tokens, 0), 0)), 0) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_output_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(GREATEST(COALESCE(settled_usage.cache_creation_input_tokens, 0), 0)),
|
||||
0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_cache_creation_tokens,
|
||||
(
|
||||
SELECT CAST(
|
||||
COALESCE(
|
||||
SUM(GREATEST(COALESCE(settled_usage.cache_read_input_tokens, 0), 0)),
|
||||
0
|
||||
) AS BIGINT
|
||||
)
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_cache_read_tokens,
|
||||
(
|
||||
SELECT MIN(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_first_finalized_at_unix_secs,
|
||||
(
|
||||
SELECT MAX(CAST(EXTRACT(EPOCH FROM settled_usage.finalized_at) AS BIGINT))
|
||||
FROM usage_billing_facts AS settled_usage
|
||||
WHERE settled_usage.created_at >= $1
|
||||
AND settled_usage.created_at < $2
|
||||
AND settled_usage.billing_status = 'settled'
|
||||
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
) AS settled_last_finalized_at_unix_secs,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN status_code >= 400
|
||||
OR lower(COALESCE(status, '')) = 'failed'
|
||||
OR error_message IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens, 0) = 0
|
||||
AND (
|
||||
COALESCE(cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(cache_creation_input_tokens_1h, 0)
|
||||
) > 0
|
||||
THEN COALESCE(cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(cache_creation_input_tokens_1h, 0)
|
||||
ELSE COALESCE(cache_creation_input_tokens, 0)
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost,
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL
|
||||
THEN GREATEST(COALESCE(response_time_ms, 0), 0)::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS response_time_sum_ms,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS response_time_samples,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
"#;
|
||||
pub(super) const UPSERT_STATS_HOURLY_SQL: &str = r#"
|
||||
INSERT INTO stats_hourly (
|
||||
id,
|
||||
hour_utc,
|
||||
total_requests,
|
||||
cache_hit_total_requests,
|
||||
cache_hit_requests,
|
||||
completed_total_requests,
|
||||
completed_cache_hit_requests,
|
||||
completed_input_tokens,
|
||||
completed_cache_creation_tokens,
|
||||
completed_cache_read_tokens,
|
||||
completed_total_input_context,
|
||||
completed_cache_creation_cost,
|
||||
completed_cache_read_cost,
|
||||
settled_total_cost,
|
||||
settled_total_requests,
|
||||
settled_input_tokens,
|
||||
settled_output_tokens,
|
||||
settled_cache_creation_tokens,
|
||||
settled_cache_read_tokens,
|
||||
settled_first_finalized_at_unix_secs,
|
||||
settled_last_finalized_at_unix_secs,
|
||||
success_requests,
|
||||
error_requests,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_creation_tokens,
|
||||
cache_read_tokens,
|
||||
total_cost,
|
||||
actual_total_cost,
|
||||
response_time_sum_ms,
|
||||
response_time_samples,
|
||||
avg_response_time_ms,
|
||||
is_complete,
|
||||
aggregated_at,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8,
|
||||
$9, $10, $11, $12, $13, $14, $15, $16,
|
||||
$17, $18, $19, $20, $21, $22, $23, $24,
|
||||
$25, $26, $27, $28, $29, $30, $31, $32,
|
||||
$33, $34, $35, $36
|
||||
)
|
||||
ON CONFLICT (hour_utc)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
cache_hit_total_requests = EXCLUDED.cache_hit_total_requests,
|
||||
cache_hit_requests = EXCLUDED.cache_hit_requests,
|
||||
completed_total_requests = EXCLUDED.completed_total_requests,
|
||||
completed_cache_hit_requests = EXCLUDED.completed_cache_hit_requests,
|
||||
completed_input_tokens = EXCLUDED.completed_input_tokens,
|
||||
completed_cache_creation_tokens = EXCLUDED.completed_cache_creation_tokens,
|
||||
completed_cache_read_tokens = EXCLUDED.completed_cache_read_tokens,
|
||||
completed_total_input_context = EXCLUDED.completed_total_input_context,
|
||||
completed_cache_creation_cost = EXCLUDED.completed_cache_creation_cost,
|
||||
completed_cache_read_cost = EXCLUDED.completed_cache_read_cost,
|
||||
settled_total_cost = EXCLUDED.settled_total_cost,
|
||||
settled_total_requests = EXCLUDED.settled_total_requests,
|
||||
settled_input_tokens = EXCLUDED.settled_input_tokens,
|
||||
settled_output_tokens = EXCLUDED.settled_output_tokens,
|
||||
settled_cache_creation_tokens = EXCLUDED.settled_cache_creation_tokens,
|
||||
settled_cache_read_tokens = EXCLUDED.settled_cache_read_tokens,
|
||||
settled_first_finalized_at_unix_secs = EXCLUDED.settled_first_finalized_at_unix_secs,
|
||||
settled_last_finalized_at_unix_secs = EXCLUDED.settled_last_finalized_at_unix_secs,
|
||||
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,
|
||||
response_time_sum_ms = EXCLUDED.response_time_sum_ms,
|
||||
response_time_samples = EXCLUDED.response_time_samples,
|
||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||
is_complete = EXCLUDED.is_complete,
|
||||
aggregated_at = EXCLUDED.aggregated_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
pub(super) const UPSERT_STATS_HOURLY_USER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
user_id,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN status_code >= 400
|
||||
OR lower(COALESCE(status, '')) = 'failed'
|
||||
OR error_message IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens, 0) = 0
|
||||
AND (
|
||||
COALESCE(cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(cache_creation_input_tokens_1h, 0)
|
||||
) > 0
|
||||
THEN COALESCE(cache_creation_input_tokens_5m, 0)
|
||||
+ COALESCE(cache_creation_input_tokens_1h, 0)
|
||||
ELSE COALESCE(cache_creation_input_tokens, 0)
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost,
|
||||
CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0)
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS DOUBLE PRECISION
|
||||
) AS settled_total_cost,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS settled_total_requests,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN GREATEST(COALESCE(input_tokens, 0), 0)
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS settled_input_tokens,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN GREATEST(COALESCE(output_tokens, 0), 0)
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS settled_output_tokens,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN GREATEST(COALESCE(cache_creation_input_tokens, 0), 0)
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS settled_cache_creation_tokens,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
THEN GREATEST(COALESCE(cache_read_input_tokens, 0), 0)
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS settled_cache_read_tokens,
|
||||
MIN(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
AND finalized_at IS NOT NULL
|
||||
THEN CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT)
|
||||
ELSE NULL
|
||||
END
|
||||
) AS settled_first_finalized_at_unix_secs,
|
||||
MAX(
|
||||
CASE
|
||||
WHEN billing_status = 'settled'
|
||||
AND COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) > 0
|
||||
AND finalized_at IS NOT NULL
|
||||
THEN CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT)
|
||||
ELSE NULL
|
||||
END
|
||||
) AS settled_last_finalized_at_unix_secs,
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL
|
||||
THEN GREATEST(COALESCE(response_time_ms, 0), 0)::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS response_time_sum_ms,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS response_time_samples
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY user_id
|
||||
)
|
||||
INSERT INTO stats_hourly_user (
|
||||
id,
|
||||
hour_utc,
|
||||
user_id,
|
||||
total_requests,
|
||||
success_requests,
|
||||
error_requests,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_creation_tokens,
|
||||
cache_read_tokens,
|
||||
total_cost,
|
||||
actual_total_cost,
|
||||
settled_total_cost,
|
||||
settled_total_requests,
|
||||
settled_input_tokens,
|
||||
settled_output_tokens,
|
||||
settled_cache_creation_tokens,
|
||||
settled_cache_read_tokens,
|
||||
settled_first_finalized_at_unix_secs,
|
||||
settled_last_finalized_at_unix_secs,
|
||||
response_time_sum_ms,
|
||||
response_time_samples,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-user:', aggregated.user_id, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.user_id,
|
||||
aggregated.total_requests,
|
||||
GREATEST(aggregated.total_requests - aggregated.error_requests, 0),
|
||||
aggregated.error_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.actual_total_cost,
|
||||
aggregated.settled_total_cost,
|
||||
aggregated.settled_total_requests,
|
||||
aggregated.settled_input_tokens,
|
||||
aggregated.settled_output_tokens,
|
||||
aggregated.settled_cache_creation_tokens,
|
||||
aggregated.settled_cache_read_tokens,
|
||||
aggregated.settled_first_finalized_at_unix_secs,
|
||||
aggregated.settled_last_finalized_at_unix_secs,
|
||||
aggregated.response_time_sum_ms,
|
||||
aggregated.response_time_samples,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
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,
|
||||
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,
|
||||
settled_total_cost = EXCLUDED.settled_total_cost,
|
||||
settled_total_requests = EXCLUDED.settled_total_requests,
|
||||
settled_input_tokens = EXCLUDED.settled_input_tokens,
|
||||
settled_output_tokens = EXCLUDED.settled_output_tokens,
|
||||
settled_cache_creation_tokens = EXCLUDED.settled_cache_creation_tokens,
|
||||
settled_cache_read_tokens = EXCLUDED.settled_cache_read_tokens,
|
||||
settled_first_finalized_at_unix_secs = EXCLUDED.settled_first_finalized_at_unix_secs,
|
||||
settled_last_finalized_at_unix_secs = EXCLUDED.settled_last_finalized_at_unix_secs,
|
||||
response_time_sum_ms = EXCLUDED.response_time_sum_ms,
|
||||
response_time_samples = EXCLUDED.response_time_samples,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
pub(super) const UPSERT_STATS_HOURLY_MODEL_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL
|
||||
THEN GREATEST(COALESCE(response_time_ms, 0), 0)::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS response_time_sum_ms,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS response_time_samples,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND model IS NOT NULL
|
||||
AND model <> ''
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY model
|
||||
)
|
||||
INSERT INTO stats_hourly_model (
|
||||
id,
|
||||
hour_utc,
|
||||
model,
|
||||
total_requests,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_cost,
|
||||
response_time_sum_ms,
|
||||
response_time_samples,
|
||||
avg_response_time_ms,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-model:', aggregated.model, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.model,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.response_time_sum_ms,
|
||||
aggregated.response_time_samples,
|
||||
aggregated.avg_response_time_ms,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
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,
|
||||
response_time_sum_ms = EXCLUDED.response_time_sum_ms,
|
||||
response_time_samples = EXCLUDED.response_time_samples,
|
||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
pub(super) const UPSERT_STATS_HOURLY_USER_MODEL_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
user_id,
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL
|
||||
THEN GREATEST(COALESCE(response_time_ms, 0), 0)::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS response_time_sum_ms,
|
||||
CAST(COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN response_time_ms IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT) AS response_time_samples
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
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
|
||||
)
|
||||
INSERT INTO stats_hourly_user_model (
|
||||
id,
|
||||
hour_utc,
|
||||
user_id,
|
||||
model,
|
||||
total_requests,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_cost,
|
||||
response_time_sum_ms,
|
||||
response_time_samples,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-user-model:', aggregated.user_id, ':', aggregated.model, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.user_id,
|
||||
aggregated.model,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.response_time_sum_ms,
|
||||
aggregated.response_time_samples,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
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,
|
||||
response_time_sum_ms = EXCLUDED.response_time_sum_ms,
|
||||
response_time_samples = EXCLUDED.response_time_samples,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
pub(super) const UPSERT_STATS_HOURLY_PROVIDER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
provider_name,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage_billing_facts AS usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND provider_name IS NOT NULL
|
||||
AND provider_name <> ''
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY provider_name
|
||||
)
|
||||
INSERT INTO stats_hourly_provider (
|
||||
id,
|
||||
hour_utc,
|
||||
provider_name,
|
||||
total_requests,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_cost,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-provider:', aggregated.provider_name, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.provider_name,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
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
|
||||
"#;
|
||||
379
crates/aether-data/src/backend/stats/sqlite.rs
Normal file
379
crates/aether-data/src/backend/stats/sqlite.rs
Normal file
@@ -0,0 +1,379 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
|
||||
use crate::backend::stats_common::{stats_id, unix_ms, unix_secs, utc_from_unix_secs};
|
||||
use crate::backend::SqliteBackend;
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::{
|
||||
DataLayerError, StatsDailyAggregationInput, StatsDailyAggregationSummary,
|
||||
StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
|
||||
};
|
||||
|
||||
impl SqliteBackend {
|
||||
pub async fn aggregate_stats_hourly(
|
||||
&self,
|
||||
input: &StatsHourlyAggregationInput,
|
||||
) -> Result<Option<StatsHourlyAggregationSummary>, 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<Option<StatsDailyAggregationSummary>, 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<Utc>,
|
||||
) -> Result<Option<i64>, DataLayerError> {
|
||||
let latest_hour: Option<i64> =
|
||||
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<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT MIN(CAST(created_at_unix_ms / 3600000 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(unix_ms(search_from)?)
|
||||
.bind(unix_ms(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<Utc>,
|
||||
) -> Result<Option<i64>, DataLayerError> {
|
||||
let latest_day: Option<i64> =
|
||||
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<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT MIN(CAST(created_at_unix_ms / 86400000 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(unix_ms(search_from)?)
|
||||
.bind(unix_ms(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_category IS NOT NULL AND error_category <> '')
|
||||
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,
|
||||
COALESCE(SUM(total_cost_usd), 0.0) AS total_cost,
|
||||
COALESCE(SUM(actual_total_cost_usd), 0.0) AS actual_total_cost,
|
||||
COALESCE(AVG(response_time_ms), 0.0) 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<Utc>,
|
||||
) -> Result<StatsHourlyAggregationSummary, DataLayerError> {
|
||||
let start_ms = unix_ms(hour_utc_unix_secs)?;
|
||||
let end_ms = unix_ms(hour_utc_unix_secs + 3600)?;
|
||||
let aggregated_at_unix_secs = unix_secs(aggregated_at);
|
||||
let mut tx = pool.begin().await.map_sql_err()?;
|
||||
let row = sqlx::query(SQLITE_STATS_AGGREGATE_SQL)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.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::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("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 = sqlite_group_count(&mut tx, "user_id", start_ms, end_ms).await?;
|
||||
let user_model_rows = sqlite_group_count(&mut tx, "user_id, model", start_ms, end_ms).await?;
|
||||
let model_rows = sqlite_group_count(&mut tx, "model", start_ms, end_ms).await?;
|
||||
let provider_rows = sqlite_group_count(&mut tx, "provider_name", start_ms, end_ms).await?;
|
||||
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<Utc>,
|
||||
) -> Result<StatsDailyAggregationSummary, DataLayerError> {
|
||||
let start_ms = unix_ms(day_start_unix_secs)?;
|
||||
let end_ms = unix_ms(day_start_unix_secs + 86_400)?;
|
||||
let aggregated_at_unix_secs = unix_secs(aggregated_at);
|
||||
let mut tx = pool.begin().await.map_sql_err()?;
|
||||
let row = sqlx::query(SQLITE_STATS_AGGREGATE_SQL)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let total_requests: i64 = row.try_get("total_requests").map_sql_err()?;
|
||||
let error_requests: i64 = row.try_get("error_requests").map_sql_err()?;
|
||||
let unique_models = sqlite_group_count(&mut tx, "model", start_ms, end_ms).await? as i64;
|
||||
let unique_providers =
|
||||
sqlite_group_count(&mut tx, "provider_name", start_ms, end_ms).await? as i64;
|
||||
|
||||
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, 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::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("total_cost").map_sql_err()?)
|
||||
.bind(row.try_get::<f64, _>("actual_total_cost").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("avg_response_time_ms")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.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 = usize::try_from(unique_models).unwrap_or(usize::MAX);
|
||||
let provider_rows = usize::try_from(unique_providers).unwrap_or(usize::MAX);
|
||||
let api_key_rows = sqlite_group_count(&mut tx, "api_key_id", start_ms, end_ms).await?;
|
||||
let error_rows = sqlite_error_group_count(&mut tx, start_ms, end_ms).await?;
|
||||
let user_rows = sqlite_group_count(&mut tx, "user_id", start_ms, end_ms).await?;
|
||||
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 sqlite_group_count(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
group_columns: &str,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let not_empty = group_columns
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.map(|column| format!("{column} IS NOT NULL AND {column} <> ''"))
|
||||
.collect::<Vec<_>>()
|
||||
.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_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn sqlite_error_group_count(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let count: i64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COUNT(*)
|
||||
FROM (
|
||||
SELECT 1
|
||||
FROM "usage"
|
||||
WHERE created_at_unix_ms >= ?
|
||||
AND created_at_unix_ms < ?
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
AND (
|
||||
status = 'failed'
|
||||
OR status_code >= 400
|
||||
OR (error_category IS NOT NULL AND error_category <> '')
|
||||
)
|
||||
GROUP BY COALESCE(NULLIF(error_category, ''), 'unknown_error'), provider_name, model
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(start_ms)
|
||||
.bind(end_ms)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(usize::try_from(count.max(0)).unwrap_or(usize::MAX))
|
||||
}
|
||||
33
crates/aether-data/src/backend/stats_common.rs
Normal file
33
crates/aether-data/src/backend/stats_common.rs
Normal file
@@ -0,0 +1,33 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::DataLayerError;
|
||||
|
||||
pub(crate) fn unix_secs(value: DateTime<Utc>) -> i64 {
|
||||
value.timestamp().max(0)
|
||||
}
|
||||
|
||||
pub(crate) fn unix_ms(value: i64) -> Result<i64, DataLayerError> {
|
||||
value.checked_mul(1000).ok_or_else(|| {
|
||||
DataLayerError::InvalidInput(format!("timestamp overflow while converting {value} to ms"))
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn utc_from_unix_secs(
|
||||
value: i64,
|
||||
field_name: &str,
|
||||
) -> Result<DateTime<Utc>, DataLayerError> {
|
||||
DateTime::<Utc>::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()
|
||||
}
|
||||
502
crates/aether-data/src/backend/system.rs
Normal file
502
crates/aether-data/src/backend/system.rs
Normal file
@@ -0,0 +1,502 @@
|
||||
use futures_util::TryStreamExt;
|
||||
use sqlx::Row;
|
||||
|
||||
use super::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::error::{SqlResultExt, SqlxResultExt};
|
||||
use crate::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const POSTGRES_FIND_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
SELECT value
|
||||
FROM system_configs
|
||||
WHERE key = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const POSTGRES_UPSERT_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
INSERT INTO system_configs (id, key, value, description, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW())
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value,
|
||||
description = COALESCE(EXCLUDED.description, system_configs.description),
|
||||
updated_at = NOW()
|
||||
RETURNING value
|
||||
"#;
|
||||
|
||||
const POSTGRES_LIST_SYSTEM_CONFIG_ENTRIES_SQL: &str = r#"
|
||||
SELECT
|
||||
key,
|
||||
value,
|
||||
description,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
FROM system_configs
|
||||
ORDER BY key ASC
|
||||
"#;
|
||||
|
||||
const POSTGRES_UPSERT_SYSTEM_CONFIG_ENTRY_SQL: &str = r#"
|
||||
INSERT INTO system_configs (id, key, value, description, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW())
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value,
|
||||
description = COALESCE(EXCLUDED.description, system_configs.description),
|
||||
updated_at = NOW()
|
||||
RETURNING
|
||||
key,
|
||||
value,
|
||||
description,
|
||||
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const POSTGRES_DELETE_SYSTEM_CONFIG_VALUE_SQL: &str = r#"
|
||||
DELETE FROM system_configs
|
||||
WHERE key = $1
|
||||
"#;
|
||||
|
||||
const POSTGRES_READ_ADMIN_SYSTEM_STATS_SQL: &str = r#"
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM users) AS total_users,
|
||||
(SELECT COUNT(*) FROM users WHERE is_active IS TRUE) AS active_users,
|
||||
(SELECT COUNT(*) FROM api_keys) AS total_api_keys,
|
||||
(SELECT COUNT(*) FROM usage) AS total_requests
|
||||
"#;
|
||||
|
||||
const MYSQL_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
|
||||
"#;
|
||||
|
||||
const SQLITE_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
|
||||
"#;
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
let row = sqlx::query(POSTGRES_FIND_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(key)
|
||||
.fetch_optional(self.pool())
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.map(|row| row.try_get("value"))
|
||||
.transpose()
|
||||
.map_postgres_err()
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, DataLayerError> {
|
||||
let row = sqlx::query(POSTGRES_UPSERT_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.bind(description)
|
||||
.fetch_one(self.pool())
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.try_get("value").map_postgres_err()
|
||||
}
|
||||
|
||||
pub async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
||||
let mut rows = sqlx::query(POSTGRES_LIST_SYSTEM_CONFIG_ENTRIES_SQL).fetch(self.pool());
|
||||
let mut entries = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
entries.push(StoredSystemConfigEntry {
|
||||
key: row.try_get("key").map_postgres_err()?,
|
||||
value: row.try_get("value").map_postgres_err()?,
|
||||
description: row.try_get("description").map_postgres_err()?,
|
||||
updated_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
||||
.map_postgres_err()?
|
||||
.map(|value| value.max(0) as u64),
|
||||
});
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config_entry(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<StoredSystemConfigEntry, DataLayerError> {
|
||||
let row = sqlx::query(POSTGRES_UPSERT_SYSTEM_CONFIG_ENTRY_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.bind(description)
|
||||
.fetch_one(self.pool())
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(StoredSystemConfigEntry {
|
||||
key: row.try_get("key").map_postgres_err()?,
|
||||
value: row.try_get("value").map_postgres_err()?,
|
||||
description: row.try_get("description").map_postgres_err()?,
|
||||
updated_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
||||
.map_postgres_err()?
|
||||
.map(|value| value.max(0) as u64),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn delete_system_config_value(&self, key: &str) -> Result<bool, DataLayerError> {
|
||||
let result = sqlx::query(POSTGRES_DELETE_SYSTEM_CONFIG_VALUE_SQL)
|
||||
.bind(key)
|
||||
.execute(self.pool())
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
pub async fn read_admin_system_stats(&self) -> Result<AdminSystemStats, DataLayerError> {
|
||||
let row = sqlx::query(POSTGRES_READ_ADMIN_SYSTEM_STATS_SQL)
|
||||
.fetch_one(self.pool())
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
postgres_admin_system_stats(row)
|
||||
}
|
||||
}
|
||||
|
||||
impl MysqlBackend {
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, 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 upsert_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, DataLayerError> {
|
||||
Ok(self
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await?
|
||||
.value)
|
||||
}
|
||||
|
||||
pub async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, 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::<Option<i64>, _>("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<StoredSystemConfigEntry, DataLayerError> {
|
||||
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
|
||||
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<bool, DataLayerError> {
|
||||
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<AdminSystemStats, DataLayerError> {
|
||||
let row = sqlx::query(MYSQL_READ_ADMIN_SYSTEM_STATS_SQL)
|
||||
.fetch_one(self.pool())
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
mysql_admin_system_stats(row)
|
||||
}
|
||||
}
|
||||
|
||||
impl SqliteBackend {
|
||||
pub async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, 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 upsert_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, DataLayerError> {
|
||||
Ok(self
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await?
|
||||
.value)
|
||||
}
|
||||
|
||||
pub async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, 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::<Option<i64>, _>("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<StoredSystemConfigEntry, DataLayerError> {
|
||||
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<bool, DataLayerError> {
|
||||
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<AdminSystemStats, DataLayerError> {
|
||||
let row = sqlx::query(SQLITE_READ_ADMIN_SYSTEM_STATS_SQL)
|
||||
.fetch_one(self.pool())
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlite_admin_system_stats(row)
|
||||
}
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
chrono::Utc::now().timestamp().max(0) as u64
|
||||
}
|
||||
|
||||
fn serialize_json_value(value: &serde_json::Value) -> Result<String, DataLayerError> {
|
||||
serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid system config JSON value: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_json_value(value: String) -> Result<serde_json::Value, DataLayerError> {
|
||||
serde_json::from_str(&value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid system config JSON value: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn postgres_admin_system_stats(
|
||||
row: sqlx::postgres::PgRow,
|
||||
) -> Result<AdminSystemStats, DataLayerError> {
|
||||
Ok(AdminSystemStats {
|
||||
total_users: row
|
||||
.try_get::<i64, _>("total_users")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64,
|
||||
active_users: row
|
||||
.try_get::<i64, _>("active_users")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64,
|
||||
total_api_keys: row
|
||||
.try_get::<i64, _>("total_api_keys")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64,
|
||||
total_requests: row
|
||||
.try_get::<i64, _>("total_requests")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64,
|
||||
})
|
||||
}
|
||||
|
||||
fn mysql_admin_system_stats(
|
||||
row: sqlx::mysql::MySqlRow,
|
||||
) -> Result<AdminSystemStats, DataLayerError> {
|
||||
Ok(AdminSystemStats {
|
||||
total_users: row.try_get::<i64, _>("total_users").map_sql_err()?.max(0) as u64,
|
||||
active_users: row.try_get::<i64, _>("active_users").map_sql_err()?.max(0) as u64,
|
||||
total_api_keys: row
|
||||
.try_get::<i64, _>("total_api_keys")
|
||||
.map_sql_err()?
|
||||
.max(0) as u64,
|
||||
total_requests: row
|
||||
.try_get::<i64, _>("total_requests")
|
||||
.map_sql_err()?
|
||||
.max(0) as u64,
|
||||
})
|
||||
}
|
||||
|
||||
fn sqlite_admin_system_stats(
|
||||
row: sqlx::sqlite::SqliteRow,
|
||||
) -> Result<AdminSystemStats, DataLayerError> {
|
||||
Ok(AdminSystemStats {
|
||||
total_users: row.try_get::<i64, _>("total_users").map_sql_err()?.max(0) as u64,
|
||||
active_users: row.try_get::<i64, _>("active_users").map_sql_err()?.max(0) as u64,
|
||||
total_api_keys: row
|
||||
.try_get::<i64, _>("total_api_keys")
|
||||
.map_sql_err()?
|
||||
.max(0) as u64,
|
||||
total_requests: row
|
||||
.try_get::<i64, _>("total_requests")
|
||||
.map_sql_err()?
|
||||
.max(0) as u64,
|
||||
})
|
||||
}
|
||||
60
crates/aether-data/src/backend/transactions.rs
Normal file
60
crates/aether-data/src/backend/transactions.rs
Normal file
@@ -0,0 +1,60 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::driver::postgres::PostgresTransactionRunner;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataTransactionBackends {
|
||||
postgres: Option<PostgresTransactionRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataTransactionBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataTransactionBackends")
|
||||
.field("has_postgres", &self.postgres.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataTransactionBackends {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
postgres: postgres.map(PostgresBackend::transaction_runner),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<PostgresTransactionRunner> {
|
||||
self.postgres.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataTransactionBackends;
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::driver::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_transaction_runner_from_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let transactions = DataTransactionBackends::from_postgres(Some(&backend));
|
||||
|
||||
assert!(transactions.has_any());
|
||||
assert!(transactions.postgres().is_some());
|
||||
}
|
||||
}
|
||||
462
crates/aether-data/src/backend/wallet.rs
Normal file
462
crates/aether-data/src/backend/wallet.rs
Normal file
@@ -0,0 +1,462 @@
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::Row;
|
||||
|
||||
use crate::backend::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::error::{SqlResultExt, SqlxResultExt};
|
||||
use crate::{DataLayerError, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult};
|
||||
|
||||
const POSTGRES_UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id,
|
||||
COUNT(*) AS total_requests,
|
||||
CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS DOUBLE PRECISION) 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_billing_facts AS usage
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id IS NOT NULL
|
||||
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) >= $1
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $2
|
||||
GROUP BY usage_settlement_snapshots.wallet_id
|
||||
)
|
||||
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
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('wallet-daily-usage:', aggregated.wallet_id, ':', CAST($3 AS TEXT), ':', $4)),
|
||||
aggregated.wallet_id,
|
||||
$3,
|
||||
$4,
|
||||
aggregated.total_cost_usd,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.first_finalized_at,
|
||||
aggregated.last_finalized_at,
|
||||
$5,
|
||||
$5,
|
||||
$5
|
||||
FROM aggregated
|
||||
ON CONFLICT (wallet_id, billing_date, billing_timezone)
|
||||
DO UPDATE SET
|
||||
total_cost_usd = EXCLUDED.total_cost_usd,
|
||||
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,
|
||||
first_finalized_at = EXCLUDED.first_finalized_at,
|
||||
last_finalized_at = EXCLUDED.last_finalized_at,
|
||||
aggregated_at = EXCLUDED.aggregated_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
|
||||
const POSTGRES_DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL: &str = r#"
|
||||
DELETE FROM wallet_daily_usage_ledgers AS ledgers
|
||||
WHERE ledgers.billing_date = $1
|
||||
AND ledgers.billing_timezone = $2
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM usage_billing_facts AS usage
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id = 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) >= $3
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $4
|
||||
)
|
||||
"#;
|
||||
|
||||
const MYSQL_SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id AS wallet_id,
|
||||
COUNT(*) AS total_requests,
|
||||
COALESCE(SUM(`usage`.total_cost_usd), 0) 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
|
||||
"#;
|
||||
|
||||
const SQLITE_SELECT_WALLET_DAILY_USAGE_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id AS wallet_id,
|
||||
COUNT(*) AS total_requests,
|
||||
COALESCE(SUM("usage".total_cost_usd), 0) 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
|
||||
"#;
|
||||
|
||||
fn u64_to_i64(value: u64, field_name: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value)
|
||||
.map_err(|_| DataLayerError::InvalidInput(format!("invalid {field_name}: {value}")))
|
||||
}
|
||||
|
||||
fn unix_secs_to_utc(
|
||||
value: u64,
|
||||
field_name: &str,
|
||||
) -> Result<chrono::DateTime<chrono::Utc>, DataLayerError> {
|
||||
let value = u64_to_i64(value, field_name)?;
|
||||
chrono::DateTime::<chrono::Utc>::from_timestamp(value, 0)
|
||||
.ok_or_else(|| DataLayerError::InvalidInput(format!("invalid {field_name}: {value}")))
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
pub async fn aggregate_wallet_daily_usage(
|
||||
&self,
|
||||
input: &WalletDailyUsageAggregationInput,
|
||||
) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
|
||||
let billing_date = chrono::NaiveDate::parse_from_str(&input.billing_date, "%Y-%m-%d")
|
||||
.map_err(|err| {
|
||||
DataLayerError::InvalidInput(format!("invalid wallet billing_date: {err}"))
|
||||
})?;
|
||||
let window_start = unix_secs_to_utc(input.window_start_unix_secs, "window_start")?;
|
||||
let window_end = unix_secs_to_utc(input.window_end_unix_secs, "window_end")?;
|
||||
let aggregated_at = unix_secs_to_utc(input.aggregated_at_unix_secs, "aggregated_at")?;
|
||||
let mut tx = self.pool().begin().await.map_postgres_err()?;
|
||||
|
||||
let aggregated_wallets = sqlx::query(POSTGRES_UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
|
||||
.bind(window_start)
|
||||
.bind(window_end)
|
||||
.bind(billing_date)
|
||||
.bind(input.billing_timezone.as_str())
|
||||
.bind(aggregated_at)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
|
||||
let deleted_stale_ledgers =
|
||||
sqlx::query(POSTGRES_DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
|
||||
.bind(billing_date)
|
||||
.bind(input.billing_timezone.as_str())
|
||||
.bind(window_start)
|
||||
.bind(window_end)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
|
||||
tx.commit().await.map_postgres_err()?;
|
||||
Ok(WalletDailyUsageAggregationResult {
|
||||
aggregated_wallets: usize::try_from(aggregated_wallets).unwrap_or(usize::MAX),
|
||||
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl MysqlBackend {
|
||||
pub async fn aggregate_wallet_daily_usage(
|
||||
&self,
|
||||
input: &WalletDailyUsageAggregationInput,
|
||||
) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
|
||||
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(MYSQL_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::<f64, _>("total_cost_usd").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("total_requests").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<Option<i64>, _>("first_finalized_at")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<Option<i64>, _>("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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl SqliteBackend {
|
||||
pub async fn aggregate_wallet_daily_usage(
|
||||
&self,
|
||||
input: &WalletDailyUsageAggregationInput,
|
||||
) -> Result<WalletDailyUsageAggregationResult, DataLayerError> {
|
||||
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(SQLITE_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::<f64, _>("total_cost_usd").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("total_requests").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens").map_sql_err()?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens").map_sql_err()?)
|
||||
.bind(
|
||||
row.try_get::<Option<i64>, _>("first_finalized_at")
|
||||
.map_sql_err()?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<Option<i64>, _>("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),
|
||||
})
|
||||
}
|
||||
}
|
||||
58
crates/aether-data/src/backend/workers.rs
Normal file
58
crates/aether-data/src/backend/workers.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::RedisBackend;
|
||||
use crate::driver::redis::{RedisStreamRunner, RedisStreamRunnerConfig};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataWorkerBackends {
|
||||
redis: Option<RedisStreamRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataWorkerBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataWorkerBackends")
|
||||
.field("has_redis", &self.redis.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataWorkerBackends {
|
||||
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
|
||||
Ok(Self {
|
||||
redis: redis
|
||||
.map(|backend| backend.stream_runner(RedisStreamRunnerConfig::default()))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn redis(&self) -> Option<RedisStreamRunner> {
|
||||
self.redis.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.redis.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataWorkerBackends;
|
||||
use crate::backend::RedisBackend;
|
||||
use crate::driver::redis::RedisClientConfig;
|
||||
|
||||
#[test]
|
||||
fn builds_redis_stream_runner_from_backend() {
|
||||
let backend = RedisBackend::from_config(RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
})
|
||||
.expect("redis backend should build");
|
||||
|
||||
let workers =
|
||||
DataWorkerBackends::from_redis(Some(&backend)).expect("worker backends should build");
|
||||
|
||||
assert!(workers.has_any());
|
||||
assert!(workers.redis().is_some());
|
||||
}
|
||||
}
|
||||
258
crates/aether-data/src/backend/write.rs
Normal file
258
crates/aether-data/src/backend/write.rs
Normal file
@@ -0,0 +1,258 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{MysqlBackend, PostgresBackend, SqliteBackend};
|
||||
use crate::repository::announcements::AnnouncementWriteRepository;
|
||||
use crate::repository::auth::AuthApiKeyWriteRepository;
|
||||
use crate::repository::auth_modules::AuthModuleWriteRepository;
|
||||
use crate::repository::candidates::RequestCandidateWriteRepository;
|
||||
use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
|
||||
use crate::repository::global_models::GlobalModelWriteRepository;
|
||||
use crate::repository::management_tokens::ManagementTokenWriteRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
||||
use crate::repository::quota::ProviderQuotaWriteRepository;
|
||||
use crate::repository::settlement::SettlementWriteRepository;
|
||||
use crate::repository::usage::UsageWriteRepository;
|
||||
use crate::repository::video_tasks::VideoTaskWriteRepository;
|
||||
use crate::repository::wallet::WalletWriteRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataWriteRepositories {
|
||||
announcements: Option<Arc<dyn AnnouncementWriteRepository>>,
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyWriteRepository>>,
|
||||
auth_modules: Option<Arc<dyn AuthModuleWriteRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
||||
gemini_file_mappings: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
|
||||
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||
settlement: Option<Arc<dyn SettlementWriteRepository>>,
|
||||
usage: Option<Arc<dyn UsageWriteRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskWriteRepository>>,
|
||||
wallets: Option<Arc<dyn WalletWriteRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataWriteRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataWriteRepositories")
|
||||
.field("has_announcements", &self.announcements.is_some())
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_auth_modules", &self.auth_modules.is_some())
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mappings",
|
||||
&self.gemini_file_mappings.is_some(),
|
||||
)
|
||||
.field("has_global_models", &self.global_models.is_some())
|
||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_settlement", &self.settlement.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
.field("has_wallets", &self.wallets.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataWriteRepositories {
|
||||
pub(crate) fn from_backends(
|
||||
postgres: Option<&PostgresBackend>,
|
||||
mysql: Option<&MysqlBackend>,
|
||||
sqlite: Option<&SqliteBackend>,
|
||||
) -> Self {
|
||||
Self {
|
||||
announcements: postgres
|
||||
.map(PostgresBackend::announcement_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::announcement_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::announcement_write_repository)),
|
||||
auth_api_keys: postgres
|
||||
.map(PostgresBackend::auth_api_key_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_api_key_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_api_key_write_repository)),
|
||||
auth_modules: postgres
|
||||
.map(PostgresBackend::auth_module_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::auth_module_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::auth_module_write_repository)),
|
||||
request_candidates: postgres
|
||||
.map(PostgresBackend::request_candidate_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::request_candidate_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::request_candidate_write_repository)),
|
||||
gemini_file_mappings: postgres
|
||||
.map(PostgresBackend::gemini_file_mapping_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::gemini_file_mapping_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::gemini_file_mapping_write_repository)),
|
||||
global_models: postgres
|
||||
.map(PostgresBackend::global_model_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::global_model_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::global_model_write_repository)),
|
||||
management_tokens: postgres
|
||||
.map(PostgresBackend::management_token_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::management_token_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::management_token_write_repository)),
|
||||
oauth_providers: postgres
|
||||
.map(PostgresBackend::oauth_provider_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
|
||||
proxy_nodes: postgres
|
||||
.map(PostgresBackend::proxy_node_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::proxy_node_write_repository)),
|
||||
provider_catalog: postgres
|
||||
.map(PostgresBackend::provider_catalog_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_catalog_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_catalog_write_repository)),
|
||||
provider_quotas: postgres
|
||||
.map(PostgresBackend::provider_quota_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_quota_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_quota_write_repository)),
|
||||
settlement: postgres
|
||||
.map(PostgresBackend::settlement_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::settlement_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::settlement_write_repository)),
|
||||
usage: postgres
|
||||
.map(PostgresBackend::usage_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::usage_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::usage_write_repository)),
|
||||
video_tasks: postgres
|
||||
.map(PostgresBackend::video_task_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::video_task_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::video_task_write_repository)),
|
||||
wallets: postgres
|
||||
.map(PostgresBackend::wallet_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::wallet_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::wallet_write_repository)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self::from_backends(postgres, None, None)
|
||||
}
|
||||
|
||||
pub fn announcements(&self) -> Option<Arc<dyn AnnouncementWriteRepository>> {
|
||||
self.announcements.clone()
|
||||
}
|
||||
|
||||
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyWriteRepository>> {
|
||||
self.auth_api_keys.clone()
|
||||
}
|
||||
|
||||
pub fn auth_modules(&self) -> Option<Arc<dyn AuthModuleWriteRepository>> {
|
||||
self.auth_modules.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageWriteRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
|
||||
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateWriteRepository>> {
|
||||
self.request_candidates.clone()
|
||||
}
|
||||
|
||||
pub fn gemini_file_mappings(&self) -> Option<Arc<dyn GeminiFileMappingWriteRepository>> {
|
||||
self.gemini_file_mappings.clone()
|
||||
}
|
||||
|
||||
pub fn global_models(&self) -> Option<Arc<dyn GlobalModelWriteRepository>> {
|
||||
self.global_models.clone()
|
||||
}
|
||||
|
||||
pub fn management_tokens(&self) -> Option<Arc<dyn ManagementTokenWriteRepository>> {
|
||||
self.management_tokens.clone()
|
||||
}
|
||||
|
||||
pub fn oauth_providers(&self) -> Option<Arc<dyn OAuthProviderWriteRepository>> {
|
||||
self.oauth_providers.clone()
|
||||
}
|
||||
|
||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
|
||||
pub fn provider_quotas(&self) -> Option<Arc<dyn ProviderQuotaWriteRepository>> {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogWriteRepository>> {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
|
||||
pub fn settlement(&self) -> Option<Arc<dyn SettlementWriteRepository>> {
|
||||
self.settlement.clone()
|
||||
}
|
||||
|
||||
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskWriteRepository>> {
|
||||
self.video_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn wallets(&self) -> Option<Arc<dyn WalletWriteRepository>> {
|
||||
self.wallets.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.announcements.is_some()
|
||||
|| self.auth_api_keys.is_some()
|
||||
|| self.auth_modules.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.gemini_file_mappings.is_some()
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.settlement.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|| self.wallets.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataWriteRepositories;
|
||||
use crate::backend::PostgresBackend;
|
||||
use crate::driver::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_write_repositories_from_postgres_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
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,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let write = DataWriteRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(write.has_any());
|
||||
assert!(write.announcements().is_some());
|
||||
assert!(write.auth_api_keys().is_some());
|
||||
assert!(write.auth_modules().is_some());
|
||||
assert!(write.request_candidates().is_some());
|
||||
assert!(write.gemini_file_mappings().is_some());
|
||||
assert!(write.global_models().is_some());
|
||||
assert!(write.management_tokens().is_some());
|
||||
assert!(write.oauth_providers().is_some());
|
||||
assert!(write.proxy_nodes().is_some());
|
||||
assert!(write.provider_catalog().is_some());
|
||||
assert!(write.provider_quotas().is_some());
|
||||
assert!(write.settlement().is_some());
|
||||
assert!(write.usage().is_some());
|
||||
assert!(write.video_tasks().is_some());
|
||||
assert!(write.wallets().is_some());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user