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:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions

View File

@@ -0,0 +1,20 @@
[package]
name = "aether-data"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared data access contracts and config for Aether Rust services"
[dependencies]
aether-cache.workspace = true
async-trait.workspace = true
futures-util.workspace = true
redis.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true
uuid.workspace = true

View File

@@ -0,0 +1,66 @@
use std::fmt;
use super::PostgresBackend;
use crate::postgres::{PostgresLeaseRunner, PostgresLeaseRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataLeaseBackends {
postgres: Option<PostgresLeaseRunner>,
}
impl fmt::Debug for DataLeaseBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataLeaseBackends")
.field("has_postgres", &self.postgres.is_some())
.finish()
}
}
impl DataLeaseBackends {
pub(crate) fn from_postgres(
postgres: Option<&PostgresBackend>,
) -> Result<Self, DataLayerError> {
Ok(Self {
postgres: postgres
.map(|backend| backend.lease_runner(PostgresLeaseRunnerConfig::default()))
.transpose()?,
})
}
pub fn postgres(&self) -> Option<PostgresLeaseRunner> {
self.postgres.clone()
}
pub fn has_any(&self) -> bool {
self.postgres.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLeaseBackends;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_postgres_lease_runner_from_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let leases =
DataLeaseBackends::from_postgres(Some(&backend)).expect("lease backends should build");
assert!(leases.has_any());
assert!(leases.postgres().is_some());
}
}

View File

@@ -0,0 +1,58 @@
use std::fmt;
use super::RedisBackend;
use crate::redis::{RedisLockRunner, RedisLockRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataLockBackends {
redis: Option<RedisLockRunner>,
}
impl fmt::Debug for DataLockBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataLockBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataLockBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.lock_runner(RedisLockRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisLockRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLockBackends;
use crate::backends::RedisBackend;
use crate::redis::RedisClientConfig;
#[test]
fn builds_redis_lock_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let locks =
DataLockBackends::from_redis(Some(&backend)).expect("lock backends should build");
assert!(locks.has_any());
assert!(locks.redis().is_some());
}
}

View File

@@ -0,0 +1,195 @@
mod leases;
mod locks;
mod postgres;
mod read;
mod redis;
mod transactions;
mod workers;
mod write;
pub use leases::DataLeaseBackends;
pub use locks::DataLockBackends;
pub use postgres::PostgresBackend;
pub use read::DataReadRepositories;
pub use redis::RedisBackend;
pub use transactions::DataTransactionBackends;
pub use workers::DataWorkerBackends;
pub use write::DataWriteRepositories;
use crate::{DataLayerConfig, DataLayerError};
#[derive(Debug, Clone, Default)]
pub struct DataBackends {
config: DataLayerConfig,
postgres: Option<PostgresBackend>,
redis: Option<RedisBackend>,
leases: DataLeaseBackends,
locks: DataLockBackends,
read: DataReadRepositories,
transactions: DataTransactionBackends,
workers: DataWorkerBackends,
write: DataWriteRepositories,
}
impl DataBackends {
pub fn from_config(config: DataLayerConfig) -> Result<Self, DataLayerError> {
config.validate()?;
let postgres = config
.postgres
.clone()
.map(PostgresBackend::from_config)
.transpose()?;
let redis = config
.redis
.clone()
.map(RedisBackend::from_config)
.transpose()?;
let leases = DataLeaseBackends::from_postgres(postgres.as_ref())?;
let locks = DataLockBackends::from_redis(redis.as_ref())?;
let read = DataReadRepositories::from_postgres(postgres.as_ref());
let transactions = DataTransactionBackends::from_postgres(postgres.as_ref());
let workers = DataWorkerBackends::from_redis(redis.as_ref())?;
let write = DataWriteRepositories::from_postgres(postgres.as_ref());
Ok(Self {
config,
postgres,
redis,
leases,
locks,
read,
transactions,
workers,
write,
})
}
pub fn config(&self) -> &DataLayerConfig {
&self.config
}
pub fn postgres(&self) -> Option<&PostgresBackend> {
self.postgres.as_ref()
}
pub fn redis(&self) -> Option<&RedisBackend> {
self.redis.as_ref()
}
pub fn read(&self) -> &DataReadRepositories {
&self.read
}
pub fn leases(&self) -> &DataLeaseBackends {
&self.leases
}
pub fn locks(&self) -> &DataLockBackends {
&self.locks
}
pub fn transactions(&self) -> &DataTransactionBackends {
&self.transactions
}
pub fn workers(&self) -> &DataWorkerBackends {
&self.workers
}
pub fn write(&self) -> &DataWriteRepositories {
&self.write
}
pub fn has_runtime_backends(&self) -> bool {
self.postgres.is_some()
|| self.redis.is_some()
|| self.leases.has_any()
|| self.locks.has_any()
|| self.read.has_any()
|| self.transactions.has_any()
|| self.workers.has_any()
|| self.write.has_any()
}
}
#[cfg(test)]
mod tests {
use super::DataBackends;
use crate::{postgres::PostgresPoolConfig, DataLayerConfig};
#[test]
fn builds_empty_backends_from_default_config() {
let backends = DataBackends::from_config(DataLayerConfig::default())
.expect("empty config should be accepted");
assert!(!backends.has_runtime_backends());
assert!(backends.postgres().is_none());
assert!(backends.redis().is_none());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_none());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.read().request_candidates().is_none());
assert!(backends.read().provider_catalog().is_none());
assert!(backends.read().usage().is_none());
assert!(backends.read().video_tasks().is_none());
assert!(backends.read().shadow_results().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.workers().redis().is_none());
assert!(backends.write().shadow_results().is_none());
}
#[tokio::test]
async fn builds_postgres_backend_from_config() {
let backends = DataBackends::from_config(DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
})
.expect("postgres backend should build");
assert!(backends.has_runtime_backends());
assert!(backends.postgres().is_some());
assert!(backends.leases().postgres().is_some());
assert!(backends.read().auth_api_keys().is_some());
assert!(backends.read().request_candidates().is_some());
assert!(backends.read().provider_catalog().is_some());
assert!(backends.read().usage().is_some());
assert!(backends.read().video_tasks().is_some());
assert!(backends.read().shadow_results().is_some());
assert!(backends.transactions().postgres().is_some());
assert!(backends.write().shadow_results().is_some());
assert!(backends.config().postgres.is_some());
}
#[test]
fn builds_redis_backend_from_config() {
let backends = DataBackends::from_config(DataLayerConfig {
postgres: None,
redis: Some(crate::redis::RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
})
.expect("redis backend should build");
assert!(backends.has_runtime_backends());
assert!(backends.postgres().is_none());
assert!(backends.redis().is_some());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_some());
assert!(backends.workers().redis().is_some());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.config().redis.is_some());
}
}

