mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
refactor: 拆分 gateway 单体为独立 crate，新增 systemd 部署方案
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、 usage-runtime、video-tasks-core 模块提取为独立 crate；重构 gateway 内部模块结构（state/router/cache/data/query 等）；移除大量遗留模块 文件；新增 systemd 二进制部署骨架及相关文档；更新前端 usage 相关 API 和组件。
This commit is contained in:
fawney19
2360
apps/aether-gateway/src/data/state/auth.rs
Normal file
2360
apps/aether-gateway/src/data/state/auth.rs
Normal file
File diff suppressed because it is too large
Load Diff
473
apps/aether-gateway/src/data/state/catalog.rs
Normal file
473
apps/aether-gateway/src/data/state/catalog.rs
Normal file
@@ -0,0 +1,473 @@
|
||||
use super::{
|
||||
DataLayerError, GatewayDataState, GeminiFileMappingListQuery, GeminiFileMappingStats,
|
||||
ProviderCatalogKeyListQuery, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
|
||||
StoredProviderCatalogProvider, StoredRequestCandidate, UpsertGeminiFileMappingRecord,
|
||||
UpsertRequestCandidateRecord,
|
||||
};
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) async fn list_request_candidates_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => repository.list_by_request_id(request_id).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_request_candidates_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => repository.list_by_provider_id(provider_id, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_request_candidates(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => repository.list_recent(limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_finalized_request_candidates_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_finalized_by_endpoint_ids_since(endpoint_ids, since_unix_secs, limit)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_finalized_request_candidate_statuses_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.count_finalized_statuses_by_endpoint_ids_since(endpoint_ids, since_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn aggregate_finalized_request_candidate_timeline_by_endpoint_ids_since(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
since_unix_secs: u64,
|
||||
until_unix_secs: u64,
|
||||
segments: u32,
|
||||
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.aggregate_finalized_timeline_by_endpoint_ids_since(
|
||||
endpoint_ids,
|
||||
since_unix_secs,
|
||||
until_unix_secs,
|
||||
segments,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_request_candidate(
|
||||
&self,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_writer {
|
||||
Some(repository) => repository.upsert(candidate).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_request_candidates_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
match &self.request_candidate_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.delete_created_before(created_before_unix_secs, limit)
|
||||
.await
|
||||
}
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn touch_auth_api_key_last_used(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.auth_api_key_writer {
|
||||
Some(repository) => repository.touch_last_used_at(api_key_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_gemini_file_mapping(
|
||||
&self,
|
||||
record: UpsertGeminiFileMappingRecord,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
match &self.gemini_file_mapping_writer {
|
||||
Some(repository) => repository.upsert(record).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_gemini_file_mappings(
|
||||
&self,
|
||||
query: &GeminiFileMappingListQuery,
|
||||
) -> Result<StoredGeminiFileMappingListPage, DataLayerError> {
|
||||
match &self.gemini_file_mapping_reader {
|
||||
Some(repository) => repository.list_mappings(query).await,
|
||||
None => Ok(StoredGeminiFileMappingListPage {
|
||||
items: Vec::new(),
|
||||
total: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_gemini_file_mappings(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<GeminiFileMappingStats, DataLayerError> {
|
||||
match &self.gemini_file_mapping_reader {
|
||||
Some(repository) => repository.summarize_mappings(now_unix_secs).await,
|
||||
None => Ok(GeminiFileMappingStats {
|
||||
total_mappings: 0,
|
||||
active_mappings: 0,
|
||||
expired_mappings: 0,
|
||||
by_mime_type: Vec::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_gemini_file_mapping_by_file_name(
|
||||
&self,
|
||||
file_name: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.gemini_file_mapping_writer {
|
||||
Some(repository) => repository.delete_by_file_name(file_name).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_gemini_file_mapping_by_id(
|
||||
&self,
|
||||
mapping_id: &str,
|
||||
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
|
||||
match &self.gemini_file_mapping_writer {
|
||||
Some(repository) => repository.delete_by_id(mapping_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_expired_gemini_file_mappings(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
match &self.gemini_file_mapping_writer {
|
||||
Some(repository) => repository.delete_expired_before(now_unix_secs).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn cache_set_string_with_ttl(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &str,
|
||||
ttl_seconds: u64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let Some(runner) = self.kv_runner() else {
|
||||
return Ok(());
|
||||
};
|
||||
runner.setex(key, value, Some(ttl_seconds)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn cache_delete_key(&self, key: &str) -> Result<(), DataLayerError> {
|
||||
let Some(runner) = self.kv_runner() else {
|
||||
return Ok(());
|
||||
};
|
||||
let _deleted = runner.del(key).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_providers_by_ids(provider_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_providers(
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_providers(active_only).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_endpoints_by_ids(endpoint_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_endpoints_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_keys_by_ids(key_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_keys_by_provider_ids(provider_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_keys_page(query).await,
|
||||
None => Ok(StoredProviderCatalogKeyPage {
|
||||
items: Vec::new(),
|
||||
total: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_key_stats_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_oauth_credentials(
|
||||
&self,
|
||||
key_id: &str,
|
||||
encrypted_api_key: &str,
|
||||
encrypted_auth_config: Option<&str>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.update_key_oauth_credentials(
|
||||
key_id,
|
||||
encrypted_api_key,
|
||||
encrypted_auth_config,
|
||||
expires_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.create_key(key).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
shift_existing_priorities_from: Option<i32>,
|
||||
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository
|
||||
.create_provider(provider, shift_existing_priorities_from)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.update_provider(provider).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_provider_catalog_provider(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.delete_provider(provider_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
endpoint_ids: &[String],
|
||||
key_ids: &[String],
|
||||
) -> Result<(), DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.create_endpoint(endpoint).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_endpoint(
|
||||
&self,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.update_endpoint(endpoint).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_provider_catalog_endpoint(
|
||||
&self,
|
||||
endpoint_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.delete_endpoint(endpoint_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.update_key(key).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_provider_catalog_key(
|
||||
&self,
|
||||
key_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.delete_key(key_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
|
||||
&self,
|
||||
key_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => repository.clear_key_oauth_invalid_marker(key_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
is_active: bool,
|
||||
health_by_format: Option<&serde_json::Value>,
|
||||
circuit_breaker_by_format: Option<&serde_json::Value>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.provider_catalog_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.update_key_health_state(
|
||||
key_id,
|
||||
is_active,
|
||||
health_by_format,
|
||||
circuit_breaker_by_format,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
430
apps/aether-gateway/src/data/state/core.rs
Normal file
430
apps/aether-gateway/src/data/state/core.rs
Normal file
@@ -0,0 +1,430 @@
|
||||
use aether_data::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner};
|
||||
use aether_data::{DataBackends, DataLayerError};
|
||||
|
||||
use super::{GatewayDataConfig, GatewayDataState, StoredSystemConfigEntry};
|
||||
|
||||
fn current_system_config_updated_at_unix_secs() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) fn disabled() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub(crate) fn from_config(config: GatewayDataConfig) -> Result<Self, DataLayerError> {
|
||||
if !config.is_enabled() {
|
||||
return Ok(Self {
|
||||
config,
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
});
|
||||
}
|
||||
|
||||
let backends = DataBackends::from_config(config.to_data_layer_config())?;
|
||||
let auth_api_key_reader = backends.read().auth_api_keys();
|
||||
let auth_api_key_writer = backends.write().auth_api_keys();
|
||||
let auth_module_reader = backends.read().auth_modules();
|
||||
let auth_module_writer = backends.write().auth_modules();
|
||||
let announcement_reader = backends.read().announcements();
|
||||
let announcement_writer = backends.write().announcements();
|
||||
let management_token_reader = backends.read().management_tokens();
|
||||
let management_token_writer = backends.write().management_tokens();
|
||||
let oauth_provider_reader = backends.read().oauth_providers();
|
||||
let oauth_provider_writer = backends.write().oauth_providers();
|
||||
let proxy_node_reader = backends.read().proxy_nodes();
|
||||
let proxy_node_writer = backends.write().proxy_nodes();
|
||||
let billing_reader = backends.read().billing();
|
||||
let gemini_file_mapping_reader = backends.read().gemini_file_mappings();
|
||||
let global_model_reader = backends.read().global_models();
|
||||
let global_model_writer = backends.write().global_models();
|
||||
let minimal_candidate_selection_reader = backends.read().minimal_candidate_selection();
|
||||
let request_candidate_reader = backends.read().request_candidates();
|
||||
let request_candidate_writer = backends.write().request_candidates();
|
||||
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
|
||||
let provider_catalog_reader = backends.read().provider_catalog();
|
||||
let provider_catalog_writer = backends.write().provider_catalog();
|
||||
let provider_quota_reader = backends.read().provider_quotas();
|
||||
let provider_quota_writer = backends.write().provider_quotas();
|
||||
let usage_reader = backends.read().usage();
|
||||
let usage_writer = backends.write().usage();
|
||||
let user_reader = backends.read().users();
|
||||
let usage_worker_runner = backends.workers().redis();
|
||||
let video_task_reader = backends.read().video_tasks();
|
||||
let video_task_writer = backends.write().video_tasks();
|
||||
let wallet_reader = backends.read().wallets();
|
||||
let wallet_writer = backends.write().wallets();
|
||||
let settlement_writer = backends.write().settlement();
|
||||
let shadow_result_reader = backends.read().shadow_results();
|
||||
let shadow_result_writer = backends.write().shadow_results();
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
backends: Some(backends),
|
||||
auth_api_key_reader,
|
||||
auth_api_key_writer,
|
||||
auth_module_reader,
|
||||
auth_module_writer,
|
||||
announcement_reader,
|
||||
announcement_writer,
|
||||
management_token_reader,
|
||||
management_token_writer,
|
||||
oauth_provider_reader,
|
||||
oauth_provider_writer,
|
||||
proxy_node_reader,
|
||||
proxy_node_writer,
|
||||
billing_reader,
|
||||
gemini_file_mapping_reader,
|
||||
gemini_file_mapping_writer,
|
||||
global_model_reader,
|
||||
global_model_writer,
|
||||
minimal_candidate_selection_reader,
|
||||
request_candidate_reader,
|
||||
request_candidate_writer,
|
||||
provider_catalog_reader,
|
||||
provider_catalog_writer,
|
||||
provider_quota_reader,
|
||||
provider_quota_writer,
|
||||
usage_reader,
|
||||
usage_writer,
|
||||
user_reader,
|
||||
user_preferences: None,
|
||||
usage_worker_runner,
|
||||
video_task_reader,
|
||||
video_task_writer,
|
||||
wallet_reader,
|
||||
wallet_writer,
|
||||
settlement_writer,
|
||||
shadow_result_reader,
|
||||
shadow_result_writer,
|
||||
system_config_values: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn has_backends(&self) -> bool {
|
||||
self.backends.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_api_key_reader(&self) -> bool {
|
||||
self.auth_api_key_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_api_key_writer(&self) -> bool {
|
||||
self.auth_api_key_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_module_writer(&self) -> bool {
|
||||
self.auth_module_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_announcement_reader(&self) -> bool {
|
||||
self.announcement_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_announcement_writer(&self) -> bool {
|
||||
self.announcement_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_management_token_reader(&self) -> bool {
|
||||
self.management_token_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_management_token_writer(&self) -> bool {
|
||||
self.management_token_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_gemini_file_mapping_reader(&self) -> bool {
|
||||
self.gemini_file_mapping_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_gemini_file_mapping_writer(&self) -> bool {
|
||||
self.gemini_file_mapping_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_global_model_reader(&self) -> bool {
|
||||
self.global_model_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_global_model_writer(&self) -> bool {
|
||||
self.global_model_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_redis_backend(&self) -> bool {
|
||||
self.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.redis())
|
||||
.is_some()
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn has_minimal_candidate_selection_reader(&self) -> bool {
|
||||
self.minimal_candidate_selection_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_request_candidate_reader(&self) -> bool {
|
||||
self.request_candidate_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_request_candidate_writer(&self) -> bool {
|
||||
self.request_candidate_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
|
||||
self.provider_catalog_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_provider_catalog_writer(&self) -> bool {
|
||||
self.provider_catalog_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_proxy_node_reader(&self) -> bool {
|
||||
self.proxy_node_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_proxy_node_writer(&self) -> bool {
|
||||
self.proxy_node_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn oauth_refresh_lock_runner(&self) -> Option<RedisLockRunner> {
|
||||
self.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.locks().redis())
|
||||
}
|
||||
|
||||
pub(crate) fn kv_runner(&self) -> Option<RedisKvRunner> {
|
||||
self.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.redis())
|
||||
.and_then(|backend| backend.kv_runner(RedisKvRunnerConfig::default()).ok())
|
||||
}
|
||||
|
||||
pub(crate) fn postgres_pool(&self) -> Option<aether_data::postgres::PostgresPool> {
|
||||
self.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
.map(|backend| backend.pool_clone())
|
||||
}
|
||||
|
||||
pub(crate) fn postgres_max_connections(&self) -> Option<u32> {
|
||||
self.config.postgres().map(|config| config.max_connections)
|
||||
}
|
||||
|
||||
pub(crate) fn has_provider_quota_writer(&self) -> bool {
|
||||
self.provider_quota_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_usage_reader(&self) -> bool {
|
||||
self.usage_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_user_reader(&self) -> bool {
|
||||
self.user_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_usage_writer(&self) -> bool {
|
||||
self.usage_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_usage_worker_runner(&self) -> bool {
|
||||
self.usage_worker_runner.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_video_task_reader(&self) -> bool {
|
||||
self.video_task_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_video_task_writer(&self) -> bool {
|
||||
self.video_task_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_wallet_reader(&self) -> bool {
|
||||
self.wallet_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_wallet_writer(&self) -> bool {
|
||||
self.wallet_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_settlement_writer(&self) -> bool {
|
||||
self.settlement_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_shadow_result_writer(&self) -> bool {
|
||||
self.shadow_result_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_shadow_result_reader(&self) -> bool {
|
||||
self.shadow_result_reader.is_some()
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn encryption_key(&self) -> Option<&str> {
|
||||
self.config.encryption_key()
|
||||
}
|
||||
|
||||
pub(crate) async fn find_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
if let Some(values) = &self.system_config_values {
|
||||
return Ok(values
|
||||
.read()
|
||||
.expect("system config values lock")
|
||||
.get(key)
|
||||
.map(|entry| entry.value.clone()));
|
||||
}
|
||||
match self
|
||||
.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
{
|
||||
Some(backend) => backend.find_system_config_value(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, DataLayerError> {
|
||||
Ok(self
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await?
|
||||
.value)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_system_config_entries(
|
||||
&self,
|
||||
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
|
||||
if let Some(values) = &self.system_config_values {
|
||||
return Ok(values
|
||||
.read()
|
||||
.expect("system config values lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect());
|
||||
}
|
||||
match self
|
||||
.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
{
|
||||
Some(backend) => backend.list_system_config_entries().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_system_config_entry(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<StoredSystemConfigEntry, DataLayerError> {
|
||||
if let Some(values) = &self.system_config_values {
|
||||
let mut values = values.write().expect("system config values lock");
|
||||
let description = description
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| values.get(key).and_then(|entry| entry.description.clone()));
|
||||
let entry = StoredSystemConfigEntry {
|
||||
key: key.to_string(),
|
||||
value: value.clone(),
|
||||
description,
|
||||
updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
|
||||
};
|
||||
values.insert(key.to_string(), entry.clone());
|
||||
return Ok(entry);
|
||||
}
|
||||
match self
|
||||
.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
{
|
||||
Some(backend) => {
|
||||
backend
|
||||
.upsert_system_config_entry(key, value, description)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredSystemConfigEntry {
|
||||
key: key.to_string(),
|
||||
value: value.clone(),
|
||||
description: description.map(ToOwned::to_owned),
|
||||
updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_system_config_value(
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
if let Some(values) = &self.system_config_values {
|
||||
return Ok(values
|
||||
.write()
|
||||
.expect("system config values lock")
|
||||
.remove(key)
|
||||
.is_some());
|
||||
}
|
||||
match self
|
||||
.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
{
|
||||
Some(backend) => backend.delete_system_config_value(key).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_admin_system_stats(
|
||||
&self,
|
||||
) -> Result<super::AdminSystemStats, DataLayerError> {
|
||||
match self
|
||||
.backends
|
||||
.as_ref()
|
||||
.and_then(|backends| backends.postgres())
|
||||
{
|
||||
Some(backend) => backend.read_admin_system_stats().await,
|
||||
None => Ok(super::AdminSystemStats::default()),
|
||||
}
|
||||
}
|
||||
}
|
||||
257
apps/aether-gateway/src/data/state/integrations.rs
Normal file
257
apps/aether-gateway/src/data/state/integrations.rs
Normal file
@@ -0,0 +1,257 @@
|
||||
use aether_billing::enrich_usage_event_with_billing;
|
||||
use aether_billing::BillingModelContextLookup;
|
||||
use aether_data::redis::RedisStreamRunner;
|
||||
use aether_data::repository::audit::RequestAuditReader;
|
||||
use aether_data::repository::auth::{
|
||||
AuthApiKeyLookupKey, ResolvedAuthApiKeySnapshotReader, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::billing::StoredBillingModelContext;
|
||||
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
|
||||
use aether_data::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
|
||||
use aether_data::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
|
||||
use aether_data::DataLayerError;
|
||||
use aether_usage_runtime::{
|
||||
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRuntimeAccess,
|
||||
UsageSettlementWriter,
|
||||
};
|
||||
use aether_video_tasks_core::StoredVideoTaskReadSide;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::GatewayDataState;
|
||||
use crate::provider_transport::ProviderTransportSnapshotSource;
|
||||
use crate::scheduler::SchedulerCandidateSelectionRowSource;
|
||||
|
||||
#[async_trait]
|
||||
impl RequestAuditReader for GatewayDataState {
|
||||
async fn find_request_usage_audit_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
GatewayDataState::find_request_usage_by_request_id(self, request_id).await
|
||||
}
|
||||
|
||||
async fn read_request_decision_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<aether_data::repository::candidates::DecisionTrace>, DataLayerError> {
|
||||
GatewayDataState::read_decision_trace(self, request_id, attempted_only).await
|
||||
}
|
||||
|
||||
async fn read_resolved_auth_api_key_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<aether_data::repository::auth::ResolvedAuthApiKeySnapshot>, DataLayerError>
|
||||
{
|
||||
GatewayDataState::read_auth_api_key_snapshot(self, user_id, api_key_id, now_unix_secs).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ResolvedAuthApiKeySnapshotReader for GatewayDataState {
|
||||
async fn find_stored_auth_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
GatewayDataState::find_auth_api_key_snapshot(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl StoredVideoTaskReadSide for GatewayDataState {
|
||||
async fn find_stored_video_task(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
GatewayDataState::find_video_task(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderTransportSnapshotSource for GatewayDataState {
|
||||
fn encryption_key(&self) -> Option<&str> {
|
||||
GatewayDataState::encryption_key(self)
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_providers_by_ids(
|
||||
&self,
|
||||
ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
GatewayDataState::list_provider_catalog_providers_by_ids(self, ids).await
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_endpoints_by_ids(
|
||||
&self,
|
||||
ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
GatewayDataState::list_provider_catalog_endpoints_by_ids(self, ids).await
|
||||
}
|
||||
|
||||
async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
GatewayDataState::list_provider_catalog_keys_by_ids(self, ids).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SchedulerCandidateSelectionRowSource for GatewayDataState {
|
||||
async fn read_minimal_candidate_selection_rows_for_api_format_and_global_model(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
self.list_minimal_candidate_selection_rows(api_format, global_model_name)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn read_minimal_candidate_selection_rows_for_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
self.list_minimal_candidate_selection_rows_for_api_format(api_format)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BillingModelContextLookup for GatewayDataState {
|
||||
async fn find_billing_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
GatewayDataState::find_billing_model_context(
|
||||
self,
|
||||
provider_id,
|
||||
provider_api_key_id,
|
||||
global_model_name,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for GatewayDataState {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
GatewayDataState::has_settlement_writer(self)
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
GatewayDataState::settle_usage(self, input).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageBillingEventEnricher for GatewayDataState {
|
||||
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError> {
|
||||
enrich_usage_event_with_billing(self, event).await
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageRuntimeAccess for GatewayDataState {
|
||||
fn has_usage_writer(&self) -> bool {
|
||||
GatewayDataState::has_usage_writer(self)
|
||||
}
|
||||
|
||||
fn has_usage_worker_runner(&self) -> bool {
|
||||
GatewayDataState::has_usage_worker_runner(self)
|
||||
}
|
||||
|
||||
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
|
||||
GatewayDataState::usage_worker_runner(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageRecordWriter for GatewayDataState {
|
||||
async fn upsert_usage_record(
|
||||
&self,
|
||||
record: UpsertUsageRecord,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
GatewayDataState::upsert_usage(self, record).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_billing::enrich_usage_event_with_billing;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::GatewayDataState;
|
||||
use crate::usage::event::{UsageEvent, UsageEventData, UsageEventType};
|
||||
|
||||
#[tokio::test]
|
||||
async fn enriches_completed_usage_event_with_billing_snapshot() {
|
||||
let state = GatewayDataState::with_billing_reader_for_tests(
|
||||
std::sync::Arc::new(
|
||||
aether_data::repository::billing::InMemoryBillingReadRepository::seed(vec![
|
||||
aether_data::repository::billing::StoredBillingModelContext::new(
|
||||
"provider-1".to_string(),
|
||||
Some("pay_as_you_go".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
Some(serde_json::json!({"openai:chat": 0.5})),
|
||||
Some(60),
|
||||
"global-model-1".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
None,
|
||||
Some(0.02),
|
||||
Some(serde_json::json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0,"cache_creation_price_per_1m":3.75,"cache_read_price_per_1m":0.30}]})),
|
||||
Some("model-1".to_string()),
|
||||
Some("gpt-5-upstream".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("billing context should build"),
|
||||
]),
|
||||
),
|
||||
);
|
||||
let mut event = UsageEvent::new(
|
||||
UsageEventType::Completed,
|
||||
"req-billing-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
input_tokens: Some(1_000),
|
||||
output_tokens: Some(500),
|
||||
cache_read_input_tokens: Some(100),
|
||||
status_code: Some(200),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
enrich_usage_event_with_billing(&state, &mut event)
|
||||
.await
|
||||
.expect("billing should succeed");
|
||||
|
||||
assert!(event.data.total_cost_usd.unwrap_or_default() > 0.0);
|
||||
assert!(event.data.actual_total_cost_usd.unwrap_or_default() > 0.0);
|
||||
assert_eq!(
|
||||
event
|
||||
.data
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("billing_snapshot"))
|
||||
.and_then(|value| value.get("status"))
|
||||
.and_then(Value::as_str),
|
||||
Some("complete")
|
||||
);
|
||||
}
|
||||
}
|
||||
272
apps/aether-gateway/src/data/state/mod.rs
Normal file
272
apps/aether-gateway/src/data/state/mod.rs
Normal file
@@ -0,0 +1,272 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use super::auth::GatewayAuthApiKeySnapshot;
|
||||
use super::candidates::{read_request_candidate_trace, RequestCandidateTrace};
|
||||
use super::config::GatewayDataConfig;
|
||||
use super::decision_trace::{read_decision_trace, DecisionTrace};
|
||||
use crate::provider_transport::{
|
||||
read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use crate::video_tasks::LocalVideoTaskReadResponse;
|
||||
use aether_data::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner, RedisStreamRunner};
|
||||
use aether_data::repository::announcements::{
|
||||
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
|
||||
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
|
||||
};
|
||||
use aether_data::repository::audit::RequestAuditBundle;
|
||||
use aether_data::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
|
||||
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::auth_modules::{
|
||||
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
|
||||
StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use aether_data::repository::billing::{BillingReadRepository, StoredBillingModelContext};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use aether_data::repository::gemini_file_mappings::{
|
||||
GeminiFileMappingListQuery, GeminiFileMappingReadRepository, GeminiFileMappingStats,
|
||||
GeminiFileMappingWriteRepository, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
|
||||
UpsertGeminiFileMappingRecord,
|
||||
};
|
||||
use aether_data::repository::global_models::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
||||
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
|
||||
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
|
||||
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
|
||||
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
use aether_data::repository::management_tokens::{
|
||||
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
|
||||
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
|
||||
StoredManagementTokenListPage, StoredManagementTokenWithUser, UpdateManagementTokenRecord,
|
||||
};
|
||||
use aether_data::repository::oauth_providers::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig,
|
||||
UpsertOAuthProviderConfigRecord,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
|
||||
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use aether_data::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use aether_data::repository::settlement::{
|
||||
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
|
||||
};
|
||||
use aether_data::repository::shadow_results::{
|
||||
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
|
||||
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
|
||||
};
|
||||
pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
|
||||
use aether_data::repository::usage::{
|
||||
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
use aether_data::repository::users::{
|
||||
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserReadRepository,
|
||||
};
|
||||
pub(crate) use aether_data::repository::users::{
|
||||
StoredUserPreferenceRecord, StoredUserSessionRecord,
|
||||
};
|
||||
use aether_data::repository::video_tasks::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
|
||||
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatusCount, VideoTaskWriteRepository,
|
||||
};
|
||||
use aether_data::repository::wallet::{
|
||||
AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminWalletLedgerQuery,
|
||||
AdminWalletListQuery, AdminWalletRefundRequestListQuery, CompleteAdminWalletRefundInput,
|
||||
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
|
||||
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
|
||||
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, FailAdminWalletRefundInput,
|
||||
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
|
||||
StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder,
|
||||
StoredAdminPaymentOrderPage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage,
|
||||
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
|
||||
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
|
||||
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
|
||||
WalletReadRepository, WalletWriteRepository,
|
||||
};
|
||||
use aether_data::{DataBackends, DataLayerError};
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct GatewayDataState {
|
||||
config: GatewayDataConfig,
|
||||
backends: Option<DataBackends>,
|
||||
auth_api_key_reader: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
auth_api_key_writer: Option<Arc<dyn AuthApiKeyWriteRepository>>,
|
||||
auth_module_reader: Option<Arc<dyn AuthModuleReadRepository>>,
|
||||
auth_module_writer: Option<Arc<dyn AuthModuleWriteRepository>>,
|
||||
announcement_reader: Option<Arc<dyn AnnouncementReadRepository>>,
|
||||
announcement_writer: Option<Arc<dyn AnnouncementWriteRepository>>,
|
||||
management_token_reader: Option<Arc<dyn ManagementTokenReadRepository>>,
|
||||
management_token_writer: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
||||
oauth_provider_reader: Option<Arc<dyn OAuthProviderReadRepository>>,
|
||||
oauth_provider_writer: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
||||
proxy_node_reader: Option<Arc<dyn ProxyNodeReadRepository>>,
|
||||
proxy_node_writer: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||
billing_reader: Option<Arc<dyn BillingReadRepository>>,
|
||||
gemini_file_mapping_reader: Option<Arc<dyn GeminiFileMappingReadRepository>>,
|
||||
gemini_file_mapping_writer: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
|
||||
global_model_reader: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||
global_model_writer: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||
minimal_candidate_selection_reader: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
||||
request_candidate_reader: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
||||
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||
usage_reader: Option<Arc<dyn UsageReadRepository>>,
|
||||
usage_writer: Option<Arc<dyn UsageWriteRepository>>,
|
||||
user_reader: Option<Arc<dyn UserReadRepository>>,
|
||||
user_preferences: Option<Arc<RwLock<BTreeMap<String, StoredUserPreferenceRecord>>>>,
|
||||
usage_worker_runner: Option<RedisStreamRunner>,
|
||||
video_task_reader: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
video_task_writer: Option<Arc<dyn VideoTaskWriteRepository>>,
|
||||
wallet_reader: Option<Arc<dyn WalletReadRepository>>,
|
||||
wallet_writer: Option<Arc<dyn WalletWriteRepository>>,
|
||||
settlement_writer: Option<Arc<dyn SettlementWriteRepository>>,
|
||||
shadow_result_reader: Option<Arc<dyn ShadowResultReadRepository>>,
|
||||
shadow_result_writer: Option<Arc<dyn ShadowResultWriteRepository>>,
|
||||
system_config_values: Option<Arc<RwLock<BTreeMap<String, StoredSystemConfigEntry>>>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for GatewayDataState {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("GatewayDataState")
|
||||
.field("config", &self.config)
|
||||
.field("has_backends", &self.backends.is_some())
|
||||
.field(
|
||||
"has_auth_api_key_reader",
|
||||
&self.auth_api_key_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_auth_api_key_writer",
|
||||
&self.auth_api_key_writer.is_some(),
|
||||
)
|
||||
.field("has_auth_module_reader", &self.auth_module_reader.is_some())
|
||||
.field("has_auth_module_writer", &self.auth_module_writer.is_some())
|
||||
.field(
|
||||
"has_announcement_reader",
|
||||
&self.announcement_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_announcement_writer",
|
||||
&self.announcement_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_management_token_reader",
|
||||
&self.management_token_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_management_token_writer",
|
||||
&self.management_token_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_oauth_provider_reader",
|
||||
&self.oauth_provider_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_oauth_provider_writer",
|
||||
&self.oauth_provider_writer.is_some(),
|
||||
)
|
||||
.field("has_proxy_node_reader", &self.proxy_node_reader.is_some())
|
||||
.field("has_proxy_node_writer", &self.proxy_node_writer.is_some())
|
||||
.field("has_billing_reader", &self.billing_reader.is_some())
|
||||
.field(
|
||||
"has_gemini_file_mapping_reader",
|
||||
&self.gemini_file_mapping_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_gemini_file_mapping_writer",
|
||||
&self.gemini_file_mapping_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_global_model_reader",
|
||||
&self.global_model_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_global_model_writer",
|
||||
&self.global_model_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_minimal_candidate_selection_reader",
|
||||
&self.minimal_candidate_selection_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_request_candidate_reader",
|
||||
&self.request_candidate_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_request_candidate_writer",
|
||||
&self.request_candidate_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_provider_catalog_reader",
|
||||
&self.provider_catalog_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_provider_catalog_writer",
|
||||
&self.provider_catalog_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_provider_quota_reader",
|
||||
&self.provider_quota_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_provider_quota_writer",
|
||||
&self.provider_quota_writer.is_some(),
|
||||
)
|
||||
.field("has_usage_reader", &self.usage_reader.is_some())
|
||||
.field("has_usage_writer", &self.usage_writer.is_some())
|
||||
.field("has_user_preferences", &self.user_preferences.is_some())
|
||||
.field(
|
||||
"has_usage_worker_runner",
|
||||
&self.usage_worker_runner.is_some(),
|
||||
)
|
||||
.field("has_video_task_reader", &self.video_task_reader.is_some())
|
||||
.field("has_video_task_writer", &self.video_task_writer.is_some())
|
||||
.field("has_wallet_reader", &self.wallet_reader.is_some())
|
||||
.field("has_wallet_writer", &self.wallet_writer.is_some())
|
||||
.field("has_settlement_writer", &self.settlement_writer.is_some())
|
||||
.field(
|
||||
"has_shadow_result_reader",
|
||||
&self.shadow_result_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_shadow_result_writer",
|
||||
&self.shadow_result_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_system_config_values",
|
||||
&self.system_config_values.is_some(),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
mod auth;
|
||||
mod catalog;
|
||||
mod core;
|
||||
mod integrations;
|
||||
mod models;
|
||||
mod runtime;
|
||||
#[cfg(test)]
|
||||
mod testing;
|
||||
257
apps/aether-gateway/src/data/state/models.rs
Normal file
257
apps/aether-gateway/src/data/state/models.rs
Normal file
@@ -0,0 +1,257 @@
|
||||
use super::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
||||
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
|
||||
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
|
||||
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredProviderActiveGlobalModel,
|
||||
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||
};
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) async fn list_minimal_candidate_selection_rows(
|
||||
&self,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
match &self.minimal_candidate_selection_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_for_exact_api_format_and_global_model(api_format, global_model_name)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format(
|
||||
&self,
|
||||
api_format: &str,
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
match &self.minimal_candidate_selection_reader {
|
||||
Some(repository) => repository.list_for_exact_api_format(api_format).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_public_global_models(
|
||||
&self,
|
||||
query: &PublicGlobalModelQuery,
|
||||
) -> Result<StoredPublicGlobalModelPage, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.list_public_models(query).await,
|
||||
None => Ok(StoredPublicGlobalModelPage {
|
||||
items: Vec::new(),
|
||||
total: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_public_global_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredPublicGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.get_public_model_by_name(model_name).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelListQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.list_public_catalog_models(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn search_public_catalog_models(
|
||||
&self,
|
||||
query: &PublicCatalogModelSearchQuery,
|
||||
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.search_public_catalog_models(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_provider_models(
|
||||
&self,
|
||||
query: &AdminProviderModelListQuery,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.list_admin_provider_models(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_global_models(
|
||||
&self,
|
||||
query: &AdminGlobalModelListQuery,
|
||||
) -> Result<StoredAdminGlobalModelPage, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.list_admin_global_models(query).await,
|
||||
None => Ok(StoredAdminGlobalModelPage {
|
||||
items: Vec::new(),
|
||||
total: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.get_admin_provider_model(provider_id, model_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_provider_available_source_models(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_admin_provider_available_source_models(provider_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_admin_global_model_by_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.get_admin_global_model_by_id(global_model_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_admin_global_model_by_name(
|
||||
&self,
|
||||
model_name: &str,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.get_admin_global_model_by_name(model_name).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_provider_models_by_global_model_id(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_admin_provider_models_by_global_model_id(global_model_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => repository.create_admin_provider_model(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_admin_provider_model(
|
||||
&self,
|
||||
record: &UpsertAdminProviderModelRecord,
|
||||
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => repository.update_admin_provider_model(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_admin_provider_model(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.delete_admin_provider_model(provider_id, model_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_admin_global_model(
|
||||
&self,
|
||||
record: &CreateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => repository.create_admin_global_model(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_admin_global_model(
|
||||
&self,
|
||||
record: &UpdateAdminGlobalModelRecord,
|
||||
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => repository.update_admin_global_model(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_admin_global_model(
|
||||
&self,
|
||||
global_model_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.global_model_writer {
|
||||
Some(repository) => repository.delete_admin_global_model(global_model_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_model_stats(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderModelStats>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => repository.list_provider_model_stats(provider_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_active_global_model_ids_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderActiveGlobalModel>, DataLayerError> {
|
||||
match &self.global_model_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_active_global_model_ids_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
958
apps/aether-gateway/src/data/state/runtime.rs
Normal file
958
apps/aether-gateway/src/data/state/runtime.rs
Normal file
@@ -0,0 +1,958 @@
|
||||
use super::{
|
||||
merge_shadow_result_sample, read_decision_trace, read_provider_transport_snapshot,
|
||||
read_request_candidate_trace, AdjustWalletBalanceInput, AdminPaymentOrderListQuery,
|
||||
AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletRefundRequestListQuery,
|
||||
AnnouncementListQuery, CompleteAdminWalletRefundInput, CreateAnnouncementRecord,
|
||||
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
|
||||
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
|
||||
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError, DecisionTrace,
|
||||
FailAdminWalletRefundInput, GatewayDataState, GatewayProviderTransportSnapshot,
|
||||
LocalVideoTaskReadResponse, ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput,
|
||||
ProcessPaymentCallbackOutcome, RecordShadowResultSample, RedisStreamRunner, RequestAuditBundle,
|
||||
RequestCandidateTrace, ShadowResultLookupKey, StoredAdminPaymentCallbackPage,
|
||||
StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminWalletLedgerPage,
|
||||
StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
|
||||
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
|
||||
StoredAdminWalletTransactionPage, StoredAnnouncement, StoredAnnouncementPage,
|
||||
StoredBillingModelContext, StoredProviderQuotaSnapshot, StoredProviderUsageSummary,
|
||||
StoredRequestUsageAudit, StoredShadowResult, StoredUsageSettlement, StoredUserAuthRecord,
|
||||
StoredUserExportRow, StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger,
|
||||
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord,
|
||||
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, VideoTaskLookupKey,
|
||||
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount, WalletLookupKey,
|
||||
WalletMutationOutcome,
|
||||
};
|
||||
use aether_video_tasks_core::read_data_backed_video_task_response;
|
||||
|
||||
fn is_missing_shadow_results_relation_error(error: &DataLayerError) -> bool {
|
||||
error
|
||||
.to_string()
|
||||
.contains("relation \"gateway_shadow_results\" does not exist")
|
||||
}
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) async fn list_announcements(
|
||||
&self,
|
||||
query: &AnnouncementListQuery,
|
||||
) -> Result<StoredAnnouncementPage, DataLayerError> {
|
||||
match &self.announcement_reader {
|
||||
Some(repository) => repository.list_announcements(query).await,
|
||||
None => Ok(StoredAnnouncementPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_announcement_by_id(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
match &self.announcement_reader {
|
||||
Some(repository) => repository.find_by_id(announcement_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_unread_active_announcements(
|
||||
&self,
|
||||
user_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match &self.announcement_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.count_unread_active_announcements(user_id, now_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_announcement(
|
||||
&self,
|
||||
record: CreateAnnouncementRecord,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
match &self.announcement_writer {
|
||||
Some(repository) => repository.create_announcement(record).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_announcement(
|
||||
&self,
|
||||
record: UpdateAnnouncementRecord,
|
||||
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
|
||||
match &self.announcement_writer {
|
||||
Some(repository) => repository.update_announcement(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_announcement(
|
||||
&self,
|
||||
announcement_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.announcement_writer {
|
||||
Some(repository) => repository.delete_announcement(announcement_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_announcement_as_read(
|
||||
&self,
|
||||
user_id: &str,
|
||||
announcement_id: &str,
|
||||
read_at_unix_secs: u64,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.announcement_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.mark_announcement_as_read(user_id, announcement_id, read_at_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_video_task(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.find(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_video_task_page(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.list_page(filter, offset, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_video_tasks(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.count(filter).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_video_tasks_by_status(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<Vec<VideoTaskStatusCount>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.count_by_status(filter).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_distinct_video_task_users(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.count_distinct_users(filter).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn top_video_task_models(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
limit: usize,
|
||||
) -> Result<Vec<VideoTaskModelCount>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.top_models(filter, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_video_tasks_created_since(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
created_since_unix_secs: u64,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.count_created_since(filter, created_since_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_video_task(
|
||||
&self,
|
||||
task: UpsertVideoTask,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_writer {
|
||||
Some(repository) => repository.upsert(task).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_active_video_task(
|
||||
&self,
|
||||
task: UpsertVideoTask,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_writer {
|
||||
Some(repository) => repository.update_if_active(task).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn claim_due_video_tasks(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
claim_until_unix_secs: u64,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.claim_due(now_unix_secs, claim_until_unix_secs, limit)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_wallet(
|
||||
&self,
|
||||
key: WalletLookupKey<'_>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.find(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_wallets_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_wallets_by_api_key_ids(api_key_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_wallets_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_wallets_by_user_ids(user_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_wallets(
|
||||
&self,
|
||||
query: &AdminWalletListQuery,
|
||||
) -> Result<StoredAdminWalletListPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_admin_wallets(query).await,
|
||||
None => Ok(StoredAdminWalletListPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_wallet_ledger(
|
||||
&self,
|
||||
query: &AdminWalletLedgerQuery,
|
||||
) -> Result<StoredAdminWalletLedgerPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_admin_wallet_ledger(query).await,
|
||||
None => Ok(StoredAdminWalletLedgerPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_wallet_refund_requests(
|
||||
&self,
|
||||
query: &AdminWalletRefundRequestListQuery,
|
||||
) -> Result<StoredAdminWalletRefundRequestPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_admin_wallet_refund_requests(query).await,
|
||||
None => Ok(StoredAdminWalletRefundRequestPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_wallet_transactions(
|
||||
&self,
|
||||
wallet_id: &str,
|
||||
limit: usize,
|
||||
offset: usize,
|
||||
) -> Result<StoredAdminWalletTransactionPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_admin_wallet_transactions(wallet_id, limit, offset)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredAdminWalletTransactionPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_wallet_today_usage(
|
||||
&self,
|
||||
wallet_id: &str,
|
||||
billing_timezone: &str,
|
||||
) -> Result<Option<StoredWalletDailyUsageLedger>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.find_wallet_today_usage(wallet_id, billing_timezone)
|
||||
.await
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_wallet_daily_usage_history(
|
||||
&self,
|
||||
wallet_id: &str,
|
||||
billing_timezone: &str,
|
||||
limit: usize,
|
||||
) -> Result<StoredWalletDailyUsageLedgerPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_wallet_daily_usage_history(wallet_id, billing_timezone, limit)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredWalletDailyUsageLedgerPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_wallet_refunds(
|
||||
&self,
|
||||
wallet_id: &str,
|
||||
limit: usize,
|
||||
offset: usize,
|
||||
) -> Result<StoredAdminWalletRefundPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_admin_wallet_refunds(wallet_id, limit, offset)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredAdminWalletRefundPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_payment_orders(
|
||||
&self,
|
||||
query: &AdminPaymentOrderListQuery,
|
||||
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.list_admin_payment_orders(query).await,
|
||||
None => Ok(StoredAdminPaymentOrderPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_payment_callbacks(
|
||||
&self,
|
||||
payment_method: Option<&str>,
|
||||
limit: usize,
|
||||
offset: usize,
|
||||
) -> Result<StoredAdminPaymentCallbackPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_admin_payment_callbacks(payment_method, limit, offset)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredAdminPaymentCallbackPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_admin_payment_order(
|
||||
&self,
|
||||
order_id: &str,
|
||||
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.find_admin_payment_order(order_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_wallet_payment_orders_by_user_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
limit: usize,
|
||||
offset: usize,
|
||||
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_wallet_payment_orders_by_user_id(user_id, limit, offset)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredAdminPaymentOrderPage::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_wallet_payment_order_by_user_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
order_id: &str,
|
||||
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.find_wallet_payment_order_by_user_id(user_id, order_id)
|
||||
.await
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_wallet_refund(
|
||||
&self,
|
||||
wallet_id: &str,
|
||||
refund_id: &str,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredAdminWalletRefund>, DataLayerError>
|
||||
{
|
||||
match &self.wallet_reader {
|
||||
Some(repository) => repository.find_wallet_refund(wallet_id, refund_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_wallet_recharge_order(
|
||||
&self,
|
||||
input: CreateWalletRechargeOrderInput,
|
||||
) -> Result<Option<CreateWalletRechargeOrderOutcome>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.create_wallet_recharge_order(input)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_wallet_refund_request(
|
||||
&self,
|
||||
input: CreateWalletRefundRequestInput,
|
||||
) -> Result<Option<CreateWalletRefundRequestOutcome>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.create_wallet_refund_request(input)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn process_payment_callback(
|
||||
&self,
|
||||
input: ProcessPaymentCallbackInput,
|
||||
) -> Result<Option<ProcessPaymentCallbackOutcome>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository.process_payment_callback(input).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn adjust_wallet_balance(
|
||||
&self,
|
||||
input: AdjustWalletBalanceInput,
|
||||
) -> Result<Option<(StoredWalletSnapshot, StoredAdminWalletTransaction)>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository.adjust_wallet_balance(input).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_manual_wallet_recharge(
|
||||
&self,
|
||||
input: CreateManualWalletRechargeInput,
|
||||
) -> Result<Option<(StoredWalletSnapshot, StoredAdminPaymentOrder)>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository.create_manual_wallet_recharge(input).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn process_admin_wallet_refund(
|
||||
&self,
|
||||
input: ProcessAdminWalletRefundInput,
|
||||
) -> Result<
|
||||
Option<
|
||||
WalletMutationOutcome<(
|
||||
StoredWalletSnapshot,
|
||||
StoredAdminWalletRefund,
|
||||
StoredAdminWalletTransaction,
|
||||
)>,
|
||||
>,
|
||||
DataLayerError,
|
||||
> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.process_admin_wallet_refund(input)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn complete_admin_wallet_refund(
|
||||
&self,
|
||||
input: CompleteAdminWalletRefundInput,
|
||||
) -> Result<Option<WalletMutationOutcome<StoredAdminWalletRefund>>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.complete_admin_wallet_refund(input)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn fail_admin_wallet_refund(
|
||||
&self,
|
||||
input: FailAdminWalletRefundInput,
|
||||
) -> Result<
|
||||
Option<
|
||||
WalletMutationOutcome<(
|
||||
StoredWalletSnapshot,
|
||||
StoredAdminWalletRefund,
|
||||
Option<StoredAdminWalletTransaction>,
|
||||
)>,
|
||||
>,
|
||||
DataLayerError,
|
||||
> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository.fail_admin_wallet_refund(input).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn expire_admin_payment_order(
|
||||
&self,
|
||||
order_id: &str,
|
||||
) -> Result<Option<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>>, DataLayerError>
|
||||
{
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.expire_admin_payment_order(order_id)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn fail_admin_payment_order(
|
||||
&self,
|
||||
order_id: &str,
|
||||
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository
|
||||
.fail_admin_payment_order(order_id)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn credit_admin_payment_order(
|
||||
&self,
|
||||
input: CreditAdminPaymentOrderInput,
|
||||
) -> Result<Option<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>>, DataLayerError>
|
||||
{
|
||||
match &self.wallet_writer {
|
||||
Some(repository) => repository.credit_admin_payment_order(input).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
match &self.settlement_writer {
|
||||
Some(repository) => repository.settle_usage(input).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn reset_due_provider_quotas(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
match &self.provider_quota_writer {
|
||||
Some(repository) => repository.reset_due(now_unix_secs).await,
|
||||
None => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_provider_quota_by_provider_id(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
|
||||
match &self.provider_quota_reader {
|
||||
Some(repository) => repository.find_by_provider_id(provider_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
|
||||
pub(crate) async fn upsert_usage(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_writer {
|
||||
Some(repository) => repository.upsert(usage).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_request_usage_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.find_by_request_id(request_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_request_usage_by_id(
|
||||
&self,
|
||||
usage_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.find_by_id(usage_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_usage_audits(
|
||||
&self,
|
||||
query: &aether_data::repository::usage::UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.list_usage_audits(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.list_recent_usage_audits(user_id, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_usage_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.summarize_total_tokens_by_api_key_ids(api_key_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(std::collections::BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_users_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_users_by_ids(user_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_export_users().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_export_users_page(
|
||||
&self,
|
||||
query: &aether_data::repository::users::UserExportListQuery,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_export_users_page(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_export_users(
|
||||
&self,
|
||||
) -> Result<aether_data::repository::users::UserExportSummary, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.summarize_export_users().await,
|
||||
None => Ok(aether_data::repository::users::UserExportSummary::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_export_user_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.find_export_user_by_id(user_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_non_admin_export_users(
|
||||
&self,
|
||||
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_non_admin_export_users().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_auth_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_user_auth_by_ids(user_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.summarize_provider_usage_since(provider_id, since_unix_secs)
|
||||
.await
|
||||
}
|
||||
None => Ok(StoredProviderUsageSummary::default()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
|
||||
self.usage_worker_runner.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn find_billing_model_context(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_api_key_id: Option<&str>,
|
||||
global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
match &self.billing_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.find_model_context(provider_id, provider_api_key_id, global_model_name)
|
||||
.await
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_candidate_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
|
||||
read_request_candidate_trace(self, request_id, attempted_only).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_decision_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<DecisionTrace>, DataLayerError> {
|
||||
read_decision_trace(self, request_id, attempted_only).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_audit(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
self.find_request_usage_by_request_id(request_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_audit_bundle(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
|
||||
aether_data::repository::audit::read_request_audit_bundle(
|
||||
self,
|
||||
request_id,
|
||||
attempted_only,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) async fn read_provider_transport_snapshot(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
) -> Result<Option<GatewayProviderTransportSnapshot>, DataLayerError> {
|
||||
read_provider_transport_snapshot(self, provider_id, endpoint_id, key_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_response(
|
||||
&self,
|
||||
route_family: Option<&str>,
|
||||
request_path: &str,
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
read_data_backed_video_task_response(self, route_family, request_path).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn write_shadow_result(
|
||||
&self,
|
||||
result: aether_data::repository::shadow_results::UpsertShadowResult,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
match &self.shadow_result_writer {
|
||||
Some(repository) => match repository.upsert(result).await {
|
||||
Ok(stored) => Ok(Some(stored)),
|
||||
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(None),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn record_shadow_result_sample(
|
||||
&self,
|
||||
sample: RecordShadowResultSample,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let Some(writer) = &self.shadow_result_writer else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let existing = match &self.shadow_result_reader {
|
||||
Some(reader) => {
|
||||
match reader
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: &sample.trace_id,
|
||||
request_fingerprint: &sample.request_fingerprint,
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(existing) => existing,
|
||||
Err(err) if is_missing_shadow_results_relation_error(&err) => return Ok(None),
|
||||
Err(err) => return Err(err),
|
||||
}
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
let merged = merge_shadow_result_sample(existing.as_ref(), sample);
|
||||
|
||||
match writer.upsert(merged).await {
|
||||
Ok(stored) => Ok(Some(stored)),
|
||||
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(None),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_shadow_results(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
match &self.shadow_result_reader {
|
||||
Some(repository) => match repository.list_recent(limit).await {
|
||||
Ok(results) => Ok(results),
|
||||
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(Vec::new()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow};
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::{is_missing_shadow_results_relation_error, GatewayDataState};
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_non_admin_export_users_from_user_reader() {
|
||||
let repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![
|
||||
StoredUserExportRow::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
"alice".to_string(),
|
||||
Some("hash".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
Some(60),
|
||||
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})),
|
||||
true,
|
||||
)
|
||||
.expect("user export row should build"),
|
||||
]));
|
||||
let state = GatewayDataState::with_user_reader_for_tests(repository);
|
||||
|
||||
let rows = state
|
||||
.list_non_admin_export_users()
|
||||
.await
|
||||
.expect("export users should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].username, "alice");
|
||||
assert!(rows[0].email_verified);
|
||||
assert_eq!(rows[0].password_hash.as_deref(), Some("hash"));
|
||||
assert_eq!(rows[0].allowed_models, Some(vec!["gpt-4.1".to_string()]));
|
||||
assert_eq!(
|
||||
rows[0].model_capability_settings,
|
||||
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_missing_shadow_results_relation_error_text() {
|
||||
let missing_relation = DataLayerError::UnexpectedValue(
|
||||
"postgres error: error returned from database: relation \"gateway_shadow_results\" does not exist"
|
||||
.to_string(),
|
||||
);
|
||||
let other = DataLayerError::UnexpectedValue(
|
||||
"postgres error: error returned from database: relation \"usage\" does not exist"
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert!(is_missing_shadow_results_relation_error(&missing_relation));
|
||||
assert!(!is_missing_shadow_results_relation_error(&other));
|
||||
}
|
||||
}
|
||||
106
apps/aether-gateway/src/data/state/testing/announcements.rs
Normal file
106
apps/aether-gateway/src/data/state/testing/announcements.rs
Normal file
@@ -0,0 +1,106 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{
|
||||
AnnouncementReadRepository, AnnouncementWriteRepository, GatewayDataConfig, GatewayDataState,
|
||||
};
|
||||
|
||||
impl GatewayDataState {
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_announcement_reader_for_tests(
|
||||
repository: Arc<dyn AnnouncementReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: Some(repository),
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_announcement_repository_for_tests<T>(repository: Arc<T>) -> Self
|
||||
where
|
||||
T: aether_data::repository::announcements::AnnouncementReadRepository
|
||||
+ aether_data::repository::announcements::AnnouncementWriteRepository
|
||||
+ 'static,
|
||||
{
|
||||
let announcement_reader: Arc<dyn AnnouncementReadRepository> = repository.clone();
|
||||
let announcement_writer: Arc<dyn AnnouncementWriteRepository> = repository;
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: Some(announcement_reader),
|
||||
announcement_writer: Some(announcement_writer),
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: None,
|
||||
video_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
2676
apps/aether-gateway/src/data/state/testing/mod.rs
Normal file
2676
apps/aether-gateway/src/data/state/testing/mod.rs
Normal file
File diff suppressed because it is too large
Load Diff
342
apps/aether-gateway/src/data/state/testing/video_tasks.rs
Normal file
342
apps/aether-gateway/src/data/state/testing/video_tasks.rs
Normal file
@@ -0,0 +1,342 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{
|
||||
AuthApiKeyReadRepository, GatewayDataConfig, GatewayDataState, ProviderCatalogReadRepository,
|
||||
RequestCandidateReadRepository, RequestCandidateWriteRepository, VideoTaskReadRepository,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) fn with_video_task_reader_for_tests(
|
||||
repository: Arc<dyn VideoTaskReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(repository),
|
||||
video_task_writer: None,
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_auth_and_video_task_repository_for_tests<T>(
|
||||
auth_repository: Arc<dyn AuthApiKeyReadRepository>,
|
||||
repository: Arc<T>,
|
||||
) -> Self
|
||||
where
|
||||
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
|
||||
{
|
||||
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
|
||||
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: Some(auth_repository),
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_repository_for_tests<T>(repository: Arc<T>) -> Self
|
||||
where
|
||||
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
|
||||
{
|
||||
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
|
||||
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_repository_and_provider_transport_for_tests<T>(
|
||||
repository: Arc<T>,
|
||||
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
|
||||
encryption_key: impl Into<String>,
|
||||
) -> Self
|
||||
where
|
||||
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
|
||||
{
|
||||
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
|
||||
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled().with_encryption_key(encryption_key),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: None,
|
||||
request_candidate_writer: None,
|
||||
provider_catalog_reader: Some(provider_catalog_repository),
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_and_request_candidate_repository_for_tests<T, U>(
|
||||
repository: Arc<T>,
|
||||
request_candidate_repository: Arc<U>,
|
||||
) -> Self
|
||||
where
|
||||
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
|
||||
U: aether_data::repository::candidates::RequestCandidateRepository + 'static,
|
||||
{
|
||||
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
|
||||
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
|
||||
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
|
||||
request_candidate_repository.clone();
|
||||
let request_candidate_writer: Arc<dyn RequestCandidateWriteRepository> =
|
||||
request_candidate_repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: Some(request_candidate_reader),
|
||||
request_candidate_writer: Some(request_candidate_writer),
|
||||
provider_catalog_reader: None,
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_provider_transport_and_request_candidate_repository_for_tests<
|
||||
T,
|
||||
U,
|
||||
V,
|
||||
>(
|
||||
repository: Arc<T>,
|
||||
provider_catalog_repository: Arc<U>,
|
||||
request_candidate_repository: Arc<V>,
|
||||
encryption_key: impl Into<String>,
|
||||
) -> Self
|
||||
where
|
||||
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
|
||||
U: ProviderCatalogReadRepository + 'static,
|
||||
V: aether_data::repository::candidates::RequestCandidateRepository + 'static,
|
||||
{
|
||||
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
|
||||
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
|
||||
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
|
||||
request_candidate_repository.clone();
|
||||
let request_candidate_writer: Arc<dyn RequestCandidateWriteRepository> =
|
||||
request_candidate_repository;
|
||||
let provider_catalog_reader: Arc<dyn ProviderCatalogReadRepository> =
|
||||
provider_catalog_repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled().with_encryption_key(encryption_key),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
auth_api_key_writer: None,
|
||||
auth_module_reader: None,
|
||||
auth_module_writer: None,
|
||||
announcement_reader: None,
|
||||
announcement_writer: None,
|
||||
management_token_reader: None,
|
||||
management_token_writer: None,
|
||||
oauth_provider_reader: None,
|
||||
oauth_provider_writer: None,
|
||||
proxy_node_reader: None,
|
||||
proxy_node_writer: None,
|
||||
billing_reader: None,
|
||||
gemini_file_mapping_reader: None,
|
||||
gemini_file_mapping_writer: None,
|
||||
global_model_reader: None,
|
||||
global_model_writer: None,
|
||||
minimal_candidate_selection_reader: None,
|
||||
request_candidate_reader: Some(request_candidate_reader),
|
||||
request_candidate_writer: Some(request_candidate_writer),
|
||||
provider_catalog_reader: Some(provider_catalog_reader),
|
||||
provider_catalog_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
user_preferences: None,
|
||||
usage_worker_runner: None,
|
||||
video_task_reader: Some(video_task_reader),
|
||||
video_task_writer: Some(video_task_writer),
|
||||
wallet_reader: None,
|
||||
wallet_writer: None,
|
||||
settlement_writer: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
system_config_values: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user