mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
feat: 扩展 Rust gateway 全功能模块,新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层:announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构,新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
112
crates/aether-data/src/repository/quota/sql.rs
Normal file
112
crates/aether-data/src/repository/quota/sql.rs
Normal file
@@ -0,0 +1,112 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_PROVIDER_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id AS provider_id,
|
||||
CAST(billing_type AS TEXT) AS billing_type,
|
||||
CAST(monthly_quota_usd AS DOUBLE PRECISION) AS monthly_quota_usd,
|
||||
CAST(COALESCE(monthly_used_usd, 0) AS DOUBLE PRECISION) AS monthly_used_usd,
|
||||
quota_reset_day,
|
||||
CAST(EXTRACT(EPOCH FROM quota_last_reset_at) AS BIGINT) AS quota_last_reset_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM quota_expires_at) AS BIGINT) AS quota_expires_at_unix_secs,
|
||||
is_active
|
||||
FROM providers
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const RESET_DUE_SQL: &str = r#"
|
||||
UPDATE providers
|
||||
SET
|
||||
monthly_used_usd = 0,
|
||||
quota_last_reset_at = TO_TIMESTAMP($1::double precision),
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
billing_type = 'monthly_quota'
|
||||
AND is_active = TRUE
|
||||
AND (
|
||||
quota_last_reset_at IS NULL
|
||||
OR (EXTRACT(EPOCH FROM TO_TIMESTAMP($1::double precision)) - EXTRACT(EPOCH FROM quota_last_reset_at)) >= (quota_reset_day * 86400)
|
||||
)
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProviderQuotaRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxProviderQuotaRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaReadRepository for SqlxProviderQuotaRepository {
|
||||
async fn find_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_PROVIDER_ID_SQL)
|
||||
.bind(provider_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderQuotaWriteRepository for SqlxProviderQuotaRepository {
|
||||
async fn reset_due(&self, now_unix_secs: u64) -> Result<usize, DataLayerError> {
|
||||
let result = sqlx::query(RESET_DUE_SQL)
|
||||
.bind(i64::try_from(now_unix_secs).map_err(|_| {
|
||||
DataLayerError::InvalidInput("provider quota reset timestamp overflow".to_string())
|
||||
})?)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() as usize)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredProviderQuotaSnapshot, DataLayerError> {
|
||||
StoredProviderQuotaSnapshot::new(
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("billing_type")?,
|
||||
row.try_get("monthly_quota_usd")?,
|
||||
row.try_get("monthly_used_usd")?,
|
||||
row.try_get("quota_reset_day")?,
|
||||
row.try_get("quota_last_reset_at_unix_secs")?,
|
||||
row.try_get("quota_expires_at_unix_secs")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxProviderQuotaRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(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("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let _repository = SqlxProviderQuotaRepository::new(pool);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user