View File

@@ -0,0 +1,123 @@
use std::sync::Arc;
use crate::postgres::{
PostgresLeaseRunner, PostgresLeaseRunnerConfig, PostgresPool, PostgresPoolConfig,
PostgresPoolFactory, PostgresTransactionRunner,
};
use crate::repository::auth::{AuthApiKeyReadRepository, SqlxAuthApiKeySnapshotReadRepository};
use crate::repository::candidates::{
RequestCandidateReadRepository, SqlxRequestCandidateReadRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, SqlxProviderCatalogReadRepository,
};
use crate::repository::shadow_results::{
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
};
use crate::repository::usage::{SqlxUsageReadRepository, UsageReadRepository};
use crate::repository::video_tasks::{SqlxVideoTaskReadRepository, VideoTaskReadRepository};
use crate::DataLayerError;
#[derive(Debug, Clone)]
pub struct PostgresBackend {
config: PostgresPoolConfig,
pool: PostgresPool,
}
impl PostgresBackend {
pub fn from_config(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
let factory = PostgresPoolFactory::new(config.clone())?;
let pool = factory.connect_lazy()?;
Ok(Self { config, pool })
}
pub fn config(&self) -> &PostgresPoolConfig {
&self.config
}
pub fn pool(&self) -> &PostgresPool {
&self.pool
}
pub fn pool_clone(&self) -> PostgresPool {
self.pool.clone()
}
pub fn auth_api_key_read_repository(&self) -> Arc<dyn AuthApiKeyReadRepository> {
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
}
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
}
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
}
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
}
pub fn transaction_runner(&self) -> PostgresTransactionRunner {
PostgresTransactionRunner::new(self.pool_clone())
}
pub fn lease_runner(
&self,
config: PostgresLeaseRunnerConfig,
) -> Result<PostgresLeaseRunner, DataLayerError> {
PostgresLeaseRunner::new(self.transaction_runner(), config)
}
pub fn shadow_result_write_repository(&self) -> Arc<dyn ShadowResultWriteRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
pub fn shadow_result_read_repository(&self) -> Arc<dyn ShadowResultReadRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
}
#[cfg(test)]
mod tests {
use super::PostgresBackend;
use crate::postgres::{PostgresLeaseRunnerConfig, PostgresPoolConfig};
#[tokio::test]
async fn backend_retains_config_and_pool() {
let config = PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
};
let backend =
PostgresBackend::from_config(config.clone()).expect("backend should build lazily");
assert_eq!(backend.config(), &config);
let _pool = backend.pool();
let _pool_clone = backend.pool_clone();
let _auth_api_key_reader = backend.auth_api_key_read_repository();
let _request_candidate_reader = backend.request_candidate_read_repository();
let _provider_catalog_reader = backend.provider_catalog_read_repository();
let _usage_reader = backend.usage_read_repository();
let _video_task_reader = backend.video_task_read_repository();
let _transaction_runner = backend.transaction_runner();
let _lease_runner = backend
.lease_runner(PostgresLeaseRunnerConfig::default())
.expect("lease runner should build");
let _shadow_result_reader = backend.shadow_result_read_repository();
let _shadow_result_writer = backend.shadow_result_write_repository();
}
}

View File

@@ -0,0 +1,111 @@
use std::fmt;
use std::sync::Arc;
use super::PostgresBackend;
use crate::repository::auth::AuthApiKeyReadRepository;
use crate::repository::candidates::RequestCandidateReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::shadow_results::ShadowResultReadRepository;
use crate::repository::usage::UsageReadRepository;
use crate::repository::video_tasks::VideoTaskReadRepository;
#[derive(Clone, Default)]
pub struct DataReadRepositories {
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
usage: Option<Arc<dyn UsageReadRepository>>,
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
shadow_results: Option<Arc<dyn ShadowResultReadRepository>>,
}
impl fmt::Debug for DataReadRepositories {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataReadRepositories")
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
.field("has_request_candidates", &self.request_candidates.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_usage", &self.usage.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
.field("has_shadow_results", &self.shadow_results.is_some())
.finish()
}
}
impl DataReadRepositories {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_read_repository),
request_candidates: postgres.map(PostgresBackend::request_candidate_read_repository),
provider_catalog: postgres.map(PostgresBackend::provider_catalog_read_repository),
usage: postgres.map(PostgresBackend::usage_read_repository),
video_tasks: postgres.map(PostgresBackend::video_task_read_repository),
shadow_results: postgres.map(PostgresBackend::shadow_result_read_repository),
}
}
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyReadRepository>> {
self.auth_api_keys.clone()
}
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateReadRepository>> {
self.request_candidates.clone()
}
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogReadRepository>> {
self.provider_catalog.clone()
}
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
self.usage.clone()
}
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskReadRepository>> {
self.video_tasks.clone()
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultReadRepository>> {
self.shadow_results.clone()
}
pub fn has_any(&self) -> bool {
self.auth_api_keys.is_some()
|| self.request_candidates.is_some()
|| self.provider_catalog.is_some()
|| self.usage.is_some()
|| self.video_tasks.is_some()
|| self.shadow_results.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataReadRepositories;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_read_repositories_from_postgres_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let read = DataReadRepositories::from_postgres(Some(&backend));
assert!(read.has_any());
assert!(read.auth_api_keys().is_some());
assert!(read.request_candidates().is_some());
assert!(read.provider_catalog().is_some());
assert!(read.usage().is_some());
assert!(read.video_tasks().is_some());
assert!(read.shadow_results().is_some());
}
}

