mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-05 17:07:46 +08:00
refactor(runtime): 优化管理端摘要查询与维护聚合链路
- 为 provider catalog key 和 video task 列表增加 summary/page 查询与排序能力,减少列表场景读取重字段 - 将多处 SQL 结果读取改为流式收集,降低 `fetch_all` 的内存占用 - 把日/小时统计、钱包日用量等维护任务改为数据库侧 `CTE + upsert` 聚合 - 修复视频任务轮询更新时从本地 snapshot 回填稀疏字段,避免 `prompt` 和请求体信息丢失
This commit is contained in:
@@ -10,9 +10,9 @@ pub(crate) use http::{
|
||||
CancelVideoTaskError,
|
||||
};
|
||||
pub(crate) use query::{
|
||||
read_video_task_detail, read_video_task_page, read_video_task_stats,
|
||||
read_video_task_video_source, VideoTaskPageResponse, VideoTaskStatsResponse,
|
||||
VideoTaskVideoSource,
|
||||
read_video_task_detail, read_video_task_page, read_video_task_page_summary,
|
||||
read_video_task_stats, read_video_task_video_source, VideoTaskPageResponse,
|
||||
VideoTaskStatsResponse, VideoTaskVideoSource,
|
||||
};
|
||||
pub(crate) use runtime::{
|
||||
execute_video_task_refresh_plan, finalize_video_task_if_terminal, spawn_video_task_poller,
|
||||
|
||||
@@ -67,6 +67,34 @@ pub(crate) async fn read_video_task_page(
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_page_summary(
|
||||
state: &AppState,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
) -> Result<VideoTaskPageResponse, GatewayError> {
|
||||
let page = page.max(1);
|
||||
let page_size = page_size.clamp(1, 100);
|
||||
let total = state.count_video_tasks(filter).await?;
|
||||
let offset = page_size.saturating_mul(page.saturating_sub(1));
|
||||
let items = state
|
||||
.list_video_task_page_summary(filter, offset, page_size)
|
||||
.await?;
|
||||
let pages = if total == 0 {
|
||||
0
|
||||
} else {
|
||||
((total as usize) + page_size - 1) / page_size
|
||||
};
|
||||
|
||||
Ok(VideoTaskPageResponse {
|
||||
items,
|
||||
total,
|
||||
page,
|
||||
page_size,
|
||||
pages,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_detail(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
|
||||
@@ -267,11 +267,14 @@ fn build_successful_poll_update(
|
||||
record.format_converted = task.format_converted;
|
||||
record.model = task.model.clone().or(record.model);
|
||||
record.prompt = task.prompt.clone().or(record.prompt);
|
||||
record.original_request_body = task.original_request_body.clone();
|
||||
record.duration_seconds = task.duration_seconds;
|
||||
record.resolution = task.resolution.clone();
|
||||
record.aspect_ratio = task.aspect_ratio.clone();
|
||||
record.size = task.size.clone();
|
||||
record.original_request_body = task
|
||||
.original_request_body
|
||||
.clone()
|
||||
.or(record.original_request_body);
|
||||
record.duration_seconds = task.duration_seconds.or(record.duration_seconds);
|
||||
record.resolution = task.resolution.clone().or(record.resolution);
|
||||
record.aspect_ratio = task.aspect_ratio.clone().or(record.aspect_ratio);
|
||||
record.size = task.size.clone().or(record.size);
|
||||
record.created_at_unix_ms = task.created_at_unix_ms;
|
||||
record.submitted_at_unix_secs = task.submitted_at_unix_secs;
|
||||
record.updated_at_unix_secs = now_unix_secs;
|
||||
@@ -349,6 +352,8 @@ fn build_failed_poll_update(
|
||||
}
|
||||
|
||||
fn stored_task_to_upsert(task: &StoredVideoTask) -> UpsertVideoTask {
|
||||
let snapshot_record =
|
||||
LocalVideoTaskSnapshot::from_stored_task(task).map(|snapshot| snapshot.to_upsert_record());
|
||||
UpsertVideoTask {
|
||||
id: task.id.clone(),
|
||||
short_id: task.short_id.clone(),
|
||||
@@ -365,12 +370,36 @@ fn stored_task_to_upsert(task: &StoredVideoTask) -> UpsertVideoTask {
|
||||
provider_api_format: task.provider_api_format.clone(),
|
||||
format_converted: task.format_converted,
|
||||
model: task.model.clone(),
|
||||
prompt: task.prompt.clone(),
|
||||
original_request_body: task.original_request_body.clone(),
|
||||
duration_seconds: task.duration_seconds,
|
||||
resolution: task.resolution.clone(),
|
||||
aspect_ratio: task.aspect_ratio.clone(),
|
||||
size: task.size.clone(),
|
||||
prompt: task.prompt.clone().or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.prompt.clone())
|
||||
}),
|
||||
original_request_body: task.original_request_body.clone().or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.original_request_body.clone())
|
||||
}),
|
||||
duration_seconds: task.duration_seconds.or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.duration_seconds)
|
||||
}),
|
||||
resolution: task.resolution.clone().or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.resolution.clone())
|
||||
}),
|
||||
aspect_ratio: task.aspect_ratio.clone().or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.aspect_ratio.clone())
|
||||
}),
|
||||
size: task.size.clone().or_else(|| {
|
||||
snapshot_record
|
||||
.as_ref()
|
||||
.and_then(|record| record.size.clone())
|
||||
}),
|
||||
status: task.status,
|
||||
progress_percent: task.progress_percent,
|
||||
progress_message: task.progress_message.clone(),
|
||||
@@ -550,3 +579,154 @@ fn now_unix_secs() -> u64 {
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_failed_poll_update, stored_task_to_upsert, VideoTaskRefreshError};
|
||||
use crate::video_tasks::{
|
||||
LocalVideoTaskPersistence, LocalVideoTaskSnapshot, LocalVideoTaskStatus,
|
||||
LocalVideoTaskTransport, OpenAiVideoTaskSeed,
|
||||
};
|
||||
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
fn sample_sparse_stored_task() -> StoredVideoTask {
|
||||
let snapshot = LocalVideoTaskSnapshot::OpenAi(OpenAiVideoTaskSeed {
|
||||
local_task_id: "task-1".to_string(),
|
||||
upstream_task_id: "ext-1".to_string(),
|
||||
created_at_unix_ms: 1,
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
seconds: Some("4".to_string()),
|
||||
remixed_from_video_id: None,
|
||||
status: LocalVideoTaskStatus::Processing,
|
||||
progress_percent: 50,
|
||||
completed_at_unix_secs: None,
|
||||
expires_at_unix_secs: None,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
persistence: LocalVideoTaskPersistence {
|
||||
request_id: "request-1".to_string(),
|
||||
username: Some("user".to_string()),
|
||||
api_key_name: Some("primary".to_string()),
|
||||
client_api_format: "openai:video".to_string(),
|
||||
provider_api_format: "openai:video".to_string(),
|
||||
original_request_body: json!({
|
||||
"prompt": "hello",
|
||||
"seconds": "4",
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
"size": "1280x720"
|
||||
}),
|
||||
format_converted: false,
|
||||
},
|
||||
transport: LocalVideoTaskTransport {
|
||||
upstream_base_url: "https://example.com".to_string(),
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".to_string()),
|
||||
model_name: Some("sora-2".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
},
|
||||
});
|
||||
|
||||
StoredVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1".to_string()),
|
||||
request_id: "request-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: Some("user".to_string()),
|
||||
api_key_name: Some("primary".to_string()),
|
||||
external_task_id: Some("ext-1".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
client_api_format: Some("openai:video".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: None,
|
||||
original_request_body: None,
|
||||
duration_seconds: None,
|
||||
resolution: None,
|
||||
aspect_ratio: None,
|
||||
size: None,
|
||||
status: VideoTaskStatus::Processing,
|
||||
progress_percent: 50,
|
||||
progress_message: Some("polling".to_string()),
|
||||
retry_count: 1,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: Some(20),
|
||||
poll_count: 2,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_ms: 1,
|
||||
submitted_at_unix_secs: Some(1),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 20,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
request_metadata: Some(json!({
|
||||
"rust_local_snapshot": serde_json::to_value(snapshot)
|
||||
.expect("snapshot should serialize")
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stored_task_to_upsert_restores_sparse_fields_from_snapshot() {
|
||||
let record = stored_task_to_upsert(&sample_sparse_stored_task());
|
||||
|
||||
assert_eq!(record.prompt.as_deref(), Some("hello"));
|
||||
assert_eq!(
|
||||
record.original_request_body,
|
||||
Some(json!({
|
||||
"prompt": "hello",
|
||||
"seconds": "4",
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
"size": "1280x720"
|
||||
}))
|
||||
);
|
||||
assert_eq!(record.duration_seconds, Some(4));
|
||||
assert_eq!(record.resolution.as_deref(), Some("720p"));
|
||||
assert_eq!(record.aspect_ratio.as_deref(), Some("16:9"));
|
||||
assert_eq!(record.size.as_deref(), Some("1280x720"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_poll_update_keeps_snapshot_backed_request_body() {
|
||||
let record = build_failed_poll_update(
|
||||
&sample_sparse_stored_task(),
|
||||
&VideoTaskRefreshError {
|
||||
message: "temporary failure".to_string(),
|
||||
permanent: false,
|
||||
},
|
||||
100,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
record.original_request_body,
|
||||
Some(json!({
|
||||
"prompt": "hello",
|
||||
"seconds": "4",
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
"size": "1280x720"
|
||||
}))
|
||||
);
|
||||
assert_eq!(record.prompt.as_deref(), Some("hello"));
|
||||
assert_eq!(record.resolution.as_deref(), Some("720p"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ use aether_data::repository::auth::{
|
||||
read_resolved_auth_api_key_snapshot_by_key_hash,
|
||||
read_resolved_auth_api_key_snapshot_by_user_api_key_ids,
|
||||
};
|
||||
use futures_util::TryStreamExt;
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -1278,12 +1279,14 @@ impl GatewayDataState {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let rows = sqlx::query(LIST_USER_SESSIONS_SQL)
|
||||
let mut rows = sqlx::query(LIST_USER_SESSIONS_SQL)
|
||||
.bind(user_id)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
rows.iter().map(map_user_session_row).collect()
|
||||
.fetch(&pool);
|
||||
let mut sessions = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
sessions.push(map_user_session_row(&row)?);
|
||||
}
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
pub(crate) async fn create_user_session(
|
||||
|
||||
@@ -280,6 +280,20 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_key_summaries_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
|
||||
@@ -135,6 +135,18 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_video_task_page_summary(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.list_page_summary(filter, offset, limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn count_video_tasks(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
|
||||
@@ -315,7 +315,7 @@ pub(crate) async fn recover_all_admin_key_health(
|
||||
Vec::new()
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
|
||||
@@ -69,7 +69,7 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
|
||||
let mut health_scores_by_format = BTreeMap::<String, Vec<f64>>::new();
|
||||
if !provider_ids.is_empty() {
|
||||
let keys = state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
@@ -235,7 +235,7 @@ pub(crate) async fn build_admin_health_summary_payload(
|
||||
Vec::new()
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
|
||||
@@ -60,7 +60,9 @@ pub(super) async fn maybe_build_local_admin_video_tasks_response(
|
||||
let page_size = query_param_value(request_context.query_string(), "page_size")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(20);
|
||||
let response = state.read_video_task_page(&filter, page, page_size).await?;
|
||||
let response = state
|
||||
.read_video_task_page_summary(&filter, page, page_size)
|
||||
.await?;
|
||||
let provider_names = build_admin_video_task_provider_names(state, &response.items).await?;
|
||||
return Ok(Some(
|
||||
Json(json!({
|
||||
|
||||
+1
-1
@@ -58,7 +58,7 @@ pub(super) async fn build_admin_monitoring_provider_name_by_id_and_keys(
|
||||
Vec::new()
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await?
|
||||
};
|
||||
Ok((provider_name_by_id, keys))
|
||||
|
||||
@@ -49,7 +49,9 @@ pub(crate) use crate::handlers::admin::provider::pool::runtime::{
|
||||
};
|
||||
pub(crate) use crate::handlers::admin::provider::shared::support::AdminProviderPoolRuntimeState;
|
||||
pub(crate) use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
pub(crate) use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyListQuery;
|
||||
pub(crate) use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_pool_response(
|
||||
state: &AdminAppState<'_>,
|
||||
|
||||
@@ -3,7 +3,7 @@ use super::{
|
||||
parse_admin_pool_page, parse_admin_pool_page_size, parse_admin_pool_search,
|
||||
parse_admin_pool_status_filter, pool_payloads, pool_selection,
|
||||
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state,
|
||||
AdminProviderPoolRuntimeState, ProviderCatalogKeyListQuery,
|
||||
AdminProviderPoolRuntimeState, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
@@ -118,6 +118,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
||||
},
|
||||
offset: page_offset,
|
||||
limit: page_size,
|
||||
order: ProviderCatalogKeyListOrder::Name,
|
||||
})
|
||||
.await?;
|
||||
(key_page.items, key_page.total)
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(crate) async fn build_admin_provider_summary_payload(
|
||||
active_global_model_ids_result,
|
||||
) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.read_provider_quota_snapshot(provider_id),
|
||||
state.list_provider_model_stats(&provider_ids),
|
||||
state.list_active_global_model_ids_by_provider_ids(&provider_ids),
|
||||
@@ -204,7 +204,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
Vec::new()
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
|
||||
@@ -48,11 +48,11 @@ pub(crate) async fn build_admin_providers_payload(
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
};
|
||||
let keys = if provider_ids.is_empty() {
|
||||
let key_stats = if provider_ids.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_stats_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
@@ -68,9 +68,12 @@ pub(crate) async fn build_admin_providers_payload(
|
||||
},
|
||||
);
|
||||
let has_any_key_by_provider =
|
||||
keys.into_iter()
|
||||
.fold(BTreeSet::<String>::new(), |mut acc, key| {
|
||||
acc.insert(key.provider_id);
|
||||
key_stats
|
||||
.into_iter()
|
||||
.fold(BTreeSet::<String>::new(), |mut acc, stats| {
|
||||
if stats.total_keys > 0 {
|
||||
acc.insert(stats.provider_id);
|
||||
}
|
||||
acc
|
||||
});
|
||||
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_payload(
|
||||
@@ -15,30 +18,26 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|mut providers| providers.drain(..).next())?;
|
||||
let mut keys = state
|
||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
let key_page = state
|
||||
.list_provider_catalog_key_page(&ProviderCatalogKeyListQuery {
|
||||
provider_id: provider.id.clone(),
|
||||
search: None,
|
||||
is_active: None,
|
||||
offset: skip,
|
||||
limit,
|
||||
order: ProviderCatalogKeyListOrder::CreatedAt,
|
||||
})
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
keys.sort_by(|left, right| {
|
||||
left.internal_priority
|
||||
.cmp(&right.internal_priority)
|
||||
.then_with(|| {
|
||||
left.created_at_unix_ms
|
||||
.unwrap_or_default()
|
||||
.cmp(&right.created_at_unix_ms.unwrap_or_default())
|
||||
})
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
.ok()?;
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
Some(serde_json::Value::Array(
|
||||
keys.into_iter()
|
||||
.skip(skip)
|
||||
.take(limit)
|
||||
key_page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
state.build_admin_provider_key_response(
|
||||
&key,
|
||||
|
||||
@@ -55,6 +55,15 @@ impl<'a> AdminAppState<'a> {
|
||||
crate::async_task::read_video_task_page(self.app, filter, page, page_size).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_page_summary(
|
||||
&self,
|
||||
filter: &aether_data_contracts::repository::video_tasks::VideoTaskQueryFilter,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
) -> Result<crate::async_task::VideoTaskPageResponse, GatewayError> {
|
||||
crate::async_task::read_video_task_page_summary(self.app, filter, page, page_size).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_stats(
|
||||
&self,
|
||||
filter: &aether_data_contracts::repository::video_tasks::VideoTaskQueryFilter,
|
||||
|
||||
@@ -69,6 +69,18 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<
|
||||
Vec<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
|
||||
@@ -148,7 +148,7 @@ pub(crate) async fn build_admin_module_runtime_state(
|
||||
false
|
||||
} else {
|
||||
state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default()
|
||||
|
||||
@@ -350,7 +350,7 @@ pub(crate) async fn build_api_format_health_monitor_payload(
|
||||
let mut key_counts_by_format = BTreeMap::<String, usize>::new();
|
||||
if options.include_key_count && !provider_ids.is_empty() {
|
||||
let keys = state
|
||||
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids)
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
|
||||
@@ -41,7 +41,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
|
||||
|
||||
let (endpoints_result, keys_result) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
);
|
||||
|
||||
let endpoint_base_url_by_provider_and_format = endpoints_result
|
||||
|
||||
@@ -109,28 +109,28 @@ const DB_MAINTENANCE_HOUR: u32 = 5;
|
||||
const DB_MAINTENANCE_MINUTE: u32 = 0;
|
||||
const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
|
||||
const DB_MAINTENANCE_TABLES: &[&str] = &["usage", "request_candidates", "audit_logs"];
|
||||
const SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id,
|
||||
COUNT(usage.id) AS total_requests,
|
||||
CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
|
||||
COALESCE(SUM(usage.input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(usage.output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(usage.cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(usage.cache_read_input_tokens), 0) AS cache_read_tokens,
|
||||
MIN(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS first_finalized_at,
|
||||
MAX(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS last_finalized_at
|
||||
FROM usage
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id IS NOT NULL
|
||||
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
|
||||
AND usage.total_cost_usd > 0
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $1
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $2
|
||||
GROUP BY usage_settlement_snapshots.wallet_id
|
||||
"#;
|
||||
const UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
usage_settlement_snapshots.wallet_id,
|
||||
COUNT(usage.id) AS total_requests,
|
||||
CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
|
||||
COALESCE(SUM(usage.input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(usage.output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(usage.cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(usage.cache_read_input_tokens), 0) AS cache_read_tokens,
|
||||
MIN(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS first_finalized_at,
|
||||
MAX(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS last_finalized_at
|
||||
FROM usage
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id IS NOT NULL
|
||||
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
|
||||
AND usage.total_cost_usd > 0
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $1
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $2
|
||||
GROUP BY usage_settlement_snapshots.wallet_id
|
||||
)
|
||||
INSERT INTO wallet_daily_usage_ledgers (
|
||||
id,
|
||||
wallet_id,
|
||||
@@ -148,11 +148,23 @@ INSERT INTO wallet_daily_usage_ledgers (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1, $2, $3, $4, $5,
|
||||
$6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15
|
||||
)
|
||||
SELECT
|
||||
md5(CONCAT('wallet-daily-usage:', aggregated.wallet_id, ':', CAST($3 AS TEXT), ':', $4)),
|
||||
aggregated.wallet_id,
|
||||
$3,
|
||||
$4,
|
||||
aggregated.total_cost_usd,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.first_finalized_at,
|
||||
aggregated.last_finalized_at,
|
||||
$5,
|
||||
$5,
|
||||
$5
|
||||
FROM aggregated
|
||||
ON CONFLICT (wallet_id, billing_date, billing_timezone)
|
||||
DO UPDATE SET
|
||||
total_cost_usd = EXCLUDED.total_cost_usd,
|
||||
@@ -397,17 +409,7 @@ WHERE request_id = ANY($1)
|
||||
"#;
|
||||
const SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
request_body,
|
||||
request_body_compressed,
|
||||
response_body,
|
||||
response_body_compressed,
|
||||
provider_request_body,
|
||||
provider_request_body_compressed,
|
||||
client_response_body
|
||||
,
|
||||
client_response_body_compressed
|
||||
id
|
||||
FROM usage
|
||||
WHERE created_at < $1
|
||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||
@@ -503,6 +505,22 @@ DO UPDATE SET
|
||||
END,
|
||||
updated_at = NOW()
|
||||
"#;
|
||||
const SELECT_USAGE_BODY_COMPRESSION_ROW_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
request_body,
|
||||
request_body_compressed,
|
||||
response_body,
|
||||
response_body_compressed,
|
||||
provider_request_body,
|
||||
provider_request_body_compressed,
|
||||
client_response_body,
|
||||
client_response_body_compressed
|
||||
FROM usage
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
const UPDATE_USAGE_BODY_COMPRESSION_SQL: &str = r#"
|
||||
UPDATE usage
|
||||
SET request_body = NULL,
|
||||
@@ -695,24 +713,24 @@ DO UPDATE SET
|
||||
aggregated_at = EXCLUDED.aggregated_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND model IS NOT NULL
|
||||
AND model <> ''
|
||||
GROUP BY model
|
||||
"#;
|
||||
const UPSERT_STATS_DAILY_MODEL_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND model IS NOT NULL
|
||||
AND model <> ''
|
||||
GROUP BY model
|
||||
)
|
||||
INSERT INTO stats_daily_model (
|
||||
id,
|
||||
date,
|
||||
@@ -727,7 +745,20 @@ INSERT INTO stats_daily_model (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
SELECT
|
||||
md5(CONCAT('stats-daily-model:', aggregated.model, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.model,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.avg_response_time_ms,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (date, model)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
@@ -739,21 +770,21 @@ DO UPDATE SET
|
||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
COALESCE(provider_name, 'Unknown') AS provider_name,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY COALESCE(provider_name, 'Unknown')
|
||||
"#;
|
||||
const UPSERT_STATS_DAILY_PROVIDER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
COALESCE(provider_name, 'Unknown') AS provider_name,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY COALESCE(provider_name, 'Unknown')
|
||||
)
|
||||
INSERT INTO stats_daily_provider (
|
||||
id,
|
||||
date,
|
||||
@@ -767,7 +798,19 @@ INSERT INTO stats_daily_provider (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
|
||||
SELECT
|
||||
md5(CONCAT('stats-daily-provider:', aggregated.provider_name, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.provider_name,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.total_cost,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (date, provider_name)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
@@ -778,24 +821,34 @@ DO UPDATE SET
|
||||
total_cost = EXCLUDED.total_cost,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
api_key_id,
|
||||
MAX(api_key_name) AS api_key_name,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(CASE WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1 ELSE 0 END), 0) AS BIGINT) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND api_key_id IS NOT NULL
|
||||
GROUP BY api_key_id
|
||||
"#;
|
||||
const UPSERT_STATS_DAILY_API_KEY_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
api_key_id,
|
||||
MAX(api_key_name) AS api_key_name,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT
|
||||
) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND api_key_id IS NOT NULL
|
||||
GROUP BY api_key_id
|
||||
)
|
||||
INSERT INTO stats_daily_api_key (
|
||||
id,
|
||||
api_key_id,
|
||||
@@ -812,7 +865,22 @@ INSERT INTO stats_daily_api_key (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
|
||||
SELECT
|
||||
md5(CONCAT('stats-daily-api-key:', aggregated.api_key_id, ':', CAST($1 AS TEXT))),
|
||||
aggregated.api_key_id,
|
||||
aggregated.api_key_name,
|
||||
$1,
|
||||
aggregated.total_requests,
|
||||
GREATEST(aggregated.total_requests - aggregated.error_requests, 0),
|
||||
aggregated.error_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.total_cost,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (api_key_id, date)
|
||||
DO UPDATE SET
|
||||
api_key_name = COALESCE(EXCLUDED.api_key_name, stats_daily_api_key.api_key_name),
|
||||
@@ -830,19 +898,19 @@ const DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL: &str = r#"
|
||||
DELETE FROM stats_daily_error
|
||||
WHERE date = $1
|
||||
"#;
|
||||
const SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
error_category,
|
||||
provider_name,
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_count
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND error_category IS NOT NULL
|
||||
GROUP BY error_category, provider_name, model
|
||||
"#;
|
||||
const INSERT_STATS_DAILY_ERROR_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
error_category,
|
||||
provider_name,
|
||||
model,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_count
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND error_category IS NOT NULL
|
||||
GROUP BY error_category, provider_name, model
|
||||
)
|
||||
INSERT INTO stats_daily_error (
|
||||
id,
|
||||
date,
|
||||
@@ -853,32 +921,56 @@ INSERT INTO stats_daily_error (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
"#;
|
||||
const SELECT_ACTIVE_USER_IDS_SQL: &str = r#"
|
||||
SELECT id
|
||||
FROM users
|
||||
WHERE is_active IS TRUE
|
||||
ORDER BY id ASC
|
||||
"#;
|
||||
const SELECT_STATS_USER_DAILY_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
user_id,
|
||||
MAX(username) AS username,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(COALESCE(SUM(CASE WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1 ELSE 0 END), 0) AS BIGINT) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
md5(
|
||||
CONCAT(
|
||||
'stats-daily-error:',
|
||||
CAST($1 AS TEXT),
|
||||
':',
|
||||
aggregated.error_category,
|
||||
':',
|
||||
COALESCE(aggregated.provider_name, ''),
|
||||
':',
|
||||
COALESCE(aggregated.model, '')
|
||||
)
|
||||
),
|
||||
$1,
|
||||
aggregated.error_category,
|
||||
aggregated.provider_name,
|
||||
aggregated.model,
|
||||
aggregated.total_count,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
"#;
|
||||
const UPSERT_STATS_USER_DAILY_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
user_id,
|
||||
MAX(username) AS username,
|
||||
CAST(COUNT(id) AS BIGINT) AS total_requests,
|
||||
CAST(
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS BIGINT
|
||||
) AS error_requests,
|
||||
CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
|
||||
CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
|
||||
CAST(COALESCE(SUM(cache_creation_input_tokens), 0) AS BIGINT) AS cache_creation_tokens,
|
||||
CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
)
|
||||
INSERT INTO stats_user_daily (
|
||||
id,
|
||||
user_id,
|
||||
@@ -895,7 +987,24 @@ INSERT INTO stats_user_daily (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
|
||||
SELECT
|
||||
md5(CONCAT('stats-user-daily:', users.id, ':', CAST($1 AS TEXT))),
|
||||
users.id,
|
||||
aggregated.username,
|
||||
$1,
|
||||
COALESCE(aggregated.total_requests, 0),
|
||||
GREATEST(COALESCE(aggregated.total_requests, 0) - COALESCE(aggregated.error_requests, 0), 0),
|
||||
COALESCE(aggregated.error_requests, 0),
|
||||
COALESCE(aggregated.input_tokens, 0),
|
||||
COALESCE(aggregated.output_tokens, 0),
|
||||
COALESCE(aggregated.cache_creation_tokens, 0),
|
||||
COALESCE(aggregated.cache_read_tokens, 0),
|
||||
COALESCE(aggregated.total_cost, 0),
|
||||
$3,
|
||||
$3
|
||||
FROM users
|
||||
LEFT JOIN aggregated ON aggregated.user_id = users.id
|
||||
WHERE users.is_active IS TRUE
|
||||
ON CONFLICT (user_id, date)
|
||||
DO UPDATE SET
|
||||
username = COALESCE(EXCLUDED.username, stats_user_daily.username),
|
||||
@@ -1034,21 +1143,29 @@ DO UPDATE SET
|
||||
aggregated_at = EXCLUDED.aggregated_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_HOURLY_USER_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
user_id,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(SUM(CASE WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1 ELSE 0 END), 0) AS error_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
"#;
|
||||
const UPSERT_STATS_HOURLY_USER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
user_id,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(
|
||||
SUM(
|
||||
CASE
|
||||
WHEN status_code >= 400 OR error_message IS NOT NULL THEN 1
|
||||
ELSE 0
|
||||
END
|
||||
),
|
||||
0
|
||||
) AS error_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
)
|
||||
INSERT INTO stats_hourly_user (
|
||||
id,
|
||||
hour_utc,
|
||||
@@ -1062,7 +1179,19 @@ INSERT INTO stats_hourly_user (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-user:', aggregated.user_id, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.user_id,
|
||||
aggregated.total_requests,
|
||||
GREATEST(aggregated.total_requests - aggregated.error_requests, 0),
|
||||
aggregated.error_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (hour_utc, user_id)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
@@ -1073,20 +1202,22 @@ DO UPDATE SET
|
||||
total_cost = EXCLUDED.total_cost,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
model,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY model
|
||||
"#;
|
||||
const UPSERT_STATS_HOURLY_MODEL_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
model,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
|
||||
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND model IS NOT NULL
|
||||
AND model <> ''
|
||||
GROUP BY model
|
||||
)
|
||||
INSERT INTO stats_hourly_model (
|
||||
id,
|
||||
hour_utc,
|
||||
@@ -1099,7 +1230,18 @@ INSERT INTO stats_hourly_model (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-model:', aggregated.model, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.model,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.avg_response_time_ms,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (hour_utc, model)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
@@ -1109,19 +1251,21 @@ DO UPDATE SET
|
||||
avg_response_time_ms = EXCLUDED.avg_response_time_ms,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
"#;
|
||||
const SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL: &str = r#"
|
||||
SELECT
|
||||
provider_name,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY provider_name
|
||||
"#;
|
||||
const UPSERT_STATS_HOURLY_PROVIDER_SQL: &str = r#"
|
||||
WITH aggregated AS (
|
||||
SELECT
|
||||
provider_name,
|
||||
COUNT(id) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
|
||||
FROM usage
|
||||
WHERE created_at >= $1
|
||||
AND created_at < $2
|
||||
AND provider_name IS NOT NULL
|
||||
AND provider_name <> ''
|
||||
GROUP BY provider_name
|
||||
)
|
||||
INSERT INTO stats_hourly_provider (
|
||||
id,
|
||||
hour_utc,
|
||||
@@ -1133,7 +1277,17 @@ INSERT INTO stats_hourly_provider (
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
|
||||
SELECT
|
||||
md5(CONCAT('stats-hourly-provider:', aggregated.provider_name, ':', CAST($1 AS TEXT))),
|
||||
$1,
|
||||
aggregated.provider_name,
|
||||
aggregated.total_requests,
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.total_cost,
|
||||
$3,
|
||||
$3
|
||||
FROM aggregated
|
||||
ON CONFLICT (hour_utc, provider_name)
|
||||
DO UPDATE SET
|
||||
total_requests = EXCLUDED.total_requests,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use chrono::Utc;
|
||||
use futures_util::TryStreamExt;
|
||||
use sqlx::Row;
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
@@ -59,13 +60,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
|
||||
|
||||
loop {
|
||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||
let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
|
||||
.bind(active_statuses.clone())
|
||||
.bind(cutoff_time)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch_all(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let stale_rows = {
|
||||
let mut stale_rows_stream = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
|
||||
.bind(active_statuses.clone())
|
||||
.bind(cutoff_time)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch(&mut *tx);
|
||||
let mut stale_rows = Vec::new();
|
||||
while let Some(row) = stale_rows_stream.try_next().await.map_err(postgres_error)? {
|
||||
stale_rows.push(row);
|
||||
}
|
||||
stale_rows
|
||||
};
|
||||
if stale_rows.is_empty() {
|
||||
tx.rollback().await.map_err(postgres_error)?;
|
||||
break;
|
||||
@@ -93,14 +99,18 @@ pub(crate) async fn cleanup_stale_pending_requests_once(
|
||||
let completed_request_ids = if request_ids.is_empty() {
|
||||
HashSet::new()
|
||||
} else {
|
||||
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
|
||||
.bind(request_ids)
|
||||
.fetch_all(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.filter_map(|row| row.try_get::<String, _>("request_id").ok())
|
||||
.collect::<HashSet<_>>()
|
||||
{
|
||||
let mut completed_rows = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
|
||||
.bind(request_ids)
|
||||
.fetch(&mut *tx);
|
||||
let mut completed_request_ids = HashSet::new();
|
||||
while let Some(row) = completed_rows.try_next().await.map_err(postgres_error)? {
|
||||
if let Ok(request_id) = row.try_get::<String, _>("request_id") {
|
||||
completed_request_ids.insert(request_id);
|
||||
}
|
||||
}
|
||||
completed_request_ids
|
||||
}
|
||||
};
|
||||
let plan = plan_pending_cleanup_batch(stale_rows, &completed_request_ids, timeout_minutes);
|
||||
let now = Utc::now();
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
@@ -10,15 +8,13 @@ use aether_data_contracts::DataLayerError;
|
||||
use super::{
|
||||
postgres_error, stats_aggregation_target_day, system_config_bool, PercentileSummary,
|
||||
StatsAggregationSummary, DELETE_STATS_DAILY_ERRORS_FOR_DATE_SQL, INSERT_STATS_DAILY_ERROR_SQL,
|
||||
INSERT_STATS_SUMMARY_SQL, SELECT_ACTIVE_USER_IDS_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
|
||||
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL,
|
||||
SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
|
||||
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL, SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL,
|
||||
SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL, SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL,
|
||||
SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL, SELECT_STATS_SUMMARY_TOTALS_SQL,
|
||||
SELECT_STATS_USER_DAILY_AGGREGATES_SQL, UPDATE_STATS_SUMMARY_SQL,
|
||||
UPSERT_STATS_DAILY_API_KEY_SQL, UPSERT_STATS_DAILY_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL,
|
||||
UPSERT_STATS_DAILY_SQL, UPSERT_STATS_USER_DAILY_SQL,
|
||||
INSERT_STATS_SUMMARY_SQL, SELECT_EXISTING_STATS_SUMMARY_ID_SQL,
|
||||
SELECT_STATS_DAILY_AGGREGATE_SQL, SELECT_STATS_DAILY_FALLBACK_COUNT_SQL,
|
||||
SELECT_STATS_DAILY_FIRST_BYTE_PERCENTILES_SQL,
|
||||
SELECT_STATS_DAILY_RESPONSE_TIME_PERCENTILES_SQL, SELECT_STATS_SUMMARY_ENTITY_COUNTS_SQL,
|
||||
SELECT_STATS_SUMMARY_TOTALS_SQL, UPDATE_STATS_SUMMARY_SQL, UPSERT_STATS_DAILY_API_KEY_SQL,
|
||||
UPSERT_STATS_DAILY_MODEL_SQL, UPSERT_STATS_DAILY_PROVIDER_SQL, UPSERT_STATS_DAILY_SQL,
|
||||
UPSERT_STATS_USER_DAILY_SQL,
|
||||
};
|
||||
|
||||
pub(super) async fn perform_stats_aggregation_once(
|
||||
@@ -226,31 +222,15 @@ async fn upsert_stats_daily_model_rows(
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows = sqlx::query(SELECT_STATS_DAILY_MODEL_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
sqlx::query(UPSERT_STATS_DAILY_MODEL_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(day_start_utc)
|
||||
.bind(row.try_get::<String, _>("model")?)
|
||||
.bind(row.try_get::<i64, _>("total_requests")?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
|
||||
.bind(row.try_get::<f64, _>("total_cost")?)
|
||||
.bind(row.try_get::<f64, _>("avg_response_time_ms")?)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(rows.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_provider_rows(
|
||||
@@ -259,30 +239,15 @@ async fn upsert_stats_daily_provider_rows(
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows = sqlx::query(SELECT_STATS_DAILY_PROVIDER_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
sqlx::query(UPSERT_STATS_DAILY_PROVIDER_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(day_start_utc)
|
||||
.bind(row.try_get::<String, _>("provider_name")?)
|
||||
.bind(row.try_get::<i64, _>("total_requests")?)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
|
||||
.bind(row.try_get::<f64, _>("total_cost")?)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(rows.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_daily_api_key_rows(
|
||||
@@ -291,35 +256,15 @@ async fn upsert_stats_daily_api_key_rows(
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let rows = sqlx::query(SELECT_STATS_DAILY_API_KEY_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
let total_requests = row.try_get::<i64, _>("total_requests")?;
|
||||
let error_requests = row.try_get::<i64, _>("error_requests")?;
|
||||
sqlx::query(UPSERT_STATS_DAILY_API_KEY_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(row.try_get::<String, _>("api_key_id")?)
|
||||
.bind(row.try_get::<Option<String>, _>("api_key_name")?)
|
||||
.bind(day_start_utc)
|
||||
.bind(total_requests)
|
||||
.bind(total_requests.saturating_sub(error_requests))
|
||||
.bind(error_requests)
|
||||
.bind(row.try_get::<i64, _>("input_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("output_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_creation_tokens")?)
|
||||
.bind(row.try_get::<i64, _>("cache_read_tokens")?)
|
||||
.bind(row.try_get::<f64, _>("total_cost")?)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(rows.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn refresh_stats_daily_error_rows(
|
||||
@@ -332,27 +277,15 @@ async fn refresh_stats_daily_error_rows(
|
||||
.bind(day_start_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
let rows = sqlx::query(SELECT_STATS_DAILY_ERROR_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
sqlx::query(INSERT_STATS_DAILY_ERROR_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(day_start_utc)
|
||||
.bind(row.try_get::<String, _>("error_category")?)
|
||||
.bind(row.try_get::<Option<String>, _>("provider_name")?)
|
||||
.bind(row.try_get::<Option<String>, _>("model")?)
|
||||
.bind(row.try_get::<i64, _>("total_count")?)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(rows.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_user_daily_rows(
|
||||
@@ -361,87 +294,15 @@ async fn upsert_stats_user_daily_rows(
|
||||
day_end_utc: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, sqlx::Error> {
|
||||
let active_user_ids = sqlx::query(SELECT_ACTIVE_USER_IDS_SQL)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| row.try_get::<String, _>("id"))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
if active_user_ids.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let aggregated_rows = sqlx::query(SELECT_STATS_USER_DAILY_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
|
||||
.bind(day_start_utc)
|
||||
.bind(day_end_utc)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
let mut aggregated_by_user = HashMap::with_capacity(aggregated_rows.len());
|
||||
for row in aggregated_rows {
|
||||
let user_id = row.try_get::<String, _>("user_id")?;
|
||||
aggregated_by_user.insert(user_id, row);
|
||||
}
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
for user_id in &active_user_ids {
|
||||
let aggregated = aggregated_by_user.get(user_id);
|
||||
let total_requests = aggregated
|
||||
.map(|row| row.try_get::<i64, _>("total_requests"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
let error_requests = aggregated
|
||||
.map(|row| row.try_get::<i64, _>("error_requests"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
sqlx::query(UPSERT_STATS_USER_DAILY_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(user_id)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<Option<String>, _>("username"))
|
||||
.transpose()?
|
||||
.flatten(),
|
||||
)
|
||||
.bind(day_start_utc)
|
||||
.bind(total_requests)
|
||||
.bind(total_requests.saturating_sub(error_requests))
|
||||
.bind(error_requests)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<i64, _>("input_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<i64, _>("output_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<i64, _>("cache_creation_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<i64, _>("cache_read_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.bind(
|
||||
aggregated
|
||||
.map(|row| row.try_get::<f64, _>("total_cost"))
|
||||
.transpose()?
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(active_user_ids.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn refresh_stats_summary_row(
|
||||
|
||||
@@ -7,9 +7,8 @@ use aether_data_contracts::DataLayerError;
|
||||
|
||||
use super::{
|
||||
stats_hourly_aggregation_target_hour, system_config_bool, SELECT_STATS_HOURLY_AGGREGATE_SQL,
|
||||
SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL, SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL,
|
||||
SELECT_STATS_HOURLY_USER_AGGREGATES_SQL, UPSERT_STATS_HOURLY_MODEL_SQL,
|
||||
UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL, UPSERT_STATS_HOURLY_USER_SQL,
|
||||
UPSERT_STATS_HOURLY_MODEL_SQL, UPSERT_STATS_HOURLY_PROVIDER_SQL, UPSERT_STATS_HOURLY_SQL,
|
||||
UPSERT_STATS_HOURLY_USER_SQL,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
@@ -115,51 +114,16 @@ async fn upsert_stats_hourly_user_rows(
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let rows = sqlx::query(SELECT_STATS_HOURLY_USER_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.fetch_all(&mut **tx)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
let user_id = row
|
||||
.try_get::<String, _>("user_id")
|
||||
.map_err(postgres_error)?;
|
||||
let total_requests = row
|
||||
.try_get::<i64, _>("total_requests")
|
||||
.map_err(postgres_error)?;
|
||||
let error_requests = row
|
||||
.try_get::<i64, _>("error_requests")
|
||||
.map_err(postgres_error)?;
|
||||
let success_requests = total_requests.saturating_sub(error_requests);
|
||||
sqlx::query(UPSERT_STATS_HOURLY_USER_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(hour_utc)
|
||||
.bind(user_id)
|
||||
.bind(total_requests)
|
||||
.bind(success_requests)
|
||||
.bind(error_requests)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("input_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("output_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("total_cost")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
}
|
||||
|
||||
Ok(rows.len())
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_model_rows(
|
||||
@@ -168,54 +132,16 @@ async fn upsert_stats_hourly_model_rows(
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let rows = sqlx::query(SELECT_STATS_HOURLY_MODEL_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.fetch_all(&mut **tx)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let mut inserted = 0usize;
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
let model = row
|
||||
.try_get::<Option<String>, _>("model")
|
||||
.map_err(postgres_error)?;
|
||||
let Some(model) = model.filter(|value| !value.is_empty()) else {
|
||||
continue;
|
||||
};
|
||||
sqlx::query(UPSERT_STATS_HOURLY_MODEL_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(hour_utc)
|
||||
.bind(model)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("total_requests")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("input_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("output_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("total_cost")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("avg_response_time_ms")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
inserted += 1;
|
||||
}
|
||||
|
||||
Ok(inserted)
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
async fn upsert_stats_hourly_provider_rows(
|
||||
@@ -224,50 +150,16 @@ async fn upsert_stats_hourly_provider_rows(
|
||||
hour_end: DateTime<Utc>,
|
||||
now_utc: DateTime<Utc>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let rows = sqlx::query(SELECT_STATS_HOURLY_PROVIDER_AGGREGATES_SQL)
|
||||
let rows_affected = sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
|
||||
.bind(hour_utc)
|
||||
.bind(hour_end)
|
||||
.fetch_all(&mut **tx)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let mut inserted = 0usize;
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
|
||||
for row in &rows {
|
||||
let provider_name = row
|
||||
.try_get::<Option<String>, _>("provider_name")
|
||||
.map_err(postgres_error)?;
|
||||
let Some(provider_name) = provider_name.filter(|value| !value.is_empty()) else {
|
||||
continue;
|
||||
};
|
||||
sqlx::query(UPSERT_STATS_HOURLY_PROVIDER_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(hour_utc)
|
||||
.bind(provider_name)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("total_requests")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("input_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("output_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<f64, _>("total_cost")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
inserted += 1;
|
||||
}
|
||||
|
||||
Ok(inserted)
|
||||
Ok(usize::try_from(rows_affected).unwrap_or(usize::MAX))
|
||||
}
|
||||
|
||||
fn postgres_error(error: sqlx::Error) -> DataLayerError {
|
||||
|
||||
@@ -29,7 +29,7 @@ use super::{
|
||||
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
||||
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
||||
UsageCleanupSettings, DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL,
|
||||
SELECT_STALE_PENDING_USAGE_BATCH_SQL, SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL,
|
||||
SELECT_STALE_PENDING_USAGE_BATCH_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
|
||||
UPDATE_FAILED_VOID_STALE_USAGE_SQL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
};
|
||||
@@ -82,9 +82,9 @@ async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
|
||||
#[test]
|
||||
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
|
||||
assert!(
|
||||
SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL.contains("JOIN usage_settlement_snapshots")
|
||||
UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL.contains("JOIN usage_settlement_snapshots")
|
||||
);
|
||||
assert!(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL
|
||||
assert!(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL
|
||||
.contains("usage_settlement_snapshots.wallet_id"));
|
||||
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL.contains("JOIN usage_settlement_snapshots"));
|
||||
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL
|
||||
|
||||
@@ -6,6 +6,7 @@ use aether_data_contracts::repository::usage::{
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use chrono::{DateTime, Utc};
|
||||
use flate2::{write::GzEncoder, Compression};
|
||||
use futures_util::TryStreamExt;
|
||||
use serde_json::{Map, Value};
|
||||
use sqlx::Row;
|
||||
use tracing::warn;
|
||||
@@ -21,7 +22,8 @@ use super::{
|
||||
DISABLE_EXPIRED_API_KEY_SQL, EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE,
|
||||
NULLIFY_REQUEST_CANDIDATE_API_KEY_BATCH_SQL, NULLIFY_USAGE_API_KEY_BATCH_SQL,
|
||||
SELECT_EXPIRED_ACTIVE_API_KEYS_SQL, SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL,
|
||||
SELECT_USAGE_HEADER_BATCH_SQL, SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL,
|
||||
SELECT_USAGE_BODY_COMPRESSION_ROW_SQL, SELECT_USAGE_HEADER_BATCH_SQL,
|
||||
SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL,
|
||||
SELECT_USAGE_STALE_BODY_BATCH_SQL, UPDATE_USAGE_BODY_COMPRESSION_SQL,
|
||||
UPDATE_USAGE_REQUEST_METADATA_SQL, UPSERT_USAGE_BODY_BLOB_SQL,
|
||||
UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
|
||||
@@ -211,23 +213,20 @@ async fn cleanup_usage_header_fields(
|
||||
|
||||
let mut total_cleaned = 0usize;
|
||||
loop {
|
||||
let rows = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
||||
let mut stream = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
.fetch(pool);
|
||||
let mut rows = Vec::new();
|
||||
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||
rows.push(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
});
|
||||
}
|
||||
if rows.is_empty() {
|
||||
break;
|
||||
}
|
||||
@@ -279,23 +278,20 @@ async fn cleanup_usage_stale_body_fields(
|
||||
|
||||
let mut total_cleaned = 0usize;
|
||||
loop {
|
||||
let rows = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
|
||||
let mut stream = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
.fetch(pool);
|
||||
let mut rows = Vec::new();
|
||||
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||
rows.push(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
});
|
||||
}
|
||||
if rows.is_empty() {
|
||||
break;
|
||||
}
|
||||
@@ -354,53 +350,59 @@ async fn compress_usage_body_fields(
|
||||
let mut no_progress_count = 0usize;
|
||||
let batch_size = batch_size.clamp(1, 25);
|
||||
loop {
|
||||
let rows = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
|
||||
let mut stream = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCompressionRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
request_body: row
|
||||
.try_get::<Option<Value>, _>("request_body")
|
||||
.map_err(postgres_error)?,
|
||||
request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
response_body: row
|
||||
.try_get::<Option<Value>, _>("response_body")
|
||||
.map_err(postgres_error)?,
|
||||
response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body: row
|
||||
.try_get::<Option<Value>, _>("provider_request_body")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body: row
|
||||
.try_get::<Option<Value>, _>("client_response_body")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("client_response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
if rows.is_empty() {
|
||||
.fetch(pool);
|
||||
let mut ids = Vec::new();
|
||||
while let Some(row) = stream.try_next().await.map_err(postgres_error)? {
|
||||
ids.push(row.try_get::<String, _>("id").map_err(postgres_error)?);
|
||||
}
|
||||
if ids.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut batch_success = 0usize;
|
||||
for row in rows {
|
||||
for id in ids {
|
||||
let row = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_ROW_SQL)
|
||||
.bind(&id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let Some(row) = row else {
|
||||
continue;
|
||||
};
|
||||
let row = UsageBodyCompressionRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
request_body: row
|
||||
.try_get::<Option<Value>, _>("request_body")
|
||||
.map_err(postgres_error)?,
|
||||
request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
response_body: row
|
||||
.try_get::<Option<Value>, _>("response_body")
|
||||
.map_err(postgres_error)?,
|
||||
response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body: row
|
||||
.try_get::<Option<Value>, _>("provider_request_body")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body: row
|
||||
.try_get::<Option<Value>, _>("client_response_body")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("client_response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
};
|
||||
let detached = build_usage_body_externalization(&row)?;
|
||||
if detached.refs.any_present() {
|
||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||
@@ -638,12 +640,9 @@ async fn cleanup_expired_api_keys(
|
||||
pool: &aether_data::postgres::PostgresPool,
|
||||
auto_delete_expired_keys: bool,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let mut expired_keys = sqlx::query(SELECT_EXPIRED_ACTIVE_API_KEYS_SQL).fetch(pool);
|
||||
let mut cleaned = 0usize;
|
||||
for row in &expired_keys {
|
||||
while let Some(row) = expired_keys.try_next().await.map_err(postgres_error)? {
|
||||
let api_key_id = row.try_get::<String, _>("id").map_err(postgres_error)?;
|
||||
let key = ExpiredApiKeyRow {
|
||||
id: api_key_id.as_str(),
|
||||
|
||||
@@ -1,14 +1,11 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
use super::{
|
||||
maintenance_timezone, wallet_daily_usage_aggregation_target,
|
||||
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL,
|
||||
UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
|
||||
DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL, UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -43,60 +40,16 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
|
||||
};
|
||||
|
||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||
let rows = sqlx::query(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL)
|
||||
let aggregated_wallets = sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
|
||||
.bind(target.window_start_utc)
|
||||
.bind(target.window_end_utc)
|
||||
.fetch_all(&mut *tx)
|
||||
.bind(target.billing_date)
|
||||
.bind(target.billing_timezone.as_str())
|
||||
.bind(now_utc)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
for row in &rows {
|
||||
sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(
|
||||
row.try_get::<String, _>("wallet_id")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(target.billing_date)
|
||||
.bind(target.billing_timezone.as_str())
|
||||
.bind(
|
||||
row.try_get::<f64, _>("total_cost_usd")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("total_requests")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("input_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("output_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_creation_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<i64, _>("cache_read_tokens")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<Option<DateTime<Utc>>, _>("first_finalized_at")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(
|
||||
row.try_get::<Option<DateTime<Utc>>, _>("last_finalized_at")
|
||||
.map_err(postgres_error)?,
|
||||
)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.bind(now_utc)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
}
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
|
||||
let deleted_stale_ledgers = sqlx::query(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL)
|
||||
.bind(target.billing_date)
|
||||
@@ -112,7 +65,7 @@ pub(super) async fn perform_wallet_daily_usage_aggregation_once(
|
||||
Ok(WalletDailyUsageAggregationSummary {
|
||||
billing_date: target.billing_date,
|
||||
billing_timezone: target.billing_timezone,
|
||||
aggregated_wallets: rows.len(),
|
||||
aggregated_wallets: usize::try_from(aggregated_wallets).unwrap_or(usize::MAX),
|
||||
deleted_stale_ledgers: usize::try_from(deleted_stale_ledgers).unwrap_or(usize::MAX),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ use crate::{
|
||||
AdminBillingRuleWriteInput, GatewayError, LocalMutationOutcome,
|
||||
};
|
||||
use aether_data::postgres::PostgresPool;
|
||||
use futures_util::TryStreamExt;
|
||||
use sqlx::Row;
|
||||
|
||||
fn internal(err: impl ToString) -> GatewayError {
|
||||
@@ -136,7 +137,7 @@ WHERE ($1::TEXT IS NULL OR task_type = $1)
|
||||
)?;
|
||||
|
||||
let offset = u64::from(page.saturating_sub(1) * page_size);
|
||||
let rows = sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
@@ -162,16 +163,13 @@ LIMIT $4
|
||||
.bind(is_enabled)
|
||||
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
||||
.bind(i64::from(page_size))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(internal)?;
|
||||
.fetch(pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_err(internal)? {
|
||||
items.push(admin_billing_rule_from_row(&row)?);
|
||||
}
|
||||
|
||||
Ok((
|
||||
rows.iter()
|
||||
.map(admin_billing_rule_from_row)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
total,
|
||||
))
|
||||
Ok((items, total))
|
||||
}
|
||||
|
||||
pub(crate) async fn find_admin_billing_rule(
|
||||
@@ -374,7 +372,7 @@ WHERE ($1::TEXT IS NULL OR api_format = $1)
|
||||
)?;
|
||||
|
||||
let offset = u64::from(page.saturating_sub(1) * page_size);
|
||||
let rows = sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
@@ -406,16 +404,13 @@ LIMIT $6
|
||||
.bind(is_enabled)
|
||||
.bind(i64::try_from(offset).map_err(|err| GatewayError::Internal(err.to_string()))?)
|
||||
.bind(i64::from(page_size))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(internal)?;
|
||||
.fetch(pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_err(internal)? {
|
||||
items.push(admin_billing_collector_from_row(&row)?);
|
||||
}
|
||||
|
||||
Ok((
|
||||
rows.iter()
|
||||
.map(admin_billing_collector_from_row)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
total,
|
||||
))
|
||||
Ok((items, total))
|
||||
}
|
||||
|
||||
pub(crate) async fn find_admin_billing_collector(
|
||||
|
||||
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
||||
|
||||
use aether_data::postgres::PostgresPool;
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures_util::TryStreamExt;
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::Row;
|
||||
|
||||
@@ -36,7 +37,7 @@ WHERE a.created_at >= $1
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("admin audit logs count failed: {err}")))?;
|
||||
|
||||
let rows = sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
a.id,
|
||||
@@ -64,21 +65,24 @@ LIMIT $4 OFFSET $5
|
||||
.bind(event_type)
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("admin audit logs read failed: {err}")))?;
|
||||
.fetch(pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.try_next()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("admin audit logs read failed: {err}")))?
|
||||
{
|
||||
items.push(admin_audit_log_row_to_json(row));
|
||||
}
|
||||
|
||||
Ok((
|
||||
rows.into_iter().map(admin_audit_log_row_to_json).collect(),
|
||||
usize::try_from(total.max(0)).unwrap_or(usize::MAX),
|
||||
))
|
||||
Ok((items, usize::try_from(total.max(0)).unwrap_or(usize::MAX)))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_suspicious_activities(
|
||||
pool: &PostgresPool,
|
||||
cutoff_time: DateTime<Utc>,
|
||||
) -> Result<Vec<Value>, GatewayError> {
|
||||
let rows = sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
@@ -102,13 +106,15 @@ LIMIT 100
|
||||
"login_failed",
|
||||
"request_rate_limited",
|
||||
])
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
.fetch(pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_err(|err| {
|
||||
GatewayError::Internal(format!("admin suspicious activities read failed: {err}"))
|
||||
})?;
|
||||
})? {
|
||||
items.push(admin_suspicious_row_to_json(row));
|
||||
}
|
||||
|
||||
Ok(rows.into_iter().map(admin_suspicious_row_to_json).collect())
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
pub(crate) async fn read_admin_user_behavior_event_counts(
|
||||
@@ -116,7 +122,7 @@ pub(crate) async fn read_admin_user_behavior_event_counts(
|
||||
user_id: &str,
|
||||
cutoff_time: DateTime<Utc>,
|
||||
) -> Result<BTreeMap<String, u64>, GatewayError> {
|
||||
let rows = sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT event_type, COUNT(*)::bigint AS count
|
||||
FROM audit_logs
|
||||
@@ -127,22 +133,25 @@ GROUP BY event_type
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(cutoff_time)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("admin user behavior read failed: {err}")))?;
|
||||
.fetch(pool);
|
||||
let mut counts = BTreeMap::new();
|
||||
while let Some(row) = rows
|
||||
.try_next()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("admin user behavior read failed: {err}")))?
|
||||
{
|
||||
let Ok(event_type) = row.try_get::<String, _>("event_type") else {
|
||||
continue;
|
||||
};
|
||||
let count = row
|
||||
.try_get::<i64, _>("count")
|
||||
.ok()
|
||||
.and_then(|value| u64::try_from(value.max(0)).ok())
|
||||
.unwrap_or(0);
|
||||
counts.insert(event_type, count);
|
||||
}
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.filter_map(|row| {
|
||||
let event_type = row.try_get::<String, _>("event_type").ok()?;
|
||||
let count = row
|
||||
.try_get::<i64, _>("count")
|
||||
.ok()
|
||||
.and_then(|value| u64::try_from(value.max(0)).ok())
|
||||
.unwrap_or(0);
|
||||
Some((event_type, count))
|
||||
})
|
||||
.collect())
|
||||
Ok(counts)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_audit_logs(
|
||||
@@ -176,7 +185,7 @@ WHERE user_id = $1
|
||||
}
|
||||
};
|
||||
|
||||
let rows = match sqlx::query(
|
||||
let mut rows = sqlx::query(
|
||||
r#"
|
||||
SELECT id, event_type, description, ip_address, status_code, created_at
|
||||
FROM audit_logs
|
||||
@@ -192,21 +201,17 @@ LIMIT $4 OFFSET $5
|
||||
.bind(event_type)
|
||||
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
||||
.bind(i64::try_from(offset).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.fetch(pool);
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.try_next()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("user audit logs read failed: {err}")))?
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"user audit logs read failed: {err}"
|
||||
)))
|
||||
}
|
||||
};
|
||||
items.push(user_audit_log_row_to_json(row));
|
||||
}
|
||||
|
||||
Ok((
|
||||
rows.into_iter().map(user_audit_log_row_to_json).collect(),
|
||||
total,
|
||||
))
|
||||
Ok((items, total))
|
||||
}
|
||||
|
||||
fn admin_audit_log_row_to_json(row: sqlx::postgres::PgRow) -> Value {
|
||||
|
||||
@@ -416,6 +416,16 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
|
||||
self.data
|
||||
.list_provider_catalog_key_summaries_by_provider_ids(provider_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
|
||||
@@ -120,6 +120,18 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_video_task_page_summary(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, GatewayError> {
|
||||
self.data
|
||||
.list_video_task_page_summary(filter, offset, limit)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn count_video_tasks(
|
||||
&self,
|
||||
filter: &VideoTaskQueryFilter,
|
||||
|
||||
Reference in New Issue
Block a user