mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate,完善并发门控与审计系统
新增 crate: - aether-runtime: 服务运行时基础设施(并发门控、分布式并发、指标、队列、优雅关闭、tracing) - aether-cache: 通用 TTL 缓存与命名空间抽象 - aether-data: 数据访问层(PostgreSQL/Redis 后端、repository 模式) - aether-http: HTTP 客户端封装(重试、配置) - aether-testkit: 集成测试工具集(gateway/executor/hub/proxy fixture、等待、负载测试) gateway 扩展: - 引入 audit 模块(shadow 执行审计、决策链路追踪、请求审计 bundle) - 引入 cache 模块(AuthContext 缓存、direct-plan bypass 缓存) - 引入 data 模块(auth/candidates/config/usage/video_tasks 数据访问) - 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控 - 新增本地 auth 拒绝、过载响应构建器 - 补充 control/auth_cache/video/concurrency 集成测试 aether-proxy 扩展: - AppState 集成 stream_gate / distributed_stream_gate 并发门控 - 新增 ProxyAdmissionError 及准入拒绝流程 - stream_handler 补充门控饱和/不可用场景测试 - 配置与注册客户端逻辑完善 aether-hub 扩展: - main.rs 引入运行时初始化、指标端点、健康检查 - local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
133
crates/aether-data/src/repository/auth/memory.rs
Normal file
133
crates/aether-data/src/repository/auth/memory.rs
Normal file
@@ -0,0 +1,133 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryAuthApiKeyIndex {
|
||||
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
|
||||
by_key_hash: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryAuthApiKeySnapshotRepository {
|
||||
index: RwLock<MemoryAuthApiKeyIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryAuthApiKeySnapshotRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
|
||||
{
|
||||
let mut by_api_key_id = BTreeMap::new();
|
||||
let mut by_key_hash = BTreeMap::new();
|
||||
for (key_hash, snapshot) in items {
|
||||
if let Some(key_hash) = key_hash {
|
||||
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
|
||||
}
|
||||
by_api_key_id.insert(snapshot.api_key_id.clone(), snapshot);
|
||||
}
|
||||
Self {
|
||||
index: RwLock::new(MemoryAuthApiKeyIndex {
|
||||
by_api_key_id,
|
||||
by_key_hash,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(match key {
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => index
|
||||
.by_key_hash
|
||||
.get(key_hash)
|
||||
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
|
||||
.cloned(),
|
||||
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
|
||||
index.by_api_key_id.get(api_key_id).cloned()
|
||||
}
|
||||
AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
} => index
|
||||
.by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|snapshot| snapshot.user_id == user_id)
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAuthApiKeySnapshotRepository;
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(200),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("snapshot should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_auth_snapshot_by_all_supported_keys() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_snapshot("key-1", "user-1"),
|
||||
)]);
|
||||
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::KeyHash("hash-1"))
|
||||
.await
|
||||
.expect("find by hash should succeed")
|
||||
.is_some());
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1"))
|
||||
.await
|
||||
.expect("find by api key id should succeed")
|
||||
.is_some());
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id: "user-1",
|
||||
api_key_id: "key-1",
|
||||
})
|
||||
.await
|
||||
.expect("find by user/api key ids should succeed")
|
||||
.is_some());
|
||||
}
|
||||
}
|
||||
9
crates/aether-data/src/repository/auth/mod.rs
Normal file
9
crates/aether-data/src/repository/auth/mod.rs
Normal file
@@ -0,0 +1,9 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryAuthApiKeySnapshotRepository;
|
||||
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
|
||||
pub use types::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
202
crates/aether-data/src/repository/auth/sql.rs
Normal file
202
crates/aether-data/src/repository/auth/sql.rs
Normal file
@@ -0,0 +1,202 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_KEY_HASH_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.key_hash = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_API_KEY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_USER_API_KEY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.id = $1 AND users.id = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxAuthApiKeySnapshotReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxAuthApiKeySnapshotReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let row = match key {
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => {
|
||||
sqlx::query(FIND_BY_KEY_HASH_SQL)
|
||||
.bind(key_hash)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
|
||||
sqlx::query(FIND_BY_API_KEY_ID_SQL)
|
||||
.bind(api_key_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
} => {
|
||||
sqlx::query(FIND_BY_USER_API_KEY_IDS_SQL)
|
||||
.bind(api_key_id)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
|
||||
row.as_ref().map(map_auth_api_key_snapshot_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyReadRepository for SqlxAuthApiKeySnapshotReadRepository {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
Self::find_api_key_snapshot(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_auth_api_key_snapshot_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredAuthApiKeySnapshot, DataLayerError> {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("email")?,
|
||||
row.try_get("user_role")?,
|
||||
row.try_get("user_auth_source")?,
|
||||
row.try_get("user_is_active")?,
|
||||
row.try_get("user_is_deleted")?,
|
||||
row.try_get("user_allowed_providers")?,
|
||||
row.try_get("user_allowed_api_formats")?,
|
||||
row.try_get("user_allowed_models")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("api_key_is_active")?,
|
||||
row.try_get("api_key_is_locked")?,
|
||||
row.try_get("api_key_is_standalone")?,
|
||||
row.try_get("api_key_rate_limit")?,
|
||||
row.try_get("api_key_concurrent_limit")?,
|
||||
row.try_get("api_key_expires_at_unix_secs")?,
|
||||
row.try_get("api_key_allowed_providers")?,
|
||||
row.try_get("api_key_allowed_api_formats")?,
|
||||
row.try_get("api_key_allowed_models")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxAuthApiKeySnapshotReadRepository;
|
||||
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 = SqlxAuthApiKeySnapshotReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
225
crates/aether-data/src/repository/auth/types.rs
Normal file
225
crates/aether-data/src/repository/auth/types.rs
Normal file
@@ -0,0 +1,225 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAuthApiKeySnapshot {
|
||||
pub user_id: String,
|
||||
pub username: String,
|
||||
pub email: Option<String>,
|
||||
pub user_role: String,
|
||||
pub user_auth_source: String,
|
||||
pub user_is_active: bool,
|
||||
pub user_is_deleted: bool,
|
||||
pub user_allowed_providers: Option<Vec<String>>,
|
||||
pub user_allowed_api_formats: Option<Vec<String>>,
|
||||
pub user_allowed_models: Option<Vec<String>>,
|
||||
pub api_key_id: String,
|
||||
pub api_key_name: Option<String>,
|
||||
pub api_key_is_active: bool,
|
||||
pub api_key_is_locked: bool,
|
||||
pub api_key_is_standalone: bool,
|
||||
pub api_key_rate_limit: Option<i32>,
|
||||
pub api_key_concurrent_limit: Option<i32>,
|
||||
pub api_key_expires_at_unix_secs: Option<u64>,
|
||||
pub api_key_allowed_providers: Option<Vec<String>>,
|
||||
pub api_key_allowed_api_formats: Option<Vec<String>>,
|
||||
pub api_key_allowed_models: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl StoredAuthApiKeySnapshot {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
user_id: String,
|
||||
username: String,
|
||||
email: Option<String>,
|
||||
user_role: String,
|
||||
user_auth_source: String,
|
||||
user_is_active: bool,
|
||||
user_is_deleted: bool,
|
||||
user_allowed_providers: Option<serde_json::Value>,
|
||||
user_allowed_api_formats: Option<serde_json::Value>,
|
||||
user_allowed_models: Option<serde_json::Value>,
|
||||
api_key_id: String,
|
||||
api_key_name: Option<String>,
|
||||
api_key_is_active: bool,
|
||||
api_key_is_locked: bool,
|
||||
api_key_is_standalone: bool,
|
||||
api_key_rate_limit: Option<i32>,
|
||||
api_key_concurrent_limit: Option<i32>,
|
||||
api_key_expires_at_unix_secs: Option<i64>,
|
||||
api_key_allowed_providers: Option<serde_json::Value>,
|
||||
api_key_allowed_api_formats: Option<serde_json::Value>,
|
||||
api_key_allowed_models: Option<serde_json::Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
Ok(Self {
|
||||
user_id,
|
||||
username,
|
||||
email,
|
||||
user_role,
|
||||
user_auth_source,
|
||||
user_is_active,
|
||||
user_is_deleted,
|
||||
user_allowed_providers: parse_string_list(
|
||||
user_allowed_providers,
|
||||
"users.allowed_providers",
|
||||
)?,
|
||||
user_allowed_api_formats: parse_string_list(
|
||||
user_allowed_api_formats,
|
||||
"users.allowed_api_formats",
|
||||
)?,
|
||||
user_allowed_models: parse_string_list(user_allowed_models, "users.allowed_models")?,
|
||||
api_key_id,
|
||||
api_key_name,
|
||||
api_key_is_active,
|
||||
api_key_is_locked,
|
||||
api_key_is_standalone,
|
||||
api_key_rate_limit,
|
||||
api_key_concurrent_limit,
|
||||
api_key_expires_at_unix_secs: api_key_expires_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid api_keys.expires_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?,
|
||||
api_key_allowed_providers: parse_string_list(
|
||||
api_key_allowed_providers,
|
||||
"api_keys.allowed_providers",
|
||||
)?,
|
||||
api_key_allowed_api_formats: parse_string_list(
|
||||
api_key_allowed_api_formats,
|
||||
"api_keys.allowed_api_formats",
|
||||
)?,
|
||||
api_key_allowed_models: parse_string_list(
|
||||
api_key_allowed_models,
|
||||
"api_keys.allowed_models",
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_currently_usable(&self, now_unix_secs: u64) -> bool {
|
||||
if !self.user_is_active || self.user_is_deleted {
|
||||
return false;
|
||||
}
|
||||
if !self.api_key_is_active {
|
||||
return false;
|
||||
}
|
||||
if self.api_key_is_locked && !self.api_key_is_standalone {
|
||||
return false;
|
||||
}
|
||||
if let Some(expires_at_unix_secs) = self.api_key_expires_at_unix_secs {
|
||||
if expires_at_unix_secs < now_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AuthApiKeyLookupKey<'a> {
|
||||
KeyHash(&'a str),
|
||||
ApiKeyId(&'a str),
|
||||
UserApiKeyIds {
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AuthApiKeyReadRepository: Send + Sync {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<serde_json::Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let array = value.as_array().ok_or_else(|| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("{field_name} is not a JSON array"))
|
||||
})?;
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
items.push(item.to_string());
|
||||
}
|
||||
Ok(Some(items))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredAuthApiKeySnapshot;
|
||||
|
||||
#[test]
|
||||
fn rejects_non_array_allowed_providers() {
|
||||
assert!(StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!({"bad": true})),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_non_standalone_key_is_not_usable() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(100),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert!(!snapshot.is_currently_usable(101));
|
||||
}
|
||||
}
|
||||
148
crates/aether-data/src/repository/candidates/memory.rs
Normal file
148
crates/aether-data/src/repository/candidates/memory.rs
Normal file
@@ -0,0 +1,148 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{RequestCandidateReadRepository, StoredRequestCandidate};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryRequestCandidateRepository {
|
||||
by_id: RwLock<BTreeMap<String, StoredRequestCandidate>>,
|
||||
}
|
||||
|
||||
impl InMemoryRequestCandidateRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRequestCandidate>,
|
||||
{
|
||||
let mut by_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_id.insert(item.id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_id: RwLock::new(by_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.filter(|row| row.request_id == request_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| {
|
||||
left.candidate_index
|
||||
.cmp(&right.candidate_index)
|
||||
.then(left.retry_index.cmp(&right.retry_index))
|
||||
.then(left.created_at_unix_secs.cmp(&right.created_at_unix_secs))
|
||||
});
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
|
||||
rows.truncate(limit);
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryRequestCandidateRepository;
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
|
||||
fn sample_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
created_at_unix_secs: i64,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
RequestCandidateStatus::Success,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
created_at_unix_secs,
|
||||
Some(created_at_unix_secs),
|
||||
Some(created_at_unix_secs + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_request_candidates_by_request_id_in_candidate_order() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-2", "req-1", 200),
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-3", "req-2", 300),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_by_request_id("req-1")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].request_id, "req-1");
|
||||
assert_eq!(rows[1].request_id, "req-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_recent_request_candidates_in_descending_created_order() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-2", "req-2", 200),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_recent(10)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].id, "cand-2");
|
||||
assert_eq!(rows[1].id, "cand-1");
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/candidates/mod.rs
Normal file
10
crates/aether-data/src/repository/candidates/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryRequestCandidateRepository;
|
||||
pub use sql::SqlxRequestCandidateReadRepository;
|
||||
pub use types::{
|
||||
RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateStatus,
|
||||
StoredRequestCandidate,
|
||||
};
|
||||
189
crates/aether-data/src/repository/candidates/sql.rs
Normal file
189
crates/aether-data/src/repository/candidates/sql.rs
Normal file
@@ -0,0 +1,189 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_BY_REQUEST_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
WHERE request_id = $1
|
||||
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxRequestCandidateReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxRequestCandidateReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_BY_REQUEST_ID_SQL)
|
||||
.bind(request_id)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid recent request candidate limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_by_request_id(self, request_id).await
|
||||
}
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_recent(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_request_candidate_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
let status =
|
||||
RequestCandidateStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
|
||||
StoredRequestCandidate::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("candidate_index")?,
|
||||
row.try_get("retry_index")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("endpoint_id")?,
|
||||
row.try_get("key_id")?,
|
||||
status,
|
||||
row.try_get("skip_reason")?,
|
||||
row.try_get("is_cached")?,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_type")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("latency_ms")?,
|
||||
row.try_get("concurrent_requests")?,
|
||||
row.try_get("extra_data")?,
|
||||
row.try_get("required_capabilities")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("started_at_unix_secs")?,
|
||||
row.try_get("finished_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxRequestCandidateReadRepository;
|
||||
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 = SqlxRequestCandidateReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
289
crates/aether-data/src/repository/candidates/types.rs
Normal file
289
crates/aether-data/src/repository/candidates/types.rs
Normal file
@@ -0,0 +1,289 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RequestCandidateStatus {
|
||||
Available,
|
||||
Unused,
|
||||
Pending,
|
||||
Streaming,
|
||||
Success,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl RequestCandidateStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"available" => Ok(Self::Available),
|
||||
"unused" => Ok(Self::Unused),
|
||||
"pending" => Ok(Self::Pending),
|
||||
"streaming" => Ok(Self::Streaming),
|
||||
"success" => Ok(Self::Success),
|
||||
"failed" => Ok(Self::Failed),
|
||||
"cancelled" => Ok(Self::Cancelled),
|
||||
"skipped" => Ok(Self::Skipped),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported request_candidates.status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_attempted(self, started_at_unix_secs: Option<u64>) -> bool {
|
||||
match self {
|
||||
Self::Available | Self::Unused | Self::Skipped => false,
|
||||
Self::Pending => started_at_unix_secs.is_some(),
|
||||
Self::Streaming | Self::Success | Self::Failed | Self::Cancelled => true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestCandidate {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub candidate_index: u32,
|
||||
pub retry_index: u32,
|
||||
pub provider_id: Option<String>,
|
||||
pub endpoint_id: Option<String>,
|
||||
pub key_id: Option<String>,
|
||||
pub status: RequestCandidateStatus,
|
||||
pub skip_reason: Option<String>,
|
||||
pub is_cached: bool,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_type: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub concurrent_requests: Option<u32>,
|
||||
pub extra_data: Option<serde_json::Value>,
|
||||
pub required_capabilities: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub started_at_unix_secs: Option<u64>,
|
||||
pub finished_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredRequestCandidate {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
candidate_index: i32,
|
||||
retry_index: i32,
|
||||
provider_id: Option<String>,
|
||||
endpoint_id: Option<String>,
|
||||
key_id: Option<String>,
|
||||
status: RequestCandidateStatus,
|
||||
skip_reason: Option<String>,
|
||||
is_cached: bool,
|
||||
status_code: Option<i32>,
|
||||
error_type: Option<String>,
|
||||
error_message: Option<String>,
|
||||
latency_ms: Option<i32>,
|
||||
concurrent_requests: Option<i32>,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
required_capabilities: Option<serde_json::Value>,
|
||||
created_at_unix_secs: i64,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
finished_at_unix_secs: Option<i64>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let candidate_index = u32::try_from(candidate_index).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.candidate_index: {candidate_index}"
|
||||
))
|
||||
})?;
|
||||
let retry_index = u32::try_from(retry_index).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.retry_index: {retry_index}"
|
||||
))
|
||||
})?;
|
||||
let status_code = status_code
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.status_code: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let latency_ms = latency_ms
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.latency_ms: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let concurrent_requests = concurrent_requests
|
||||
.map(|value| {
|
||||
u32::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.concurrent_requests: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let started_at_unix_secs = started_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.started_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let finished_at_unix_secs = finished_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.finished_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RequestCandidateReadRepository: Send + Sync {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait RequestCandidateRepository: RequestCandidateReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> RequestCandidateRepository for T where T: RequestCandidateReadRepository + Send + Sync {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RequestCandidateStatus, StoredRequestCandidate};
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
assert_eq!(
|
||||
RequestCandidateStatus::from_database("streaming").expect("status should parse"),
|
||||
RequestCandidateStatus::Streaming
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(RequestCandidateStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_candidate_index() {
|
||||
assert!(StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
-1,
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_created_at() {
|
||||
assert!(StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
-1,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_without_started_at_is_not_attempted() {
|
||||
assert!(!RequestCandidateStatus::Pending.is_attempted(None));
|
||||
assert!(RequestCandidateStatus::Pending.is_attempted(Some(1)));
|
||||
}
|
||||
}
|
||||
6
crates/aether-data/src/repository/mod.rs
Normal file
6
crates/aether-data/src/repository/mod.rs
Normal file
@@ -0,0 +1,6 @@
|
||||
pub mod auth;
|
||||
pub mod candidates;
|
||||
pub mod provider_catalog;
|
||||
pub mod shadow_results;
|
||||
pub mod usage;
|
||||
pub mod video_tasks;
|
||||
157
crates/aether-data/src/repository/provider_catalog/memory.rs
Normal file
157
crates/aether-data/src/repository/provider_catalog/memory.rs
Normal file
@@ -0,0 +1,157 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryProviderCatalogIndex {
|
||||
providers: BTreeMap<String, StoredProviderCatalogProvider>,
|
||||
endpoints: BTreeMap<String, StoredProviderCatalogEndpoint>,
|
||||
keys: BTreeMap<String, StoredProviderCatalogKey>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryProviderCatalogReadRepository {
|
||||
index: RwLock<MemoryProviderCatalogIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryProviderCatalogReadRepository {
|
||||
pub fn seed(
|
||||
providers: Vec<StoredProviderCatalogProvider>,
|
||||
endpoints: Vec<StoredProviderCatalogEndpoint>,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
) -> Self {
|
||||
Self {
|
||||
index: RwLock::new(MemoryProviderCatalogIndex {
|
||||
providers: providers
|
||||
.into_iter()
|
||||
.map(|provider| (provider.id.clone(), provider))
|
||||
.collect(),
|
||||
endpoints: endpoints
|
||||
.into_iter()
|
||||
.map(|endpoint| (endpoint.id.clone(), endpoint))
|
||||
.collect(),
|
||||
keys: keys.into_iter().map(|key| (key.id.clone(), key)).collect(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(provider_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.providers.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(endpoint_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.endpoints.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(key_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.keys.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryProviderCatalogReadRepository;
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
id.to_string(),
|
||||
format!("provider-{id}"),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_endpoint(id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_key(id: &str, provider_id: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_provider_catalog_items_by_id() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![sample_endpoint("endpoint-1", "provider-1")],
|
||||
vec![sample_key("key-1", "provider-1")],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_providers_by_ids(&["provider-1".to_string()])
|
||||
.await
|
||||
.expect("providers should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_endpoints_by_ids(&["endpoint-1".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/provider_catalog/mod.rs
Normal file
10
crates/aether-data/src/repository/provider_catalog/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryProviderCatalogReadRepository;
|
||||
pub use sql::SqlxProviderCatalogReadRepository;
|
||||
pub use types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
209
crates/aether-data/src/repository/provider_catalog/sql.rs
Normal file
209
crates/aether-data/src/repository/provider_catalog/sql.rs
Normal file
@@ -0,0 +1,209 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||
|
||||
use super::types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_PROVIDERS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
website,
|
||||
provider_type
|
||||
FROM providers
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
const LIST_ENDPOINTS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
provider_id,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
is_active
|
||||
FROM provider_endpoints
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
const LIST_KEYS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
provider_id,
|
||||
name,
|
||||
auth_type,
|
||||
capabilities,
|
||||
is_active
|
||||
FROM provider_api_keys
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProviderCatalogReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxProviderCatalogReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_PROVIDERS_BY_IDS_PREFIX,
|
||||
provider_ids,
|
||||
" ORDER BY name ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_provider_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_ENDPOINTS_BY_IDS_PREFIX,
|
||||
endpoint_ids,
|
||||
" ORDER BY api_format ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_endpoint_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
if key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_KEYS_BY_IDS_PREFIX,
|
||||
key_ids,
|
||||
" ORDER BY name ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
Self::list_providers_by_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
Self::list_endpoints_by_ids(self, endpoint_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
Self::list_keys_by_ids(self, key_ids).await
|
||||
}
|
||||
}
|
||||
|
||||
fn build_list_query<'a>(
|
||||
prefix: &'static str,
|
||||
ids: &'a [String],
|
||||
suffix: &'static str,
|
||||
) -> QueryBuilder<'a, Postgres> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(prefix);
|
||||
let mut separated = builder.separated(", ");
|
||||
for id in ids {
|
||||
separated.push_bind(id);
|
||||
}
|
||||
separated.push_unseparated(")");
|
||||
builder.push(suffix);
|
||||
builder
|
||||
}
|
||||
|
||||
fn map_provider_row(row: &PgRow) -> Result<StoredProviderCatalogProvider, DataLayerError> {
|
||||
StoredProviderCatalogProvider::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("name")?,
|
||||
row.try_get("website")?,
|
||||
row.try_get("provider_type")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_endpoint_row(row: &PgRow) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("api_format")?,
|
||||
row.try_get("api_family")?,
|
||||
row.try_get("endpoint_kind")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
StoredProviderCatalogKey::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("name")?,
|
||||
row.try_get("auth_type")?,
|
||||
row.try_get("capabilities")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxProviderCatalogReadRepository;
|
||||
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 = SqlxProviderCatalogReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
175
crates/aether-data/src/repository/provider_catalog/types.rs
Normal file
175
crates/aether-data/src/repository/provider_catalog/types.rs
Normal file
@@ -0,0 +1,175 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogProvider {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub website: Option<String>,
|
||||
pub provider_type: String,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogProvider {
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
website: Option<String>,
|
||||
provider_type: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"providers.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"providers.provider_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
website,
|
||||
provider_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogEndpoint {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub api_format: String,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogEndpoint {
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
api_format: String,
|
||||
api_family: Option<String>,
|
||||
endpoint_kind: Option<String>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if api_format.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_endpoints.api_format is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogKey {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub name: String,
|
||||
pub auth_type: String,
|
||||
pub capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogKey {
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
name: String,
|
||||
auth_type: String,
|
||||
capabilities: Option<serde_json::Value>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_api_keys.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if auth_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_api_keys.auth_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
name,
|
||||
auth_type,
|
||||
capabilities,
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderCatalogReadRepository: Send + Sync {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_provider_name() {
|
||||
assert!(StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
"custom".to_string(),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_endpoint_api_format() {
|
||||
assert!(StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_key_auth_type() {
|
||||
assert!(StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
170
crates/aether-data/src/repository/shadow_results/memory.rs
Normal file
170
crates/aether-data/src/repository/shadow_results/memory.rs
Normal file
@@ -0,0 +1,170 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ShadowResultLookupKey, ShadowResultReadRepository, ShadowResultWriteRepository,
|
||||
StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryShadowResultRepository {
|
||||
results: RwLock<BTreeMap<(String, String), StoredShadowResult>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultReadRepository for InMemoryShadowResultRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let results = self.results.read().expect("shadow result repository lock");
|
||||
Ok(match key {
|
||||
ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
} => results
|
||||
.get(&(trace_id.to_string(), request_fingerprint.to_string()))
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut results = self
|
||||
.results
|
||||
.read()
|
||||
.expect("shadow result repository lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
results.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
|
||||
results.truncate(limit);
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultWriteRepository for InMemoryShadowResultRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
let stored = result.into_stored();
|
||||
let mut results = self.results.write().expect("shadow result repository lock");
|
||||
results.insert(
|
||||
(stored.trace_id.clone(), stored.request_fingerprint.clone()),
|
||||
stored.clone(),
|
||||
);
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryShadowResultRepository;
|
||||
use crate::repository::shadow_results::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultWriteRepository, UpsertShadowResult,
|
||||
};
|
||||
|
||||
fn sample_result(
|
||||
trace_id: &str,
|
||||
request_fingerprint: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> UpsertShadowResult {
|
||||
UpsertShadowResult {
|
||||
trace_id: trace_id.to_string(),
|
||||
request_fingerprint: request_fingerprint.to_string(),
|
||||
request_id: Some(format!("req-{trace_id}")),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: Some("cand-1".to_string()),
|
||||
rust_result_digest: Some("rust-digest".to_string()),
|
||||
python_result_digest: Some("python-digest".to_string()),
|
||||
match_status: ShadowResultMatchStatus::Match,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
|
||||
updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_result_by_trace_and_fingerprint() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: "trace-1",
|
||||
request_fingerprint: "fp-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_recent_returns_results_in_descending_update_order() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_result("trace-2", "fp-2", 200))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let recent = repo
|
||||
.list_recent(10)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
assert_eq!(recent.len(), 2);
|
||||
assert_eq!(recent[0].trace_id, "trace-2");
|
||||
assert_eq!(recent[1].trace_id, "trace-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_replaces_existing_shadow_result() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(UpsertShadowResult {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-trace-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: Some("cand-2".to_string()),
|
||||
rust_result_digest: Some("rust-digest-2".to_string()),
|
||||
python_result_digest: Some("python-digest-2".to_string()),
|
||||
match_status: ShadowResultMatchStatus::Mismatch,
|
||||
status_code: Some(502),
|
||||
error_message: Some("mismatch".to_string()),
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 200,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let stored = repo
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: "trace-1",
|
||||
request_fingerprint: "fp-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.request_id.as_deref(), Some("req-trace-1"));
|
||||
assert_eq!(stored.candidate_id.as_deref(), Some("cand-2"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Mismatch);
|
||||
}
|
||||
}
|
||||
12
crates/aether-data/src/repository/shadow_results/mod.rs
Normal file
12
crates/aether-data/src/repository/shadow_results/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
mod memory;
|
||||
mod record;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryShadowResultRepository;
|
||||
pub use record::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
|
||||
pub use sql::SqlxShadowResultRepository;
|
||||
pub use types::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultRepository, ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
188
crates/aether-data/src/repository/shadow_results/record.rs
Normal file
188
crates/aether-data/src/repository/shadow_results/record.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
use super::types::{ShadowResultMatchStatus, StoredShadowResult, UpsertShadowResult};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ShadowResultSampleOrigin {
|
||||
Rust,
|
||||
Python,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RecordShadowResultSample {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub origin: ShadowResultSampleOrigin,
|
||||
pub result_digest: String,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub recorded_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
pub fn merge_shadow_result_sample(
|
||||
existing: Option<&StoredShadowResult>,
|
||||
sample: RecordShadowResultSample,
|
||||
) -> UpsertShadowResult {
|
||||
let RecordShadowResultSample {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
origin,
|
||||
result_digest,
|
||||
status_code,
|
||||
error_message,
|
||||
recorded_at_unix_secs,
|
||||
} = sample;
|
||||
|
||||
let (rust_result_digest, python_result_digest) = match origin {
|
||||
ShadowResultSampleOrigin::Rust => (
|
||||
Some(result_digest),
|
||||
existing.and_then(|row| row.python_result_digest.clone()),
|
||||
),
|
||||
ShadowResultSampleOrigin::Python => (
|
||||
existing.and_then(|row| row.rust_result_digest.clone()),
|
||||
Some(result_digest),
|
||||
),
|
||||
};
|
||||
|
||||
let match_status = resolve_match_status(
|
||||
rust_result_digest.as_deref(),
|
||||
python_result_digest.as_deref(),
|
||||
);
|
||||
|
||||
UpsertShadowResult {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id: request_id.or_else(|| existing.and_then(|row| row.request_id.clone())),
|
||||
route_family: route_family.or_else(|| existing.and_then(|row| row.route_family.clone())),
|
||||
route_kind: route_kind.or_else(|| existing.and_then(|row| row.route_kind.clone())),
|
||||
candidate_id: candidate_id.or_else(|| existing.and_then(|row| row.candidate_id.clone())),
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code: status_code.or(existing.and_then(|row| row.status_code)),
|
||||
error_message: resolve_error_message(existing, error_message, match_status),
|
||||
created_at_unix_secs: existing
|
||||
.map(|row| row.created_at_unix_secs)
|
||||
.unwrap_or(recorded_at_unix_secs),
|
||||
updated_at_unix_secs: recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_match_status(
|
||||
rust_result_digest: Option<&str>,
|
||||
python_result_digest: Option<&str>,
|
||||
) -> ShadowResultMatchStatus {
|
||||
match (rust_result_digest, python_result_digest) {
|
||||
(Some(rust_digest), Some(python_digest)) if rust_digest == python_digest => {
|
||||
ShadowResultMatchStatus::Match
|
||||
}
|
||||
(Some(_), Some(_)) => ShadowResultMatchStatus::Mismatch,
|
||||
_ => ShadowResultMatchStatus::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_error_message(
|
||||
existing: Option<&StoredShadowResult>,
|
||||
error_message: Option<String>,
|
||||
match_status: ShadowResultMatchStatus,
|
||||
) -> Option<String> {
|
||||
if match_status == ShadowResultMatchStatus::Mismatch {
|
||||
error_message
|
||||
.or_else(|| existing.and_then(|row| row.error_message.clone()))
|
||||
.or_else(|| Some("shadow result digest mismatch".to_string()))
|
||||
} else {
|
||||
error_message.or_else(|| existing.and_then(|row| row.error_message.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
|
||||
use crate::repository::shadow_results::{ShadowResultMatchStatus, UpsertShadowResult};
|
||||
|
||||
fn rust_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
|
||||
RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Rust,
|
||||
result_digest: result_digest.to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn python_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
|
||||
RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Python,
|
||||
result_digest: result_digest.to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn stored(upsert: UpsertShadowResult) -> crate::repository::shadow_results::StoredShadowResult {
|
||||
upsert.into_stored()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_pending_until_both_samples_exist() {
|
||||
let merged = merge_shadow_result_sample(None, rust_sample("digest-1", 100));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Pending);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert!(merged.python_result_digest.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn marks_match_when_rust_and_python_digests_are_equal() {
|
||||
let existing = stored(merge_shadow_result_sample(
|
||||
None,
|
||||
rust_sample("digest-1", 100),
|
||||
));
|
||||
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-1", 200));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Match);
|
||||
assert_eq!(merged.created_at_unix_secs, 100);
|
||||
assert_eq!(merged.updated_at_unix_secs, 200);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn marks_mismatch_when_rust_and_python_digests_differ() {
|
||||
let existing = stored(merge_shadow_result_sample(
|
||||
None,
|
||||
rust_sample("digest-1", 100),
|
||||
));
|
||||
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-2", 200));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Mismatch);
|
||||
assert_eq!(
|
||||
merged.error_message.as_deref(),
|
||||
Some("shadow result digest mismatch")
|
||||
);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-2"));
|
||||
}
|
||||
}
|
||||
283
crates/aether-data/src/repository/shadow_results/sql.rs
Normal file
283
crates/aether-data/src/repository/shadow_results/sql.rs
Normal file
@@ -0,0 +1,283 @@
|
||||
use async_trait::async_trait;
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_TRACE_FINGERPRINT_SQL: &str = r#"
|
||||
SELECT
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM gateway_shadow_results
|
||||
WHERE trace_id = $1 AND request_fingerprint = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_SQL: &str = r#"
|
||||
SELECT
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM gateway_shadow_results
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT $1
|
||||
"#;
|
||||
|
||||
const UPSERT_SQL: &str = r#"
|
||||
INSERT INTO gateway_shadow_results (
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
created_at,
|
||||
updated_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
TO_TIMESTAMP($11::double precision),
|
||||
TO_TIMESTAMP($12::double precision)
|
||||
)
|
||||
ON CONFLICT (trace_id, request_fingerprint)
|
||||
DO UPDATE SET
|
||||
route_family = EXCLUDED.route_family,
|
||||
route_kind = EXCLUDED.route_kind,
|
||||
candidate_id = EXCLUDED.candidate_id,
|
||||
rust_result_digest = EXCLUDED.rust_result_digest,
|
||||
python_result_digest = EXCLUDED.python_result_digest,
|
||||
match_status = EXCLUDED.match_status,
|
||||
status_code = EXCLUDED.status_code,
|
||||
error_message = EXCLUDED.error_message,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
RETURNING
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxShadowResultRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxShadowResultRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.tx_runner
|
||||
}
|
||||
|
||||
pub async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
match key {
|
||||
ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
} => {
|
||||
self.find_by_trace_fingerprint(trace_id, request_fingerprint)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn find_by_trace_fingerprint(
|
||||
&self,
|
||||
trace_id: &str,
|
||||
request_fingerprint: &str,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_TRACE_FINGERPRINT_SQL)
|
||||
.bind(trace_id)
|
||||
.bind(request_fingerprint)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_shadow_result_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid recent shadow result limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_shadow_result_row).collect()
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(UPSERT_SQL)
|
||||
.bind(&result.trace_id)
|
||||
.bind(&result.request_fingerprint)
|
||||
.bind(&result.route_family)
|
||||
.bind(&result.route_kind)
|
||||
.bind(&result.candidate_id)
|
||||
.bind(&result.rust_result_digest)
|
||||
.bind(&result.python_result_digest)
|
||||
.bind(match_status_to_database(result.match_status))
|
||||
.bind(result.status_code.map(i32::from))
|
||||
.bind(&result.error_message)
|
||||
.bind(result.created_at_unix_secs as f64)
|
||||
.bind(result.updated_at_unix_secs as f64)
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
map_shadow_result_row(&row)
|
||||
}) as BoxFuture<'_, Result<StoredShadowResult, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultReadRepository for SqlxShadowResultRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
Self::find(self, key).await
|
||||
}
|
||||
|
||||
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
Self::list_recent(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultWriteRepository for SqlxShadowResultRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
Self::upsert(self, result).await
|
||||
}
|
||||
}
|
||||
|
||||
fn match_status_to_database(status: ShadowResultMatchStatus) -> &'static str {
|
||||
match status {
|
||||
ShadowResultMatchStatus::Pending => "pending",
|
||||
ShadowResultMatchStatus::Match => "match",
|
||||
ShadowResultMatchStatus::Mismatch => "mismatch",
|
||||
ShadowResultMatchStatus::Error => "error",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_shadow_result_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
let match_status =
|
||||
ShadowResultMatchStatus::from_database(row.try_get::<String, _>("match_status")?.as_str())?;
|
||||
StoredShadowResult::new(
|
||||
row.try_get("trace_id")?,
|
||||
row.try_get("request_fingerprint")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("route_family")?,
|
||||
row.try_get("route_kind")?,
|
||||
row.try_get("candidate_id")?,
|
||||
row.try_get("rust_result_digest")?,
|
||||
row.try_get("python_result_digest")?,
|
||||
match_status,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxShadowResultRepository;
|
||||
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 = SqlxShadowResultRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
let _ = repository.transaction_runner();
|
||||
}
|
||||
}
|
||||
227
crates/aether-data/src/repository/shadow_results/types.rs
Normal file
227
crates/aether-data/src/repository/shadow_results/types.rs
Normal file
@@ -0,0 +1,227 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum ShadowResultMatchStatus {
|
||||
Pending,
|
||||
Match,
|
||||
Mismatch,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl ShadowResultMatchStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"pending" => Ok(Self::Pending),
|
||||
"match" => Ok(Self::Match),
|
||||
"mismatch" => Ok(Self::Mismatch),
|
||||
"error" => Ok(Self::Error),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported gateway_shadow_results.match_status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredShadowResult {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub rust_result_digest: Option<String>,
|
||||
pub python_result_digest: Option<String>,
|
||||
pub match_status: ShadowResultMatchStatus,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl StoredShadowResult {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
trace_id: String,
|
||||
request_fingerprint: String,
|
||||
request_id: Option<String>,
|
||||
route_family: Option<String>,
|
||||
route_kind: Option<String>,
|
||||
candidate_id: Option<String>,
|
||||
rust_result_digest: Option<String>,
|
||||
python_result_digest: Option<String>,
|
||||
match_status: ShadowResultMatchStatus,
|
||||
status_code: Option<i32>,
|
||||
error_message: Option<String>,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let status_code = status_code
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid status_code: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpsertShadowResult {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub rust_result_digest: Option<String>,
|
||||
pub python_result_digest: Option<String>,
|
||||
pub match_status: ShadowResultMatchStatus,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl UpsertShadowResult {
|
||||
pub fn into_stored(self) -> StoredShadowResult {
|
||||
StoredShadowResult {
|
||||
trace_id: self.trace_id,
|
||||
request_fingerprint: self.request_fingerprint,
|
||||
request_id: self.request_id,
|
||||
route_family: self.route_family,
|
||||
route_kind: self.route_kind,
|
||||
candidate_id: self.candidate_id,
|
||||
rust_result_digest: self.rust_result_digest,
|
||||
python_result_digest: self.python_result_digest,
|
||||
match_status: self.match_status,
|
||||
status_code: self.status_code,
|
||||
error_message: self.error_message,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ShadowResultLookupKey<'a> {
|
||||
TraceFingerprint {
|
||||
trace_id: &'a str,
|
||||
request_fingerprint: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ShadowResultReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ShadowResultWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait ShadowResultRepository:
|
||||
ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> ShadowResultRepository for T where
|
||||
T: ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ShadowResultMatchStatus, StoredShadowResult};
|
||||
|
||||
#[test]
|
||||
fn parses_match_status_from_database_text() {
|
||||
assert_eq!(
|
||||
ShadowResultMatchStatus::from_database("match").expect("status should parse"),
|
||||
ShadowResultMatchStatus::Match
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(ShadowResultMatchStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_numeric_fields() {
|
||||
assert!(StoredShadowResult::new(
|
||||
"trace-1".to_string(),
|
||||
"fp-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
ShadowResultMatchStatus::Pending,
|
||||
Some(-1),
|
||||
None,
|
||||
1,
|
||||
1,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_updated_at_values() {
|
||||
assert!(StoredShadowResult::new(
|
||||
"trace-1".to_string(),
|
||||
"fp-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
ShadowResultMatchStatus::Pending,
|
||||
Some(200),
|
||||
None,
|
||||
1,
|
||||
-1,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
107
crates/aether-data/src/repository/usage/memory.rs
Normal file
107
crates/aether-data/src/repository/usage/memory.rs
Normal file
@@ -0,0 +1,107 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUsageReadRepository {
|
||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||
}
|
||||
|
||||
impl InMemoryUsageReadRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRequestUsageAudit>,
|
||||
{
|
||||
let mut by_request_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_request_id.insert(item.request_id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.get(request_id)
|
||||
.cloned())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryUsageReadRepository;
|
||||
use crate::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
|
||||
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
Some("gpt-4.1-mini".to_string()),
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
false,
|
||||
100,
|
||||
50,
|
||||
150,
|
||||
0.12,
|
||||
0.18,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(420),
|
||||
Some(120),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_secs + 1,
|
||||
Some(created_at_unix_secs + 2),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_usage_by_request_id() {
|
||||
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage("req-1", 100),
|
||||
sample_usage("req-2", 200),
|
||||
]);
|
||||
|
||||
let usage = repository
|
||||
.find_by_request_id("req-2")
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("usage should exist");
|
||||
|
||||
assert_eq!(usage.request_id, "req-2");
|
||||
assert_eq!(usage.total_tokens, 150);
|
||||
}
|
||||
}
|
||||
7
crates/aether-data/src/repository/usage/mod.rs
Normal file
7
crates/aether-data/src/repository/usage/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryUsageReadRepository;
|
||||
pub use sql::SqlxUsageReadRepository;
|
||||
pub use types::{StoredRequestUsageAudit, UsageReadRepository, UsageRepository};
|
||||
150
crates/aether-data/src/repository/usage/sql.rs
Normal file
150
crates/aether-data/src/repository/usage/sql.rs
Normal file
@@ -0,0 +1,150 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_REQUEST_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxUsageReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxUsageReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_REQUEST_ID_SQL)
|
||||
.bind(request_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for SqlxUsageReadRepository {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_request_id(self, request_id).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
StoredRequestUsageAudit::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("provider_name")?,
|
||||
row.try_get("model")?,
|
||||
row.try_get("target_model")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("provider_endpoint_id")?,
|
||||
row.try_get("provider_api_key_id")?,
|
||||
row.try_get("request_type")?,
|
||||
row.try_get("api_format")?,
|
||||
row.try_get("api_family")?,
|
||||
row.try_get("endpoint_kind")?,
|
||||
row.try_get("endpoint_api_format")?,
|
||||
row.try_get("provider_api_family")?,
|
||||
row.try_get("provider_endpoint_kind")?,
|
||||
row.try_get("has_format_conversion")?,
|
||||
row.try_get("is_stream")?,
|
||||
row.try_get("input_tokens")?,
|
||||
row.try_get("output_tokens")?,
|
||||
row.try_get("total_tokens")?,
|
||||
row.try_get("total_cost_usd")?,
|
||||
row.try_get("actual_total_cost_usd")?,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("error_category")?,
|
||||
row.try_get("response_time_ms")?,
|
||||
row.try_get("first_byte_time_ms")?,
|
||||
row.try_get("status")?,
|
||||
row.try_get("billing_status")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("finalized_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxUsageReadRepository;
|
||||
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 = SqlxUsageReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
304
crates/aether-data/src/repository/usage/types.rs
Normal file
304
crates/aether-data/src/repository/usage/types.rs
Normal file
@@ -0,0 +1,304 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestUsageAudit {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub provider_name: String,
|
||||
pub model: String,
|
||||
pub target_model: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub provider_endpoint_id: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub request_type: Option<String>,
|
||||
pub api_format: Option<String>,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub endpoint_api_format: Option<String>,
|
||||
pub provider_api_family: Option<String>,
|
||||
pub provider_endpoint_kind: Option<String>,
|
||||
pub has_format_conversion: bool,
|
||||
pub is_stream: bool,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub total_cost_usd: f64,
|
||||
pub actual_total_cost_usd: f64,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_category: Option<String>,
|
||||
pub response_time_ms: Option<u64>,
|
||||
pub first_byte_time_ms: Option<u64>,
|
||||
pub status: String,
|
||||
pub billing_status: String,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredRequestUsageAudit {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
provider_name: String,
|
||||
model: String,
|
||||
target_model: Option<String>,
|
||||
provider_id: Option<String>,
|
||||
provider_endpoint_id: Option<String>,
|
||||
provider_api_key_id: Option<String>,
|
||||
request_type: Option<String>,
|
||||
api_format: Option<String>,
|
||||
api_family: Option<String>,
|
||||
endpoint_kind: Option<String>,
|
||||
endpoint_api_format: Option<String>,
|
||||
provider_api_family: Option<String>,
|
||||
provider_endpoint_kind: Option<String>,
|
||||
has_format_conversion: bool,
|
||||
is_stream: bool,
|
||||
input_tokens: i32,
|
||||
output_tokens: i32,
|
||||
total_tokens: i32,
|
||||
total_cost_usd: f64,
|
||||
actual_total_cost_usd: f64,
|
||||
status_code: Option<i32>,
|
||||
error_message: Option<String>,
|
||||
error_category: Option<String>,
|
||||
response_time_ms: Option<i32>,
|
||||
first_byte_time_ms: Option<i32>,
|
||||
status: String,
|
||||
billing_status: String,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
finalized_at_unix_secs: Option<i64>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.request_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.provider_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if model.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.model is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if billing_status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.billing_status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
if !actual_total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.actual_total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
has_format_conversion,
|
||||
is_stream,
|
||||
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
|
||||
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
|
||||
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code: parse_u16(status_code, "usage.status_code")?,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms: parse_optional_u64(response_time_ms, "usage.response_time_ms")?,
|
||||
first_byte_time_ms: parse_optional_u64(first_byte_time_ms, "usage.first_byte_time_ms")?,
|
||||
status,
|
||||
billing_status,
|
||||
created_at_unix_secs: parse_timestamp(
|
||||
created_at_unix_secs,
|
||||
"usage.created_at_unix_secs",
|
||||
)?,
|
||||
updated_at_unix_secs: parse_timestamp(
|
||||
updated_at_unix_secs,
|
||||
"usage.updated_at_unix_secs",
|
||||
)?,
|
||||
finalized_at_unix_secs: finalized_at_unix_secs
|
||||
.map(|value| parse_timestamp(value, "usage.finalized_at_unix_secs"))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadRepository: Send + Sync {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + Send + Sync {}
|
||||
|
||||
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_u64(
|
||||
value: Option<i32>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<u64>, crate::DataLayerError> {
|
||||
value
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_u16(value: Option<i32>, field_name: &str) -> Result<Option<u16>, crate::DataLayerError> {
|
||||
value
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredRequestUsageAudit;
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_request_id() {
|
||||
assert!(StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
0.1,
|
||||
0.1,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_token_count() {
|
||||
assert!(StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
-1,
|
||||
20,
|
||||
30,
|
||||
0.1,
|
||||
0.1,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
232
crates/aether-data/src/repository/video_tasks/memory.rs
Normal file
232
crates/aether-data/src/repository/video_tasks/memory.rs
Normal file
@@ -0,0 +1,232 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryVideoTaskIndex {
|
||||
by_id: BTreeMap<String, StoredVideoTask>,
|
||||
short_to_id: BTreeMap<String, String>,
|
||||
user_external_to_id: BTreeMap<(String, String), String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryVideoTaskRepository {
|
||||
index: RwLock<MemoryVideoTaskIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryVideoTaskRepository {
|
||||
fn store_locked(index: &mut MemoryVideoTaskIndex, task: StoredVideoTask) -> StoredVideoTask {
|
||||
if let Some(previous) = index.by_id.insert(task.id.clone(), task.clone()) {
|
||||
if let Some(short_id) = previous.short_id {
|
||||
index.short_to_id.remove(&short_id);
|
||||
}
|
||||
if let (Some(user_id), Some(external_task_id)) =
|
||||
(previous.user_id, previous.external_task_id)
|
||||
{
|
||||
index
|
||||
.user_external_to_id
|
||||
.remove(&(user_id, external_task_id));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(short_id) = &task.short_id {
|
||||
index.short_to_id.insert(short_id.clone(), task.id.clone());
|
||||
}
|
||||
if let (Some(user_id), Some(external_task_id)) = (&task.user_id, &task.external_task_id) {
|
||||
index
|
||||
.user_external_to_id
|
||||
.insert((user_id.clone(), external_task_id.clone()), task.id.clone());
|
||||
}
|
||||
|
||||
task
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let index = self.index.read().expect("video task repository lock");
|
||||
Ok(match key {
|
||||
VideoTaskLookupKey::Id(id) => index.by_id.get(id).cloned(),
|
||||
VideoTaskLookupKey::ShortId(short_id) => index
|
||||
.short_to_id
|
||||
.get(short_id)
|
||||
.and_then(|id| index.by_id.get(id))
|
||||
.cloned(),
|
||||
VideoTaskLookupKey::UserExternal {
|
||||
user_id,
|
||||
external_task_id,
|
||||
} => index
|
||||
.user_external_to_id
|
||||
.get(&(user_id.to_string(), external_task_id.to_string()))
|
||||
.and_then(|id| index.by_id.get(id))
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut tasks = self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| task.status.is_active())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
tasks.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
|
||||
tasks.truncate(limit);
|
||||
Ok(tasks)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskWriteRepository for InMemoryVideoTaskRepository {
|
||||
async fn upsert(&self, task: UpsertVideoTask) -> Result<StoredVideoTask, DataLayerError> {
|
||||
let mut index = self.index.write().expect("video task repository lock");
|
||||
Ok(Self::store_locked(&mut index, task.into_stored()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryVideoTaskRepository;
|
||||
use crate::repository::video_tasks::{
|
||||
UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
fn sample_task(
|
||||
id: &str,
|
||||
status: VideoTaskStatus,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> UpsertVideoTask {
|
||||
UpsertVideoTask {
|
||||
id: id.to_string(),
|
||||
short_id: Some(format!("short-{id}")),
|
||||
user_id: Some("user-1".to_string()),
|
||||
external_task_id: Some(format!("ext-{id}")),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status,
|
||||
progress_percent: 0,
|
||||
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
|
||||
updated_at_unix_secs,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_task_by_all_supported_lookup_keys() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::Id("task-1"))
|
||||
.await
|
||||
.expect("find by id should succeed")
|
||||
.is_some());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1"))
|
||||
.await
|
||||
.expect("find by short id should succeed")
|
||||
.is_some());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::UserExternal {
|
||||
user_id: "user-1",
|
||||
external_task_id: "ext-task-1",
|
||||
})
|
||||
.await
|
||||
.expect("find by user/external should succeed")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_active_only_returns_active_tasks_in_descending_update_order() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Completed, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 200))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-3", VideoTaskStatus::Queued, 150))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let active = repo
|
||||
.list_active(10)
|
||||
.await
|
||||
.expect("list active should succeed");
|
||||
assert_eq!(active.len(), 2);
|
||||
assert_eq!(active[0].id, "task-2");
|
||||
assert_eq!(active[1].id, "task-3");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_replaces_secondary_indexes() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
repo.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1b".to_string()),
|
||||
user_id: Some("user-2".to_string()),
|
||||
external_task_id: Some("ext-task-1b".to_string()),
|
||||
provider_api_format: Some("gemini:video".to_string()),
|
||||
model: Some("veo-3".to_string()),
|
||||
prompt: Some("remix".to_string()),
|
||||
size: Some("720p".to_string()),
|
||||
status: VideoTaskStatus::Processing,
|
||||
progress_percent: 50,
|
||||
created_at_unix_secs: 150,
|
||||
updated_at_unix_secs: 200,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1"))
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_none());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::UserExternal {
|
||||
user_id: "user-1",
|
||||
external_task_id: "ext-task-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_none());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1b"))
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_some());
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/video_tasks/mod.rs
Normal file
10
crates/aether-data/src/repository/video_tasks/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryVideoTaskRepository;
|
||||
pub use sql::SqlxVideoTaskReadRepository;
|
||||
pub use types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskRepository, VideoTaskStatus, VideoTaskWriteRepository,
|
||||
};
|
||||
259
crates/aether-data/src/repository/video_tasks/sql.rs
Normal file
259
crates/aether-data/src/repository/video_tasks/sql.rs
Normal file
@@ -0,0 +1,259 @@
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::repository::video_tasks::{
|
||||
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_SHORT_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE short_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_USER_EXTERNAL_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE user_id = $1 AND external_task_id = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_ACTIVE_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE status = ANY($1)
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT $2
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxVideoTaskReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxVideoTaskReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match key {
|
||||
VideoTaskLookupKey::Id(id) => self.find_by_id(id).await,
|
||||
VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await,
|
||||
VideoTaskLookupKey::UserExternal {
|
||||
user_id,
|
||||
external_task_id,
|
||||
} => self.find_by_user_external(user_id, external_task_id).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn find_by_id(&self, id: &str) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_ID_SQL)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_short_id(
|
||||
&self,
|
||||
short_id: &str,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_SHORT_ID_SQL)
|
||||
.bind(short_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_user_external(
|
||||
&self,
|
||||
user_id: &str,
|
||||
external_task_id: &str,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_USER_EXTERNAL_SQL)
|
||||
.bind(user_id)
|
||||
.bind(external_task_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let active_statuses = vec!["pending", "submitted", "queued", "processing"];
|
||||
let rows = sqlx::query(LIST_ACTIVE_SQL)
|
||||
.bind(active_statuses)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}"))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_video_task_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskReadRepository for SqlxVideoTaskReadRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
Self::find(self, key).await
|
||||
}
|
||||
|
||||
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
Self::list_active(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_video_task_row(row: &sqlx::postgres::PgRow) -> Result<StoredVideoTask, DataLayerError> {
|
||||
let status = VideoTaskStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
|
||||
StoredVideoTask::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("short_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("external_task_id")?,
|
||||
row.try_get("provider_api_format")?,
|
||||
row.try_get("model")?,
|
||||
row.try_get("prompt")?,
|
||||
row.try_get("size")?,
|
||||
status,
|
||||
row.try_get("progress_percent")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("error_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("video_url")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxVideoTaskReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
use crate::repository::video_tasks::{VideoTaskLookupKey, VideoTaskReadRepository};
|
||||
|
||||
#[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 = SqlxVideoTaskReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_trait_delegates_to_sqlx_repository() {
|
||||
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 = SqlxVideoTaskReadRepository::new(pool);
|
||||
let _ = VideoTaskReadRepository::find(&repository, VideoTaskLookupKey::Id("task-1")).await;
|
||||
}
|
||||
}
|
||||
277
crates/aether-data/src/repository/video_tasks/types.rs
Normal file
277
crates/aether-data/src/repository/video_tasks/types.rs
Normal file
@@ -0,0 +1,277 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum VideoTaskStatus {
|
||||
Pending,
|
||||
Submitted,
|
||||
Queued,
|
||||
Processing,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Expired,
|
||||
Deleted,
|
||||
}
|
||||
|
||||
impl VideoTaskStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"pending" => Ok(Self::Pending),
|
||||
"submitted" => Ok(Self::Submitted),
|
||||
"queued" => Ok(Self::Queued),
|
||||
"processing" => Ok(Self::Processing),
|
||||
"completed" => Ok(Self::Completed),
|
||||
"failed" => Ok(Self::Failed),
|
||||
"cancelled" => Ok(Self::Cancelled),
|
||||
"expired" => Ok(Self::Expired),
|
||||
"deleted" => Ok(Self::Deleted),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported video_tasks.status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_active(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Pending | Self::Submitted | Self::Queued | Self::Processing
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
}
|
||||
|
||||
impl StoredVideoTask {
|
||||
pub fn new(
|
||||
id: String,
|
||||
short_id: Option<String>,
|
||||
user_id: Option<String>,
|
||||
external_task_id: Option<String>,
|
||||
provider_api_format: Option<String>,
|
||||
model: Option<String>,
|
||||
prompt: Option<String>,
|
||||
size: Option<String>,
|
||||
status: VideoTaskStatus,
|
||||
progress_percent: i32,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
error_code: Option<String>,
|
||||
error_message: Option<String>,
|
||||
video_url: Option<String>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid progress_percent: {progress_percent}"
|
||||
))
|
||||
})?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpsertVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
}
|
||||
|
||||
impl UpsertVideoTask {
|
||||
pub fn into_stored(self) -> StoredVideoTask {
|
||||
StoredVideoTask {
|
||||
id: self.id,
|
||||
short_id: self.short_id,
|
||||
user_id: self.user_id,
|
||||
external_task_id: self.external_task_id,
|
||||
provider_api_format: self.provider_api_format,
|
||||
model: self.model,
|
||||
prompt: self.prompt,
|
||||
size: self.size,
|
||||
status: self.status,
|
||||
progress_percent: self.progress_percent,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
error_code: self.error_code,
|
||||
error_message: self.error_message,
|
||||
video_url: self.video_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum VideoTaskLookupKey<'a> {
|
||||
Id(&'a str),
|
||||
ShortId(&'a str),
|
||||
UserExternal {
|
||||
user_id: &'a str,
|
||||
external_task_id: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn list_active(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskWriteRepository: Send + Sync {
|
||||
async fn upsert(&self, task: UpsertVideoTask)
|
||||
-> Result<StoredVideoTask, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait VideoTaskRepository:
|
||||
VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> VideoTaskRepository for T where
|
||||
T: VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredVideoTask, VideoTaskStatus};
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
assert_eq!(
|
||||
VideoTaskStatus::from_database("processing").expect("status should parse"),
|
||||
VideoTaskStatus::Processing
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(VideoTaskStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_numeric_fields() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
-1,
|
||||
1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_updated_at_values() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
1,
|
||||
-1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_created_at_values() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
-1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user