View File

@@ -0,0 +1,76 @@
use crate::redis::{
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisLockRunner,
RedisLockRunnerConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use crate::DataLayerError;
#[derive(Debug, Clone)]
pub struct RedisBackend {
config: RedisClientConfig,
client: RedisClient,
}
impl RedisBackend {
pub fn from_config(config: RedisClientConfig) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config.clone())?;
let client = factory.connect_lazy()?;
Ok(Self { config, client })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn client_clone(&self) -> RedisClient {
self.client.clone()
}
pub fn keyspace(&self) -> RedisKeyspace {
self.config.keyspace()
}
pub fn lock_runner(
&self,
config: RedisLockRunnerConfig,
) -> Result<RedisLockRunner, DataLayerError> {
RedisLockRunner::new(self.client_clone(), self.keyspace(), config)
}
pub fn stream_runner(
&self,
config: RedisStreamRunnerConfig,
) -> Result<RedisStreamRunner, DataLayerError> {
RedisStreamRunner::new(self.client_clone(), self.keyspace(), config)
}
}
#[cfg(test)]
mod tests {
use super::RedisBackend;
use crate::redis::{RedisClientConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig};
#[test]
fn backend_retains_config_client_and_shared_runners() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let backend = RedisBackend::from_config(config.clone()).expect("backend should build");
assert_eq!(backend.config(), &config);
assert_eq!(backend.keyspace().key("audit"), "aether:audit");
let _client_ref = backend.client();
let _client_clone = backend.client_clone();
let _lock_runner = backend
.lock_runner(RedisLockRunnerConfig::default())
.expect("lock runner should build");
let _stream_runner = backend
.stream_runner(RedisStreamRunnerConfig::default())
.expect("stream runner should build");
}
}

View File

@@ -0,0 +1,60 @@
use std::fmt;
use super::PostgresBackend;
use crate::postgres::PostgresTransactionRunner;
#[derive(Clone, Default)]
pub struct DataTransactionBackends {
postgres: Option<PostgresTransactionRunner>,
}
impl fmt::Debug for DataTransactionBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataTransactionBackends")
.field("has_postgres", &self.postgres.is_some())
.finish()
}
}
impl DataTransactionBackends {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
postgres: postgres.map(PostgresBackend::transaction_runner),
}
}
pub fn postgres(&self) -> Option<PostgresTransactionRunner> {
self.postgres.clone()
}
pub fn has_any(&self) -> bool {
self.postgres.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataTransactionBackends;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_postgres_transaction_runner_from_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let transactions = DataTransactionBackends::from_postgres(Some(&backend));
assert!(transactions.has_any());
assert!(transactions.postgres().is_some());
}
}

View File

@@ -0,0 +1,58 @@
use std::fmt;
use super::RedisBackend;
use crate::redis::{RedisStreamRunner, RedisStreamRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataWorkerBackends {
redis: Option<RedisStreamRunner>,
}
impl fmt::Debug for DataWorkerBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataWorkerBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataWorkerBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.stream_runner(RedisStreamRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisStreamRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataWorkerBackends;
use crate::backends::RedisBackend;
use crate::redis::RedisClientConfig;
#[test]
fn builds_redis_stream_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let workers =
DataWorkerBackends::from_redis(Some(&backend)).expect("worker backends should build");
assert!(workers.has_any());
assert!(workers.redis().is_some());
}
}

View File

@@ -0,0 +1,61 @@
use std::fmt;
use std::sync::Arc;
use super::PostgresBackend;
use crate::repository::shadow_results::ShadowResultWriteRepository;
#[derive(Clone, Default)]
pub struct DataWriteRepositories {
shadow_results: Option<Arc<dyn ShadowResultWriteRepository>>,
}
impl fmt::Debug for DataWriteRepositories {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataWriteRepositories")
.field("has_shadow_results", &self.shadow_results.is_some())
.finish()
}
}
impl DataWriteRepositories {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
shadow_results: postgres.map(PostgresBackend::shadow_result_write_repository),
}
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultWriteRepository>> {
self.shadow_results.clone()
}
pub fn has_any(&self) -> bool {
self.shadow_results.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataWriteRepositories;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_shadow_result_writer_from_postgres_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let write = DataWriteRepositories::from_postgres(Some(&backend));
assert!(write.has_any());
assert!(write.shadow_results().is_some());
}
}

View File

@@ -0,0 +1,74 @@
use crate::postgres::PostgresPoolConfig;
use crate::redis::RedisClientConfig;
use crate::DataLayerError;
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct DataLayerConfig {
pub postgres: Option<PostgresPoolConfig>,
pub redis: Option<RedisClientConfig>,
}
impl DataLayerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if let Some(postgres) = &self.postgres {
postgres.validate()?;
}
if let Some(redis) = &self.redis {
redis.validate()?;
}
Ok(())
}
pub fn has_persistent_backends(&self) -> bool {
self.postgres.is_some() || self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLayerConfig;
use crate::postgres::PostgresPoolConfig;
use crate::redis::RedisClientConfig;
#[test]
fn validates_nested_backend_configs() {
let config = DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 2,
max_connections: 8,
acquire_timeout_ms: 1_500,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: Some(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
};
assert!(config.validate().is_ok());
assert!(config.has_persistent_backends());
}
#[test]
fn rejects_invalid_nested_backend_configs() {
let config = DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: String::new(),
min_connections: 4,
max_connections: 2,
acquire_timeout_ms: 1_500,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
};
assert!(config.validate().is_err());
}
}

View File

@@ -0,0 +1,20 @@
#[derive(Debug, thiserror::Error)]
pub enum DataLayerError {
#[error("invalid configuration: {0}")]
InvalidConfiguration(String),
#[error("invalid input: {0}")]
InvalidInput(String),
#[error("postgres error: {0}")]
Postgres(#[from] sqlx::Error),
#[error("redis error: {0}")]
Redis(#[from] redis::RedisError),
#[error("operation timed out: {0}")]
TimedOut(String),
#[error("unexpected database value: {0}")]
UnexpectedValue(String),
}

View File

@@ -0,0 +1,14 @@
pub mod backends;
mod config;
mod error;
pub mod postgres;
pub mod redis;
pub mod repository;
pub use backends::{
DataBackends, DataLeaseBackends, DataLockBackends, DataReadRepositories,
DataTransactionBackends, DataWorkerBackends, DataWriteRepositories, PostgresBackend,
RedisBackend,
};
pub use config::DataLayerConfig;
pub use error::DataLayerError;

View File

@@ -0,0 +1,417 @@
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
use crate::DataLayerError;
use futures_util::FutureExt;
use sqlx::query_scalar;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PostgresLeaseClaimOptions {
pub batch_size: usize,
pub lease_ms: u64,
}
impl PostgresLeaseClaimOptions {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease batch_size must be positive".to_string(),
));
}
if self.lease_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lease_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PostgresLeaseClaimSpec {
pub table: &'static str,
pub id_column: &'static str,
pub lease_owner_column: &'static str,
pub lease_expires_at_column: &'static str,
pub eligibility_predicate_sql: &'static str,
pub order_by_sql: &'static str,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresLeaseRunnerConfig {
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresLeaseRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct PostgresLeaseRunner {
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
}
impl PostgresLeaseRunner {
pub fn new(
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
transaction_runner,
config,
})
}
pub fn config(&self) -> PostgresLeaseRunnerConfig {
self.config
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.transaction_runner
}
pub async fn claim_ids(
&self,
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
let sql = build_postgres_lease_claim_sql(spec, options)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(options.lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(owner)
.bind(lease_ms)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
pub async fn release_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_release_sql(spec)?;
let owner = owner.trim().to_string();
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
pub async fn renew_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
lease_ms: u64,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if lease_ms == 0 {
return Err(DataLayerError::InvalidInput(
"postgres lease lease_ms must be positive".to_string(),
));
}
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_renew_sql(spec)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.bind(lease_ms)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
}
pub fn build_postgres_lease_claim_sql(
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
) -> Result<String, DataLayerError> {
options.validate()?;
validate_lease_spec(spec)?;
Ok(format!(
"WITH claimable AS (\
SELECT {id_column} \
FROM {table} \
WHERE ({eligibility_predicate_sql}) \
AND ({lease_expires_at_column} IS NULL OR {lease_expires_at_column} <= NOW()) \
ORDER BY {order_by_sql} \
FOR UPDATE SKIP LOCKED \
LIMIT {batch_size}\
) \
UPDATE {table} AS target \
SET {lease_owner_column} = $1, \
{lease_expires_at_column} = NOW() + ($2::bigint * INTERVAL '1 millisecond') \
FROM claimable \
WHERE target.{id_column} = claimable.{id_column} \
RETURNING target.{id_column}",
id_column = spec.id_column,
table = spec.table,
eligibility_predicate_sql = spec.eligibility_predicate_sql,
lease_expires_at_column = spec.lease_expires_at_column,
order_by_sql = spec.order_by_sql,
batch_size = options.batch_size,
lease_owner_column = spec.lease_owner_column,
))
}
pub fn build_postgres_lease_release_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_owner_column} = NULL, \
{lease_expires_at_column} = NULL \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
pub fn build_postgres_lease_renew_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_expires_at_column} = NOW() + ($3::bigint * INTERVAL '1 millisecond') \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
fn validate_lease_spec(spec: &PostgresLeaseClaimSpec) -> Result<(), DataLayerError> {
for (field, value) in [
("table", spec.table),
("id_column", spec.id_column),
("lease_owner_column", spec.lease_owner_column),
("lease_expires_at_column", spec.lease_expires_at_column),
("eligibility_predicate_sql", spec.eligibility_predicate_sql),
("order_by_sql", spec.order_by_sql),
] {
if value.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(format!(
"postgres lease {field} cannot be empty"
)));
}
}
Ok(())
}
fn validate_lease_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"postgres lease owner cannot be empty".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory, PostgresTransactionRunner};
fn sample_spec() -> PostgresLeaseClaimSpec {
PostgresLeaseClaimSpec {
table: "video_tasks",
id_column: "id",
lease_owner_column: "lease_owner",
lease_expires_at_column: "lease_expires_at",
eligibility_predicate_sql: "status IN ('submitted', 'processing')",
order_by_sql: "updated_at ASC",
}
}
#[test]
fn builds_skip_locked_claim_sql() {
let sql = build_postgres_lease_claim_sql(
&sample_spec(),
PostgresLeaseClaimOptions {
batch_size: 16,
lease_ms: 15_000,
},
)
.expect("claim SQL should build");
assert!(sql.contains("FOR UPDATE SKIP LOCKED"));
assert!(sql.contains("LIMIT 16"));
assert!(sql.contains("lease_owner = $1"));
assert!(sql.contains("NOW() + ($2::bigint * INTERVAL '1 millisecond')"));
}
#[test]
fn builds_release_sql() {
let sql =
build_postgres_lease_release_sql(&sample_spec()).expect("release SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("lease_expires_at = NULL"));
}
#[test]
fn builds_renew_sql() {
let sql = build_postgres_lease_renew_sql(&sample_spec()).expect("renew SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("NOW() + ($3::bigint * INTERVAL '1 millisecond')"));
}
#[tokio::test]
async fn lease_runner_reuses_transaction_runner() {
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("lazy pool should build");
let transaction_runner = PostgresTransactionRunner::new(pool);
let lease_runner = PostgresLeaseRunner::new(
transaction_runner.clone(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
},
)
.expect("lease runner should build");
assert_eq!(
lease_runner.config(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
}
);
let _runner_ref = lease_runner.transaction_runner();
}
#[tokio::test]
async fn release_and_renew_empty_ids_short_circuit() {
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("lazy pool should build");
let runner = PostgresLeaseRunner::new(
PostgresTransactionRunner::new(pool),
PostgresLeaseRunnerConfig::default(),
)
.expect("lease runner should build");
assert!(runner
.release_ids(&sample_spec(), &[], "worker-1")
.await
.expect("empty release should succeed")
.is_empty());
assert!(runner
.renew_ids(&sample_spec(), &[], "worker-1", 5_000)
.await
.expect("empty renew should succeed")
.is_empty());
}
}

View File

@@ -0,0 +1,15 @@
mod lease;
mod pool;
mod tx;
mod types;
pub use lease::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
pub use pool::{PostgresPool, PostgresPoolConfig, PostgresPoolFactory};
pub use tx::{
PostgresTransaction, PostgresTransactionOptions, PostgresTransactionRunner, TransactionMode,
};
pub use types::DatabaseRecordId;

View File

@@ -0,0 +1,122 @@
use crate::DataLayerError;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode};
use sqlx::PgPool;
use std::str::FromStr;
use std::time::Duration;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct PostgresPoolConfig {
pub database_url: String,
pub min_connections: u32,
pub max_connections: u32,
pub acquire_timeout_ms: u64,
pub idle_timeout_ms: u64,
pub max_lifetime_ms: u64,
pub statement_cache_capacity: usize,
pub require_ssl: bool,
}
impl Default for PostgresPoolConfig {
fn default() -> Self {
Self {
database_url: String::new(),
min_connections: 1,
max_connections: 20,
acquire_timeout_ms: 5_000,
idle_timeout_ms: 60_000,
max_lifetime_ms: 30 * 60_000,
statement_cache_capacity: 100,
require_ssl: false,
}
}
}
impl PostgresPoolConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.database_url.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"postgres database_url cannot be empty".to_string(),
));
}
if self.min_connections > self.max_connections {
return Err(DataLayerError::InvalidConfiguration(
"postgres min_connections cannot exceed max_connections".to_string(),
));
}
if self.statement_cache_capacity == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_cache_capacity must be positive".to_string(),
));
}
Ok(())
}
pub fn connect_options(&self) -> Result<PgConnectOptions, DataLayerError> {
self.validate()?;
let ssl_mode = if self.require_ssl {
PgSslMode::Require
} else {
PgSslMode::Prefer
};
let options = PgConnectOptions::from_str(self.database_url.trim()).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid postgres database_url: {err}"))
})?;
Ok(options
.ssl_mode(ssl_mode)
.statement_cache_capacity(self.statement_cache_capacity))
}
}
pub type PostgresPool = PgPool;
#[derive(Debug, Clone)]
pub struct PostgresPoolFactory {
config: PostgresPoolConfig,
}
impl PostgresPoolFactory {
pub fn new(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &PostgresPoolConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<PostgresPool, DataLayerError> {
let options = self.config.connect_options()?;
Ok(PgPoolOptions::new()
.min_connections(self.config.min_connections)
.max_connections(self.config.max_connections)
.acquire_timeout(Duration::from_millis(self.config.acquire_timeout_ms))
.idle_timeout(Duration::from_millis(self.config.idle_timeout_ms))
.max_lifetime(Duration::from_millis(self.config.max_lifetime_ms))
.connect_lazy_with(options))
}
}
#[cfg(test)]
mod tests {
use super::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn factory_builds_lazy_pool_from_valid_config() {
let config = PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
};
let factory = PostgresPoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
}
}

View File

@@ -0,0 +1,202 @@
use futures_util::future::BoxFuture;
use sqlx::{Postgres, Transaction};
use crate::postgres::PostgresPool;
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TransactionMode {
ReadOnly,
#[default]
ReadWrite,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresTransactionOptions {
pub mode: TransactionMode,
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresTransactionOptions {
pub fn read_only() -> Self {
Self {
mode: TransactionMode::ReadOnly,
..Self::default()
}
}
pub fn read_write() -> Self {
Self {
mode: TransactionMode::ReadWrite,
..Self::default()
}
}
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
pub type PostgresTransaction = Transaction<'static, Postgres>;
#[derive(Debug, Clone)]
pub struct PostgresTransactionRunner {
pool: PostgresPool,
}
impl PostgresTransactionRunner {
pub fn new(pool: PostgresPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PostgresPool {
&self.pool
}
pub async fn begin(
&self,
options: PostgresTransactionOptions,
) -> Result<PostgresTransaction, DataLayerError> {
options.validate()?;
let mut tx = self.pool.begin().await?;
for statement in build_transaction_setup_statements(options) {
sqlx::query(statement.as_str()).execute(&mut *tx).await?;
}
Ok(tx)
}
pub async fn run<T, F>(
&self,
options: PostgresTransactionOptions,
f: F,
) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
let mut tx = self.begin(options).await?;
match f(&mut tx).await {
Ok(value) => {
tx.commit().await?;
Ok(value)
}
Err(err) => {
let _ = tx.rollback().await;
Err(err)
}
}
}
pub async fn run_read_only<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_only(), f).await
}
pub async fn run_read_write<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_write(), f).await
}
}
pub(crate) fn build_transaction_setup_statements(
options: PostgresTransactionOptions,
) -> Vec<String> {
let mut statements = Vec::new();
if options.mode == TransactionMode::ReadOnly {
statements.push("SET TRANSACTION READ ONLY".to_string());
}
if let Some(statement_timeout_ms) = options.statement_timeout_ms {
statements.push(format!(
"SET LOCAL statement_timeout = {}",
statement_timeout_ms
));
}
if let Some(lock_timeout_ms) = options.lock_timeout_ms {
statements.push(format!("SET LOCAL lock_timeout = {}", lock_timeout_ms));
}
statements
}
#[cfg(test)]
mod tests {
use super::{
build_transaction_setup_statements, PostgresTransactionOptions, PostgresTransactionRunner,
TransactionMode,
};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
fn validates_transaction_options() {
assert!(PostgresTransactionOptions {
statement_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
assert!(PostgresTransactionOptions {
lock_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
}
#[test]
fn builds_expected_setup_statements() {
let statements = build_transaction_setup_statements(PostgresTransactionOptions {
mode: TransactionMode::ReadOnly,
statement_timeout_ms: Some(1_500),
lock_timeout_ms: Some(750),
});
assert_eq!(
statements,
vec![
"SET TRANSACTION READ ONLY".to_string(),
"SET LOCAL statement_timeout = 1500".to_string(),
"SET LOCAL lock_timeout = 750".to_string(),
]
);
}
#[tokio::test]
async fn runner_reuses_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("lazy pool should build");
let runner = PostgresTransactionRunner::new(pool.clone());
let _pool_ref = runner.pool();
}
}

View File

@@ -0,0 +1,2 @@
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DatabaseRecordId(pub String);

View File

@@ -0,0 +1,68 @@
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
pub type RedisClient = redis::Client;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
pub url: String,
pub key_prefix: Option<String>,
}
impl RedisClientConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
let raw = self.url.trim();
if raw.is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"redis url cannot be empty".to_string(),
));
}
url::Url::parse(raw).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}"))
})?;
Ok(())
}
pub fn keyspace(&self) -> RedisKeyspace {
RedisKeyspace::new(self.key_prefix.as_deref())
}
}
#[derive(Debug, Clone)]
pub struct RedisClientFactory {
config: RedisClientConfig,
}
impl RedisClientFactory {
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
Ok(RedisClient::open(self.config.url.clone())?)
}
}
#[cfg(test)]
mod tests {
use super::{RedisClientConfig, RedisClientFactory};
#[test]
fn factory_builds_lazy_client_from_valid_config() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let factory = RedisClientFactory::new(config.clone()).expect("factory should build");
assert_eq!(factory.config(), &config);
let _client = factory
.connect_lazy()
.expect("lazy redis client should build");
}
}

View File

@@ -0,0 +1,307 @@
use std::future::Future;
use std::time::Duration;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisLockKey(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisLockLease {
pub key: RedisLockKey,
pub owner: String,
pub token: String,
pub ttl_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisLockRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub default_ttl_ms: u64,
}
impl Default for RedisLockRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
default_ttl_ms: 15_000,
}
}
}
impl RedisLockRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis lock command_timeout_ms must be positive".to_string(),
));
}
if self.default_ttl_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis lock default_ttl_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisLockRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
}
impl RedisLockRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisLockRunnerConfig {
self.config
}
pub async fn try_acquire(
&self,
key: &RedisLockKey,
owner: &str,
ttl_ms: Option<u64>,
) -> Result<Option<RedisLockLease>, DataLayerError> {
validate_owner(owner)?;
validate_key(key)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
let token = format!("{owner}:{}", Uuid::new_v4());
self.run_with_timeout("redis lock acquire", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let status = redis::cmd("SET")
.arg(&key.0)
.arg(&token)
.arg("NX")
.arg("PX")
.arg(ttl_ms)
.query_async::<Option<String>>(&mut connection)
.await?;
Ok(status.map(|_| RedisLockLease {
key: key.clone(),
owner: owner.to_string(),
token,
ttl_ms,
}))
})
.await
}
pub async fn release(&self, lease: &RedisLockLease) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
self.run_with_timeout("redis lock release", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let deleted = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('del', KEYS[1]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.invoke_async::<i32>(&mut connection)
.await?;
Ok(deleted > 0)
})
.await
}
pub async fn renew(
&self,
lease: &RedisLockLease,
ttl_ms: Option<u64>,
) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
self.run_with_timeout("redis lock renew", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let renewed = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('pexpire', KEYS[1], ARGV[2]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.arg(ttl_ms)
.invoke_async::<i32>(&mut connection)
.await?;
Ok(renewed > 0)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
let ttl_ms = ttl_ms.unwrap_or(self.config.default_ttl_ms);
if ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(ttl_ms)
}
}
fn validate_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock owner cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_key(key: &RedisLockKey) -> Result<(), DataLayerError> {
if key.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock key cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
validate_key(&lease.key)?;
validate_owner(&lease.owner)?;
if lease.token.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock token cannot be empty".to_string(),
));
}
if lease.ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
use crate::redis::{RedisClientConfig, RedisClientFactory};
fn sample_runner() -> RedisLockRunner {
let client = RedisClientFactory::new(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisLockRunner::new(
client,
RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}
.keyspace(),
RedisLockRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_runner_config() {
assert!(RedisLockRunnerConfig {
command_timeout_ms: Some(0),
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisLockRunnerConfig {
default_ttl_ms: 0,
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisLockRunnerConfig::default());
assert_eq!(runner.keyspace().lock_key("poller").0, "aether:lock:poller");
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_owner_or_lease_before_network() {
let runner = sample_runner();
assert!(runner
.try_acquire(&RedisLockKey("aether:lock:poller".to_string()), "", None)
.await
.is_err());
assert!(runner
.release(&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: String::new(),
ttl_ms: 1_000,
})
.await
.is_err());
assert!(runner
.renew(
&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: "token-1".to_string(),
ttl_ms: 1_000,
},
Some(0),
)
.await
.is_err());
}
}

View File

@@ -0,0 +1,12 @@
mod client;
mod lock;
mod namespace;
mod stream;
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
pub use namespace::RedisKeyspace;
pub use stream::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
};

View File

@@ -0,0 +1,43 @@
use aether_cache::CacheKeyNamespace;
use crate::redis::{RedisLockKey, RedisStreamName};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisKeyspace {
namespace: CacheKeyNamespace,
}
impl RedisKeyspace {
pub fn new(prefix: Option<&str>) -> Self {
let normalized = prefix.unwrap_or_default().trim().trim_matches(':');
Self {
namespace: CacheKeyNamespace::new(normalized),
}
}
pub fn key(&self, raw_key: &str) -> String {
self.namespace.key(raw_key)
}
pub fn lock_key(&self, raw_key: &str) -> RedisLockKey {
RedisLockKey(self.namespace.child("lock").key(raw_key))
}
pub fn stream_name(&self, raw_name: &str) -> RedisStreamName {
RedisStreamName(self.namespace.child("stream").key(raw_name))
}
}
#[cfg(test)]
mod tests {
use super::RedisKeyspace;
#[test]
fn composes_prefixed_lock_and_stream_names() {
let keyspace = RedisKeyspace::new(Some("aether"));
assert_eq!(keyspace.key("auth:user"), "aether:auth:user");
assert_eq!(keyspace.lock_key("poller").0, "aether:lock:poller");
assert_eq!(keyspace.stream_name("audit").0, "aether:stream:audit");
}
}

View File

@@ -0,0 +1,658 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::time::Duration;
use redis::from_redis_value;
use redis::streams::StreamReadReply;
use redis::Value as RedisValue;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisStreamName(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerGroup(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerName(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamEntry {
pub id: String,
pub fields: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamReclaimResult {
pub next_start_id: String,
pub entries: Vec<RedisStreamEntry>,
pub deleted_ids: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamReclaimConfig {
pub min_idle_ms: u64,
pub count: usize,
}
impl Default for RedisStreamReclaimConfig {
fn default() -> Self {
Self {
min_idle_ms: 60_000,
count: 32,
}
}
}
impl RedisStreamReclaimConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.min_idle_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim min_idle_ms must be positive".to_string(),
));
}
if self.count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub read_block_ms: Option<u64>,
pub read_count: usize,
}
impl Default for RedisStreamRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
read_block_ms: Some(1_000),
read_count: 32,
}
}
}
impl RedisStreamRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream command_timeout_ms must be positive".to_string(),
));
}
if matches!(self.read_block_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_block_ms must be positive".to_string(),
));
}
if self.read_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisStreamRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
}
impl RedisStreamRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisStreamRunnerConfig {
self.config
}
pub async fn ensure_consumer_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
start_id: &str,
) -> Result<(), DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_stream_position(start_id)?;
self.run_with_timeout("redis stream ensure consumer group", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let result = redis::cmd("XGROUP")
.arg("CREATE")
.arg(&stream.0)
.arg(&group.0)
.arg(start_id)
.arg("MKSTREAM")
.query_async::<String>(&mut connection)
.await;
match result {
Ok(_) => Ok(()),
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
Err(err) => Err(DataLayerError::Redis(err)),
}
})
.await
}
pub async fn append_fields(
&self,
stream: &RedisStreamName,
fields: &BTreeMap<String, String>,
) -> Result<String, DataLayerError> {
validate_stream_name(stream)?;
if fields.is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream fields cannot be empty".to_string(),
));
}
self.run_with_timeout("redis stream append", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XADD");
command.arg(&stream.0).arg("*");
for (key, value) in fields {
command.arg(key).arg(value);
}
Ok(command.query_async::<String>(&mut connection).await?)
})
.await
}
pub async fn append_json(
&self,
stream: &RedisStreamName,
field: &str,
payload: &serde_json::Value,
) -> Result<String, DataLayerError> {
if field.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream json field cannot be empty".to_string(),
));
}
let mut fields = BTreeMap::new();
fields.insert(
field.to_string(),
serde_json::to_string(payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to serialize redis stream payload: {err}"
))
})?,
);
self.append_fields(stream, &fields).await
}
pub async fn read_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
self.run_with_timeout("redis stream read group", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XREADGROUP");
command
.arg("GROUP")
.arg(&group.0)
.arg(&consumer.0)
.arg("COUNT")
.arg(self.config.read_count);
if let Some(block_ms) = self.config.read_block_ms {
command.arg("BLOCK").arg(block_ms);
}
command.arg("STREAMS").arg(&stream.0).arg(">");
let reply = command
.query_async::<StreamReadReply>(&mut connection)
.await?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
.ok()
.map(|text| (field, text))
})
.collect(),
})
.collect())
})
.await
}
pub async fn ack(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<usize, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
if ids.is_empty() {
return Ok(0);
}
self.run_with_timeout("redis stream ack", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XACK");
command.arg(&stream.0).arg(&group.0);
for id in ids {
command.arg(id);
}
Ok(command.query_async::<usize>(&mut connection).await?)
})
.await
}
pub async fn claim_stale(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
start_id: &str,
config: RedisStreamReclaimConfig,
) -> Result<RedisStreamReclaimResult, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
validate_stream_position(start_id)?;
config.validate()?;
self.run_with_timeout("redis stream reclaim", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let reply = redis::cmd("XAUTOCLAIM")
.arg(&stream.0)
.arg(&group.0)
.arg(&consumer.0)
.arg(config.min_idle_ms)
.arg(start_id)
.arg("COUNT")
.arg(config.count)
.query_async::<RedisValue>(&mut connection)
.await?;
parse_reclaim_result(reply)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
}
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
if stream.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_group(group: &RedisConsumerGroup) -> Result<(), DataLayerError> {
if group.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer group cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_consumer(consumer: &RedisConsumerName) -> Result<(), DataLayerError> {
if consumer.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
if position.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream position cannot be empty".to_string(),
));
}
Ok(())
}
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim returned non-array payload".to_string(),
));
};
if parts.len() < 2 || parts.len() > 3 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim returned {} top-level fields, expected 2 or 3",
parts.len()
)));
}
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
let entries = parse_reclaim_entries(&parts[1])?;
let deleted_ids = match parts.get(2) {
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
None => Vec::new(),
};
Ok(RedisStreamReclaimResult {
next_start_id,
entries,
deleted_ids,
})
}
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
match value {
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entries payload was not an array".to_string(),
)),
}
}
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entry was not an array".to_string(),
));
};
if parts.len() != 2 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim entry had {} fields, expected 2",
parts.len()
)));
}
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
Ok(RedisStreamEntry { id, fields })
}
fn parse_string_map(
value: &RedisValue,
context: &str,
) -> Result<BTreeMap<String, String>, DataLayerError> {
match value {
RedisValue::Array(values) => {
if values.len() % 2 != 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"{context} expected an even number of field elements, got {}",
values.len()
)));
}
let mut fields = BTreeMap::new();
for pair in values.chunks(2) {
let key = parse_string_value(&pair[0], context)?;
let value = parse_string_value(&pair[1], context)?;
fields.insert(key, value);
}
Ok(fields)
}
RedisValue::Map(entries) => entries
.iter()
.map(|(key, value)| {
Ok((
parse_string_value(key, context)?,
parse_string_value(value, context)?,
))
})
.collect(),
RedisValue::Nil => Ok(BTreeMap::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array/map payload"
))),
}
}
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
match value {
RedisValue::Array(values) => values
.iter()
.map(|value| parse_string_value(value, context))
.collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array payload"
))),
}
}
fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataLayerError> {
from_redis_value::<String>(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"{context} was not a string-compatible redis value: {err}"
))
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use crate::redis::{RedisClientConfig, RedisClientFactory};
use redis::Value as RedisValue;
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_stream_runner_config() {
assert!(RedisStreamRunnerConfig {
command_timeout_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_block_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_count: 0,
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn validates_reclaim_config() {
assert!(RedisStreamReclaimConfig {
min_idle_ms: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamReclaimConfig {
count: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisStreamRunnerConfig::default());
assert_eq!(
runner.keyspace().stream_name("audit").0,
"aether:stream:audit"
);
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_inputs_before_network() {
let runner = sample_runner();
let stream = RedisStreamName("aether:stream:audit".to_string());
let group = RedisConsumerGroup("audit-workers".to_string());
let consumer = RedisConsumerName("worker-1".to_string());
assert!(runner
.ensure_consumer_group(&stream, &group, "")
.await
.is_err());
assert!(runner
.append_fields(&stream, &BTreeMap::new())
.await
.is_err());
assert!(runner
.append_json(&stream, "", &serde_json::json!({"ok": true}))
.await
.is_err());
assert!(runner
.read_group(&stream, &group, &RedisConsumerName(String::new()))
.await
.is_err());
assert_eq!(
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
0
);
assert!(runner
.claim_stale(
&stream,
&group,
&consumer,
"",
RedisStreamReclaimConfig::default()
)
.await
.is_err());
let _ = consumer;
}
#[test]
fn parses_reclaim_result_with_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::Array(vec![
RedisValue::BulkString(b"payload".to_vec()),
RedisValue::BulkString(br#"{"ok":true}"#.to_vec()),
RedisValue::BulkString(b"kind".to_vec()),
RedisValue::BulkString(b"shadow".to_vec()),
]),
])]),
RedisValue::Array(vec![RedisValue::BulkString(b"1709999999999-0".to_vec())]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: vec![super::RedisStreamEntry {
id: "1710000000000-0".to_string(),
fields: BTreeMap::from([
("kind".to_string(), "shadow".to_string()),
("payload".to_string(), r#"{"ok":true}"#.to_string()),
]),
}],
deleted_ids: vec!["1709999999999-0".to_string()],
}
);
}
#[test]
fn parses_reclaim_result_without_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: Vec::new(),
deleted_ids: Vec::new(),
}
);
}
}

View 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());
}
}

View 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,
};

View 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();
}
}

View 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));
}
}

View 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");
}
}

View 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,
};

View 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();
}
}

View 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)));
}
}

View 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;

View 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
);
}
}

View 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,
};

View 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();
}
}

View 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());
}
}

View 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);
}
}

View 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,
};

View 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"));
}
}

View 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();
}
}

View 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());
}
}

View 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);
}
}

View 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};

View 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();
}
}

View 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());
}
}

View 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());
}
}

View 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,
};

View 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;
}
}

View 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());
}
}